Add query engine support for mutable transactions (#3524)

# Description of Changes

<!-- Please describe your change, mention any related tickets, and so on
here. -->

The query engine can now execute queries using mutable transactions
because this patch implements the `Datastore` trait for `MutTxId`.

Note this is both a refactor and a feature patch. It is a refactor
because the `Datastore` trait was updated to allow for mutable
transactions.

# API and ABI breaking changes

<!-- If this is an API or ABI breaking change, please apply the
corresponding GitHub label. -->

None

# Expected complexity level and risk

1.5

<!--
How complicated do you think these changes are? Grade on a scale from 1
to 5,
where 1 is a trivial change, and 5 is a deep-reaching and complex
change.

This complexity rating applies not only to the complexity apparent in
the diff,
but also to its interactions with existing and future code.

If you answered more than a 2, explain what is complex about the PR,
and what other components it interacts with in potentially concerning
ways. -->

# Testing

<!-- Describe any testing you've done, and any testing you'd like your
reviewers to do,
so that you're confident that all the changes work as expected! -->

- [x] Current tests pass after refactor
- [ ] Use `MutTxId` for query execution

Author: joshua-spacetime

crates/core/src/subscription/mod.rs

@@ -174,8 +174,7 @@ where
         .filter(|(_, plan)| {
             // Since subscriptions only support selects and inner joins,
             // we filter out any plans that read from an empty table.
-            plan.table_ids()
-                .all(|table_id| tx.table(table_id).is_some_and(|t| t.row_count > 0))
+            plan.table_ids().all(|table_id| tx.row_count(table_id) > 0)
         })
         .map(|(sql, plan)| (sql, plan, plan.subscribed_table_id(), plan.subscribed_table_name()))
         .map(|(sql, plan, table_id, table_name)| {

crates/core/src/subscription/tx.rs

@@ -10,7 +10,7 @@ use smallvec::SmallVec;
 use spacetimedb_execution::{Datastore, DeltaStore, Row};
 use spacetimedb_lib::{query::Delta, AlgebraicValue, ProductValue};
 use spacetimedb_primitives::{IndexId, TableId};
-use spacetimedb_table::{blob_store::BlobStore, table::Table};
+use spacetimedb_table::table::{IndexScanRangeIter, TableScanIter};
 
 use spacetimedb_datastore::{
     locking_tx_datastore::{state_view::StateView, TxId},
@@ -119,12 +119,31 @@ impl<'a> From<&'a TxId> for DeltaTx<'a> {
 }
 
 impl Datastore for DeltaTx<'_> {
-    fn table(&self, table_id: TableId) -> Option<&Table> {
-        self.tx.table(table_id)
+    type TableIter<'a>
+        = TableScanIter<'a>
+    where
+        Self: 'a;
+
+    type IndexIter<'a>
+        = IndexScanRangeIter<'a>
+    where
+        Self: 'a;
+
+    fn row_count(&self, table_id: TableId) -> u64 {
+        self.tx.row_count(table_id)
     }
 
-    fn blob_store(&self) -> &dyn BlobStore {
-        self.tx.blob_store()
+    fn table_scan<'a>(&'a self, table_id: TableId) -> anyhow::Result<Self::TableIter<'a>> {
+        self.tx.table_scan(table_id)
+    }
+
+    fn index_scan<'a>(
+        &'a self,
+        table_id: TableId,
+        index_id: IndexId,
+        range: &impl RangeBounds<AlgebraicValue>,
+    ) -> anyhow::Result<Self::IndexIter<'a>> {
+        self.tx.index_scan(table_id, index_id, range)
     }
 }
 

crates/datastore/src/locking_tx_datastore/datastore.rs

@@ -2662,7 +2662,8 @@ mod tests {
         let table_id = datastore.create_table_mut_tx(&mut tx, table_schema)?;
         let index_id = datastore.index_id_from_name_mut_tx(&tx, "index")?.unwrap();
         let find_row_by_key = |tx: &MutTxId, key: u32| {
-            tx.index_scan_point(table_id, index_id, &key.into())
+            let key: AlgebraicValue = key.into();
+            tx.index_scan(table_id, index_id, &key)
                 .unwrap()
                 .map(|row| row.pointer())
                 .collect::<Vec<_>>()
@@ -3441,7 +3442,8 @@ mod tests {
         let mut tx = begin_mut_tx(&datastore);
         assert_eq!(tx.get_schema(table_id).unwrap().columns, columns_original);
         let index_key_types = |tx: &MutTxId| {
-            tx.table(table_id)
+            tx.committed_state_write_lock
+                .get_table(table_id)
                 .unwrap()
                 .indexes
                 .values()
@@ -3570,7 +3572,7 @@ mod tests {
             .map(|row| row.to_product_value())
             .collect::<Vec<_>>();
         assert_eq!(rows, old_rows, "Rows shouldn't be changed if rolledback");
-        let table = tx.table(rollback_table_id);
+        let table = tx.table_name(rollback_table_id);
         assert!(table.is_none(), "new table shouldn't be created if rolledback");
 
         // Add column and actually commit this time.

crates/datastore/src/locking_tx_datastore/mut_tx.rs

@@ -165,12 +165,43 @@ impl MutTxId {
 }
 
 impl Datastore for MutTxId {
-    fn blob_store(&self) -> &dyn BlobStore {
-        &self.committed_state_write_lock.blob_store
+    type TableIter<'a>
+        = IterMutTx<'a>
+    where
+        Self: 'a;
+
+    type IndexIter<'a>
+        = IndexScanRanged<'a>
+    where
+        Self: 'a;
+
+    fn row_count(&self, table_id: TableId) -> u64 {
+        self.table_row_count(table_id).unwrap_or_default()
+    }
+
+    fn table_scan<'a>(&'a self, table_id: TableId) -> anyhow::Result<Self::TableIter<'a>> {
+        Ok(self.iter(table_id)?)
     }
 
-    fn table(&self, table_id: TableId) -> Option<&Table> {
-        self.committed_state_write_lock.get_table(table_id)
+    fn index_scan<'a>(
+        &'a self,
+        table_id: TableId,
+        index_id: IndexId,
+        range: &impl RangeBounds<AlgebraicValue>,
+    ) -> anyhow::Result<Self::IndexIter<'a>> {
+        // Extract the table id, and commit/tx indices.
+        let (_, commit_index, tx_index) = self
+            .get_table_and_index(index_id)
+            .ok_or_else(|| IndexError::NotFound(index_id))?;
+
+        // Get an index seek iterator for the tx and committed state.
+        let tx_iter = tx_index.map(|i| i.seek_range(range));
+        let commit_iter = commit_index.seek_range(range);
+
+        let dt = self.tx_state.get_delete_table(table_id);
+        let iter = combine_range_index_iters(dt, tx_iter, commit_iter);
+
+        Ok(iter)
     }
 }
 

crates/datastore/src/locking_tx_datastore/tx.rs

@@ -9,11 +9,10 @@ use crate::locking_tx_datastore::state_view::IterTx;
 use spacetimedb_durability::TxOffset;
 use spacetimedb_execution::Datastore;
 use spacetimedb_lib::metrics::ExecutionMetrics;
-use spacetimedb_primitives::{ColList, TableId};
+use spacetimedb_primitives::{ColList, IndexId, TableId};
 use spacetimedb_sats::AlgebraicValue;
 use spacetimedb_schema::schema::TableSchema;
-use spacetimedb_table::blob_store::BlobStore;
-use spacetimedb_table::table::Table;
+use spacetimedb_table::table::{IndexScanRangeIter, TableScanIter};
 use std::sync::Arc;
 use std::{future, num::NonZeroU64};
 use std::{
@@ -33,12 +32,44 @@ pub struct TxId {
 }
 
 impl Datastore for TxId {
-    fn blob_store(&self) -> &dyn BlobStore {
-        &self.committed_state_shared_lock.blob_store
+    type TableIter<'a>
+        = TableScanIter<'a>
+    where
+        Self: 'a;
+
+    type IndexIter<'a>
+        = IndexScanRangeIter<'a>
+    where
+        Self: 'a;
+
+    fn row_count(&self, table_id: TableId) -> u64 {
+        self.committed_state_shared_lock
+            .table_row_count(table_id)
+            .unwrap_or_default()
     }
 
-    fn table(&self, table_id: TableId) -> Option<&Table> {
-        self.committed_state_shared_lock.get_table(table_id)
+    fn table_scan<'a>(&'a self, table_id: TableId) -> anyhow::Result<Self::TableIter<'a>> {
+        self.committed_state_shared_lock
+            .get_table(table_id)
+            .map(|table| table.scan_rows(&self.committed_state_shared_lock.blob_store))
+            .ok_or_else(|| anyhow::anyhow!("TableId `{table_id}` does not exist"))
+    }
+
+    fn index_scan<'a>(
+        &'a self,
+        table_id: TableId,
+        index_id: IndexId,
+        range: &impl RangeBounds<AlgebraicValue>,
+    ) -> anyhow::Result<Self::IndexIter<'a>> {
+        self.committed_state_shared_lock
+            .get_table(table_id)
+            .ok_or_else(|| anyhow::anyhow!("TableId `{table_id}` does not exist"))
+            .and_then(|table| {
+                table
+                    .get_index_by_id_with_table(&self.committed_state_shared_lock.blob_store, index_id)
+                    .map(|i| i.seek_range(range))
+                    .ok_or_else(|| anyhow::anyhow!("IndexId `{index_id}` does not exist"))
+            })
     }
 }