Mongo ORM: extendable Collection class + collections registry (ADR 175)

docs/architecture docs/adrs/ADR 175 - Shared ORM Collection interface.md

@@ -2,6 +2,8 @@
 
 > **Note (later supersession):** this ADR was written before [ADR 183 — Aggregation pipeline only, never find API](ADR%20183%20-%20Aggregation%20pipeline%20only,%20never%20find%20API.md). On the Mongo family, ORM reads compile to aggregation pipelines (`AggregateCommand`) only — `FindCommand` is not a peer compilation target. Wherever this ADR pairs `FindCommand / AggregateCommand`, read `AggregateCommand` only.
 
+> **Status update:** the custom-collection half of the Mongo target has landed: `@prisma-next/mongo-orm` exports an extendable `Collection` class, `mongoOrm({ collections })` / `mongo({ collections })` register subclasses by model name, and chaining preserves the subclass (clone-through-`this.constructor`, as in the SQL ORM). The shared-interface *extraction* (a framework-level `Collection` both families implement) remains future work per "spike then extract".
+
 ## At a glance
 
 The same chaining API works for both families. The consumer doesn't know (or care) whether the data lives in Postgres or MongoDB:

packages/2-mongo-family/5-query-builders/orm/src/collection.ts

@@ -32,6 +32,7 @@ import {
 } from '@prisma-next/mongo-query-ast/execution';
 import type { MongoValue } from '@prisma-next/mongo-value';
 import { MongoParamRef } from '@prisma-next/mongo-value';
+import { blindCast } from '@prisma-next/utils/casts';
 import type { MongoIncludeExpr } from './collection-state';
 import { emptyCollectionState, type MongoCollectionState } from './collection-state';
 import { compileMongoQuery } from './compile';
@@ -154,7 +155,34 @@ function resolveCollectionName(model: MongoModelDefinition, modelName: string):
   return model.storage.collection ?? modelName;
 }
 
-class MongoCollectionImpl<
+type MongoCollectionConstructor<
+  TContract extends MongoContractWithTypeMaps<MongoContract, AnyMongoTypeMaps>,
+  ModelName extends string & keyof MongoModelsMap<TContract>,
+  TIncludes extends MongoIncludeSpec<TContract, ModelName>,
+  TVariant extends string,
+> = new (
+  contract: TContract,
+  modelName: ModelName,
+  executor: MongoQueryExecutor,
+) => Collection<TContract, ModelName, TIncludes, TVariant>;
+
+/**
+ * The extendable Mongo ORM collection ([ADR 175](../../../../../docs/architecture%20docs/adrs/ADR%20175%20-%20Shared%20ORM%20Collection%20interface.md)).
+ * Subclass it to add domain methods that start chains, mirroring the SQL ORM's
+ * `Collection`:
+ *
+ * ```typescript
+ * class UserRepository extends Collection<Contract, 'User'> {
+ *   byName(name: string) { return this.where({ name }); }
+ * }
+ * ```
+ *
+ * Register subclasses via `mongoOrm({ collections: { User: UserRepository } })` (keyed
+ * by model name). Chaining methods clone through `this.constructor`, so a chain started
+ * from a subclass stays an instance of that subclass. Subclasses must keep the base
+ * constructor signature `(contract, modelName, executor)`.
+ */
+export class Collection<
   TContract extends MongoContractWithTypeMaps<MongoContract, AnyMongoTypeMaps>,
   ModelName extends string & keyof MongoModelsMap<TContract>,
   TIncludes extends MongoIncludeSpec<TContract, ModelName> = NoIncludes,
@@ -270,12 +298,7 @@ class MongoCollectionImpl<
 
     return this.#clone({
       includes: [...this.#state.includes, includeExpr],
-    }) as unknown as MongoCollectionImpl<
-      TContract,
-      ModelName,
-      TIncludes & Record<K, true>,
-      TVariant
-    >;
+    }) as unknown as Collection<TContract, ModelName, TIncludes & Record<K, true>, TVariant>;
   }
 
   orderBy(
@@ -804,14 +827,16 @@ class MongoCollectionImpl<
     return { ...data, [model.discriminator.field]: variantEntry.value };
   }
 
+  // Clones instantiate through `this.constructor` so chains started from a custom
+  // subclass keep returning that subclass (same pattern as the SQL ORM Collection).
   #clone(
     overrides: Partial<MongoCollectionState>,
-  ): MongoCollectionImpl<TContract, ModelName, TIncludes, TVariant> {
-    const instance = new MongoCollectionImpl<TContract, ModelName, TIncludes, TVariant>(
-      this.#contract,
-      this.#modelName,
-      this.#executor,
-    );
+  ): Collection<TContract, ModelName, TIncludes, TVariant> {
+    const Ctor = blindCast<
+      MongoCollectionConstructor<TContract, ModelName, TIncludes, TVariant>,
+      'this.constructor is this Collection (sub)class; subclasses keep the base constructor signature'
+    >(this.constructor);
+    const instance = new Ctor(this.#contract, this.#modelName, this.#executor);
     instance.#state = { ...this.#state, ...overrides };
     instance.#collectionName = this.#collectionName;
     instance.#variantName = this.#variantName;
@@ -821,12 +846,12 @@ class MongoCollectionImpl<
   #cloneWithVariant<VNew extends string>(
     overrides: Partial<MongoCollectionState>,
     variantName: string,
-  ): MongoCollectionImpl<TContract, ModelName, TIncludes, VNew> {
-    const instance = new MongoCollectionImpl<TContract, ModelName, TIncludes, VNew>(
-      this.#contract,
-      this.#modelName,
-      this.#executor,
-    );
+  ): Collection<TContract, ModelName, TIncludes, VNew> {
+    const Ctor = blindCast<
+      MongoCollectionConstructor<TContract, ModelName, TIncludes, VNew>,
+      'this.constructor is this Collection (sub)class; subclasses keep the base constructor signature'
+    >(this.constructor);
+    const instance = new Ctor(this.#contract, this.#modelName, this.#executor);
     instance.#state = { ...this.#state, ...overrides };
     instance.#collectionName = this.#collectionName;
     instance.#variantName = variantName;
@@ -842,5 +867,5 @@ export function createMongoCollection<
   modelName: ModelName,
   executor: MongoQueryExecutor,
 ): MongoCollection<TContract, ModelName> {
-  return new MongoCollectionImpl(contract, modelName, executor);
+  return new Collection(contract, modelName, executor);
 }

packages/2-mongo-family/5-query-builders/orm/src/exports/index.ts

@@ -1,7 +1,7 @@
 export type { MongoQueryPlan } from '@prisma-next/mongo-query-ast/execution';
 export type { SimplifyDeep } from '@prisma-next/utils/simplify-deep';
 export type { MongoCollection } from '../collection';
-export { createMongoCollection } from '../collection';
+export { Collection, createMongoCollection } from '../collection';
 export { compileMongoQuery } from '../compile';
 export type { MongoQueryExecutor } from '../executor';
 export type {
@@ -13,7 +13,7 @@ export type {
   UpdateOperator,
 } from '../field-accessor';
 export { createFieldAccessor } from '../field-accessor';
-export type { MongoOrmClient, MongoOrmOptions } from '../mongo-orm';
+export type { AnyMongoCollectionClass, MongoOrmClient, MongoOrmOptions } from '../mongo-orm';
 export { mongoOrm } from '../mongo-orm';
 export type { MongoRawClient } from '../mongo-raw';
 export { mongoRaw } from '../mongo-raw';

packages/2-mongo-family/5-query-builders/orm/src/mongo-orm.ts

@@ -9,36 +9,70 @@ import type { MongoCollection } from './collection';
 import { createMongoCollection } from './collection';
 import type { MongoQueryExecutor } from './executor';
 
-export interface MongoOrmOptions<TContract extends MongoContract> {
+export type AnyMongoCollectionClass = new (...args: never[]) => object;
+
+export interface MongoOrmOptions<
+  TContract extends MongoContract,
+  Collections extends Partial<Record<string, AnyMongoCollectionClass>> = Record<never, never>,
+> {
   readonly contract: TContract;
   readonly executor: MongoQueryExecutor;
+  /**
+   * Custom `Collection` subclasses keyed by model name (not root/collection name),
+   * mirroring the SQL ORM's `orm({ collections })`. Roots whose model has no entry
+   * get the base collection.
+   */
+  readonly collections?: Collections;
 }
 
+type CustomCollectionForModel<
+  Collections extends Partial<Record<string, AnyMongoCollectionClass>>,
+  ModelName extends string,
+> = ModelName extends keyof Collections
+  ? Collections[ModelName] extends AnyMongoCollectionClass
+    ? InstanceType<Collections[ModelName]>
+    : never
+  : never;
+
+type RootCollection<
+  TContract extends MongoContractWithTypeMaps<MongoContract, AnyMongoTypeMaps>,
+  Collections extends Partial<Record<string, AnyMongoCollectionClass>>,
+  RootName extends keyof TContract['roots'] & string,
+> = [CustomCollectionForModel<Collections, RootModelName<TContract, RootName>>] extends [never]
+  ? MongoCollection<TContract, RootModelName<TContract, RootName>>
+  : CustomCollectionForModel<Collections, RootModelName<TContract, RootName>>;
+
 export type MongoOrmClient<
   TContract extends MongoContractWithTypeMaps<MongoContract, AnyMongoTypeMaps>,
+  Collections extends Partial<Record<string, AnyMongoCollectionClass>> = Record<never, never>,
 > = {
-  readonly [K in keyof TContract['roots'] & string]: MongoCollection<
-    TContract,
-    RootModelName<TContract, K>
-  >;
+  readonly [K in keyof TContract['roots'] & string]: RootCollection<TContract, Collections, K>;
 };
 
 export function mongoOrm<
   TContract extends MongoContractWithTypeMaps<MongoContract, AnyMongoTypeMaps>,
->(options: MongoOrmOptions<TContract>): MongoOrmClient<TContract> {
-  const { contract, executor } = options;
+  Collections extends Partial<Record<string, AnyMongoCollectionClass>> = Record<never, never>,
+>(options: MongoOrmOptions<TContract, Collections>): MongoOrmClient<TContract, Collections> {
+  const { contract, executor, collections } = options;
   const client: Record<string, unknown> = {};
 
   for (const [rootName, rootRef] of Object.entries(contract.roots)) {
-    client[rootName] = createMongoCollection(
-      contract,
-      blindCast<
-        RootModelName<TContract, typeof rootName & keyof TContract['roots'] & string>,
-        'roots entries are CrossReferences; rootRef.model is a valid RootModelName for this contract'
-      >(rootRef.model),
-      executor,
-    );
+    const modelName = blindCast<
+      RootModelName<TContract, typeof rootName & keyof TContract['roots'] & string>,
+      'roots entries are CrossReferences; rootRef.model is a valid RootModelName for this contract'
+    >(rootRef.model);
+    const CustomCtor = collections?.[rootRef.model];
+    client[rootName] = CustomCtor
+      ? new (blindCast<
+          new (
+            contract: TContract,
+            modelName: string,
+            executor: MongoQueryExecutor,
+          ) => object,
+          'a registered collection class is a Collection subclass constructor'
+        >(CustomCtor))(contract, modelName, executor)
+      : createMongoCollection(contract, modelName, executor);
   }
 
-  return client as MongoOrmClient<TContract>;
+  return client as MongoOrmClient<TContract, Collections>;
 }

packages/2-mongo-family/5-query-builders/orm/test/custom-collection.test-d.ts

@@ -0,0 +1,40 @@
+import { describe, expectTypeOf, it } from 'vitest';
+import type { Contract } from '../../../1-foundation/mongo-contract/test/fixtures/orm-contract';
+import ormContractJson from '../../../1-foundation/mongo-contract/test/fixtures/orm-contract.json';
+import type { MongoCollection } from '../src/collection';
+import { Collection } from '../src/collection';
+import type { MongoQueryExecutor } from '../src/executor';
+import { mongoOrm } from '../src/mongo-orm';
+
+const contract = ormContractJson as unknown as Contract;
+declare const executor: MongoQueryExecutor;
+
+class UserRepository extends Collection<Contract, 'User'> {
+  byName(name: string) {
+    return this.where({ name });
+  }
+}
+
+describe('custom collection typing', () => {
+  it('a registered root is typed as the subclass', () => {
+    const client = mongoOrm({ contract, executor, collections: { User: UserRepository } });
+    expectTypeOf(client.users).toEqualTypeOf<UserRepository>();
+    expectTypeOf(client.users.byName).toBeFunction();
+  });
+
+  it('unregistered roots keep the base collection type', () => {
+    const client = mongoOrm({ contract, executor, collections: { User: UserRepository } });
+    expectTypeOf(client.tasks).toEqualTypeOf<MongoCollection<Contract, 'Task'>>();
+  });
+
+  it('no collections option keeps every root on the base type', () => {
+    const client = mongoOrm({ contract, executor });
+    expectTypeOf(client.users).toEqualTypeOf<MongoCollection<Contract, 'User'>>();
+  });
+
+  it('domain methods return a chainable collection', async () => {
+    const client = mongoOrm({ contract, executor, collections: { User: UserRepository } });
+    const first = await client.users.byName('Alice').take(1).first();
+    expectTypeOf(first).not.toBeAny();
+  });
+});

packages/2-mongo-family/5-query-builders/orm/test/custom-collection.test.ts

@@ -0,0 +1,130 @@
+import { AsyncIterableResult } from '@prisma-next/framework-components/runtime';
+import type {
+  MongoMatchStage,
+  MongoQueryPlan,
+  MongoSortStage,
+} from '@prisma-next/mongo-query-ast/execution';
+import { describe, expect, it } from 'vitest';
+import type { Contract } from '../../../1-foundation/mongo-contract/test/fixtures/orm-contract';
+import ormContractJson from '../../../1-foundation/mongo-contract/test/fixtures/orm-contract.json';
+import { Collection, createMongoCollection } from '../src/collection';
+import type { MongoQueryExecutor } from '../src/executor';
+import { mongoOrm } from '../src/mongo-orm';
+
+const contract = ormContractJson as unknown as Contract;
+
+function createMockExecutor(...responses: unknown[][]): MongoQueryExecutor & {
+  lastPlan: MongoQueryPlan | undefined;
+  readonly lastStages: ReadonlyArray<unknown> | undefined;
+} {
+  let callIndex = 0;
+  const mock = {
+    lastPlan: undefined as MongoQueryPlan | undefined,
+    get lastStages(): ReadonlyArray<unknown> | undefined {
+      const cmd = mock.lastPlan?.command;
+      if (cmd?.kind === 'aggregate') return cmd.pipeline as ReadonlyArray<unknown>;
+      return undefined;
+    },
+    execute<Row>(plan: MongoQueryPlan<Row>): AsyncIterableResult<Row> {
+      mock.lastPlan = plan as MongoQueryPlan;
+      const data = responses[callIndex] ?? [];
+      callIndex++;
+      async function* gen(): AsyncGenerator<Row> {
+        for (const row of data) yield row as Row;
+      }
+      return new AsyncIterableResult(gen());
+    },
+  };
+  return mock;
+}
+
+class UserRepository extends Collection<Contract, 'User'> {
+  byName(name: string) {
+    return this.where({ name });
+  }
+
+  newestFirst() {
+    return this.orderBy({ _id: -1 });
+  }
+}
+
+class TaskRepository extends Collection<Contract, 'Task'> {
+  bugs() {
+    return this.variant('Bug');
+  }
+}
+
+describe('custom collection subclasses', () => {
+  it('mongoOrm exposes a registered subclass on its root', () => {
+    const executor = createMockExecutor();
+    const client = mongoOrm({ contract, executor, collections: { User: UserRepository } });
+    expect(client.users).toBeInstanceOf(UserRepository);
+  });
+
+  it('roots without a registered subclass get the base collection', () => {
+    const executor = createMockExecutor();
+    const client = mongoOrm({ contract, executor, collections: { User: UserRepository } });
+    expect(client.tasks).toBeInstanceOf(Collection);
+    expect(client.tasks).not.toBeInstanceOf(UserRepository);
+  });
+
+  it('a domain method compiles the same plan as the base collection', async () => {
+    const repoExecutor = createMockExecutor([]);
+    const baseExecutor = createMockExecutor([]);
+    const client = mongoOrm({
+      contract,
+      executor: repoExecutor,
+      collections: { User: UserRepository },
+    });
+    const base = createMongoCollection(contract, 'User', baseExecutor);
+
+    for await (const _row of (client.users as UserRepository).byName('Alice').all()) {
+      // drain
+    }
+    for await (const _row of base.where({ name: 'Alice' }).all()) {
+      // drain
+    }
+
+    const repoMatch = repoExecutor.lastStages?.[0] as MongoMatchStage;
+    const baseMatch = baseExecutor.lastStages?.[0] as MongoMatchStage;
+    expect(repoMatch).toEqual(baseMatch);
+  });
+
+  it('chaining preserves the subclass', () => {
+    const executor = createMockExecutor();
+    const client = mongoOrm({ contract, executor, collections: { User: UserRepository } });
+    const users = client.users as UserRepository;
+    expect(users.byName('Alice')).toBeInstanceOf(UserRepository);
+    expect(users.newestFirst()).toBeInstanceOf(UserRepository);
+    expect(users.byName('Alice').take(5).skip(1).select('name')).toBeInstanceOf(UserRepository);
+  });
+
+  it('variant() preserves the subclass', () => {
+    const executor = createMockExecutor();
+    const client = mongoOrm({ contract, executor, collections: { Task: TaskRepository } });
+    const tasks = client.tasks as TaskRepository;
+    expect(tasks.bugs()).toBeInstanceOf(TaskRepository);
+  });
+
+  it('a directly constructed subclass executes queries', async () => {
+    const executor = createMockExecutor([{ _id: '1', name: 'Alice' }]);
+    const repo = new UserRepository(contract, 'User', executor);
+    const row = await repo.byName('Alice').first();
+    expect(row).toEqual({ _id: '1', name: 'Alice' });
+  });
+
+  it('domain methods compose with base chaining before execution', async () => {
+    const executor = createMockExecutor([]);
+    const client = mongoOrm({ contract, executor, collections: { User: UserRepository } });
+    const users = client.users as UserRepository;
+    for await (const _row of users.byName('Alice').orderBy({ name: 1 }).take(2).all()) {
+      // drain
+    }
+    const stages = executor.lastStages as Array<MongoMatchStage | MongoSortStage>;
+    expect(stages.map((s) => (s as { kind: string }).kind).slice(0, 3)).toEqual([
+      'match',
+      'sort',
+      'limit',
+    ]);
+  });
+});

packages/3-extensions/mongo/src/exports/runtime.ts

@@ -1,3 +1,5 @@
+export type { AnyMongoCollectionClass, MongoCollection } from '@prisma-next/mongo-orm';
+export { Collection } from '@prisma-next/mongo-orm';
 export type { MongoBinding, MongoBindingInput } from '../runtime/binding';
 export type {
   MongoClient,