diff --git a/crates/core/src/sql/compiler.rs b/crates/core/src/sql/compiler.rs index 1e74b0f9ba..2da0a70207 100644 --- a/crates/core/src/sql/compiler.rs +++ b/crates/core/src/sql/compiler.rs @@ -403,24 +403,8 @@ fn compile_update( } else { QueryExpr::new(&table.root) }; - let mut cols = Vec::with_capacity(table.root.columns.len()); - - for field in table.root.columns.iter() { - let field = FieldName::named(&table.root.table_name, &field.col_name); - if let Some(f) = assignments.get(&field) { - cols.push(f.clone()); - } else { - cols.push(FieldExpr::Name(field)); - } - } - let insert = QueryExpr::new(&table.root).with_project(&cols, None); - let insert = if let Some(filter) = selection { - compile_where(insert, &table, filter)? - } else { - insert - }; - Ok(CrudExpr::Update { insert, delete }) + Ok(CrudExpr::Update { delete, assignments }) } /// Compiles a `CREATE TABLE ...` clause diff --git a/crates/core/src/vm.rs b/crates/core/src/vm.rs index a9c460dfe7..f838dc27bd 100644 --- a/crates/core/src/vm.rs +++ b/crates/core/src/vm.rs @@ -11,7 +11,6 @@ use spacetimedb_lib::relation::{Header, MemTable, RelIter, RelValue, RowCount, T use spacetimedb_lib::table::ProductTypeMeta; use spacetimedb_primitives::{ColId, TableId}; use spacetimedb_sats::{AlgebraicValue, ProductValue}; -use spacetimedb_vm::dsl::mem_table; use spacetimedb_vm::env::EnvDb; use spacetimedb_vm::errors::ErrorVm; use spacetimedb_vm::eval::IterRows; @@ -308,16 +307,6 @@ impl<'db, 'tx> DbProgram<'db, 'tx> { } } - fn insert_query(&mut self, table: &Table, query: QueryCode) -> Result { - let result = self._eval_query(query)?; - match result { - Code::Table(result) => { - self._execute_insert(table, result.data.into_iter().map(|row| row.data).collect_vec()) - } - _ => Ok(result), - } - } - fn create_table( &mut self, table_name: &str, @@ -402,7 +391,10 @@ impl ProgramVm for DbProgram<'_, '_> { match query { CrudCode::Query(query) => self._eval_query(query), CrudCode::Insert { table, rows } => self._execute_insert(&table, rows), - CrudCode::Update { mut insert, delete } => { + CrudCode::Update { + delete, + mut assignments, + } => { let table = delete.table.clone(); let result = self._eval_query(delete)?; @@ -415,11 +407,38 @@ impl ProgramVm for DbProgram<'_, '_> { deleted.data.clone().into_iter().map(|row| row.data).collect_vec(), )?; - let to_insert = mem_table(table.head().clone(), deleted.data.into_iter().map(|row| row.data)); - insert.table = Table::MemTable(to_insert); + // Replace the columns in the matched rows with the assigned + // values. No typechecking is performed here, nor that all + // assignments are consumed. + let exprs: Vec> = table + .head() + .fields + .iter() + .map(|col| assignments.remove(&col.field)) + .collect(); + let insert_rows = deleted + .data + .into_iter() + .map(|row| { + let elements = row + .data + .elements + .into_iter() + .zip(&exprs) + .map(|(val, expr)| { + if let Some(FieldExpr::Value(assigned)) = expr { + assigned.clone() + } else { + val + } + }) + .collect(); + + ProductValue { elements } + }) + .collect_vec(); - let result = self.insert_query(&table, insert)?; - Ok(result) + self._execute_insert(&table, insert_rows) } CrudCode::Delete { query } => { let result = self.delete_query(query)?; diff --git a/crates/vm/src/eval.rs b/crates/vm/src/eval.rs index c9dc2242b6..0c2eead5c3 100644 --- a/crates/vm/src/eval.rs +++ b/crates/vm/src/eval.rs @@ -101,11 +101,10 @@ fn build_typed(p: &mut P, node: Expr) -> ExprOpt { } ExprOpt::Crud(Box::new(CrudExprOpt::Insert { source, rows })) } - CrudExpr::Update { insert, delete } => { - let insert = build_query_opt(insert); + CrudExpr::Update { delete, assignments } => { let delete = build_query_opt(delete); - ExprOpt::Crud(Box::new(CrudExprOpt::Update { insert, delete })) + ExprOpt::Crud(Box::new(CrudExprOpt::Update { delete, assignments })) } CrudExpr::Delete { query } => { let query = build_query_opt(query); @@ -270,10 +269,9 @@ fn compile(p: &mut P, node: ExprOpt) -> Result { }; Code::Crud(q) } - CrudExprOpt::Update { insert, delete } => { - let insert = compile_query(insert); + CrudExprOpt::Update { delete, assignments } => { let delete = compile_query(delete); - Code::Crud(CrudCode::Update { insert, delete }) + Code::Crud(CrudCode::Update { delete, assignments }) } CrudExprOpt::Delete { query } => { let query = compile_query(query); diff --git a/crates/vm/src/expr.rs b/crates/vm/src/expr.rs index a4d2e989aa..859fa2d549 100644 --- a/crates/vm/src/expr.rs +++ b/crates/vm/src/expr.rs @@ -400,7 +400,7 @@ pub enum Crud { Drop(DbType), } -#[derive(Debug, Clone, Eq, PartialEq, PartialOrd, Ord)] +#[derive(Debug, Clone, Eq, PartialEq)] pub enum CrudExpr { Query(QueryExpr), Insert { @@ -408,8 +408,8 @@ pub enum CrudExpr { rows: Vec>, }, Update { - insert: QueryExpr, delete: QueryExpr, + assignments: HashMap, }, Delete { query: QueryExpr, @@ -1081,8 +1081,8 @@ pub enum CrudExprOpt { rows: Vec, }, Update { - insert: QueryExprOpt, delete: QueryExprOpt, + assignments: HashMap, }, Delete { query: QueryExprOpt, @@ -1343,8 +1343,8 @@ pub enum CrudCode { rows: Vec, }, Update { - insert: QueryCode, delete: QueryCode, + assignments: HashMap, }, Delete { query: QueryCode, @@ -1576,11 +1576,13 @@ mod tests { #[test] fn test_auth_crud_code_update() { - let mut qc = query_codes().into_iter(); - let insert = qc.next().unwrap(); - let delete = qc.next().unwrap(); - let crud = CrudCode::Update { insert, delete }; - assert_owner_required(crud); + for qc in query_codes() { + let crud = CrudCode::Update { + delete: qc, + assignments: Default::default(), + }; + assert_owner_required(crud); + } } #[test] diff --git a/crates/vm/src/typecheck.rs b/crates/vm/src/typecheck.rs index 2e7aed6ad2..f5b973d480 100644 --- a/crates/vm/src/typecheck.rs +++ b/crates/vm/src/typecheck.rs @@ -136,7 +136,7 @@ pub(crate) fn check_types(env: &mut EnvTy, ast: &ExprOpt) -> Result Ok(ty_source(source)), - CrudExprOpt::Update { insert, .. } => Ok(ty_source(&insert.source)), + CrudExprOpt::Update { delete, .. } => Ok(ty_source(&delete.source)), CrudExprOpt::Delete { query } => Ok(ty_source(&query.source)), CrudExprOpt::CreateTable { columns, .. } => Ok(AlgebraicType::Product(columns.columns.clone()).into()), CrudExprOpt::Drop { .. } => {