[Do Not Merge] PoC for shell-api autocomplete type definitions

packages/shell-api/api-extractor.json

@@ -0,0 +1,50 @@
+{
+  "$schema": "https://developer.microsoft.com/json-schemas/api-extractor/v7/api-extractor.schema.json",
+  "mainEntryPointFilePath": "<projectFolder>/lib/index.d.ts",
+  "apiReport": {
+    "enabled": false
+  },
+  "docModel": {
+    "enabled": false
+  },
+  "bundledPackages": [
+    "@mongosh/service-provider-core",
+    "@mongosh/types",
+    "mongodb",
+    "bson"
+  ],
+  "dtsRollup": {
+    "enabled": true,
+    "untrimmedFilePath": "",
+    "publicTrimmedFilePath": "<projectFolder>/lib/api-raw.d.ts"
+  },
+  "tsdocMetadata": {
+    "enabled": false
+  },
+  "newlineKind": "lf",
+  "messages": {
+    "compilerMessageReporting": {
+      "default": {
+        "logLevel": "error"
+      }
+    },
+    "extractorMessageReporting": {
+      "default": {
+        "logLevel": "error"
+      },
+      "ae-internal-missing-underscore": {
+        "logLevel": "none",
+        "addToApiReportFile": false
+      },
+      "ae-forgotten-export": {
+        "logLevel": "error",
+        "addToApiReportFile": false
+      }
+    },
+    "tsdocMessageReporting": {
+      "default": {
+        "logLevel": "none"
+      }
+    }
+  }
+}

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

@@ -0,0 +1,148 @@
+import * as babel from '@babel/core';
+import type * as BabelTypes from '@babel/types';
+import { promises as fs } from 'fs';
+import path from 'path';
+import { signatures } from '../';
+
+function applyAsyncRewriterChanges() {
+  return ({
+    types: t,
+  }: {
+    types: typeof BabelTypes;
+  }): babel.PluginObj<{
+    processedMethods: [string, string, BabelTypes.ClassBody][];
+  }> => {
+    return {
+      pre() {
+        this.processedMethods = [];
+      },
+      post() {
+        for (const className of Object.keys(signatures)) {
+          for (const methodName of Object.keys(
+            signatures[className].attributes ?? {}
+          )) {
+            if (
+              signatures[className].attributes?.[methodName].returnsPromise &&
+              !signatures[className].attributes?.[methodName].inherited
+            ) {
+              if (
+                !this.processedMethods.find(
+                  ([cls, method]) => cls === className && method === methodName
+                )
+              ) {
+                console.error(
+                  `Expected to find and transpile type for @returnsPromise-annotated method ${className}.${methodName}`
+                );
+              }
+            }
+          }
+        }
+      },
+      visitor: {
+        TSDeclareMethod(path) {
+          if ('isMongoshAsyncRewrittenMethod' in path.node) return;
+
+          if (path.parent.type !== 'ClassBody') return;
+          if (path.parentPath.parent.type !== 'ClassDeclaration') return;
+          const classId = path.parentPath.parent.id;
+          if (classId?.type !== 'Identifier') return;
+          const className = classId.name;
+          if (path.node.key.type !== 'Identifier') return;
+          const methodName = path.node.key.name;
+
+          if (
+            this.processedMethods.find(
+              ([cls, method, classBody]) =>
+                cls === className &&
+                method === methodName &&
+                classBody !== path.parent
+            )
+          ) {
+            throw new Error(`Duplicate method: ${className}.${methodName}`);
+          }
+          this.processedMethods.push([className, methodName, path.parent]);
+
+          if (!signatures[className]?.attributes?.[methodName]?.returnsPromise)
+            return;
+
+          const { returnType } = path.node;
+          if (returnType?.type !== 'TSTypeAnnotation') return;
+          if (returnType.typeAnnotation.type !== 'TSTypeReference') return;
+          if (returnType.typeAnnotation.typeName.type !== 'Identifier') return;
+          if (returnType.typeAnnotation.typeName.name !== 'Promise') return;
+          if (!returnType.typeAnnotation.typeParameters?.params.length) return;
+          path.replaceWith({
+            ...path.node,
+            returnType: {
+              ...returnType,
+              typeAnnotation:
+                returnType.typeAnnotation.typeParameters.params[0],
+            },
+            isMongoshAsyncRewrittenMethod: true,
+          });
+        },
+      },
+    };
+  };
+}
+
+async function main() {
+  const apiRaw = await fs.readFile(
+    path.resolve(__dirname, '..', 'lib', 'api-raw.d.ts'),
+    'utf8'
+  );
+  const result = babel.transformSync(apiRaw, {
+    code: true,
+    ast: false,
+    configFile: false,
+    babelrc: false,
+    browserslistConfigFile: false,
+    compact: false,
+    sourceType: 'module',
+    plugins: [applyAsyncRewriterChanges()],
+    parserOpts: {
+      plugins: ['typescript'],
+    },
+  });
+  let code = result?.code ?? '';
+  code += `
+// REPLACEME
+type MongodbServerSchema = {
+  admin: {},
+  config: {},
+  test: {
+    test: {
+      schema: {
+        _id: ObjectId;
+        foo: number;
+      }
+    },
+    with: { schema: never },
+    'with.dots': {
+      schema: {
+        _id: ObjectId;
+        bar: string;
+      }
+    }
+  }
+}
+// REPLACEME
+
+declare global {
+  // second argument optional
+  var db: Database<MongodbServerSchema, MongodbServerSchema['test']>;
+
+  var use: (collection: StringKey<MongodbServerSchema>) => void;
+}
+`;
+  await fs.writeFile(
+    path.resolve(__dirname, '..', 'lib', 'api-processed.d.ts'),
+    code
+  );
+}
+
+main().catch((err) =>
+  process.nextTick(() => {
+    throw err;
+  })
+);

packages/shell-api/package.json

@@ -13,6 +13,7 @@
   },
   "scripts": {
     "compile": "tsc -p tsconfig.json",
+    "api-generate": "api-extractor run ; ts-node bin/api-postprocess.ts",
     "pretest": "npm run compile",
     "eslint": "eslint",
     "lint": "npm run eslint . && npm run prettier -- --check .",
@@ -48,6 +49,7 @@
     "mongodb-redact": "^0.2.2"
   },
   "devDependencies": {
+    "@microsoft/api-extractor": "^7.39.3",
     "@mongodb-js/eslint-config-mongosh": "^1.0.0",
     "@mongodb-js/prettier-config-devtools": "^1.0.1",
     "@mongodb-js/tsconfig-mongosh": "^1.0.0",

packages/shell-api/src/collection.ts

@@ -22,6 +22,10 @@ import type {
   FindAndModifyMethodShellOptions,
   RemoveShellOptions,
   MapReduceShellOptions,
+  GenericCollectionSchema,
+  GenericDatabaseSchema,
+  GenericServerSideSchema,
+  StringKey,
 } from './helpers';
 import {
   adaptAggregateOptions,
@@ -88,13 +92,37 @@ 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 StringKey<D> as k extends `${N}.${infer S}` ? S : never]: Collection<
+    M,
+    D,
+    D[k],
+    k
+  >;
+};
+
 @shellApiClassDefault
 @addSourceToResults
-export default class Collection extends ShellApiWithMongoClass {
-  _mongo: Mongo;
-  _database: Database;
-  _name: string;
-  constructor(mongo: Mongo, database: Database, name: string) {
+export class CollectionImpl<
+  M extends GenericServerSideSchema = GenericServerSideSchema,
+  D extends GenericDatabaseSchema = M[keyof M],
+  C extends GenericCollectionSchema = D[keyof D],
+  N extends StringKey<D> = StringKey<D>
+> extends ShellApiWithMongoClass {
+  _mongo: Mongo<M>;
+  _database: Database<M, 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;
@@ -513,7 +541,7 @@ export default class Collection extends ShellApiWithMongoClass {
     query: Document = {},
     projection?: Document,
     options: FindOptions = {}
-  ): Promise<Document | null> {
+  ): Promise<C['schema'] | null> {
     if (projection) {
       options.projection = projection;
     }
@@ -1406,7 +1434,7 @@ export default class Collection extends ShellApiWithMongoClass {
    * @return {Database}
    */
   @returnType('Database')
-  getDB(): Database {
+  getDB(): Database<M, D> {
     this._emitCollectionApiCall('getDB');
     return this._database;
   }
@@ -1417,7 +1445,7 @@ export default class Collection extends ShellApiWithMongoClass {
    * @return {Mongo}
    */
   @returnType('Mongo')
-  getMongo(): Mongo {
+  getMongo(): Mongo<M> {
     this._emitCollectionApiCall('getMongo');
     return this._mongo;
   }
@@ -1547,9 +1575,9 @@ export default class Collection extends ShellApiWithMongoClass {
     return `${this._database._name}.${this._name}`;
   }
 
-  getName(): string {
+  getName(): N {
     this._emitCollectionApiCall('getName');
-    return `${this._name}`;
+    return this._name;
   }
 
   @returnsPromise
@@ -1596,7 +1624,7 @@ export default class Collection extends ShellApiWithMongoClass {
   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);
   }
 
   /**
@@ -1754,7 +1782,7 @@ export default class Collection extends ShellApiWithMongoClass {
     }
 
     const ns = `${this._database._name}.${this._name}`;
-    const config = this._mongo.getDB('config');
+    const config = this._mongo.getDB('config' as StringKey<M>);
     if (collStats[0].shard) {
       result.shards = shardStats;
     }
@@ -1964,7 +1992,7 @@ export default class Collection extends ShellApiWithMongoClass {
       true,
       await this._database._baseOptions()
     );
-    return new Bulk(this, innerBulk, true);
+    return new Bulk(this._typeLaunder(), innerBulk, true);
   }
 
   @returnsPromise
@@ -1978,14 +2006,14 @@ export default class Collection extends ShellApiWithMongoClass {
       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
@@ -2061,7 +2089,7 @@ export default class Collection extends ShellApiWithMongoClass {
     this._emitCollectionApiCall('getShardDistribution', {});
 
     const result = {} as Document;
-    const config = this._mongo.getDB('config');
+    const config = this._mongo.getDB('config' as StringKey<M>);
     const ns = `${this._database._name}.${this._name}`;
 
     const configCollectionsInfo = await config
@@ -2234,15 +2262,15 @@ export default class Collection extends ShellApiWithMongoClass {
   @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])
@@ -2396,3 +2424,5 @@ export default class Collection extends ShellApiWithMongoClass {
     );
   }
 }
+
+export default Collection;