Add spacetimedb::sql::query_debug_info::QueryDebugInfo (#560)

Add spacetimedb::sql::query_debug_info::QueryDebugInfo

---------

Signed-off-by: james gilles <jameshgilles@gmail.com>

crates/core/src/execution_context.rs

@@ -1,15 +1,24 @@
 use derive_more::Display;
 use spacetimedb_lib::Address;
 
+use crate::sql::query_debug_info::QueryDebugInfo;
+
 /// Represents the context under which a database runtime method is executed.
 /// In particular it provides details about the currently executing txn to runtime operations.
 /// More generally it acts as a container for information that database operations may require to function correctly.
 #[derive(Default)]
 pub struct ExecutionContext<'a> {
-    // The database on which a transaction is being executed.
+    /// The database on which a transaction is being executed.
     database: Address,
-    // The reducer from which the current transaction originated.
+    /// The reducer from which the current transaction originated.
+    /// Note: this will never be set at the same time as `query`.
     reducer: Option<&'a str>,
+    /// The SQL query being executed, if any.
+    /// Note: this will never be set at the same time as `reducer`.
+    /// It is also NOT guaranteed to be set, even if txn_type == Sql.
+    /// This is because some transactions tagged "SQL" don't exactly correspond
+    /// to any particular query.
+    query_debug_info: Option<&'a QueryDebugInfo>,
     // The type of transaction that is being executed.
     txn_type: TransactionType,
 }
@@ -37,15 +46,17 @@ impl<'a> ExecutionContext<'a> {
         Self {
             database,
             reducer: Some(name),
+            query_debug_info: None,
             txn_type: TransactionType::Reducer,
         }
     }
 
     /// Returns an [ExecutionContext] for a sql or subscription transaction.
-    pub fn sql(database: Address) -> Self {
+    pub fn sql(database: Address, query_debug_info: Option<&'a QueryDebugInfo>) -> Self {
         Self {
             database,
             reducer: None,
+            query_debug_info,
             txn_type: TransactionType::Sql,
         }
     }
@@ -55,6 +66,7 @@ impl<'a> ExecutionContext<'a> {
         Self {
             database,
             reducer: None,
+            query_debug_info: None,
             txn_type: TransactionType::Internal,
         }
     }
@@ -72,6 +84,13 @@ impl<'a> ExecutionContext<'a> {
         self.reducer
     }
 
+    /// Returns the debug info for the query being executed.
+    /// Returns [None] if this is not a sql context.
+    #[inline]
+    pub fn query_debug_info(&self) -> Option<&QueryDebugInfo> {
+        self.query_debug_info
+    }
+
     /// Returns the type of transaction that is being executed.
     #[inline]
     pub fn txn_type(&self) -> TransactionType {

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

@@ -27,6 +27,7 @@ use crate::host::{
 };
 use crate::identity::Identity;
 use crate::sql;
+use crate::sql::query_debug_info::QueryDebugInfo;
 use crate::subscription::module_subscription_actor::ModuleSubscriptionManager;
 use crate::util::{const_unwrap, ResultInspectExt};
 use crate::worker_metrics::WORKER_METRICS;
@@ -249,21 +250,24 @@ impl<T: WasmModule> Module for WasmModuleHostActor<T> {
         let auth = AuthCtx::new(self.database_instance_context.identity, caller_identity);
         // TODO(jgilles): make this a read-only TX when those get added
 
-        db.with_read_only(&ExecutionContext::sql(db.address()), |tx| {
-            log::debug!("One-off query: {query}");
-
-            let compiled = sql::compiler::compile_sql(db, tx, &query)?
-                .into_iter()
-                .map(|expr| {
-                    if matches!(expr, CrudExpr::Query { .. }) {
-                        Ok(expr)
-                    } else {
-                        Err(anyhow!("One-off queries are not allowed to modify the database"))
-                    }
-                })
-                .collect::<Result<_, _>>()?;
-            sql::execute::execute_sql(db, tx, compiled, auth)
-        })
+        db.with_read_only(
+            &ExecutionContext::sql(db.address(), Some(&QueryDebugInfo::from_source(&query))),
+            |tx| {
+                log::debug!("One-off query: {query}");
+
+                let compiled = sql::compiler::compile_sql(db, tx, &query)?
+                    .into_iter()
+                    .map(|expr| {
+                        if matches!(expr, CrudExpr::Query { .. }) {
+                            Ok(expr)
+                        } else {
+                            Err(anyhow!("One-off queries are not allowed to modify the database"))
+                        }
+                    })
+                    .collect::<Result<_, _>>()?;
+                sql::execute::execute_sql(db, tx, compiled, Some(&QueryDebugInfo::from_source(&query)), auth)
+            },
+        )
     }
 
     fn clear_table(&self, table_name: String) -> Result<(), anyhow::Error> {

crates/core/src/sql/execute.rs

@@ -13,6 +13,8 @@ use crate::execution_context::ExecutionContext;
 use crate::sql::compiler::compile_sql;
 use crate::vm::DbProgram;
 
+use super::query_debug_info::QueryDebugInfo;
+
 pub struct StmtResult {
     pub schema: ProductType,
     pub rows: Vec<ProductValue>,
@@ -32,7 +34,8 @@ pub fn execute(
     info!(sql = sql_text);
     if let Some((database_instance_context, _)) = db_inst_ctx_controller.get(database_instance_id) {
         let db = &database_instance_context.relational_db;
-        let ctx = ExecutionContext::sql(db.address());
+        let info = QueryDebugInfo::from_source(&sql_text);
+        let ctx = ExecutionContext::sql(db.address(), Some(&info));
         db.with_auto_commit(&ctx, |tx| {
             run(&database_instance_context.relational_db, tx, &sql_text, auth)
         })
@@ -62,9 +65,10 @@ pub fn execute_single_sql(
     db: &RelationalDB,
     tx: &mut MutTxId,
     ast: CrudExpr,
+    query_debug_info: Option<&QueryDebugInfo>,
     auth: AuthCtx,
 ) -> Result<Vec<MemTable>, DBError> {
-    let ctx = ExecutionContext::sql(db.address());
+    let ctx = ExecutionContext::sql(db.address(), query_debug_info);
     let p = &mut DbProgram::new(&ctx, db, tx, auth);
     let q = Expr::Crud(Box::new(ast));
 
@@ -79,10 +83,11 @@ pub fn execute_sql(
     db: &RelationalDB,
     tx: &mut MutTxId,
     ast: Vec<CrudExpr>,
+    query_debug_info: Option<&QueryDebugInfo>,
     auth: AuthCtx,
 ) -> Result<Vec<MemTable>, DBError> {
     let total = ast.len();
-    let ctx = ExecutionContext::sql(db.address());
+    let ctx = ExecutionContext::sql(db.address(), query_debug_info);
     let p = &mut DbProgram::new(&ctx, db, tx, auth);
     let q = Expr::Block(ast.into_iter().map(|x| Expr::Crud(Box::new(x))).collect());
 
@@ -95,7 +100,7 @@ pub fn execute_sql(
 #[tracing::instrument(skip_all)]
 pub fn run(db: &RelationalDB, tx: &mut MutTxId, sql_text: &str, auth: AuthCtx) -> Result<Vec<MemTable>, DBError> {
     let ast = compile_sql(db, tx, sql_text)?;
-    execute_sql(db, tx, ast, auth)
+    execute_sql(db, tx, ast, Some(&QueryDebugInfo::from_source(sql_text)), auth)
 }
 
 #[cfg(test)]

crates/core/src/sql/mod.rs

@@ -1,3 +1,4 @@
 pub mod ast;
 pub mod compiler;
 pub mod execute;
+pub mod query_debug_info;

crates/core/src/sql/query_debug_info.rs

@@ -0,0 +1,22 @@
+use std::sync::Arc;
+
+/// Used to store the source of a SQL query in a way that can be cheaply cloned,
+/// without proliferating lifetimes everywhere.
+///
+/// TODO: if CrudExpr ever gets refactored, this should probably be attached to those.
+/// That would be a large refactoring though. It would be nice if we could get
+/// more precise spans from sqlparser. We could stick all sorts of other things in here too.
+#[derive(Clone, Debug, Eq, Ord, PartialEq, PartialOrd)]
+pub struct QueryDebugInfo(Arc<str>);
+
+impl QueryDebugInfo {
+    /// Create a new [QueryDebugInfo] from the given source text.
+    pub fn from_source<T: AsRef<str>>(source: T) -> Self {
+        Self(source.as_ref().into())
+    }
+
+    /// Get the source text of the query, if available.
+    pub fn source(&self) -> &str {
+        &self.0
+    }
+}

crates/core/src/subscription/module_subscription_actor.rs

@@ -163,10 +163,14 @@ impl ModuleSubscriptionActor {
         sender: ClientConnectionSender,
         subscription: Subscribe,
     ) -> Result<(), DBError> {
-        //Split logic to properly handle `Error` + `Tx`
+        // Split logic to properly handle `Error` + `Tx`
         let mut tx = self.relational_db.begin_tx();
+
         let result = self._add_subscription(sender, subscription, &mut tx).await;
-        let ctx = ExecutionContext::sql(self.relational_db.address());
+
+        // Note: the missing QueryDebugInfo here is only used for finishing the transaction;
+        // all of the relevant queries already executed, with debug info, in _add_subscription
+        let ctx = ExecutionContext::sql(self.relational_db.address(), None);
         self.relational_db.finish_tx(&ctx, tx, result)
     }
 
@@ -214,7 +218,7 @@ impl ModuleSubscriptionActor {
         //Split logic to properly handle `Error` + `Tx`
         let mut tx = self.relational_db.begin_tx();
         let result = self._broadcast_commit_event(event, &mut tx).await;
-        let ctx = ExecutionContext::sql(self.relational_db.address());
+        let ctx = ExecutionContext::sql(self.relational_db.address(), None);
         self.relational_db.finish_tx(&ctx, tx, result)
     }
 }

crates/core/src/subscription/query.rs

@@ -4,7 +4,8 @@ use crate::error::{DBError, SubscriptionError};
 use crate::host::module_host::DatabaseTableUpdate;
 use crate::sql::compiler::compile_sql;
 use crate::sql::execute::execute_single_sql;
-use crate::subscription::subscription::QuerySet;
+use crate::sql::query_debug_info::QueryDebugInfo;
+use crate::subscription::subscription::{QuerySet, SupportedQuery};
 use spacetimedb_lib::identity::AuthCtx;
 use spacetimedb_sats::relation::{Column, FieldName, MemTable, RelValue};
 use spacetimedb_sats::AlgebraicType;
@@ -69,9 +70,10 @@ pub(crate) fn run_query(
     db: &RelationalDB,
     tx: &mut MutTxId,
     query: &QueryExpr,
+    query_debug_info: Option<&QueryDebugInfo>,
     auth: AuthCtx,
 ) -> Result<Vec<MemTable>, DBError> {
-    execute_single_sql(db, tx, CrudExpr::Query(query.clone()), auth)
+    execute_single_sql(db, tx, CrudExpr::Query(query.clone()), query_debug_info, auth)
 }
 
 // TODO: It's semantically wrong to `SUBSCRIBE_TO_ALL_QUERY`
@@ -123,8 +125,13 @@ pub fn compile_read_only_query(
         }
     }
 
+    let info = QueryDebugInfo::from_source(input);
+
     if !queries.is_empty() {
-        Ok(queries.into_iter().map(TryFrom::try_from).collect::<Result<_, _>>()?)
+        Ok(queries
+            .into_iter()
+            .map(|query| SupportedQuery::new(query, Some(info.clone())))
+            .collect::<Result<_, _>>()?)
     } else {
         Err(SubscriptionError::Empty.into())
     }
@@ -331,7 +338,7 @@ mod tests {
         data: &DatabaseTableUpdate,
     ) -> ResultTest<()> {
         let q = to_mem_table(q.clone(), data);
-        let result = run_query(db, tx, &q, AuthCtx::for_testing())?;
+        let result = run_query(db, tx, &q, None, AuthCtx::for_testing())?;
 
         assert_eq!(
             Some(table.as_without_table_name()),
@@ -803,6 +810,7 @@ mod tests {
             &db,
             &mut tx,
             &q,
+            None,
             AuthCtx::new(Identity::__dummy(), Identity::from_byte_array([1u8; 32])),
         ) {
             Ok(_) => {
@@ -925,7 +933,7 @@ mod tests {
         let qset = compile_read_only_query(&db, &tx, &AuthCtx::for_testing(), sql_query)?;
 
         for q in qset {
-            let result = run_query(&db, &mut tx, q.as_expr(), AuthCtx::for_testing())?;
+            let result = run_query(&db, &mut tx, q.as_expr(), None, AuthCtx::for_testing())?;
             assert_eq!(result.len(), 1, "Join query did not return any rows");
         }