[Rust] Primary keys for query builder views (#4572)

# Description of Changes

Query builder views return a subset of rows from a physical table. If
that table has a primary key, then so should the view. What this means
concretely is that the view should expose the same api as the table,
specifically as it relates to the primary key column.

With that in mind, this patch commits the following changes:
1. Annotates `ViewDef` with a `primary_key`
2. Updates the return type of query builder views in the raw module def
to a special product type
3. Adds an index for the primary key on the view's backing table
4. Updates the query planner to use this index
5. Updates rust client codegen to generate `on_update` for such views

# API and ABI breaking changes

None

Old `impl Query` views compiled with an older version of SpacetimeDB
will continue to work as they did before - without a primary key.

# Expected complexity level and risk

3

# Testing

- [x] New rust sdk integration suite exercising `on_update` for PK views
and semijoin scenarios
- [x] Smoketests for PK views and semijoin scenarios

Author: joshua-spacetime

Cargo.lock

@@ -50,12 +50,14 @@ members = [
   "modules/sdk-test-connect-disconnect",
   "modules/sdk-test-procedure",
   "modules/sdk-test-view",
+  "modules/sdk-test-view-pk",
   "modules/sdk-test-event-table",
   "sdks/rust/tests/test-client",
   "sdks/rust/tests/test-counter",
   "sdks/rust/tests/connect_disconnect_client",
   "sdks/rust/tests/procedure-client",
   "sdks/rust/tests/view-client",
+  "sdks/rust/tests/view-pk-client",
   "sdks/rust/tests/event-table-client",
   "tools/ci",
   "tools/upgrade-version",

Cargo.toml

@@ -77,27 +77,27 @@ error[E0425]: cannot find type `T` in this scope
 201 | fn view_nonexistent_table(ctx: &ViewContext) -> impl Query<T> {
     |                                                            ^ not found in this scope
 
-error[E0277]: the trait bound `ViewKind<ReducerContext>: ViewKindTrait` is not satisfied
+error[E0277]: the trait bound `spacetimedb::rt::ViewKind<ReducerContext>: ViewKindTrait` is not satisfied
    --> tests/ui/views.rs:106:1
     |
 106 | #[view(accessor = view_def_wrong_context, public)]
-    | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ the trait `ViewKindTrait` is not implemented for `ViewKind<ReducerContext>`
+    | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ the trait `ViewKindTrait` is not implemented for `spacetimedb::rt::ViewKind<ReducerContext>`
     |
 help: the following other types implement trait `ViewKindTrait`
    --> src/rt.rs
     |
     | impl ViewKindTrait for ViewKind<ViewContext> {
-    | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ `ViewKind<ViewContext>`
+    | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ `spacetimedb::rt::ViewKind<ViewContext>`
 ...
     | impl ViewKindTrait for ViewKind<AnonymousViewContext> {
-    | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ `ViewKind<AnonymousViewContext>`
+    | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ `spacetimedb::rt::ViewKind<AnonymousViewContext>`
     = note: this error originates in the attribute macro `view` (in Nightly builds, run with -Z macro-backtrace for more info)
 
 error[E0276]: impl has stricter requirements than trait
    --> tests/ui/views.rs:106:1
     |
 106 | #[view(accessor = view_def_wrong_context, public)]
-    | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ impl has extra requirement `ViewKind<ReducerContext>: ViewKindTrait`
+    | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ impl has extra requirement `spacetimedb::rt::ViewKind<ReducerContext>: ViewKindTrait`
     |
     = note: this error originates in the attribute macro `view` (in Nightly builds, run with -Z macro-backtrace for more info)
 

crates/bindings/tests/ui/views.stderr

@@ -1411,9 +1411,19 @@ impl __sdk::InModule for DbUpdate {{
                     }
                     for view in iter_views(module) {
                         let field_name = table_method_name(&view.accessor_name);
+                        let with_updates = view
+                            .primary_key
+                            .map(|col| {
+                                let pk_field = view.return_columns[col.idx()]
+                                    .accessor_name
+                                    .deref()
+                                    .to_case(Case::Snake);
+                                format!(".with_updates_by_pk(|row| &row.{pk_field})")
+                            })
+                            .unwrap_or_default();
                         writeln!(
                             out,
-                            "diff.{field_name} = cache.apply_diff_to_table::<{}>({:?}, &self.{field_name});",
+                            "diff.{field_name} = cache.apply_diff_to_table::<{}>({:?}, &self.{field_name}){with_updates};",
                             type_ref_name(module, view.product_type_ref),
                             view.name.deref(),
                         );

crates/codegen/src/rust.rs

@@ -341,6 +341,7 @@ impl From<PhysicalPlan> for PipelinedExecutor {
                     lhs,
                     rhs,
                     rhs_index,
+                    rhs_prefix,
                     rhs_field,
                     unique,
                     lhs_field,
@@ -352,6 +353,7 @@ impl From<PhysicalPlan> for PipelinedExecutor {
                 lhs: Box::new(Self::from(*lhs)),
                 rhs_table: rhs.table_id,
                 rhs_index,
+                rhs_prefix,
                 rhs_field,
                 lhs_field,
                 unique,
@@ -362,6 +364,7 @@ impl From<PhysicalPlan> for PipelinedExecutor {
                     lhs,
                     rhs,
                     rhs_index,
+                    rhs_prefix,
                     rhs_field,
                     unique,
                     lhs_field,
@@ -373,6 +376,7 @@ impl From<PhysicalPlan> for PipelinedExecutor {
                 lhs: Box::new(Self::from(*lhs)),
                 rhs_table: rhs.table_id,
                 rhs_index,
+                rhs_prefix,
                 rhs_field,
                 rhs_delta,
                 lhs_field,
@@ -923,6 +927,16 @@ fn combine_prefix_and_last(prefix: Vec<(ColId, AlgebraicValue)>, last: Algebraic
     }
 }
 
+fn combine_probe_prefix_and_last(prefix: &[AlgebraicValue], last: AlgebraicValue) -> AlgebraicValue {
+    if prefix.is_empty() {
+        last
+    } else {
+        AlgebraicValue::product(ProductValue::from_iter(
+            prefix.iter().cloned().chain(std::iter::once(last)),
+        ))
+    }
+}
+
 impl PipelinedIxScanEq {
     /// We don't know statically if an index scan will return rows
     pub fn is_empty(&self, _: &impl DeltaStore) -> bool {
@@ -969,6 +983,8 @@ pub struct PipelinedIxJoin {
     pub rhs_table: TableId,
     /// The rhs index
     pub rhs_index: IndexId,
+    /// Constant prefix values for multi-column index probes.
+    pub rhs_prefix: Vec<AlgebraicValue>,
     /// The rhs join field
     pub rhs_field: ColId,
     /// The lhs join field
@@ -996,7 +1012,7 @@ impl PipelinedIxJoin {
         let mut bytes_scanned = 0;
 
         let iter_rhs = |u: &Tuple, lhs_field: &TupleField, bytes_scanned: &mut usize| -> Result<_> {
-            let key = project(u, lhs_field, bytes_scanned);
+            let key = combine_probe_prefix_and_last(&self.rhs_prefix, project(u, lhs_field, bytes_scanned));
             Ok(tx
                 .index_scan_point(self.rhs_table, self.rhs_index, &key)?
                 .map(Row::Ptr)
@@ -1135,6 +1151,8 @@ pub struct PipelinedIxDeltaJoin {
     pub rhs_delta: Delta,
     /// The rhs index
     pub rhs_index: IndexId,
+    /// Constant prefix values for multi-column index probes.
+    pub rhs_prefix: Vec<AlgebraicValue>,
     /// The rhs join field
     pub rhs_field: ColId,
     /// The lhs join field
@@ -1177,13 +1195,10 @@ impl PipelinedIxDeltaJoin {
                 lhs.execute(tx, metrics, &mut |u| {
                     n += 1;
                     index_seeks += 1;
+                    let key =
+                        combine_probe_prefix_and_last(&self.rhs_prefix, project(&u, lhs_field, &mut bytes_scanned));
                     if tx
-                        .index_scan_point_for_delta(
-                            self.rhs_table,
-                            self.rhs_index,
-                            self.rhs_delta,
-                            &project(&u, lhs_field, &mut bytes_scanned),
-                        )
+                        .index_scan_point_for_delta(self.rhs_table, self.rhs_index, self.rhs_delta, &key)
                         .next()
                         .is_some()
                     {
@@ -1203,13 +1218,10 @@ impl PipelinedIxDeltaJoin {
                 lhs.execute(tx, metrics, &mut |u| {
                     n += 1;
                     index_seeks += 1;
+                    let key =
+                        combine_probe_prefix_and_last(&self.rhs_prefix, project(&u, lhs_field, &mut bytes_scanned));
                     if let Some(v) = tx
-                        .index_scan_point_for_delta(
-                            self.rhs_table,
-                            self.rhs_index,
-                            self.rhs_delta,
-                            &project(&u, lhs_field, &mut bytes_scanned),
-                        )
+                        .index_scan_point_for_delta(self.rhs_table, self.rhs_index, self.rhs_delta, &key)
                         .next()
                         .map(Tuple::Row)
                     {
@@ -1229,13 +1241,10 @@ impl PipelinedIxDeltaJoin {
                 lhs.execute(tx, metrics, &mut |u| {
                     n += 1;
                     index_seeks += 1;
+                    let key =
+                        combine_probe_prefix_and_last(&self.rhs_prefix, project(&u, lhs_field, &mut bytes_scanned));
                     if let Some(v) = tx
-                        .index_scan_point_for_delta(
-                            self.rhs_table,
-                            self.rhs_index,
-                            self.rhs_delta,
-                            &project(&u, lhs_field, &mut bytes_scanned),
-                        )
+                        .index_scan_point_for_delta(self.rhs_table, self.rhs_index, self.rhs_delta, &key)
                         .next()
                         .map(Tuple::Row)
                     {
@@ -1256,13 +1265,10 @@ impl PipelinedIxDeltaJoin {
                 lhs.execute(tx, metrics, &mut |u| {
                     n += 1;
                     index_seeks += 1;
+                    let key =
+                        combine_probe_prefix_and_last(&self.rhs_prefix, project(&u, lhs_field, &mut bytes_scanned));
                     for _ in 0..tx
-                        .index_scan_point_for_delta(
-                            self.rhs_table,
-                            self.rhs_index,
-                            self.rhs_delta,
-                            &project(&u, lhs_field, &mut bytes_scanned),
-                        )
+                        .index_scan_point_for_delta(self.rhs_table, self.rhs_index, self.rhs_delta, &key)
                         .count()
                     {
                         f(u.clone())?;
@@ -1281,13 +1287,10 @@ impl PipelinedIxDeltaJoin {
                 lhs.execute(tx, metrics, &mut |u| {
                     n += 1;
                     index_seeks += 1;
+                    let key =
+                        combine_probe_prefix_and_last(&self.rhs_prefix, project(&u, lhs_field, &mut bytes_scanned));
                     for v in tx
-                        .index_scan_point_for_delta(
-                            self.rhs_table,
-                            self.rhs_index,
-                            self.rhs_delta,
-                            &project(&u, lhs_field, &mut bytes_scanned),
-                        )
+                        .index_scan_point_for_delta(self.rhs_table, self.rhs_index, self.rhs_delta, &key)
                         .map(Tuple::Row)
                     {
                         f(v)?;
@@ -1306,13 +1309,10 @@ impl PipelinedIxDeltaJoin {
                 lhs.execute(tx, metrics, &mut |u| {
                     n += 1;
                     index_seeks += 1;
+                    let key =
+                        combine_probe_prefix_and_last(&self.rhs_prefix, project(&u, lhs_field, &mut bytes_scanned));
                     for v in tx
-                        .index_scan_point_for_delta(
-                            self.rhs_table,
-                            self.rhs_index,
-                            self.rhs_delta,
-                            &project(&u, lhs_field, &mut bytes_scanned),
-                        )
+                        .index_scan_point_for_delta(self.rhs_table, self.rhs_index, self.rhs_delta, &key)
                         .map(Tuple::Row)
                     {
                         f(u.clone().join(v.clone()))?;

crates/execution/src/pipelined.rs

@@ -4,3 +4,4 @@ pub mod attr;
 pub mod auth;
 pub mod default_element_ordering;
 pub mod raw_def;
+pub mod view;

crates/lib/src/db/mod.rs

@@ -26,6 +26,7 @@ use crate::db::raw_def::v10::RawConstraintDefV10;
 use crate::db::raw_def::v10::RawScopedTypeNameV10;
 use crate::db::raw_def::v10::RawSequenceDefV10;
 use crate::db::raw_def::v10::RawTypeDefV10;
+use crate::db::view::extract_view_return_product_type_ref;
 
 /// A not-yet-validated `sql`.
 pub type RawSql = Box<str>;
@@ -124,16 +125,7 @@ impl RawModuleDefV9 {
     fn type_ref_for_view(&self, view_name: &str) -> Option<AlgebraicTypeRef> {
         self.find_view_def(view_name)
             .map(|view_def| &view_def.return_type)
-            .and_then(|return_type| {
-                return_type
-                    .as_option()
-                    .and_then(|inner| inner.clone().into_ref().ok())
-                    .or_else(|| {
-                        return_type
-                            .as_array()
-                            .and_then(|inner| inner.elem_ty.clone().into_ref().ok())
-                    })
-            })
+            .and_then(|return_type| extract_view_return_product_type_ref(return_type).map(|(ref_, _)| ref_))
     }
 
     /// Find and return the product type ref for a table or view in this module def

crates/lib/src/db/raw_def/v9.rs

@@ -0,0 +1,32 @@
+use spacetimedb_sats::{AlgebraicType, AlgebraicTypeRef};
+
+pub const QUERY_VIEW_RETURN_TAG: &str = "__query__";
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum ViewKind {
+    Procedural,
+    Query,
+}
+
+pub fn extract_view_return_product_type_ref(return_type: &AlgebraicType) -> Option<(AlgebraicTypeRef, ViewKind)> {
+    // Query-builder views (`Query<T>`) are encoded as: { __query__: T }.
+    if let Some(product) = return_type.as_product()
+        && product.elements.len() == 1
+        && product.elements[0].name.as_deref() == Some(QUERY_VIEW_RETURN_TAG)
+        && let Some(product_type_ref) = product.elements[0].algebraic_type.as_ref().copied()
+    {
+        return Some((product_type_ref, ViewKind::Query));
+    }
+
+    return_type
+        .as_option()
+        .and_then(AlgebraicType::as_ref)
+        .or_else(|| {
+            return_type
+                .as_array()
+                .map(|array_type| array_type.elem_ty.as_ref())
+                .and_then(AlgebraicType::as_ref)
+        })
+        .copied()
+        .map(|product_type_ref| (product_type_ref, ViewKind::Procedural))
+}

crates/lib/src/db/view.rs

@@ -575,6 +575,7 @@ impl PhysicalPlan {
                     rhs,
                     rhs_label,
                     rhs_index,
+                    rhs_prefix,
                     rhs_field,
                     unique,
                     lhs_field,
@@ -587,6 +588,7 @@ impl PhysicalPlan {
                     rhs,
                     rhs_label,
                     rhs_index,
+                    rhs_prefix,
                     rhs_field,
                     unique,
                     lhs_field,
@@ -1252,6 +1254,8 @@ pub struct IxJoin {
     pub rhs_label: Label,
     /// The index id
     pub rhs_index: IndexId,
+    /// Optional constant prefix values for multi-column index probes.
+    pub rhs_prefix: Vec<AlgebraicValue>,
     /// The index field
     pub rhs_field: ColId,
     /// Is the index a unique constraint index?
@@ -1498,7 +1502,7 @@ mod tests {
     ) -> TableOrViewSchema {
         TableOrViewSchema::from(Arc::new(TableSchema::new(
             table_id,
-            TableName::for_test(table_name),
+            TableName::new(Identifier::for_test(table_name)),
             None,
             columns
                 .iter()