fixup: add index signature support

Author: addaleax

packages/shell-api/bin/api-postprocess.ts

@@ -116,6 +116,13 @@ type MongodbServerSchema = {
         _id: ObjectId;
         foo: number;
       }
+    },
+    with: { schema: never },
+    'with.dots': {
+      schema: {
+        _id: ObjectId;
+        bar: string;
+      }
     }
   }
 }

packages/shell-api/src/collection.ts

@@ -92,17 +92,36 @@ import PlanCache from './plan-cache';
 import ChangeStreamCursor from './change-stream-cursor';
 import { ShellApiErrors } from './error-codes';
 
+export type Collection<
+  M extends GenericServerSideSchema = GenericServerSideSchema,
+  D extends GenericDatabaseSchema = M[keyof M],
+  C extends GenericCollectionSchema = D[keyof D],
+  N extends StringKey<D> = StringKey<D>
+> = CollectionImpl<M, D, C, N> & {
+  [k in keyof D as k extends `${N}.${infer S}` ? S : never]: Collection<
+    M,
+    D,
+    D[k]
+  >;
+};
+
 @shellApiClassDefault
 @addSourceToResults
-export class Collection<
+export class CollectionImpl<
   M extends GenericServerSideSchema = GenericServerSideSchema,
   D extends GenericDatabaseSchema = M[keyof M],
-  C extends GenericCollectionSchema = D[keyof D]
+  C extends GenericCollectionSchema = D[keyof D],
+  N extends StringKey<D> = StringKey<D>
 > extends ShellApiWithMongoClass {
   _mongo: Mongo<M>;
   _database: Database<M, D>;
-  _name: StringKey<D>;
-  constructor(mongo: Mongo<M>, database: Database<M, D>, name: StringKey<D>) {
+  _name: N;
+
+  _typeLaunder(): Collection<M, D> {
+    return this as Collection<M, D>;
+  }
+
+  constructor(mongo: Mongo<M>, database: Database<M, D>, name: N) {
     super();
     this._mongo = mongo;
     this._database = database;
@@ -1555,9 +1574,9 @@ export class Collection<
     return `${this._database._name}.${this._name}`;
   }
 
-  getName(): string {
+  getName(): N {
     this._emitCollectionApiCall('getName');
-    return `${this._name}`;
+    return this._name;
   }
 
   @returnsPromise
@@ -1604,7 +1623,7 @@ export class Collection<
   explain(verbosity: ExplainVerbosityLike = 'queryPlanner'): Explainable {
     verbosity = validateExplainableVerbosity(verbosity);
     this._emitCollectionApiCall('explain', { verbosity });
-    return new Explainable(this._mongo, this, verbosity);
+    return new Explainable(this._mongo, this._typeLaunder(), verbosity);
   }
 
   /**
@@ -1972,7 +1991,7 @@ export class Collection<
       true,
       await this._database._baseOptions()
     );
-    return new Bulk(this, innerBulk, true);
+    return new Bulk(this._typeLaunder(), innerBulk, true);
   }
 
   @returnsPromise
@@ -1986,14 +2005,14 @@ export class Collection<
       false,
       await this._database._baseOptions()
     );
-    return new Bulk(this, innerBulk);
+    return new Bulk(this._typeLaunder(), innerBulk);
   }
 
   @returnType('PlanCache')
   @apiVersions([])
   getPlanCache(): PlanCache {
     this._emitCollectionApiCall('getPlanCache');
-    return new PlanCache(this);
+    return new PlanCache(this._typeLaunder());
   }
 
   @returnsPromise
@@ -2242,15 +2261,15 @@ export class Collection<
   @apiVersions([1])
   async hideIndex(index: string | Document): Promise<Document> {
     this._emitCollectionApiCall('hideIndex');
-    return setHideIndex(this, index, true);
+    return setHideIndex(this._typeLaunder(), index, true);
   }
 
   @serverVersions(['4.4.0', ServerVersions.latest])
   @returnsPromise
   @apiVersions([1])
   async unhideIndex(index: string | Document): Promise<Document> {
     this._emitCollectionApiCall('unhideIndex');
-    return setHideIndex(this, index, false);
+    return setHideIndex(this._typeLaunder(), index, false);
   }
 
   @serverVersions(['7.0.0', ServerVersions.latest])

packages/shell-api/src/database.ts

@@ -1,5 +1,6 @@
 import type Mongo from './mongo';
-import Collection from './collection';
+import type Collection from './collection';
+import { CollectionImpl } from './collection';
 import {
   returnsPromise,
   returnType,
@@ -70,26 +71,37 @@ type AuthDoc = {
   mechanism?: string;
 };
 
+export type Database<
+  M extends GenericServerSideSchema = GenericServerSideSchema,
+  D extends GenericDatabaseSchema = GenericDatabaseSchema
+> = DatabaseImpl<M, D> & {
+  [k in StringKey<D>]: Collection<M, D, D[k], k>;
+};
+
 @shellApiClassDefault
-export default class Database<
+export class DatabaseImpl<
   M extends GenericServerSideSchema = GenericServerSideSchema,
-  D extends GenericDatabaseSchema = M[keyof M]
+  D extends GenericDatabaseSchema = GenericDatabaseSchema
 > extends ShellApiWithMongoClass {
   _mongo: Mongo<M>;
   _name: StringKey<M>;
-  _collections: Record<string, Collection>;
+  _collections: Record<string, Collection<M, D>>;
   _session: Session | undefined;
   _cachedCollectionNames: StringKey<D>[] = [];
   _cachedHello: Document | null = null;
 
+  _typeLaunder(): Database<M, D> {
+    return this as Database<M, D>;
+  }
+
   constructor(mongo: Mongo<M>, name: StringKey<M>, session?: Session) {
     super();
     this._mongo = mongo;
     this._name = name;
-    const collections: Record<string, Collection> = Object.create(null);
+    const collections: Record<string, Collection<M, D>> = Object.create(null);
     this._collections = collections;
     this._session = session;
-    const proxy = new Proxy(this, {
+    const proxy = new Proxy(this._typeLaunder(), {
       get: (target, prop): any => {
         if (prop in target) {
           return (target as any)[prop];
@@ -104,7 +116,11 @@ export default class Database<
         }
 
         if (!collections[prop]) {
-          collections[prop] = new Collection(mongo, proxy, prop);
+          collections[prop] = new CollectionImpl(
+            mongo,
+            proxy,
+            prop
+          )._typeLaunder();
         }
 
         return collections[prop];
@@ -449,13 +465,13 @@ export default class Database<
     assertArgsDefinedType([db], ['string'], 'Database.getSiblingDB');
     this._emitDatabaseApiCall('getSiblingDB', { db });
     if (this._session) {
-      return this._session.getDatabase(db);
+      return this._session.getDatabase(db) as Database<M, M[K]>;
     }
     return this._mongo._getDb(db);
   }
 
   @returnType('Collection')
-  getCollection<K extends StringKey<D>>(coll: K): Collection<M, D, D[K]> {
+  getCollection<K extends StringKey<D>>(coll: K): Collection<M, D, D[K], K> {
     assertArgsDefinedType([coll], ['string'], 'Database.getColl');
     this._emitDatabaseApiCall('getCollection', { coll });
     if (!isValidCollectionName(coll)) {
@@ -465,13 +481,17 @@ export default class Database<
       );
     }
 
-    const collections: Record<string, Collection> = this._collections;
+    const collections: Record<string, Collection<M, D>> = this._collections;
 
     if (!collections[coll]) {
-      collections[coll] = new Collection(this._mongo, this, coll);
+      collections[coll] = new CollectionImpl(
+        this._mongo,
+        this._typeLaunder(),
+        coll
+      )._typeLaunder();
     }
 
-    return collections[coll];
+    return collections[coll] as Collection<M, D, D[K], K>;
   }
 
   @returnsPromise
@@ -1471,7 +1491,7 @@ export default class Database<
   async printShardingStatus(verbose = false): Promise<CommandResult> {
     this._emitDatabaseApiCall('printShardingStatus', { verbose });
     const result = await getPrintableShardStatus(
-      await getConfigDB(this),
+      await getConfigDB(this._typeLaunder()),
       verbose
     );
     return new CommandResult('StatsResult', result);
@@ -1760,3 +1780,5 @@ export default class Database<
     });
   }
 }
+
+export default Database;

packages/shell-api/src/mongo.ts

@@ -19,31 +19,31 @@ import {
   topologies,
   deprecated,
 } from './decorators';
-import type {
-  ChangeStreamOptions,
-  ClientSessionOptions,
-  CommandOperationOptions,
-  Document,
-  ListDatabasesOptions,
-  ReadConcernLevel,
-  ReadPreference,
-  ReadPreferenceLike,
-  ReadPreferenceMode,
-  ServiceProvider,
-  TransactionOptions,
-  MongoClientOptions,
-  AutoEncryptionOptions as SPAutoEncryption,
-  ServerApi,
-  ServerApiVersion,
-  WriteConcern,
+import {
+  type ChangeStreamOptions,
+  type ClientSessionOptions,
+  type CommandOperationOptions,
+  type Document,
+  type ListDatabasesOptions,
+  type ReadConcernLevel,
+  type ReadPreference,
+  type ReadPreferenceLike,
+  type ReadPreferenceMode,
+  type ServiceProvider,
+  type TransactionOptions,
+  type MongoClientOptions,
+  type AutoEncryptionOptions as SPAutoEncryption,
+  type ServerApi,
+  type ServerApiVersion,
+  type WriteConcern,
 } from '@mongosh/service-provider-core';
 import type { ConnectionInfo } from '@mongosh/arg-parser';
 import {
   mapCliToDriver,
   generateConnectionInfoFromCliArgs,
 } from '@mongosh/arg-parser';
-import type Collection from './collection';
-import Database from './database';
+import type { Database } from './database';
+import { DatabaseImpl } from './database';
 import type ShellInstanceState from './shell-instance-state';
 import { CommandResult } from './result';
 import { redactURICredentials } from '@mongosh/history';
@@ -62,6 +62,7 @@ import { KeyVault, ClientEncryption } from './field-level-encryption';
 import { ShellApiErrors } from './error-codes';
 import type { LogEntry } from './log-entry';
 import { parseAnyLogEntry } from './log-entry';
+import type { Collection } from './collection';
 
 /* Utility, inverse of Readonly<T> */
 type Mutable<T> = {
@@ -262,9 +263,9 @@ export default class Mongo<
     }
 
     if (!(name in this._databases)) {
-      this._databases[name] = new Database(this, name);
+      this._databases[name] = new DatabaseImpl(this, name)._typeLaunder();
     }
-    return this._databases[name];
+    return this._databases[name] as Database<M, M[K]>;
   }
 
   @returnType('Database')
@@ -286,7 +287,11 @@ export default class Mongo<
         CommonErrors.InvalidArgument
       );
     }
-    return this._getDb(db as StringKey<M>).getCollection(coll);
+    return this._getDb(db as StringKey<M>).getCollection(coll) as Collection<
+      M,
+      M[KD],
+      M[KD][KC]
+    >;
   }
 
   getURI(): string {

packages/shell-api/src/session.ts

@@ -14,21 +14,25 @@ import type {
 } from '@mongosh/service-provider-core';
 import { asPrintable } from './enums';
 import type Mongo from './mongo';
-import Database from './database';
+import type Database from './database';
+import { DatabaseImpl } from './database';
 import { CommonErrors, MongoshInvalidInputError } from '@mongosh/errors';
+import type { GenericServerSideSchema, StringKey } from './helpers';
 import { assertArgsDefinedType, isValidDatabaseName } from './helpers';
 
 @shellApiClassDefault
 @classPlatforms(['CLI'])
-export default class Session extends ShellApiWithMongoClass {
+export default class Session<
+  M extends GenericServerSideSchema = GenericServerSideSchema
+> extends ShellApiWithMongoClass {
   public id: ServerSessionId | undefined;
   public _session: ClientSession;
   public _options: ClientSessionOptions;
-  public _mongo: Mongo;
-  private _databases: Record<string, Database>;
+  public _mongo: Mongo<M>;
+  private _databases: Record<string, Database<M>>;
 
   constructor(
-    mongo: Mongo,
+    mongo: Mongo<M>,
     options: ClientSessionOptions,
     session: ClientSession
   ) {
@@ -47,7 +51,7 @@ export default class Session extends ShellApiWithMongoClass {
     return this._session.id;
   }
 
-  getDatabase(name: string): Database {
+  getDatabase<K extends StringKey<M>>(name: K): Database<M, M[K]> {
     assertArgsDefinedType([name], ['string'], 'Session.getDatabase');
 
     if (!isValidDatabaseName(name)) {
@@ -58,9 +62,13 @@ export default class Session extends ShellApiWithMongoClass {
     }
 
     if (!(name in this._databases)) {
-      this._databases[name] = new Database(this._mongo, name, this);
+      this._databases[name] = new DatabaseImpl(
+        this._mongo,
+        name,
+        this
+      )._typeLaunder();
     }
-    return this._databases[name];
+    return this._databases[name] as Database<M, M[K]>;
   }
 
   advanceOperationTime(ts: TimestampType): void {