Materialize views for subscriptions

Author: joshua-spacetime

crates/core/src/subscription/module_subscription_actor.rs

@@ -31,10 +31,11 @@ use spacetimedb_client_api_messages::websocket::{
 use spacetimedb_datastore::db_metrics::DB_METRICS;
 use spacetimedb_datastore::execution_context::{Workload, WorkloadType};
 use spacetimedb_datastore::locking_tx_datastore::datastore::TxMetrics;
-use spacetimedb_datastore::locking_tx_datastore::TxId;
-use spacetimedb_datastore::traits::TxData;
+use spacetimedb_datastore::locking_tx_datastore::{MutTxId, TxId};
+use spacetimedb_datastore::traits::{IsolationLevel, TxData};
 use spacetimedb_durability::TxOffset;
 use spacetimedb_execution::pipelined::{PipelinedProject, ViewProject};
+use spacetimedb_expr::expr::CollectViews;
 use spacetimedb_lib::identity::AuthCtx;
 use spacetimedb_lib::metrics::ExecutionMetrics;
 use spacetimedb_lib::Identity;
@@ -368,7 +369,7 @@ impl ModuleSubscriptions {
         let hash = QueryHash::from_string(&sql, auth.caller, false);
         let hash_with_param = QueryHash::from_string(&sql, auth.caller, true);
 
-        let (tx, tx_offset) = self.begin_tx(Workload::Subscribe);
+        let (mut_tx, _) = self.begin_mut_tx(Workload::Subscribe);
 
         let existing_query = {
             let guard = self.subscriptions.read();
@@ -378,7 +379,7 @@ impl ModuleSubscriptions {
         let query = return_on_err_with_sql!(
             existing_query.map(Ok).unwrap_or_else(|| compile_query_with_hashes(
                 &auth,
-                &tx,
+                &*mut_tx,
                 &sql,
                 hash,
                 hash_with_param
@@ -388,6 +389,9 @@ impl ModuleSubscriptions {
             send_err_msg
         );
 
+        let mut_tx = ScopeGuard::<MutTxId, _>::into_inner(mut_tx);
+        let (tx, tx_offset) = self.materialize_views_and_downgrade_tx(mut_tx, &query, auth.caller)?;
+
         let (table_rows, metrics) = return_on_err_with_sql!(
             self.evaluate_initial_subscription(sender.clone(), query.clone(), &tx, &auth, TableUpdateType::Subscribe),
             query.sql(),
@@ -605,7 +609,7 @@ impl ModuleSubscriptions {
         queries: &[Box<str>],
         num_queries: usize,
         metrics: &SubscriptionMetrics,
-    ) -> Result<(Vec<Arc<Plan>>, AuthCtx, TxId, HistogramTimer), DBError> {
+    ) -> Result<(Vec<Arc<Plan>>, AuthCtx, MutTxId, HistogramTimer), DBError> {
         let mut subscribe_to_all_tables = false;
         let mut plans = Vec::with_capacity(num_queries);
         let mut query_hashes = Vec::with_capacity(num_queries);
@@ -624,7 +628,7 @@ impl ModuleSubscriptions {
         let auth = AuthCtx::new(self.owner_identity, sender);
 
         // We always get the db lock before the subscription lock to avoid deadlocks.
-        let (tx, _tx_offset) = self.begin_tx(Workload::Subscribe);
+        let (mut_tx, _tx_offset) = self.begin_mut_tx(Workload::Subscribe);
 
         let compile_timer = metrics.compilation_time.start_timer();
 
@@ -637,9 +641,14 @@ impl ModuleSubscriptions {
 
         if subscribe_to_all_tables {
             plans.extend(
-                super::subscription::get_all(&self.relational_db, &tx, &auth)?
-                    .into_iter()
-                    .map(Arc::new),
+                super::subscription::get_all(
+                    |relational_db, tx| relational_db.get_all_tables_mut(tx).map(|schemas| schemas.into_iter()),
+                    &self.relational_db,
+                    &*mut_tx,
+                    &auth,
+                )?
+                .into_iter()
+                .map(Arc::new),
             );
         }
 
@@ -652,7 +661,7 @@ impl ModuleSubscriptions {
                 plans.push(unit);
             } else {
                 plans.push(Arc::new(
-                    compile_query_with_hashes(&auth, &tx, sql, hash, hash_with_param).map_err(|err| {
+                    compile_query_with_hashes(&auth, &*mut_tx, sql, hash, hash_with_param).map_err(|err| {
                         DBError::WithSql {
                             error: Box::new(DBError::Other(err.into())),
                             sql: sql.into(),
@@ -666,7 +675,7 @@ impl ModuleSubscriptions {
         // How many queries in this subscription are not cached?
         metrics.num_new_queries_subscribed.inc_by(new_queries);
 
-        Ok((plans, auth, scopeguard::ScopeGuard::into_inner(tx), compile_timer))
+        Ok((plans, auth, ScopeGuard::<MutTxId, _>::into_inner(mut_tx), compile_timer))
     }
 
     /// Send a message to a client connection.
@@ -731,7 +740,7 @@ impl ModuleSubscriptions {
         // How many queries make up this subscription?
         subscription_metrics.num_queries_subscribed.inc_by(num_queries as _);
 
-        let (queries, auth, tx, compile_timer) = return_on_err!(
+        let (queries, auth, mut_tx, compile_timer) = return_on_err!(
             self.compile_queries(
                 sender.id.identity,
                 &request.query_strings,
@@ -741,7 +750,7 @@ impl ModuleSubscriptions {
             send_err_msg,
             None
         );
-        let (tx, tx_offset) = self.guard_tx(tx, <_>::default());
+        let (mut_tx, _) = self.guard_mut_tx(mut_tx, <_>::default());
 
         // We minimize locking so that other clients can add subscriptions concurrently.
         // We are protected from race conditions with broadcasts, because we have the db lock,
@@ -761,6 +770,9 @@ impl ModuleSubscriptions {
         // Record how long it took to compile the subscription
         drop(compile_timer);
 
+        let mut_tx = ScopeGuard::<MutTxId, _>::into_inner(mut_tx);
+        let (tx, tx_offset) = self.materialize_views_and_downgrade_tx(mut_tx, &queries, auth.caller)?;
+
         let Ok((update, metrics)) =
             self.evaluate_queries(sender.clone(), &queries, &tx, &auth, TableUpdateType::Subscribe)
         else {
@@ -827,13 +839,14 @@ impl ModuleSubscriptions {
         // How many queries make up this subscription?
         subscription_metrics.num_queries_subscribed.inc_by(num_queries as _);
 
-        let (queries, auth, tx, compile_timer) = self.compile_queries(
+        let (queries, auth, mut_tx, compile_timer) = self.compile_queries(
             sender.id.identity,
             &subscription.query_strings,
             num_queries,
             &subscription_metrics,
         )?;
-        let (tx, tx_offset) = self.guard_tx(tx, <_>::default());
+
+        let (tx, tx_offset) = self.materialize_views_and_downgrade_tx(mut_tx, &queries, auth.caller)?;
 
         check_row_limit(
             &queries,
@@ -983,13 +996,35 @@ impl ModuleSubscriptions {
         }))
     }
 
+    /// Materialize views for `sender`, collected from `view_collector`.
+    fn materialize_views_and_downgrade_tx(
+        &self,
+        mut tx: MutTxId,
+        view_collector: &impl CollectViews,
+        sender: Identity,
+    ) -> Result<(TxGuard<impl FnOnce(TxId) + '_>, TransactionOffset), DBError> {
+        tx.materialize_views(view_collector, sender)?;
+        let (tx_data, tx_metrics_mut, tx) = tx.commit_downgrade(Workload::Subscribe);
+        let opts = GuardTxOptions::from_mut(tx_data, tx_metrics_mut);
+        Ok(self.guard_tx(tx, opts))
+    }
+
     /// Helper that starts a new read transaction, and guards it using
     /// [`Self::guard_tx`] with the default configuration.
-    fn begin_tx(&self, workload: Workload) -> (ScopeGuard<TxId, impl FnOnce(TxId) + '_>, TransactionOffset) {
+    fn begin_tx(&self, workload: Workload) -> (TxGuard<impl FnOnce(TxId) + '_>, TransactionOffset) {
         self.guard_tx(self.relational_db.begin_tx(workload), <_>::default())
     }
 
-    /// Helper wrapping `tx` in a scopegard, with a configurable drop fn.
+    /// Helper that starts a new mutable transaction, and guards it using
+    /// [`Self::guard_mut_tx`] with the default configuration.
+    fn begin_mut_tx(&self, workload: Workload) -> (MutTxGuard<impl FnOnce(MutTxId) + '_>, TransactionOffset) {
+        self.guard_mut_tx(
+            self.relational_db.begin_mut_tx(IsolationLevel::Serializable, workload),
+            <_>::default(),
+        )
+    }
+
+    /// Helper wrapping a [`TxId`] in a scopegard, with a configurable drop fn.
     ///
     /// By default, `tx` is released when the returned [`ScopeGuard`] is dropped,
     /// and reports the transaction metrics via [`RelationalDB::report_tx_metrics`].
@@ -1004,27 +1039,42 @@ impl ModuleSubscriptions {
     /// If another receiver of the transaction offset is needed, its sending
     /// side can be passed in as `extra_tx_offset_sender`. It will be sent the
     /// offset as well.
-    fn guard_tx(
-        &self,
-        tx: TxId,
-        GuardTxOptions {
-            extra_tx_offset_sender,
-            tx_data,
-            tx_metrics_mut,
-        }: GuardTxOptions,
-    ) -> (ScopeGuard<TxId, impl FnOnce(TxId) + '_>, TransactionOffset) {
+    fn guard_tx(&self, tx: TxId, opts: GuardTxOptions) -> (TxGuard<impl FnOnce(TxId) + '_>, TransactionOffset) {
         let (offset_tx, offset_rx) = oneshot::channel();
         let guard = scopeguard::guard(tx, |tx| {
             let (tx_offset, tx_metrics, reducer) = self.relational_db.release_tx(tx);
             log::trace!("read tx released with offset {tx_offset}");
             let _ = offset_tx.send(tx_offset);
-            if let Some(extra) = extra_tx_offset_sender {
+            if let Some(extra) = opts.extra_tx_offset_sender {
                 let _ = extra.send(tx_offset);
             }
             self.relational_db
-                .report_tx_metrics(reducer, tx_data, tx_metrics_mut, Some(tx_metrics));
+                .report_tx_metrics(reducer, opts.tx_data, opts.tx_metrics_mut, Some(tx_metrics));
         });
+        (guard, offset_rx)
+    }
 
+    /// The same as [`Self::guard_tx`] but for mutable transactions.
+    ///
+    /// By default, `tx` is committed when the returned [`ScopeGuard`] is dropped,
+    /// and reports the transaction metrics via [`RelationalDB::report_tx_metrics`].
+    fn guard_mut_tx(
+        &self,
+        tx: MutTxId,
+        opts: GuardTxOptions,
+    ) -> (MutTxGuard<impl FnOnce(MutTxId) + '_>, TransactionOffset) {
+        let (offset_tx, offset_rx) = oneshot::channel();
+        let guard = scopeguard::guard(tx, |tx| {
+            if let Ok(Some((tx_offset, tx_data, tx_metrics_mut, reducer))) = self.relational_db.commit_tx(tx) {
+                log::trace!("mutable tx committed with offset {tx_offset}");
+                let _ = offset_tx.send(tx_offset);
+                if let Some(extra) = opts.extra_tx_offset_sender {
+                    let _ = extra.send(tx_offset);
+                }
+                self.relational_db
+                    .report_tx_metrics(reducer, Some(Arc::new(tx_data)), Some(tx_metrics_mut), None);
+            }
+        });
         (guard, offset_rx)
     }
 }
@@ -1052,10 +1102,24 @@ impl GuardTxOptions {
             tx_metrics_mut: tx_metrics_mut.into(),
         }
     }
+
+    fn from_mut(tx_data: TxData, tx_metrics_mut: TxMetrics) -> Self {
+        Self {
+            extra_tx_offset_sender: None,
+            tx_data: Some(Arc::new(tx_data)),
+            tx_metrics_mut: tx_metrics_mut.into(),
+        }
+    }
 }
 
 pub struct WriteConflict;
 
+/// A [`ScopeGuard`] for [`TxId`]
+type TxGuard<F> = ScopeGuard<TxId, F>;
+
+/// A [`ScopeGuard`] for [`MutTxId`]
+type MutTxGuard<F> = ScopeGuard<MutTxId, F>;
+
 #[cfg(test)]
 mod tests {
     use super::{AssertTxFn, ModuleSubscriptions};

crates/core/src/subscription/module_subscription_manager.rs

@@ -22,9 +22,10 @@ use spacetimedb_client_api_messages::websocket::{
 use spacetimedb_data_structures::map::{Entry, IntMap};
 use spacetimedb_datastore::locking_tx_datastore::state_view::StateView;
 use spacetimedb_durability::TxOffset;
+use spacetimedb_expr::expr::CollectViews;
 use spacetimedb_lib::metrics::ExecutionMetrics;
 use spacetimedb_lib::{AlgebraicValue, ConnectionId, Identity, ProductValue};
-use spacetimedb_primitives::{ColId, IndexId, TableId};
+use spacetimedb_primitives::{ColId, IndexId, TableId, ViewId};
 use spacetimedb_subscription::{JoinEdge, SubscriptionPlan, TableName};
 use std::collections::BTreeMap;
 use std::fmt::Debug;
@@ -53,6 +54,14 @@ pub struct Plan {
     plans: Vec<SubscriptionPlan>,
 }
 
+impl CollectViews for Plan {
+    fn collect_views(&self, views: &mut std::collections::HashSet<ViewId>) {
+        for plan in &self.plans {
+            plan.collect_views(views);
+        }
+    }
+}
+
 impl Plan {
     /// Create a new subscription plan to be cached
     pub fn new(plans: Vec<SubscriptionPlan>, hash: QueryHash, text: String) -> Self {

crates/core/src/subscription/query.rs

@@ -5,6 +5,8 @@ use crate::sql::compiler::compile_sql;
 use crate::subscription::subscription::SupportedQuery;
 use once_cell::sync::Lazy;
 use regex::Regex;
+use spacetimedb_datastore::locking_tx_datastore::state_view::StateView;
+use spacetimedb_execution::Datastore;
 use spacetimedb_lib::identity::AuthCtx;
 use spacetimedb_subscription::SubscriptionPlan;
 use spacetimedb_vm::expr::{self, Crud, CrudExpr, QueryExpr};
@@ -93,7 +95,7 @@ pub fn compile_read_only_query(auth: &AuthCtx, tx: &Tx, input: &str) -> Result<P
 
 /// Compile a string into a single read-only query.
 /// This returns an error if the string has multiple queries or mutations.
-pub fn compile_query_with_hashes(
+pub fn compile_query_with_hashes<Tx: Datastore + StateView>(
     auth: &AuthCtx,
     tx: &Tx,
     input: &str,

crates/core/src/subscription/subscription.rs

@@ -34,6 +34,7 @@ use anyhow::Context;
 use itertools::Either;
 use spacetimedb_client_api_messages::websocket::Compression;
 use spacetimedb_data_structures::map::HashSet;
+use spacetimedb_datastore::locking_tx_datastore::state_view::StateView;
 use spacetimedb_datastore::locking_tx_datastore::TxId;
 use spacetimedb_lib::db::auth::{StAccess, StTableType};
 use spacetimedb_lib::identity::AuthCtx;
@@ -42,13 +43,13 @@ use spacetimedb_primitives::TableId;
 use spacetimedb_sats::ProductValue;
 use spacetimedb_schema::def::error::AuthError;
 use spacetimedb_schema::relation::DbTable;
+use spacetimedb_schema::schema::TableSchema;
 use spacetimedb_subscription::SubscriptionPlan;
 use spacetimedb_vm::expr::{self, AuthAccess, IndexJoin, Query, QueryExpr, SourceExpr, SourceProvider, SourceSet};
 use spacetimedb_vm::rel_ops::RelOps;
 use spacetimedb_vm::relation::{MemTable, RelValue};
 use std::hash::Hash;
 use std::iter;
-use std::ops::Deref;
 use std::sync::Arc;
 use std::time::Duration;
 
@@ -611,11 +612,18 @@ impl AuthAccess for ExecutionSet {
 /// Queries all the [`StTableType::User`] tables *right now*
 /// and turns them into [`QueryExpr`],
 /// the moral equivalent of `SELECT * FROM table`.
-pub(crate) fn get_all(relational_db: &RelationalDB, tx: &Tx, auth: &AuthCtx) -> Result<Vec<Plan>, DBError> {
-    Ok(relational_db
-        .get_all_tables(tx)?
-        .iter()
-        .map(Deref::deref)
+pub(crate) fn get_all<T, F, I>(
+    get_all_tables: F,
+    relational_db: &RelationalDB,
+    tx: &T,
+    auth: &AuthCtx,
+) -> Result<Vec<Plan>, DBError>
+where
+    T: StateView,
+    F: Fn(&RelationalDB, &T) -> Result<I, DBError>,
+    I: Iterator<Item = Arc<TableSchema>>,
+{
+    Ok(get_all_tables(relational_db, tx)?
         .filter(|t| t.table_type == StTableType::User && (auth.is_owner() || t.table_access == StAccess::Public))
         .map(|schema| {
             let sql = format!("SELECT * FROM {}", schema.table_name);
@@ -648,6 +656,8 @@ pub(crate) fn legacy_get_all(
     tx: &Tx,
     auth: &AuthCtx,
 ) -> Result<Vec<SupportedQuery>, DBError> {
+    use std::ops::Deref;
+
     Ok(relational_db
         .get_all_tables(tx)?
         .iter()