Making tx mut so we can auto-create params on compilation

Author: mamcx

crates/bench/benches/subscription.rs

@@ -126,10 +126,10 @@ fn eval(c: &mut Criterion) {
     // A benchmark runner for the new query engine
     let bench_query = |c: &mut Criterion, name, sql| {
         c.bench_function(name, |b| {
-            let tx = raw.db.begin_tx(Workload::Subscribe);
+            let mut tx = raw.db.begin_tx(Workload::Subscribe);
             let auth = AuthCtx::for_testing();
-            let schema_viewer = &SchemaViewer::new(&tx, &auth);
-            let (plans, table_id, table_name, _) = compile_subscription(sql, schema_viewer, &auth).unwrap();
+            let mut schema_viewer = SchemaViewer::new(&mut tx, &auth);
+            let (plans, table_id, table_name, _) = compile_subscription(sql, &mut schema_viewer, &auth).unwrap();
             let plans = plans
                 .into_iter()
                 .map(|plan| plan.optimize(&auth).unwrap())
@@ -155,8 +155,8 @@ fn eval(c: &mut Criterion) {
 
     let bench_eval = |c: &mut Criterion, name, sql| {
         c.bench_function(name, |b| {
-            let tx = raw.db.begin_tx(Workload::Update);
-            let query = compile_read_only_queryset(&raw.db, &AuthCtx::for_testing(), &tx, sql).unwrap();
+            let mut tx = raw.db.begin_tx(Workload::Update);
+            let query = compile_read_only_queryset(&raw.db, &AuthCtx::for_testing(), &mut tx, sql).unwrap();
             let query: ExecutionSet = query.into();
 
             b.iter(|| {
@@ -207,11 +207,11 @@ fn eval(c: &mut Criterion) {
         // A passthru executed independently of the database.
         let select_lhs = "select * from footprint";
         let select_rhs = "select * from location";
-        let tx = &raw.db.begin_tx(Workload::Update);
-        let query_lhs = compile_read_only_queryset(&raw.db, &AuthCtx::for_testing(), tx, select_lhs).unwrap();
-        let query_rhs = compile_read_only_queryset(&raw.db, &AuthCtx::for_testing(), tx, select_rhs).unwrap();
+        let mut tx = raw.db.begin_tx(Workload::Update);
+        let query_lhs = compile_read_only_queryset(&raw.db, &AuthCtx::for_testing(), &mut tx, select_lhs).unwrap();
+        let query_rhs = compile_read_only_queryset(&raw.db, &AuthCtx::for_testing(), &mut tx, select_rhs).unwrap();
         let query = ExecutionSet::from_iter(query_lhs.into_iter().chain(query_rhs));
-        let tx = &tx.into();
+        let tx = &(&mut tx).into();
 
         b.iter(|| drop(black_box(query.eval_incr_for_test(&raw.db, tx, &update, None))))
     });
@@ -226,10 +226,10 @@ fn eval(c: &mut Criterion) {
             from footprint join location on footprint.entity_id = location.entity_id \
             where location.chunk_index = {chunk_index}"
         );
-        let tx = &raw.db.begin_tx(Workload::Update);
-        let query = compile_read_only_queryset(&raw.db, &AuthCtx::for_testing(), tx, &join).unwrap();
+        let mut tx = raw.db.begin_tx(Workload::Update);
+        let query = compile_read_only_queryset(&raw.db, &AuthCtx::for_testing(), &mut tx, &join).unwrap();
         let query: ExecutionSet = query.into();
-        let tx = &tx.into();
+        let tx = &(&mut tx).into();
 
         b.iter(|| drop(black_box(query.eval_incr_for_test(&raw.db, tx, &update, None))));
     });

crates/core/src/estimation.rs

@@ -181,8 +181,8 @@ mod tests {
     }
 
     fn num_rows_for(db: &RelationalDB, sql: &str) -> u64 {
-        let tx = begin_tx(db);
-        match &*compile_sql(db, &AuthCtx::for_testing(), &tx, sql).expect("Failed to compile sql") {
+        let mut tx = begin_tx(db);
+        match &*compile_sql(db, &AuthCtx::for_testing(), &mut tx, sql).expect("Failed to compile sql") {
             [CrudExpr::Query(expr)] => num_rows(&tx, expr),
             exprs => panic!("unexpected result from compilation: {exprs:#?}"),
         }
@@ -191,10 +191,10 @@ mod tests {
     /// Using the new query plan
     fn new_row_estimate(db: &RelationalDB, sql: &str) -> u64 {
         let auth = AuthCtx::for_testing();
-        let tx = begin_tx(db);
-        let tx = SchemaViewer::new(&tx, &auth);
+        let mut tx = begin_tx(db);
+        let mut tx = SchemaViewer::new(&mut tx, &auth);
 
-        compile_subscription(sql, &tx, &auth)
+        compile_subscription(sql, &mut tx, &auth)
             .map(|(plans, ..)| plans)
             .expect("failed to compile sql query")
             .into_iter()

crates/core/src/host/module_host.rs

@@ -1869,7 +1869,7 @@ impl ModuleHost {
         let metrics = self
             .on_module_thread("one_off_query", move || {
                 let (tx_offset_sender, tx_offset_receiver) = oneshot::channel();
-                let tx = scopeguard::guard(db.begin_tx(Workload::Sql), |tx| {
+                let mut tx = scopeguard::guard(db.begin_tx(Workload::Sql), |tx| {
                     let (tx_offset, tx_metrics, reducer) = db.release_tx(tx);
                     let _ = tx_offset_sender.send(tx_offset);
                     db.report_read_tx_metrics(reducer, tx_metrics);
@@ -1878,7 +1878,7 @@ impl ModuleHost {
                 // We wrap the actual query in a closure so we can use ? to handle errors without making
                 // the entire transaction abort with an error.
                 let result: Result<(OneOffTable<F>, ExecutionMetrics), anyhow::Error> = (|| {
-                    let tx = SchemaViewer::new(&*tx, &auth);
+                    let mut tx = SchemaViewer::new(&mut *tx, &auth);
 
                     let (
                         // A query may compile down to several plans.
@@ -1888,7 +1888,7 @@ impl ModuleHost {
                         _,
                         table_name,
                         _,
-                    ) = compile_subscription(&query, &tx, &auth)?;
+                    ) = compile_subscription(&query, &mut tx, &auth)?;
 
                     // Optimize each fragment
                     let optimized = plans

crates/core/src/host/wasm_common/module_host_actor.rs

@@ -1126,10 +1126,10 @@ impl InstanceCommon {
 
         // Views bypass RLS, since views should enforce their own access control procedurally.
         let auth = AuthCtx::for_current(self.info.database_identity);
-        let schema_view = SchemaViewer::new(&*tx, &auth);
+        let mut schema_view = SchemaViewer::new(&mut *tx, &auth);
 
         // Compile to subscription plans.
-        let (plans, has_params) = SubscriptionPlan::compile(the_query, &schema_view, &auth)?;
+        let (plans, has_params) = SubscriptionPlan::compile(the_query, &mut schema_view, &auth)?;
         ensure!(
             !has_params,
             "parameterized SQL is not supported for view materialization yet"

crates/core/src/sql/ast.rs

@@ -5,6 +5,7 @@ use spacetimedb_data_structures::map::{HashCollectionExt as _, IntMap};
 use spacetimedb_datastore::locking_tx_datastore::state_view::StateView;
 use spacetimedb_datastore::system_tables::{StRowLevelSecurityFields, ST_ROW_LEVEL_SECURITY_ID};
 use spacetimedb_expr::check::{SchemaView, TypingResult};
+use spacetimedb_expr::errors::TypingError;
 use spacetimedb_expr::statement::compile_sql_stmt;
 use spacetimedb_lib::identity::AuthCtx;
 use spacetimedb_primitives::{ArgId, ColId, TableId};
@@ -22,7 +23,7 @@ use sqlparser::ast::{
 };
 use sqlparser::dialect::PostgreSqlDialect;
 use sqlparser::parser::Parser;
-use std::ops::Deref;
+use std::ops::{Deref, DerefMut};
 use std::sync::Arc;
 
 /// Simplify to detect features of the syntax we don't support yet
@@ -477,7 +478,7 @@ fn compile_where(table: &From, filter: Option<SqlExpr>) -> Result<Option<Selecti
 }
 
 pub struct SchemaViewer<'a, T> {
-    pub(crate) tx: &'a T,
+    tx: &'a mut T,
     auth: &'a AuthCtx,
 }
 
@@ -489,6 +490,12 @@ impl<T> Deref for SchemaViewer<'_, T> {
     }
 }
 
+impl<T> DerefMut for SchemaViewer<'_, T> {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        self.tx
+    }
+}
+
 impl<T: StateView> SchemaView for SchemaViewer<'_, T> {
     fn table_id(&self, name: &str) -> Option<TableId> {
         // Get the schema from the in-memory state instead of fetching from the database for speed
@@ -536,10 +543,15 @@ impl<T: StateView> SchemaView for SchemaViewer<'_, T> {
             })
             .collect::<anyhow::Result<_>>()
     }
+
+    fn get_or_create_params(&mut self, _params: ProductValue) -> TypingResult<ArgId> {
+        // Caller should have used `SchemaViewerMut` on crate `core`
+        Err(TypingError::ParamsReadOnly)
+    }
 }
 
 impl<'a, T> SchemaViewer<'a, T> {
-    pub fn new(tx: &'a T, auth: &'a AuthCtx) -> Self {
+    pub fn new(tx: &'a mut T, auth: &'a AuthCtx) -> Self {
         Self { tx, auth }
     }
 }
@@ -1000,13 +1012,12 @@ fn compile_statement<T: TableSchemaView + StateView>(
 pub(crate) fn compile_to_ast<T: TableSchemaView + StateView>(
     db: &RelationalDB,
     auth: &AuthCtx,
-    tx: &T,
+    tx: &mut T,
     sql_text: &str,
 ) -> Result<Vec<SqlAst>, DBError> {
     // NOTE: The following ensures compliance with the 1.0 sql api.
     // Come 1.0, it will have replaced the current compilation stack.
-    compile_sql_stmt(sql_text, &SchemaViewer::new(tx, auth), auth)?;
-
+    compile_sql_stmt(sql_text, &mut SchemaViewer::new(tx, auth), auth)?;
     let dialect = PostgreSqlDialect {};
     let ast = Parser::parse_sql(&dialect, sql_text).map_err(|error| DBError::SqlParser {
         sql: sql_text.to_string(),