[TS]: Primary keys for query builder views

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

@@ -54,7 +54,7 @@ export const isRowTypedQuery = (val: unknown): val is RowTypedQuery<any, any> =>
 export const isTypedQuery = (val: unknown): val is TableTypedQuery<any> =>
   !!val && typeof val === 'object' && QueryBrand in (val as object);
 
-export function toSql(q: Query<any>): string {
+export function toSql(q: RowTypedQuery<any, any>): string {
   return (q as unknown as { toSql(): string }).toSql();
 }
 

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

@@ -19,6 +19,7 @@ import {
   ArrayBuilder,
   OptionBuilder,
   ProductBuilder,
+  QueryTypeBuilder,
   RefBuilder,
   ResultBuilder,
   RowBuilder,
@@ -307,6 +308,10 @@ export class ModuleContext {
       return new ArrayBuilder(
         this.registerTypesRecursively(typeBuilder.element)
       ) as any;
+    } else if (typeBuilder instanceof QueryTypeBuilder) {
+      return new QueryTypeBuilder(
+        this.registerTypesRecursively(typeBuilder.row)
+      ) as any;
     } else {
       return typeBuilder as any;
     }

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

@@ -15,6 +15,8 @@ import { Uuid, type UuidAlgebraicType } from './uuid';
 // Used in codegen files
 export { type AlgebraicTypeType } from './algebraic_type';
 
+export const QUERY_VIEW_RETURN_TAG = '__query__' as const;
+
 /**
  * Helper type to extract the TypeScript type from a TypeBuilder
  */
@@ -1295,6 +1297,43 @@ export class ArrayBuilder<Element extends TypeBuilder<any, any>>
   }
 }
 
+type QueryReturnType<Row extends TypeBuilder<object, any>> = {
+  tag: 'Product';
+  value: {
+    elements: [
+      {
+        name: typeof QUERY_VIEW_RETURN_TAG;
+        algebraicType: InferSpacetimeTypeOfTypeBuilder<Row>;
+      },
+    ];
+  };
+};
+
+export class QueryTypeBuilder<
+  Row extends TypeBuilder<object, any>,
+> extends TypeBuilder<
+  readonly InferTypeOfTypeBuilder<Row>[],
+  QueryReturnType<Row>
+> {
+  readonly row: Row;
+
+  constructor(row: Row) {
+    super(
+      AlgebraicType.Product({
+        elements: [
+          {
+            name: QUERY_VIEW_RETURN_TAG,
+            get algebraicType() {
+              return row.algebraicType;
+            },
+          },
+        ],
+      }) as QueryReturnType<Row>
+    );
+    this.row = row;
+  }
+}
+
 export class ByteArrayBuilder
   extends TypeBuilder<
     Uint8Array,
@@ -3871,6 +3910,16 @@ export const t = {
     return new ArrayBuilder(e);
   },
 
+  /**
+   * Creates a return type marker for query-builder views.
+   *
+   * This encodes as the special SATS product shape `{ __query__: T }`,
+   * where `T` is the row product type returned by the query.
+   */
+  query<Row extends TypeBuilder<object, any>>(row: Row): QueryTypeBuilder<Row> {
+    return new QueryTypeBuilder(row);
+  },
+
   enum: enumImpl,
 
   /**

crates/bindings-typescript/src/server/view.test-d.ts

@@ -71,13 +71,23 @@ const spacetime = schema({
   personWithMissing,
 });
 
+const queryRetValue = t.query(person.rowType);
 const arrayRetValue = t.array(person.rowType);
 const optionalPerson = t.option(person.rowType);
 
-spacetime.anonymousView({ name: 'v1', public: true }, arrayRetValue, ctx => {
+spacetime.anonymousView({ name: 'v1', public: true }, queryRetValue, ctx => {
   return ctx.from.person.build();
 });
 
+// Legacy compatibility: query-builder views can still be declared with array return types.
+spacetime.anonymousView(
+  { name: 'v1_legacy_array_query', public: true },
+  arrayRetValue,
+  ctx => {
+    return ctx.from.person.build();
+  }
+);
+
 spacetime.anonymousView(
   { name: 'optionalPerson', public: true },
   optionalPerson,
@@ -106,7 +116,7 @@ spacetime.anonymousView(
 
 spacetime.anonymousView(
   { name: 'v2', public: true },
-  arrayRetValue,
+  queryRetValue,
   // @ts-expect-error returns a query of the wrong type.
   ctx => {
     return ctx.from.order.build();
@@ -116,7 +126,7 @@ spacetime.anonymousView(
 // For queries, we can't return rows with extra fields.
 spacetime.anonymousView(
   { name: 'v3', public: true },
-  arrayRetValue,
+  queryRetValue,
   // @ts-expect-error returns a query of the wrong type.
   ctx => {
     return ctx.from.personWithExtra.build();
@@ -126,7 +136,7 @@ spacetime.anonymousView(
 // Ideally this would fail, since we depend on the field ordering for serialization.
 spacetime.anonymousView(
   { name: 'reorderedPerson', public: true },
-  arrayRetValue,
+  queryRetValue,
   // Comment this out if we can fix the types.
   // // @ts-expect-error returns a query of the wrong type.
   ctx => {
@@ -137,21 +147,21 @@ spacetime.anonymousView(
 // Fails because it is missing a field.
 spacetime.anonymousView(
   { name: 'missingField', public: true },
-  arrayRetValue,
+  queryRetValue,
   // @ts-expect-error returns a query of the wrong type.
   ctx => {
     return ctx.from.personWithMissing.build();
   }
 );
 
-spacetime.anonymousView({ name: 'v4', public: true }, arrayRetValue, ctx => {
+spacetime.anonymousView({ name: 'v4', public: true }, queryRetValue, ctx => {
   // @ts-expect-error returns a query of the wrong type.
   const _invalid = ctx.from.person.where(row => row.id.eq('string')).build();
   const _columnEqs = ctx.from.person.where(row => row.id.eq(row.id)).build();
   return ctx.from.person.where(row => row.id.eq(5)).build();
 });
 
-spacetime.anonymousView({ name: 'v5', public: true }, arrayRetValue, ctx => {
+spacetime.anonymousView({ name: 'v5', public: true }, queryRetValue, ctx => {
   const _nonIndexedSemijoin = ctx.from.person
     .where(row => row.id.eq(5))
     // @ts-expect-error person_id is not indexed.

crates/bindings-typescript/src/server/views.ts

@@ -10,6 +10,7 @@ import type { OptionAlgebraicType } from '../lib/option';
 import type { ParamsObj } from '../lib/reducers';
 import { type UntypedSchemaDef } from '../lib/schema';
 import {
+  QueryTypeBuilder,
   RowBuilder,
   type Infer,
   type InferSpacetimeTypeOfTypeBuilder,
@@ -90,6 +91,58 @@ export type ViewOpts = {
 
 type FlattenedArray<T> = T extends readonly (infer E)[] ? E : never;
 
+type ArrayViewReturnTypeBuilder = TypeBuilder<
+  readonly object[],
+  { tag: 'Array'; value: AlgebraicTypeVariants.Product }
+>;
+
+type OptionViewReturnTypeBuilder = TypeBuilder<
+  object | undefined,
+  OptionAlgebraicType<AlgebraicTypeVariants.Product>
+>;
+
+type ProceduralViewReturnTypeBuilder =
+  | ArrayViewReturnTypeBuilder
+  | OptionViewReturnTypeBuilder;
+
+export type QueryViewReturnTypeBuilder = QueryTypeBuilder<
+  TypeBuilder<object, any>
+>;
+
+export type ViewReturnTypeBuilder =
+  | ProceduralViewReturnTypeBuilder
+  | QueryViewReturnTypeBuilder;
+
+type ExtractProductFromTypeBuilder<T extends TypeBuilder<any, any>> =
+  InferSpacetimeTypeOfTypeBuilder<T> extends { tag: 'Product'; value: infer P }
+    ? P
+    : never;
+
+type QueryReturnRow<Ret extends QueryViewReturnTypeBuilder> =
+  Ret extends QueryTypeBuilder<infer Row> ? Infer<Row> : never;
+
+type QueryReturnProduct<Ret extends QueryViewReturnTypeBuilder> =
+  Ret extends QueryTypeBuilder<infer Row>
+    ? ExtractProductFromTypeBuilder<Row>
+    : never;
+
+type ExtractArrayProduct<T extends TypeBuilder<any, any>> =
+  InferSpacetimeTypeOfTypeBuilder<T> extends { tag: 'Array'; value: infer V }
+    ? V extends { tag: 'Product'; value: infer P }
+      ? P
+      : never
+    : never;
+
+type LegacyArrayQueryReturn<Ret extends ArrayViewReturnTypeBuilder> =
+  RowTypedQuery<FlattenedArray<Infer<Ret>>, ExtractArrayProduct<Ret>>;
+
+type ViewReturn<Ret extends ViewReturnTypeBuilder> =
+  Ret extends QueryViewReturnTypeBuilder
+    ? RowTypedQuery<QueryReturnRow<Ret>, QueryReturnProduct<Ret>>
+    : Ret extends ArrayViewReturnTypeBuilder
+      ? Infer<Ret> | LegacyArrayQueryReturn<Ret>
+      : Infer<Ret>;
+
 // // If we allowed functions to return either.
 // type ViewReturn<Ret extends ViewReturnTypeBuilder> =
 //   | Infer<Ret>
@@ -99,33 +152,16 @@ export type ViewFn<
   S extends UntypedSchemaDef,
   Params extends ParamsObj,
   Ret extends ViewReturnTypeBuilder,
-> =
-  | ((ctx: ViewCtx<S>, params: InferTypeOfRow<Params>) => Infer<Ret>)
-  | ((
-      ctx: ViewCtx<S>,
-      params: InferTypeOfRow<Params>
-    ) => RowTypedQuery<FlattenedArray<Infer<Ret>>, ExtractArrayProduct<Ret>>);
+> = (ctx: ViewCtx<S>, params: InferTypeOfRow<Params>) => ViewReturn<Ret>;
 
 export type AnonymousViewFn<
   S extends UntypedSchemaDef,
   Params extends ParamsObj,
   Ret extends ViewReturnTypeBuilder,
-> =
-  | ((ctx: AnonymousViewCtx<S>, params: InferTypeOfRow<Params>) => Infer<Ret>)
-  | ((
-      ctx: AnonymousViewCtx<S>,
-      params: InferTypeOfRow<Params>
-    ) => RowTypedQuery<FlattenedArray<Infer<Ret>>, ExtractArrayProduct<Ret>>);
-
-export type ViewReturnTypeBuilder =
-  | TypeBuilder<
-      readonly object[],
-      { tag: 'Array'; value: AlgebraicTypeVariants.Product }
-    >
-  | TypeBuilder<
-      object | undefined,
-      OptionAlgebraicType<AlgebraicTypeVariants.Product>
-    >;
+> = (
+  ctx: AnonymousViewCtx<S>,
+  params: InferTypeOfRow<Params>
+) => ViewReturn<Ret>;
 
 export function registerView<
   S extends UntypedSchemaDef,
@@ -202,12 +238,3 @@ type ViewInfo<F> = {
 
 export type Views = ViewInfo<ViewFn<any, any, any>>[];
 export type AnonViews = ViewInfo<AnonymousViewFn<any, any, any>>[];
-
-// A helper to get the product type out of a type builder.
-// This is only non-never if the type builder is an array.
-type ExtractArrayProduct<T extends TypeBuilder<any, any>> =
-  InferSpacetimeTypeOfTypeBuilder<T> extends { tag: 'Array'; value: infer V }
-    ? V extends { tag: 'Product'; value: infer P }
-      ? P
-      : never
-    : never;