Add submodule 'merge' api in typescript

Author: aasoni

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

@@ -1,6 +1,7 @@
 export * from '../lib/type_builders';
 export {
   schema,
+  merge,
   type InferSchema,
   type ModuleExport,
   type ModuleSettings,

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

@@ -79,6 +79,8 @@ export class SchemaInner<
   > = new Map();
   pendingSchedules: PendingSchedule[] = [];
   mountedDispatchInfos: MountedDispatchInfo[] = [];
+  pendingMergedExports: Array<[string, ModuleExport]> = [];
+  mergedSchemas: Set<SchemaInner> = new Set();
 
   constructor(getSchemaType: (ctx: SchemaInner<S>) => S) {
     super();
@@ -185,6 +187,14 @@ export class Schema<S extends UntypedSchemaDef> implements ModuleDefaultExport {
     exports: object,
     opts?: { ignoreNonModuleExports?: boolean }
   ): RawModuleDefV10 {
+    // Register merged exports first so they are in ctx.reducers before
+    // registerModuleExports runs. Skip any already registered (e.g. if the
+    // consumer re-exports the same reducer by name).
+    for (const [name, exp] of this.#ctx.pendingMergedExports) {
+      if (!this.#ctx.functionExports.has(exp as ReducerExport<any, any>)) {
+        exp[registerExport](this.#ctx, name);
+      }
+    }
     registerModuleExports(this.#ctx, exports, {
       ignoreNonModuleExports: opts?.ignoreNonModuleExports ?? false,
     });
@@ -550,9 +560,14 @@ function isModuleExport(x: unknown): x is ModuleExport {
   );
 }
 
-/** Verify that the ModuleContext that `exp` comes from is the same as `schema` */
+/** Verify that the ModuleContext that `exp` comes from is the same as `schema`,
+ * or is a library that was merged into `schema` via merge(). */
 function checkExportContext(exp: ModuleExport, schema: SchemaInner) {
-  if (exp[exportContext] != null && exp[exportContext] !== schema) {
+  if (
+    exp[exportContext] != null &&
+    exp[exportContext] !== schema &&
+    !schema.mergedSchemas.has(exp[exportContext])
+  ) {
     throw new TypeError('multiple schemas are not supported');
   }
 }
@@ -598,7 +613,7 @@ type MountedModuleNamespace = {
   [key: string]: unknown;
 };
 
-type SchemaEntry = UntypedTableSchema | MountedModuleNamespace;
+type SchemaEntry = UntypedTableSchema | MountedModuleNamespace | ModuleExport;
 
 type ExtractTableEntries<H extends Record<string, SchemaEntry>> = {
   [K in keyof H as H[K] extends UntypedTableSchema ? K : never]: Extract<
@@ -639,6 +654,11 @@ function registerModuleExports(
       throw new TypeError('exporting something that is not a spacetime export');
     }
     checkExportContext(moduleExport, schema);
+    // Skip exports already registered via pendingMergedExports to prevent
+    // double-registration when the consumer re-exports a merged reducer by name.
+    if (schema.functionExports.has(moduleExport as ReducerExport<any, any>)) {
+      continue;
+    }
     moduleExport[registerExport](schema, name);
   }
 }
@@ -666,15 +686,31 @@ export function schema<const H extends Record<string, SchemaEntry>>(
         ctx.mountedDispatchInfos.push(dispatch);
         continue;
       }
+      if (isModuleExport(entry)) {
+        // Entry came from a merge() spread — defer registration to buildRawModuleDefV10
+        // and track the source schema so checkExportContext allows re-exports.
+        ctx.pendingMergedExports.push([accName, entry]);
+        if (entry[exportContext] != null) {
+          ctx.mergedSchemas.add(entry[exportContext]);
+        }
+        continue;
+      }
       if (!isUntypedTableSchema(entry)) {
         throw new TypeError(
-          `schema entry '${accName}' must be a table or a mounted module namespace object`
+          `schema entry '${accName}' must be a table, a mounted module namespace object, or a merge() result`
         );
       }
 
       const table = entry;
-      const tableDef = table.tableDef(ctx, accName);
-      tableSchemas[accName] = tableToSchema(accName, table, tableDef);
+      // UntypedTableSchema.tableDef is a factory fn; UntypedTableDef.tableDef is a
+      // pre-materialized RawTableDefV10 (produced by merge()). Handle both.
+      const isMaterialized = typeof (table as any).tableDef !== 'function';
+      const tableDef: RawTableDefV10 = isMaterialized
+        ? (table as unknown as UntypedTableDef).tableDef
+        : table.tableDef(ctx, accName);
+      tableSchemas[accName] = isMaterialized
+        ? (table as unknown as UntypedTableDef)
+        : tableToSchema(accName, table, tableDef);
       ctx.moduleDef.tables.push(tableDef);
       if (table.schedule) {
         ctx.pendingSchedules.push({
@@ -697,3 +733,33 @@ export function schema<const H extends Record<string, SchemaEntry>>(
 
   return new Schema(ctx);
 }
+
+/**
+ * Flattens a library's tables and named exports into a plain object suitable
+ * for spreading into `schema({...})`. The library's tables and reducers land
+ * directly in the consumer's public namespace with no namespace prefix.
+ *
+ * Multiple `merge()` calls compose without `default`-key collisions:
+ * ```ts
+ * const spacetimedb = schema({
+ *   players,
+ *   ...merge(authLib),
+ *   ...merge(utilLib),
+ * });
+ * ```
+ *
+ * Libraries with scheduled reducers or user-defined compound column types
+ * should use the namespaced mount form instead: `schema({ alias: lib })`.
+ */
+export function merge(
+  lib: MountedModuleNamespace
+): Record<string, UntypedTableDef | ModuleExport> {
+  const { default: libSchema, ...namedExports } = lib as Record<string, unknown> & {
+    default: Schema<any>;
+  };
+  const tables: Record<string, UntypedTableDef> = {};
+  for (const td of Object.values(libSchema.schemaType.tables)) {
+    tables[td.accessorName] = td;
+  }
+  return { ...tables, ...(namedExports as Record<string, ModuleExport>) };
+}

crates/bindings-typescript/tests/schema_mounts.test.ts

@@ -23,11 +23,12 @@ vi.mock('../src/server/runtime', () => ({
 
 describe('schema mounts', () => {
   let schema: typeof import('../src/server/schema').schema;
+  let merge: typeof import('../src/server/schema').merge;
   let table: typeof import('../src/lib/table').table;
   let t: typeof import('../src/lib/type_builders').t;
 
   beforeAll(async () => {
-    ({ schema } = await import('../src/server/schema'));
+    ({ schema, merge } = await import('../src/server/schema'));
     ({ table } = await import('../src/lib/table'));
     ({ t } = await import('../src/lib/type_builders'));
   });
@@ -180,4 +181,85 @@ describe('schema mounts', () => {
     // Unused variable check
     void consumerReducer;
   });
+
+  it('merge() flattens library tables and reducers into the consumer namespace', () => {
+    const sessions = table(
+      { name: 'sessions' },
+      { id: t.u64().primaryKey().autoInc() }
+    );
+    const authSchema = schema({ sessions });
+    const cleanExpiredSessions = authSchema.reducer(() => {});
+    const authLib = { default: authSchema, cleanExpiredSessions };
+
+    const players = table({ name: 'players' }, { id: t.u32().primaryKey() });
+    const consumer = schema({ players, ...merge(authLib) });
+
+    const raw = consumer.buildRawModuleDefV10({});
+
+    // No namespaced mounts — library content lives in the root module
+    const mounts = raw.sections.find(section => section.tag === 'Mounts');
+    expect(mounts?.value ?? []).toHaveLength(0);
+
+    // Both tables are in the root Tables section
+    const tables = raw.sections.find(section => section.tag === 'Tables')?.value ?? [];
+    const tableNames = tables.map((t: any) => t.sourceName);
+    expect(tableNames).toContain('players');
+    expect(tableNames).toContain('sessions');
+
+    // The merged reducer is in the root Reducers section
+    const reducers = raw.sections.find(section => section.tag === 'Reducers')?.value ?? [];
+    expect(reducers).toEqual(
+      expect.arrayContaining([
+        expect.objectContaining({ sourceName: 'cleanExpiredSessions' }),
+      ])
+    );
+  });
+
+  it('merge() allows re-exporting a merged reducer without checkExportContext error', () => {
+    const sessions = table(
+      { name: 'sessions' },
+      { id: t.u64().primaryKey().autoInc() }
+    );
+    const authSchema = schema({ sessions });
+    const cleanExpiredSessions = authSchema.reducer(() => {});
+    const authLib = { default: authSchema, cleanExpiredSessions };
+
+    const players = table({ name: 'players' }, { id: t.u32().primaryKey() });
+    const consumer = schema({ players, ...merge(authLib) });
+
+    // Simulate consumer re-exporting the merged reducer by name — should not throw
+    expect(() =>
+      consumer.buildRawModuleDefV10({ cleanExpiredSessions })
+    ).not.toThrow();
+  });
+
+  it('merge() from two libraries composes without default-key collision', () => {
+    const sessionsTable = table(
+      { name: 'sessions' },
+      { id: t.u64().primaryKey().autoInc() }
+    );
+    const authSchema = schema({ sessionsTable });
+    const authReducer = authSchema.reducer(() => {});
+    const authLib = { default: authSchema, authReducer };
+
+    const itemsTable = table({ name: 'items' }, { id: t.u32().primaryKey() });
+    const utilSchema = schema({ itemsTable });
+    const utilReducer = utilSchema.reducer(() => {});
+    const utilLib = { default: utilSchema, utilReducer };
+
+    const players = table({ name: 'players' }, { id: t.u32().primaryKey() });
+
+    // Both merges should compose cleanly
+    expect(() =>
+      schema({ players, ...merge(authLib), ...merge(utilLib) })
+    ).not.toThrow();
+
+    const consumer = schema({ players, ...merge(authLib), ...merge(utilLib) });
+    const raw = consumer.buildRawModuleDefV10({});
+
+    const reducers = raw.sections.find(section => section.tag === 'Reducers')?.value ?? [];
+    const reducerNames = reducers.map((r: any) => r.sourceName);
+    expect(reducerNames).toContain('authReducer');
+    expect(reducerNames).toContain('utilReducer');
+  });
 });