Add primary key support for procedural views to rust and ts modules

Author: joshua-spacetime

Cargo.lock

@@ -50,6 +50,7 @@ members = [
   "modules/sdk-test-view",
   "modules/sdk-test-case-conversion",
   "modules/sdk-test-view-pk",
+  "modules/sdk-test-procedural-view-pk",
   "modules/sdk-test-event-table",
   "sdks/rust/tests/test-client",
   "sdks/rust/tests/test-counter",
@@ -58,6 +59,7 @@ members = [
   "sdks/rust/tests/view-client",
   "sdks/rust/tests/case-conversion-client",
   "sdks/rust/tests/view-pk-client",
+  "sdks/rust/tests/procedural-view-pk-client",
   "sdks/rust/tests/event-table-client",
   "tools/ci",
   "tools/upgrade-version",

Cargo.toml

@@ -12,15 +12,49 @@ use crate::util::{check_duplicate_msg, match_meta};
 pub(crate) struct ViewArgs {
     name: Option<LitStr>,
     accessor: Ident,
+    primary_key: Option<ViewPrimaryKeyArg>,
     #[allow(unused)]
     public: bool,
 }
 
+/// Argument accepted by `#[view(primary_key = ...)]`.
+///
+/// Both identifier and string literal syntax is supported.
+enum ViewPrimaryKeyArg {
+    Ident(Ident),
+    Literal(LitStr),
+}
+
+impl ViewPrimaryKeyArg {
+    fn parse(input: syn::parse::ParseStream<'_>) -> syn::Result<Self> {
+        if input.peek(LitStr) {
+            input.parse().map(Self::Literal)
+        } else {
+            input.parse().map(Self::Ident)
+        }
+    }
+
+    fn name(&self) -> String {
+        match self {
+            Self::Ident(ident) => ident.unraw().to_string(),
+            Self::Literal(lit) => lit.value(),
+        }
+    }
+
+    fn ident(&self) -> Option<&Ident> {
+        match self {
+            Self::Ident(ident) => Some(ident),
+            Self::Literal(_) => None,
+        }
+    }
+}
+
 impl ViewArgs {
     /// Parse `#[view(accessor = ..., public)]` where both `name` and `public` are required.
     pub(crate) fn parse(input: TokenStream, func_ident: &Ident) -> syn::Result<Self> {
         let mut name = None;
         let mut accessor = None;
+        let mut primary_key = None;
         let mut public = None;
         syn::meta::parser(|meta| {
             match_meta!(match meta {
@@ -36,6 +70,10 @@ impl ViewArgs {
                     check_duplicate_msg(&accessor, &meta, "`accessor` already specified")?;
                     accessor = Some(meta.value()?.parse()?);
                 }
+                sym::primary_key => {
+                    check_duplicate_msg(&primary_key, &meta, "`primary_key` already specified")?;
+                    primary_key = Some(ViewPrimaryKeyArg::parse(meta.value()?)?);
+                }
             });
             Ok(())
         })
@@ -51,6 +89,7 @@ impl ViewArgs {
             .ok_or_else(|| syn::Error::new(Span::call_site(), "views must be `public`, e.g. `#[view(public)]`"))?;
         Ok(Self {
             name,
+            primary_key,
             public: true,
             accessor,
         })
@@ -74,6 +113,29 @@ fn extract_impl_query_inner(ty: &syn::Type) -> Option<&syn::Type> {
     None
 }
 
+/// If `ty` is a supported view return type, returns the row type `T`.
+fn extract_view_return_row_type(ty: &syn::Type) -> Option<&syn::Type> {
+    if let Some(inner) = extract_impl_query_inner(ty) {
+        return Some(inner);
+    }
+
+    let syn::Type::Path(path) = ty else {
+        return None;
+    };
+
+    let seg = path.path.segments.last()?;
+    if !matches!(seg.ident.to_string().as_str(), "Vec" | "Option" | "RawQuery") {
+        return None;
+    }
+    let syn::PathArguments::AngleBracketed(args) = &seg.arguments else {
+        return None;
+    };
+    let Some(syn::GenericArgument::Type(inner)) = args.args.first() else {
+        return None;
+    };
+    Some(inner)
+}
+
 pub(crate) fn view_impl(args: ViewArgs, original_function: &ItemFn) -> syn::Result<TokenStream> {
     let vis = &original_function.vis;
     let func_name = &original_function.sig.ident;
@@ -221,6 +283,24 @@ pub(crate) fn view_impl(args: ViewArgs, original_function: &ItemFn) -> syn::Resu
     };
 
     let eff_ret_ty = &effective_ret_ty;
+    let primary_key_column_name = args.primary_key.as_ref().map(ViewPrimaryKeyArg::name);
+    let primary_key_field_check = args
+        .primary_key
+        .as_ref()
+        .and_then(ViewPrimaryKeyArg::ident)
+        .zip(extract_view_return_row_type(ret_ty))
+        .map(|(primary_key, row_ty)| {
+            quote! {
+                const _: () = {
+                    fn _assert_view_primary_key_column #lt_params (__row: &#row_ty) #lt_where_clause {
+                        let _ = &__row.#primary_key;
+                    }
+                };
+            }
+        });
+    let primary_key_column_const = primary_key_column_name
+        .as_ref()
+        .map(|primary_key| quote! { const VIEW_PRIMARY_KEY_COLUMNS: &'static [&'static str] = &[#primary_key]; });
 
     Ok(quote! {
         #emitted_fn
@@ -243,6 +323,8 @@ pub(crate) fn view_impl(args: ViewArgs, original_function: &ItemFn) -> syn::Resu
             }
         };
 
+        #primary_key_field_check
+
         impl #func_name {
             fn invoke(__ctx: #ctx_ty, __args: &[u8]) -> Vec<u8> {
                 spacetimedb::rt::ViewDispatcher::<#ctx_ty>::invoke::<_, _, _>(#func_name, __ctx, __args)
@@ -266,6 +348,8 @@ pub(crate) fn view_impl(args: ViewArgs, original_function: &ItemFn) -> syn::Resu
             /// The pointer for invoking this function
             const INVOKE: Self::Invoke = #func_name::invoke;
 
+            #primary_key_column_const
+
             /// The return type of this function
             fn return_type(
                 ts: &mut impl spacetimedb::sats::typespace::TypespaceBuilder

crates/bindings-macro/src/view.rs

@@ -194,6 +194,7 @@ export class ModuleContext {
     schedules: [],
     procedures: [],
     views: [],
+    viewPrimaryKeys: [],
     lifeCycleReducers: [],
     caseConversionPolicy: { tag: 'SnakeCase' },
     explicitNames: {
@@ -220,6 +221,12 @@ export class ModuleContext {
     push(module.reducers && { tag: 'Reducers', value: module.reducers });
     push(module.procedures && { tag: 'Procedures', value: module.procedures });
     push(module.views && { tag: 'Views', value: module.views });
+    push(
+      module.viewPrimaryKeys && {
+        tag: 'ViewPrimaryKeys',
+        value: module.viewPrimaryKeys,
+      }
+    );
     push(module.schedules && { tag: 'Schedules', value: module.schedules });
     push(
       module.lifeCycleReducers && {

crates/bindings-typescript/src/lib/autogen/types.ts

@@ -38,6 +38,7 @@ import {
   type ViewFn,
   type ViewOpts,
   type ViewReturnTypeBuilder,
+  type ValidateViewPrimaryKey,
   type Views,
 } from './views';
 import type { UntypedTableDef } from '../lib/table';
@@ -347,7 +348,11 @@ export class Schema<S extends UntypedSchemaDef> implements ModuleDefaultExport {
   view<Ret extends ViewReturnTypeBuilder, F extends ViewFn<S, {}, Ret>>(
     opts: ViewOpts,
     ret: Ret,
-    fn: F
+    fn: F,
+    // Compile-time-only guard: this rest parameter is `[]` for valid return
+    // builders, but becomes a required error tuple when a returned row builder
+    // marks more than one column with `.primaryKey()`.
+    ..._: ValidateViewPrimaryKey<Ret>
   ): ViewExport<F> {
     return makeViewExport<S, {}, Ret, F>(this.#ctx, opts, {}, ret, fn);
   }
@@ -380,7 +385,15 @@ export class Schema<S extends UntypedSchemaDef> implements ModuleDefaultExport {
   anonymousView<
     Ret extends ViewReturnTypeBuilder,
     F extends AnonymousViewFn<S, {}, Ret>,
-  >(opts: ViewOpts, ret: Ret, fn: F): ViewExport<F> {
+  >(
+    opts: ViewOpts,
+    ret: Ret,
+    fn: F,
+    // Compile-time-only guard: this rest parameter is `[]` for valid return
+    // builders, but becomes a required error tuple when a returned row builder
+    // marks more than one column with `.primaryKey()`.
+    ..._: ValidateViewPrimaryKey<Ret>
+  ): ViewExport<F> {
     return makeAnonViewExport<S, {}, Ret, F>(this.#ctx, opts, {}, ret, fn);
   }
 

crates/bindings-typescript/src/lib/schema.ts

@@ -73,11 +73,24 @@ const spacetime = schema({
 
 const arrayRetValue = t.array(person.rowType);
 const optionalPerson = t.option(person.rowType);
+const multiplePrimaryKeyRows = t.array(
+  t.row('MultiplePrimaryKeyRows', {
+    id: t.u32().primaryKey(),
+    name: t.string().primaryKey(),
+  })
+);
 
 spacetime.anonymousView({ name: 'v1', public: true }, arrayRetValue, ctx => {
   return ctx.from.person.build();
 });
 
+// @ts-expect-error views can have at most one primaryKey column on the returned row type.
+spacetime.anonymousView(
+  { name: 'multiplePrimaryRows', public: true },
+  multiplePrimaryKeyRows,
+  () => []
+);
+
 spacetime.anonymousView(
   { name: 'optionalPerson', public: true },
   optionalPerson,