[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,141 @@
+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;
+      }
+    }
+  }
+}
+// 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,
@@ -90,11 +94,15 @@ import { ShellApiErrors } from './error-codes';
 
 @shellApiClassDefault
 @addSourceToResults
-export default class Collection extends ShellApiWithMongoClass {
-  _mongo: Mongo;
-  _database: Database;
-  _name: string;
-  constructor(mongo: Mongo, database: Database, name: string) {
+export class Collection<
+  M extends GenericServerSideSchema = GenericServerSideSchema,
+  D extends GenericDatabaseSchema = M[keyof M],
+  C extends GenericCollectionSchema = D[keyof D]
+> extends ShellApiWithMongoClass {
+  _mongo: Mongo<M>;
+  _database: Database<M, D>;
+  _name: StringKey<D>;
+  constructor(mongo: Mongo<M>, database: Database<M, D>, name: StringKey<D>) {
     super();
     this._mongo = mongo;
     this._database = database;
@@ -513,7 +521,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 +1414,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 +1425,7 @@ export default class Collection extends ShellApiWithMongoClass {
    * @return {Mongo}
    */
   @returnType('Mongo')
-  getMongo(): Mongo {
+  getMongo(): Mongo<M> {
     this._emitCollectionApiCall('getMongo');
     return this._mongo;
   }
@@ -1754,7 +1762,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;
     }
@@ -2061,7 +2069,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
@@ -2396,3 +2404,5 @@ export default class Collection extends ShellApiWithMongoClass {
     );
   }
 }
+
+export default Collection;

packages/shell-api/src/database.ts

@@ -11,6 +11,11 @@ import {
   ShellApiWithMongoClass,
 } from './decorators';
 import { asPrintable, ServerVersions, Topologies } from './enums';
+import type {
+  GenericDatabaseSchema,
+  GenericServerSideSchema,
+  StringKey,
+} from './helpers';
 import {
   adaptAggregateOptions,
   adaptOptions,
@@ -66,15 +71,18 @@ type AuthDoc = {
 };
 
 @shellApiClassDefault
-export default class Database extends ShellApiWithMongoClass {
-  _mongo: Mongo;
-  _name: string;
+export default class Database<
+  M extends GenericServerSideSchema = GenericServerSideSchema,
+  D extends GenericDatabaseSchema = M[keyof M]
+> extends ShellApiWithMongoClass {
+  _mongo: Mongo<M>;
+  _name: StringKey<M>;
   _collections: Record<string, Collection>;
   _session: Session | undefined;
-  _cachedCollectionNames: string[] = [];
+  _cachedCollectionNames: StringKey<D>[] = [];
   _cachedHello: Document | null = null;
 
-  constructor(mongo: Mongo, name: string, session?: Session) {
+  constructor(mongo: Mongo<M>, name: StringKey<M>, session?: Session) {
     super();
     this._mongo = mongo;
     this._name = name;
@@ -308,11 +316,11 @@ export default class Database extends ShellApiWithMongoClass {
   }
 
   @returnType('Mongo')
-  getMongo(): Mongo {
+  getMongo(): Mongo<M> {
     return this._mongo;
   }
 
-  getName(): string {
+  getName(): StringKey<M> {
     return this._name;
   }
 
@@ -323,9 +331,9 @@ export default class Database extends ShellApiWithMongoClass {
    */
   @returnsPromise
   @apiVersions([1])
-  async getCollectionNames(): Promise<string[]> {
+  async getCollectionNames(): Promise<StringKey<D>[]> {
     this._emitDatabaseApiCall('getCollectionNames');
-    return this._getCollectionNames();
+    return (await this._getCollectionNames()) as StringKey<D>[];
   }
 
   /**
@@ -437,7 +445,7 @@ export default class Database extends ShellApiWithMongoClass {
   }
 
   @returnType('Database')
-  getSiblingDB(db: string): Database {
+  getSiblingDB<K extends StringKey<M>>(db: K): Database<M, M[K]> {
     assertArgsDefinedType([db], ['string'], 'Database.getSiblingDB');
     this._emitDatabaseApiCall('getSiblingDB', { db });
     if (this._session) {
@@ -447,7 +455,7 @@ export default class Database extends ShellApiWithMongoClass {
   }
 
   @returnType('Collection')
-  getCollection(coll: string): Collection {
+  getCollection<K extends StringKey<D>>(coll: K): Collection<M, D, D[K]> {
     assertArgsDefinedType([coll], ['string'], 'Database.getColl');
     this._emitDatabaseApiCall('getCollection', { coll });
     if (!isValidCollectionName(coll)) {

packages/shell-api/src/decorators.ts

@@ -381,6 +381,7 @@ export interface TypeSignature {
   isDirectShellCommand?: boolean;
   acceptsRawInput?: boolean;
   shellCommandCompleter?: ShellCommandCompleter;
+  inherited?: boolean;
 }
 
 /**
@@ -426,6 +427,7 @@ type ClassSignature = {
       isDirectShellCommand: boolean;
       acceptsRawInput?: boolean;
       shellCommandCompleter?: ShellCommandCompleter;
+      inherited?: true;
     };
   };
 };
@@ -574,6 +576,7 @@ function shellApiClassGeneric<T extends { prototype: any }>(
         isDirectShellCommand: method.isDirectShellCommand,
         acceptsRawInput: method.acceptsRawInput,
         shellCommandCompleter: method.shellCommandCompleter,
+        inherited: true,
       };
 
       const attributeHelpKeyPrefix = `${superClassHelpKeyPrefix}.attributes.${propertyName}`;