fluidB
/
fluidb
1
0
Fork 0
Code Pull Requests Activity

Preliminary work to support https://github.com/cozodb/cozo/issues/109

Browse Source
main
Ziyang Hu 1 year ago
parent db29a65c68
commit bf325b76f5
6 changed files with 208 additions and 99 deletions
Show Stats Download Patch File Download Diff File

3
cozo-core/src/cozoscript.pest
Unescape Escape View File

@ -127,13 +127,14 @@ object_pair = {expr ~ ":" ~ expr}
list = { "[" ~ (expr ~ ",")* ~ expr? ~ "]" } list = { "[" ~ (expr ~ ",")* ~ expr? ~ "]" }
grouping = { "(" ~ expr ~ ")" } grouping = { "(" ~ expr ~ ")" }
option = _{(limit_option|offset_option|sort_option|relation_option|timeout_option|sleep_option| option = _{(limit_option|offset_option|sort_option|relation_option|timeout_option|sleep_option|returning_option|
assert_none_option|assert_some_option|disable_magic_rewrite_option) ~ ";"?} assert_none_option|assert_some_option|disable_magic_rewrite_option) ~ ";"?}
out_arg = @{var ~ ("(" ~ var ~ ")")?} out_arg = @{var ~ ("(" ~ var ~ ")")?}
disable_magic_rewrite_option = {":disable_magic_rewrite" ~ expr} disable_magic_rewrite_option = {":disable_magic_rewrite" ~ expr}
limit_option = {":limit" ~ expr} limit_option = {":limit" ~ expr}
offset_option = {":offset" ~ expr} offset_option = {":offset" ~ expr}
sort_option = {(":sort" | ":order") ~ (sort_arg ~ ",")* ~ sort_arg } sort_option = {(":sort" | ":order") ~ (sort_arg ~ ",")* ~ sort_arg }
returning_option = {":returning"}
relation_option = {relation_op ~ (compound_ident | underscore_ident) ~ table_schema?} relation_option = {relation_op ~ (compound_ident | underscore_ident) ~ table_schema?}
relation_op = _{relation_create | relation_replace | relation_insert | relation_put | relation_update | relation_rm | relation_delete | relation_ensure_not | relation_ensure } relation_op = _{relation_create | relation_replace | relation_insert | relation_put | relation_update | relation_rm | relation_delete | relation_ensure_not | relation_ensure }
relation_create = {":create"} relation_create = {":create"}

14
cozo-core/src/data/program.rs
Unescape Escape View File

@ -40,6 +40,12 @@ pub(crate) enum QueryAssertion {
AssertSome(SourceSpan), AssertSome(SourceSpan),
} }
#[derive(Debug, Copy, Clone, Eq, PartialEq)]
pub(crate) enum ReturnMutation {
NotReturning,
Returning,
}
#[derive(Clone, PartialEq, Default)] #[derive(Clone, PartialEq, Default)]
pub(crate) struct QueryOutOptions { pub(crate) struct QueryOutOptions {
pub(crate) limit: Option<usize>, pub(crate) limit: Option<usize>,
@ -47,7 +53,7 @@ pub(crate) struct QueryOutOptions {
pub(crate) timeout: Option<f64>, pub(crate) timeout: Option<f64>,
pub(crate) sleep: Option<f64>, pub(crate) sleep: Option<f64>,
pub(crate) sorters: Vec<(Symbol, SortDir)>, pub(crate) sorters: Vec<(Symbol, SortDir)>,
pub(crate) store_relation: Option<(InputRelationHandle, RelationOp)>, pub(crate) store_relation: Option<(InputRelationHandle, RelationOp, ReturnMutation)>,
pub(crate) assertion: Option<QueryAssertion>, pub(crate) assertion: Option<QueryAssertion>,
} }
@ -84,8 +90,12 @@ impl Display for QueryOutOptions {
.. ..
}, },
op, op,
return_mutation,
)) = &self.store_relation )) = &self.store_relation
{ {
if *return_mutation == ReturnMutation::Returning {
write!(f, ":returning\n")?;
}
match op { match op {
RelationOp::Create => { RelationOp::Create => {
write!(f, ":create ")?; write!(f, ":create ")?;
@ -546,7 +556,7 @@ pub(crate) struct NoEntryError;
impl InputProgram { impl InputProgram {
pub(crate) fn needs_write_lock(&self) -> Option<SmartString<LazyCompact>> { pub(crate) fn needs_write_lock(&self) -> Option<SmartString<LazyCompact>> {
if let Some((h, _)) = &self.out_opts.store_relation { if let Some((h, _, _)) = &self.out_opts.store_relation {
if !h.name.name.starts_with('_') { if !h.name.name.starts_with('_') {
Some(h.name.name.clone()) Some(h.name.name.clone())
} else { } else {

33
cozo-core/src/parse/query.rs
Unescape Escape View File

@ -22,11 +22,7 @@ use thiserror::Error;
use crate::data::aggr::{parse_aggr, Aggregation}; use crate::data::aggr::{parse_aggr, Aggregation};
use crate::data::expr::Expr; use crate::data::expr::Expr;
use crate::data::functions::{str2vld, MAX_VALIDITY_TS}; use crate::data::functions::{str2vld, MAX_VALIDITY_TS};
use crate::data::program::{ use crate::data::program::{FixedRuleApply, FixedRuleArg, InputAtom, InputInlineRule, InputInlineRulesOrFixed, InputNamedFieldRelationApplyAtom, InputProgram, InputRelationApplyAtom, InputRuleApplyAtom, QueryAssertion, QueryOutOptions, RelationOp, ReturnMutation, SearchInput, SortDir, Unification};
FixedRuleApply, FixedRuleArg, InputAtom, InputInlineRule, InputInlineRulesOrFixed,
InputNamedFieldRelationApplyAtom, InputProgram, InputRelationApplyAtom, InputRuleApplyAtom,
QueryAssertion, QueryOutOptions, RelationOp, SearchInput, SortDir, Unification,
};
use crate::data::relation::{ColType, ColumnDef, NullableColType, StoredRelationMetadata}; use crate::data::relation::{ColType, ColumnDef, NullableColType, StoredRelationMetadata};
use crate::data::symb::{Symbol, PROG_ENTRY}; use crate::data::symb::{Symbol, PROG_ENTRY};
use crate::data::value::{DataValue, ValidityTs}; use crate::data::value::{DataValue, ValidityTs};
@ -87,7 +83,7 @@ impl Diagnostic for MultipleRuleDefinitionError {
fn code<'a>(&'a self) -> Option<Box<dyn Display + 'a>> { fn code<'a>(&'a self) -> Option<Box<dyn Display + 'a>> {
Some(Box::new("parser::mult_rule_def")) Some(Box::new("parser::mult_rule_def"))
} }
fn labels(&self) -> Option<Box<dyn Iterator<Item = LabeledSpan> + '_>> { fn labels(&self) -> Option<Box<dyn Iterator<Item=LabeledSpan> + '_>> {
Some(Box::new( Some(Box::new(
self.1.iter().map(|s| LabeledSpan::new_with_span(None, s)), self.1.iter().map(|s| LabeledSpan::new_with_span(None, s)),
)) ))
@ -113,6 +109,7 @@ pub(crate) fn parse_query(
let mut disable_magic_rewrite = false; let mut disable_magic_rewrite = false;
let mut stored_relation = None; let mut stored_relation = None;
let mut returning_mutation = ReturnMutation::NotReturning;
for pair in src { for pair in src {
match pair.as_rule() { match pair.as_rule() {
@ -310,6 +307,9 @@ pub(crate) fn parse_query(
out_opts.sorters.push((Symbol::new(var, span), dir)); out_opts.sorters.push((Symbol::new(var, span), dir));
} }
} }
Rule::returning_option => {
returning_mutation = ReturnMutation::Returning;
}
Rule::relation_option => { Rule::relation_option => {
let span = pair.extract_span(); let span = pair.extract_span();
let mut args = pair.into_inner(); let mut args = pair.into_inner();
@ -389,13 +389,14 @@ pub(crate) fn parse_query(
if prog.prog.is_empty() { if prog.prog.is_empty() {
if let Some(( if let Some((
InputRelationHandle { InputRelationHandle {
key_bindings, key_bindings,
dep_bindings, dep_bindings,
.. ..
}, },
RelationOp::Create, RelationOp::Create,
)) = &prog.out_opts.store_relation _
)) = &prog.out_opts.store_relation
{ {
let mut bindings = key_bindings.clone(); let mut bindings = key_bindings.clone();
bindings.extend_from_slice(dep_bindings); bindings.extend_from_slice(dep_bindings);
@ -435,13 +436,13 @@ pub(crate) fn parse_query(
dep_bindings: vec![], dep_bindings: vec![],
span, span,
}; };
prog.out_opts.store_relation = Some((handle, op)) prog.out_opts.store_relation = Some((handle, op, returning_mutation))
} }
Some(Right(r)) => prog.out_opts.store_relation = Some(r), Some(Right((h, o))) => prog.out_opts.store_relation = Some((h, o, returning_mutation)),
} }
if prog.prog.is_empty() { if prog.prog.is_empty() {
if let Some((handle, RelationOp::Create)) = &prog.out_opts.store_relation { if let Some((handle, RelationOp::Create, _)) = &prog.out_opts.store_relation {
let mut bindings = handle.dep_bindings.clone(); let mut bindings = handle.dep_bindings.clone();
bindings.extend_from_slice(&handle.key_bindings); bindings.extend_from_slice(&handle.key_bindings);
make_empty_const_rule(&mut prog, &bindings); make_empty_const_rule(&mut prog, &bindings);

63
cozo-core/src/query/stored.rs
Unescape Escape View File

@ -44,7 +44,7 @@ impl<'a> SessionTx<'a> {
pub(crate) fn execute_relation<'s, S: Storage<'s>>( pub(crate) fn execute_relation<'s, S: Storage<'s>>(
&mut self, &mut self,
db: &Db<S>, db: &Db<S>,
res_iter: impl Iterator<Item = Tuple>, res_iter: impl Iterator<Item=Tuple>,
op: RelationOp, op: RelationOp,
meta: &InputRelationHandle, meta: &InputRelationHandle,
headers: &[Symbol], headers: &[Symbol],
@ -52,6 +52,7 @@ impl<'a> SessionTx<'a> {
callback_targets: &BTreeSet<SmartString<LazyCompact>>, callback_targets: &BTreeSet<SmartString<LazyCompact>>,
callback_collector: &mut CallbackCollector, callback_collector: &mut CallbackCollector,
propagate_triggers: bool, propagate_triggers: bool,
force_collect: &str,
) -> Result<Vec<(Vec<u8>, Vec<u8>)>> { ) -> Result<Vec<(Vec<u8>, Vec<u8>)>> {
let mut to_clear = vec![]; let mut to_clear = vec![];
let mut replaced_old_triggers = None; let mut replaced_old_triggers = None;
@ -88,7 +89,7 @@ impl<'a> SessionTx<'a> {
&db.fixed_rules.read().unwrap(), &db.fixed_rules.read().unwrap(),
cur_vld, cur_vld,
)? )?
.get_single_program()?; .get_single_program()?;
let (_, cleanups) = db let (_, cleanups) = db
.run_query( .run_query(
@ -141,17 +142,18 @@ impl<'a> SessionTx<'a> {
callback_collector, callback_collector,
propagate_triggers, propagate_triggers,
&mut to_clear, &mut to_clear,
&mut relation_store, &relation_store,
metadata, metadata,
key_bindings, key_bindings,
op == RelationOp::Delete, op == RelationOp::Delete,
force_collect,
*span, *span,
)?, )?,
RelationOp::Ensure => self.ensure_in_relation( RelationOp::Ensure => self.ensure_in_relation(
res_iter, res_iter,
headers, headers,
cur_vld, cur_vld,
&mut relation_store, &relation_store,
metadata, metadata,
key_bindings, key_bindings,
*span, *span,
@ -160,7 +162,7 @@ impl<'a> SessionTx<'a> {
res_iter, res_iter,
headers, headers,
cur_vld, cur_vld,
&mut relation_store, &relation_store,
metadata, metadata,
key_bindings, key_bindings,
*span, *span,
@ -174,9 +176,10 @@ impl<'a> SessionTx<'a> {
callback_collector, callback_collector,
propagate_triggers, propagate_triggers,
&mut to_clear, &mut to_clear,
&mut relation_store, &relation_store,
metadata, metadata,
key_bindings, key_bindings,
force_collect,
*span, *span,
)?, )?,
RelationOp::Create | RelationOp::Replace | RelationOp::Put | RelationOp::Insert => self RelationOp::Create | RelationOp::Replace | RelationOp::Put | RelationOp::Insert => self
@ -189,11 +192,12 @@ impl<'a> SessionTx<'a> {
callback_collector, callback_collector,
propagate_triggers, propagate_triggers,
&mut to_clear, &mut to_clear,
&mut relation_store, &relation_store,
metadata, metadata,
key_bindings, key_bindings,
dep_bindings, dep_bindings,
op == RelationOp::Insert, op == RelationOp::Insert,
force_collect,
*span, *span,
)?, )?,
}; };
@ -204,21 +208,22 @@ impl<'a> SessionTx<'a> {
fn put_into_relation<'s, S: Storage<'s>>( fn put_into_relation<'s, S: Storage<'s>>(
&mut self, &mut self,
db: &Db<S>, db: &Db<S>,
res_iter: impl Iterator<Item = Tuple>, res_iter: impl Iterator<Item=Tuple>,
headers: &[Symbol], headers: &[Symbol],
cur_vld: ValidityTs, cur_vld: ValidityTs,
callback_targets: &BTreeSet<SmartString<LazyCompact>>, callback_targets: &BTreeSet<SmartString<LazyCompact>>,
callback_collector: &mut CallbackCollector, callback_collector: &mut CallbackCollector,
propagate_triggers: bool, propagate_triggers: bool,
to_clear: &mut Vec<(Vec<u8>, Vec<u8>)>, to_clear: &mut Vec<(Vec<u8>, Vec<u8>)>,
relation_store: &mut RelationHandle, relation_store: &RelationHandle,
metadata: &StoredRelationMetadata, metadata: &StoredRelationMetadata,
key_bindings: &[Symbol], key_bindings: &[Symbol],
dep_bindings: &[Symbol], dep_bindings: &[Symbol],
is_insert: bool, is_insert: bool,
force_collect: &str,
span: SourceSpan, span: SourceSpan,
) -> Result<()> { ) -> Result<()> {
let is_callback_target = callback_targets.contains(&relation_store.name); let is_callback_target = callback_targets.contains(&relation_store.name) || force_collect == &relation_store.name;
if relation_store.access_level < AccessLevel::Protected { if relation_store.access_level < AccessLevel::Protected {
bail!(InsufficientAccessLevel( bail!(InsufficientAccessLevel(
@ -235,9 +240,9 @@ impl<'a> SessionTx<'a> {
headers, headers,
)?; )?;
let need_to_collect = !relation_store.is_temp let need_to_collect = !force_collect.is_empty() || (!relation_store.is_temp
&& (is_callback_target && (is_callback_target
|| (propagate_triggers && !relation_store.put_triggers.is_empty())); || (propagate_triggers && !relation_store.put_triggers.is_empty())));
let has_indices = !relation_store.indices.is_empty(); let has_indices = !relation_store.indices.is_empty();
let has_hnsw_indices = !relation_store.hnsw_indices.is_empty(); let has_hnsw_indices = !relation_store.hnsw_indices.is_empty();
let has_fts_indices = !relation_store.fts_indices.is_empty(); let has_fts_indices = !relation_store.fts_indices.is_empty();
@ -517,19 +522,20 @@ impl<'a> SessionTx<'a> {
fn update_in_relation<'s, S: Storage<'s>>( fn update_in_relation<'s, S: Storage<'s>>(
&mut self, &mut self,
db: &Db<S>, db: &Db<S>,
res_iter: impl Iterator<Item = Tuple>, res_iter: impl Iterator<Item=Tuple>,
headers: &[Symbol], headers: &[Symbol],
cur_vld: ValidityTs, cur_vld: ValidityTs,
callback_targets: &BTreeSet<SmartString<LazyCompact>>, callback_targets: &BTreeSet<SmartString<LazyCompact>>,
callback_collector: &mut CallbackCollector, callback_collector: &mut CallbackCollector,
propagate_triggers: bool, propagate_triggers: bool,
to_clear: &mut Vec<(Vec<u8>, Vec<u8>)>, to_clear: &mut Vec<(Vec<u8>, Vec<u8>)>,
relation_store: &mut RelationHandle, relation_store: &RelationHandle,
metadata: &StoredRelationMetadata, metadata: &StoredRelationMetadata,
key_bindings: &[Symbol], key_bindings: &[Symbol],
force_collect: &str,
span: SourceSpan, span: SourceSpan,
) -> Result<()> { ) -> Result<()> {
let is_callback_target = callback_targets.contains(&relation_store.name); let is_callback_target = callback_targets.contains(&relation_store.name) || force_collect == &relation_store.name;
if relation_store.access_level < AccessLevel::Protected { if relation_store.access_level < AccessLevel::Protected {
bail!(InsufficientAccessLevel( bail!(InsufficientAccessLevel(
@ -546,9 +552,9 @@ impl<'a> SessionTx<'a> {
headers, headers,
)?; )?;
let need_to_collect = !relation_store.is_temp let need_to_collect = !force_collect.is_empty() || (!relation_store.is_temp
&& (is_callback_target && (is_callback_target
|| (propagate_triggers && !relation_store.put_triggers.is_empty())); || (propagate_triggers && !relation_store.put_triggers.is_empty())));
let has_indices = !relation_store.indices.is_empty(); let has_indices = !relation_store.indices.is_empty();
let has_hnsw_indices = !relation_store.hnsw_indices.is_empty(); let has_hnsw_indices = !relation_store.hnsw_indices.is_empty();
let has_fts_indices = !relation_store.fts_indices.is_empty(); let has_fts_indices = !relation_store.fts_indices.is_empty();
@ -695,7 +701,7 @@ impl<'a> SessionTx<'a> {
&db.fixed_rules.read().unwrap(), &db.fixed_rules.read().unwrap(),
cur_vld, cur_vld,
)? )?
.get_single_program()?; .get_single_program()?;
make_const_rule( make_const_rule(
&mut program, &mut program,
@ -791,10 +797,10 @@ impl<'a> SessionTx<'a> {
fn ensure_not_in_relation( fn ensure_not_in_relation(
&mut self, &mut self,
res_iter: impl Iterator<Item = Tuple>, res_iter: impl Iterator<Item=Tuple>,
headers: &[Symbol], headers: &[Symbol],
cur_vld: ValidityTs, cur_vld: ValidityTs,
relation_store: &mut RelationHandle, relation_store: &RelationHandle,
metadata: &StoredRelationMetadata, metadata: &StoredRelationMetadata,
key_bindings: &[Symbol], key_bindings: &[Symbol],
span: SourceSpan, span: SourceSpan,
@ -838,10 +844,10 @@ impl<'a> SessionTx<'a> {
fn ensure_in_relation( fn ensure_in_relation(
&mut self, &mut self,
res_iter: impl Iterator<Item = Tuple>, res_iter: impl Iterator<Item=Tuple>,
headers: &[Symbol], headers: &[Symbol],
cur_vld: ValidityTs, cur_vld: ValidityTs,
relation_store: &mut RelationHandle, relation_store: &RelationHandle,
metadata: &StoredRelationMetadata, metadata: &StoredRelationMetadata,
key_bindings: &[Symbol], key_bindings: &[Symbol],
span: SourceSpan, span: SourceSpan,
@ -908,20 +914,21 @@ impl<'a> SessionTx<'a> {
fn remove_from_relation<'s, S: Storage<'s>>( fn remove_from_relation<'s, S: Storage<'s>>(
&mut self, &mut self,
db: &Db<S>, db: &Db<S>,
res_iter: impl Iterator<Item = Tuple>, res_iter: impl Iterator<Item=Tuple>,
headers: &[Symbol], headers: &[Symbol],
cur_vld: ValidityTs, cur_vld: ValidityTs,
callback_targets: &BTreeSet<SmartString<LazyCompact>>, callback_targets: &BTreeSet<SmartString<LazyCompact>>,
callback_collector: &mut CallbackCollector, callback_collector: &mut CallbackCollector,
propagate_triggers: bool, propagate_triggers: bool,
to_clear: &mut Vec<(Vec<u8>, Vec<u8>)>, to_clear: &mut Vec<(Vec<u8>, Vec<u8>)>,
relation_store: &mut RelationHandle, relation_store: &RelationHandle,
metadata: &StoredRelationMetadata, metadata: &StoredRelationMetadata,
key_bindings: &[Symbol], key_bindings: &[Symbol],
check_exists: bool, check_exists: bool,
force_collect: &str,
span: SourceSpan, span: SourceSpan,
) -> Result<()> { ) -> Result<()> {
let is_callback_target = callback_targets.contains(&relation_store.name); let is_callback_target = callback_targets.contains(&relation_store.name) || force_collect == relation_store.name;
if relation_store.access_level < AccessLevel::Protected { if relation_store.access_level < AccessLevel::Protected {
bail!(InsufficientAccessLevel( bail!(InsufficientAccessLevel(
@ -937,9 +944,9 @@ impl<'a> SessionTx<'a> {
headers, headers,
)?; )?;
let need_to_collect = !relation_store.is_temp let need_to_collect = !force_collect.is_empty() || (!relation_store.is_temp
&& (is_callback_target && (is_callback_target
|| (propagate_triggers && !relation_store.rm_triggers.is_empty())); || (propagate_triggers && !relation_store.rm_triggers.is_empty())));
let has_indices = !relation_store.indices.is_empty(); let has_indices = !relation_store.indices.is_empty();
let has_hnsw_indices = !relation_store.hnsw_indices.is_empty(); let has_hnsw_indices = !relation_store.hnsw_indices.is_empty();
let has_fts_indices = !relation_store.fts_indices.is_empty(); let has_fts_indices = !relation_store.fts_indices.is_empty();
@ -1028,7 +1035,7 @@ impl<'a> SessionTx<'a> {
&db.fixed_rules.read().unwrap(), &db.fixed_rules.read().unwrap(),
cur_vld, cur_vld,
)? )?
.get_single_program()?; .get_single_program()?;
make_const_rule(&mut program, "_new", k_bindings.clone(), new_tuples.clone()); make_const_rule(&mut program, "_new", k_bindings.clone(), new_tuples.clone());

156
cozo-core/src/runtime/db.rs
Unescape Escape View File

@ -34,7 +34,7 @@ use thiserror::Error;
use crate::data::functions::current_validity; use crate::data::functions::current_validity;
use crate::data::json::JsonValue; use crate::data::json::JsonValue;
use crate::data::program::{InputProgram, QueryAssertion, RelationOp}; use crate::data::program::{InputProgram, QueryAssertion, RelationOp, ReturnMutation};
use crate::data::relation::ColumnDef; use crate::data::relation::ColumnDef;
use crate::data::tuple::{Tuple, TupleT}; use crate::data::tuple::{Tuple, TupleT};
use crate::data::value::{DataValue, ValidityTs, LARGEST_UTF_CHAR}; use crate::data::value::{DataValue, ValidityTs, LARGEST_UTF_CHAR};
@ -51,9 +51,7 @@ use crate::query::ra::{
use crate::runtime::callback::{ use crate::runtime::callback::{
CallbackCollector, CallbackDeclaration, CallbackOp, EventCallbackRegistry, CallbackCollector, CallbackDeclaration, CallbackOp, EventCallbackRegistry,
}; };
use crate::runtime::relation::{ use crate::runtime::relation::{extend_tuple_from_v, AccessLevel, InsufficientAccessLevel, RelationHandle, RelationId, InputRelationHandle};
extend_tuple_from_v, AccessLevel, InsufficientAccessLevel, RelationHandle, RelationId,
};
use crate::runtime::transact::SessionTx; use crate::runtime::transact::SessionTx;
use crate::storage::temp::TempStorage; use crate::storage::temp::TempStorage;
use crate::storage::{Storage, StoreTx}; use crate::storage::{Storage, StoreTx};
@ -377,9 +375,9 @@ impl<'s, S: Storage<'s>> Db<S> {
/// ///
/// `relations` contains names of the stored relations to export. /// `relations` contains names of the stored relations to export.
pub fn export_relations<I, T>(&'s self, relations: I) -> Result<BTreeMap<String, NamedRows>> pub fn export_relations<I, T>(&'s self, relations: I) -> Result<BTreeMap<String, NamedRows>>
where where
T: AsRef<str>, T: AsRef<str>,
I: Iterator<Item = T>, I: Iterator<Item=T>,
{ {
let tx = self.transact()?; let tx = self.transact()?;
let mut ret: BTreeMap<String, NamedRows> = BTreeMap::new(); let mut ret: BTreeMap<String, NamedRows> = BTreeMap::new();
@ -689,8 +687,8 @@ impl<'s, S: Storage<'s>> Db<S> {
} }
/// Register a custom fixed rule implementation. /// Register a custom fixed rule implementation.
pub fn register_fixed_rule<R>(&self, name: String, rule_impl: R) -> Result<()> pub fn register_fixed_rule<R>(&self, name: String, rule_impl: R) -> Result<()>
where where
R: FixedRule + 'static, R: FixedRule + 'static,
{ {
match self.fixed_rules.write().unwrap().entry(name) { match self.fixed_rules.write().unwrap().entry(name) {
Entry::Vacant(ent) => { Entry::Vacant(ent) => {
@ -759,7 +757,7 @@ impl<'s, S: Storage<'s>> Db<S> {
ret.is_some() ret.is_some()
} }
pub(crate) fn obtain_relation_locks<'a, T: Iterator<Item = &'a SmartString<LazyCompact>>>( pub(crate) fn obtain_relation_locks<'a, T: Iterator<Item=&'a SmartString<LazyCompact>>>(
&'s self, &'s self,
rels: T, rels: T,
) -> Vec<Arc<ShardedLock<()>>> { ) -> Vec<Arc<ShardedLock<()>>> {
@ -831,7 +829,7 @@ impl<'s, S: Storage<'s>> Db<S> {
callback_collector: &mut CallbackCollector, callback_collector: &mut CallbackCollector,
) -> Result<NamedRows> { ) -> Result<NamedRows> {
#[allow(unused_variables)] #[allow(unused_variables)]
let sleep_opt = p.out_opts.sleep; let sleep_opt = p.out_opts.sleep;
let (q_res, q_cleanups) = let (q_res, q_cleanups) =
self.run_query(tx, p, cur_vld, callback_targets, callback_collector, true)?; self.run_query(tx, p, cur_vld, callback_targets, callback_collector, true)?;
cleanups.extend(q_cleanups); cleanups.extend(q_cleanups);
@ -970,28 +968,28 @@ impl<'s, S: Storage<'s>> Db<S> {
("fixed", json!(null), json!(null), json!(null)) ("fixed", json!(null), json!(null), json!(null))
} }
RelAlgebra::TempStore(TempStoreRA { RelAlgebra::TempStore(TempStoreRA {
storage_key, storage_key,
filters, filters,
.. ..
}) => ( }) => (
"load_mem", "load_mem",
json!(storage_key.to_string()), json!(storage_key.to_string()),
json!(null), json!(null),
json!(filters.iter().map(|f| f.to_string()).collect_vec()), json!(filters.iter().map(|f| f.to_string()).collect_vec()),
), ),
RelAlgebra::Stored(StoredRA { RelAlgebra::Stored(StoredRA {
storage, filters, .. storage, filters, ..
}) => ( }) => (
"load_stored", "load_stored",
json!(format!(":{}", storage.name)), json!(format!(":{}", storage.name)),
json!(null), json!(null),
json!(filters.iter().map(|f| f.to_string()).collect_vec()), json!(filters.iter().map(|f| f.to_string()).collect_vec()),
), ),
RelAlgebra::StoredWithValidity(StoredWithValidityRA { RelAlgebra::StoredWithValidity(StoredWithValidityRA {
storage, storage,
filters, filters,
.. ..
}) => ( }) => (
"load_stored_with_validity", "load_stored_with_validity",
json!(format!(":{}", storage.name)), json!(format!(":{}", storage.name)),
json!(null), json!(null),
@ -1030,10 +1028,10 @@ impl<'s, S: Storage<'s>> Db<S> {
("reorder", json!(null), json!(null), json!(null)) ("reorder", json!(null), json!(null), json!(null))
} }
RelAlgebra::Filter(FilteredRA { RelAlgebra::Filter(FilteredRA {
parent, parent,
filters: pred, filters: pred,
.. ..
}) => { }) => {
rel_stack.push(parent); rel_stack.push(parent);
( (
"filter", "filter",
@ -1043,12 +1041,12 @@ impl<'s, S: Storage<'s>> Db<S> {
) )
} }
RelAlgebra::Unification(UnificationRA { RelAlgebra::Unification(UnificationRA {
parent, parent,
binding, binding,
expr, expr,
is_multi, is_multi,
.. ..
}) => { }) => {
rel_stack.push(parent); rel_stack.push(parent);
( (
if *is_multi { "multi-unify" } else { "unify" }, if *is_multi { "multi-unify" } else { "unify" },
@ -1058,8 +1056,8 @@ impl<'s, S: Storage<'s>> Db<S> {
) )
} }
RelAlgebra::HnswSearch(HnswSearchRA { RelAlgebra::HnswSearch(HnswSearchRA {
hnsw_search, .. hnsw_search, ..
}) => ( }) => (
"hnsw_index", "hnsw_index",
json!(format!(":{}", hnsw_search.query.name)), json!(format!(":{}", hnsw_search.query.name)),
json!(hnsw_search.query.name), json!(hnsw_search.query.name),
@ -1353,7 +1351,7 @@ impl<'s, S: Storage<'s>> Db<S> {
let mut clean_ups = vec![]; let mut clean_ups = vec![];
// Some checks in case the query specifies mutation // Some checks in case the query specifies mutation
if let Some((meta, op)) = &input_program.out_opts.store_relation { if let Some((meta, op, _)) = &input_program.out_opts.store_relation {
if *op == RelationOp::Create { if *op == RelationOp::Create {
#[derive(Debug, Error, Diagnostic)] #[derive(Debug, Error, Diagnostic)]
#[error("Stored relation {0} conflicts with an existing one")] #[error("Stored relation {0} conflicts with an existing one")]
@ -1442,7 +1440,7 @@ impl<'s, S: Storage<'s>> Db<S> {
if let Some(tuple) = result_store.all_iter().next() { if let Some(tuple) = result_store.all_iter().next() {
#[derive(Debug, Error, Diagnostic)] #[derive(Debug, Error, Diagnostic)]
#[error( #[error(
"The query is asserted to return no result, but a tuple {0:?} is found" "The query is asserted to return no result, but a tuple {0:?} is found"
)] )]
#[diagnostic(code(eval::assert_none_failure))] #[diagnostic(code(eval::assert_none_failure))]
struct AssertNoneFailure(Tuple, #[label] SourceSpan); struct AssertNoneFailure(Tuple, #[label] SourceSpan);
@ -1475,7 +1473,7 @@ impl<'s, S: Storage<'s>> Db<S> {
} else { } else {
Right(sorted_iter) Right(sorted_iter)
}; };
if let Some((meta, relation_op)) = &out_opts.store_relation { if let Some((meta, relation_op, returning)) = &out_opts.store_relation {
let to_clear = tx let to_clear = tx
.execute_relation( .execute_relation(
self, self,
@ -1487,16 +1485,16 @@ impl<'s, S: Storage<'s>> Db<S> {
callback_targets, callback_targets,
callback_collector, callback_collector,
top_level, top_level,
if *returning == ReturnMutation::Returning {
&meta.name.name
} else {
""
},
) )
.wrap_err_with(|| format!("when executing against relation '{}'", meta.name))?; .wrap_err_with(|| format!("when executing against relation '{}'", meta.name))?;
clean_ups.extend(to_clear); clean_ups.extend(to_clear);
Ok(( let returned_rows = Self::get_returning_rows(callback_collector, meta, returning);
NamedRows::new( Ok((returned_rows, clean_ups))
vec![STATUS_STR.to_string()],
vec![vec![DataValue::from(OK_STR)]],
),
clean_ups,
))
} else { } else {
// not sorting outputs // not sorting outputs
let rows: Vec<Tuple> = sorted_iter.collect_vec(); let rows: Vec<Tuple> = sorted_iter.collect_vec();
@ -1530,7 +1528,7 @@ impl<'s, S: Storage<'s>> Db<S> {
Left(result_store.all_iter().map(|t| t.into_tuple())) Left(result_store.all_iter().map(|t| t.into_tuple()))
}; };
if let Some((meta, relation_op)) = &out_opts.store_relation { if let Some((meta, relation_op, returning)) = &out_opts.store_relation {
let to_clear = tx let to_clear = tx
.execute_relation( .execute_relation(
self, self,
@ -1542,16 +1540,17 @@ impl<'s, S: Storage<'s>> Db<S> {
callback_targets, callback_targets,
callback_collector, callback_collector,
top_level, top_level,
if *returning == ReturnMutation::Returning {
&meta.name.name
} else {
""
},
) )
.wrap_err_with(|| format!("when executing against relation '{}'", meta.name))?; .wrap_err_with(|| format!("when executing against relation '{}'", meta.name))?;
clean_ups.extend(to_clear); clean_ups.extend(to_clear);
Ok(( let returned_rows = Self::get_returning_rows(callback_collector, meta, returning);
NamedRows::new(
vec![STATUS_STR.to_string()], Ok((returned_rows, clean_ups))
vec![vec![DataValue::from(OK_STR)]],
),
clean_ups,
))
} else { } else {
let rows: Vec<Tuple> = scan.collect_vec(); let rows: Vec<Tuple> = scan.collect_vec();
@ -1568,6 +1567,59 @@ impl<'s, S: Storage<'s>> Db<S> {
} }
} }
} }
fn get_returning_rows(callback_collector: &mut CallbackCollector, meta: &InputRelationHandle, returning: &ReturnMutation) -> NamedRows {
let returned_rows = {
match returning {
ReturnMutation::NotReturning => {
NamedRows::new(
vec![STATUS_STR.to_string()],
vec![vec![DataValue::from(OK_STR)]],
)
}
ReturnMutation::Returning => {
let target_len = meta.metadata.keys.len() + meta.metadata.non_keys.len();
let mut returned_rows = Vec::new();
if let Some(collected) = callback_collector.get(&meta.name.name) {
for (kind, insertions, deletions) in collected {
let (pos_key, neg_key) = match kind {
CallbackOp::Put => { ("inserted", "replaced") }
CallbackOp::Rm => { ("requested", "deleted") }
};
for row in &insertions.rows {
let mut v = Vec::with_capacity(target_len + 1);
v.push(DataValue::from(pos_key));
v.extend_from_slice(&row[..target_len]);
while v.len() <= target_len {
v.push(DataValue::Null);
}
returned_rows.push(v);
}
for row in &deletions.rows {
let mut v = Vec::with_capacity(target_len + 1);
v.push(DataValue::from(neg_key));
v.extend_from_slice(&row[..target_len]);
while v.len() <= target_len {
v.push(DataValue::Null);
}
returned_rows.push(v);
}
}
}
let mut header = vec!["_kind".to_string()];
header.extend(meta.metadata.keys
.iter()
.chain(meta.metadata.non_keys.iter())
.map(|s| s.name.to_string()));
NamedRows::new(
header,
returned_rows,
)
}
}
};
returned_rows
}
pub(crate) fn list_running(&self) -> Result<NamedRows> { pub(crate) fn list_running(&self) -> Result<NamedRows> {
let rows = self let rows = self
.running_queries .running_queries
@ -1793,7 +1845,7 @@ impl Poison {
pub(crate) fn seconds_since_the_epoch() -> Result<f64> { pub(crate) fn seconds_since_the_epoch() -> Result<f64> {
#[cfg(not(target_arch = "wasm32"))] #[cfg(not(target_arch = "wasm32"))]
let now = SystemTime::now(); let now = SystemTime::now();
#[cfg(not(target_arch = "wasm32"))] #[cfg(not(target_arch = "wasm32"))]
return Ok(now return Ok(now
.duration_since(UNIX_EPOCH) .duration_since(UNIX_EPOCH)

38
cozo-core/src/runtime/tests.rs
Unescape Escape View File

@ -1218,6 +1218,44 @@ fn deletion() {
.run_script(r"?[x] <- [[1]] :delete a {x}", Default::default()).unwrap(); .run_script(r"?[x] <- [[1]] :delete a {x}", Default::default()).unwrap();
} }
#[test]
fn returning() {
let db = DbInstance::new("mem", "", "").unwrap();
db.run_script(":create a {x => y}", Default::default())
.unwrap();
let res = db
.run_script(
r"?[x, y] <- [[1, 2]] :insert a {x => y} ",
Default::default(),
)
.unwrap();
for row in res.into_json()["rows"].as_array().unwrap() {
println!("{}", row);
}
let res = db
.run_script(
r"?[x, y] <- [[1, 3], [2, 4]] :returning :put a {x => y} ",
Default::default(),
)
.unwrap();
println!("{:?}", res.headers);
for row in res.into_json()["rows"].as_array().unwrap() {
println!("{}", row);
}
let res = db
.run_script(
r"?[x] <- [[1], [4]] :returning :rm a {x} ",
Default::default(),
)
.unwrap();
println!("{:?}", res.headers);
for row in res.into_json()["rows"].as_array().unwrap() {
println!("{}", row);
}
}
#[test] #[test]
fn parser_corner_case() { fn parser_corner_case() {
let db = DbInstance::new("mem", "", "").unwrap(); let db = DbInstance::new("mem", "", "").unwrap();

Write Preview
Loading…
Cancel
Save