feat: allow opting out of schema sampling with config disableSchemaSampling=true MONGOSH-2370

packages/autocomplete/src/index.ts

@@ -293,7 +293,9 @@ export async function completer(
 // from https://github.com/mongodb-js/devtools-shared/commit/e4a5b00a83b19a76bdf380799a421511230168db
 function satisfiesVersion(v1: string, v2: string): boolean {
   const isGTECheck = /^\d+?\.\d+?\.\d+?$/.test(v2);
-  return semver.satisfies(v1, isGTECheck ? `>=${v2}` : v2);
+  return semver.satisfies(v1, isGTECheck ? `>=${v2}` : v2, {
+    includePrerelease: true,
+  });
 }
 
 function isAcceptable(
@@ -329,6 +331,7 @@ function isAcceptable(
       : connectionInfo.is_atlas || connectionInfo.is_local_atlas
       ? entry.env.includes(ATLAS)
       : entry.env.includes(ON_PREM));
+
   return isAcceptableVersion && isAcceptableEnvironment;
 }
 

packages/cli-repl/src/cli-repl.spec.ts

@@ -309,6 +309,7 @@ describe('CliRepl', function () {
           'maxTimeMS',
           'enableTelemetry',
           'editor',
+          'disableSchemaSampling',
           'snippetIndexSourceURLs',
           'snippetRegistryURL',
           'snippetAutoload',
@@ -2478,10 +2479,6 @@ describe('CliRepl', function () {
         // be listed
         wantWatch = true;
         wantShardDistribution = true;
-
-        // TODO: we need MQL support in mongodb-ts-autocomplete in order for it
-        // to autocomplete collection field names
-        wantQueryOperators = false;
       }
     }
 
@@ -2492,8 +2489,11 @@ describe('CliRepl', function () {
         await tick();
         input.write('\u0009');
         await waitCompletion(cliRepl.bus);
+        await tick();
       };
 
+      let docsLoadedPromise: Promise<void>;
+
       beforeEach(async function () {
         if (testServer === null) {
           cliReplOptions.shellCliOptions = { nodb: true };
@@ -2508,9 +2508,16 @@ describe('CliRepl', function () {
         if (!(testServer as any)?._opts.args?.includes('--auth')) {
           // make sure there are some collections we can autocomplete on below
           input.write('db.coll.insertOne({})\n');
-          input.write('db.movies.insertOne({})\n');
+          await waitEval(cliRepl.bus);
+          input.write('db.movies.insertOne({ year: 1 })\n');
           await waitEval(cliRepl.bus);
         }
+
+        docsLoadedPromise = new Promise<void>((resolve) => {
+          cliRepl.bus.once('mongosh:load-sample-docs-complete', () => {
+            resolve();
+          });
+        });
       });
 
       afterEach(async function () {
@@ -2521,6 +2528,29 @@ describe('CliRepl', function () {
         await cliRepl.mongoshRepl.close();
       });
 
+      it(`${
+        wantQueryOperators ? 'completes' : 'does not complete'
+      } query operators`, async function () {
+        input.write('db.movies.find({year: {$g');
+        await tabCompletion();
+
+        if (wantQueryOperators) {
+          if (process.env.USE_NEW_AUTOCOMPLETE) {
+            // wait for the documents to finish loading to be sure that the next
+            // tabCompletion() call will work
+            await docsLoadedPromise;
+          }
+        }
+
+        await tabCompletion();
+
+        if (wantQueryOperators) {
+          expect(output).to.include('db.movies.find({year: {$gte');
+        } else {
+          expect(output).to.not.include('db.movies.find({year: {$gte');
+        }
+      });
+
       it(`${
         wantWatch ? 'completes' : 'does not complete'
       } the watch method`, async function () {
@@ -2644,19 +2674,6 @@ describe('CliRepl', function () {
         expect(output).to.include('use admin');
       });
 
-      it(`${
-        wantQueryOperators ? 'completes' : 'does not complete'
-      } query operators`, async function () {
-        input.write('db.movies.find({year: {$g');
-        await tabCompletion();
-        await tabCompletion();
-        if (wantQueryOperators) {
-          expect(output).to.include('db.movies.find({year: {$gte');
-        } else {
-          expect(output).to.not.include('db.movies.find({year: {$gte');
-        }
-      });
-
       it('completes properties of shell API result types', async function () {
         if (!hasCollectionNames) return this.skip();
 

packages/cli-repl/src/mongosh-repl.spec.ts

@@ -55,6 +55,8 @@ describe('MongoshNodeRepl', function () {
   let config: Record<string, any>;
   const tmpdir = useTmpdir();
 
+  let docsLoadedPromise: Promise<void>;
+
   beforeEach(function () {
     input = new PassThrough();
     outputStream = new PassThrough();
@@ -116,6 +118,12 @@ describe('MongoshNodeRepl', function () {
       ioProvider: ioProvider,
     };
     mongoshRepl = new MongoshNodeRepl(mongoshReplOptions);
+
+    docsLoadedPromise = new Promise<void>((resolve) => {
+      mongoshRepl.bus.once('mongosh:load-sample-docs-complete', () => {
+        resolve();
+      });
+    });
   });
 
   let originalEnvVars: Record<string, string | undefined>;
@@ -379,6 +387,7 @@ describe('MongoshNodeRepl', function () {
       });
       const initialized = await mongoshRepl.initialize(serviceProvider);
       await mongoshRepl.startRepl(initialized);
+      await mongoshRepl.setConfig('disableSchemaSampling', false);
     });
 
     it('provides an editor action', async function () {
@@ -481,14 +490,37 @@ describe('MongoshNodeRepl', function () {
           await tick();
           expect(output, output).to.include('db.coll.updateOne');
         });
-        // this will eventually be supported in the new autocomplete
-        it.skip('autocompletes collection schema fields', async function () {
-          input.write('db.coll.find({');
-          await tabtab();
+        it('autocompletes collection schema fields', async function () {
+          if (!process.env.USE_NEW_AUTOCOMPLETE) {
+            // auto-completing collection field names only supported by new autocomplete
+            return this.skip();
+          }
+          input.write('db.coll.find({fo');
+          await tab();
+          await docsLoadedPromise;
+          await tab();
           await waitCompletion(bus);
           await tick();
           expect(output, output).to.include('db.coll.find({foo');
         });
+
+        it('does not autocomplete collection schema fields if disableSchemaSampling=true', async function () {
+          if (!process.env.USE_NEW_AUTOCOMPLETE) {
+            // auto-completing collection field names only supported by new autocomplete
+            return this.skip();
+          }
+          await mongoshRepl.setConfig('disableSchemaSampling', true);
+          try {
+            input.write('db.coll.find({fo');
+            await tab();
+            await tab();
+            await waitCompletion(bus);
+            await tick();
+            expect(output, output).to.not.include('db.coll.find({foo');
+          } finally {
+            await mongoshRepl.setConfig('disableSchemaSampling', false);
+          }
+        });
         it('autocompletes shell-api methods (once)', async function () {
           input.write('vers');
           await tabtab();

packages/cli-repl/src/mongosh-repl.ts

@@ -469,11 +469,6 @@ class MongoshNodeRepl implements EvaluationListener {
           }
         })(),
       ]);
-      this.bus.emit(
-        'mongosh:autocompletion-complete',
-        replResults,
-        mongoshResults
-      ); // For testing.
 
       // Sometimes the mongosh completion knows that what it is doing is right,
       // and that autocompletion based on inspecting the actual objects that
@@ -513,6 +508,8 @@ class MongoshNodeRepl implements EvaluationListener {
           results = results.filter(
             (result) => !CONTROL_CHAR_REGEXP.test(result)
           );
+          // emit here so that on nextTick the results should be output
+          this.bus.emit('mongosh:autocompletion-complete'); // For testing.
           return [results, completeOn];
         } finally {
           this.insideAutoCompleteOrGetPrompt = false;

packages/cli-repl/test/repl-helpers.ts

@@ -71,10 +71,7 @@ async function waitEval(bus: MongoshBus) {
 }
 
 async function waitCompletion(bus: MongoshBus) {
-  await Promise.race([
-    waitBus(bus, 'mongosh:autocompletion-complete'),
-    new Promise((resolve) => setTimeout(resolve, 10_000)?.unref?.()),
-  ]);
+  await waitBus(bus, 'mongosh:autocompletion-complete');
   await tick();
 }
 

packages/shell-api/src/collection.ts

@@ -2590,7 +2590,11 @@ export class Collection<
   async _getSampleDocsForCompletion(): Promise<Document[]> {
     return await Promise.race([
       (async () => {
-        return await this._getSampleDocs();
+        const result = await this._getSampleDocs();
+        this._mongo._instanceState.messageBus.emit(
+          'mongosh:load-sample-docs-complete'
+        );
+        return result;
       })(),
       (async () => {
         // 200ms should be a good compromise between giving the server a chance

packages/shell-api/src/database.ts

@@ -295,9 +295,13 @@ export class Database<
   async _getCollectionNamesForCompletion(): Promise<string[]> {
     return await Promise.race([
       (async () => {
-        return await this._getCollectionNames({
+        const result = await this._getCollectionNames({
           readPreference: 'primaryPreferred',
         });
+        this._mongo._instanceState.messageBus.emit(
+          'mongosh:load-collections-complete'
+        );
+        return result;
       })(),
       (async () => {
         // 200ms should be a good compromise between giving the server a chance

packages/shell-api/src/mongo.ts

@@ -362,9 +362,12 @@ export default class Mongo<
   async _getDatabaseNamesForCompletion(): Promise<string[]> {
     return await Promise.race([
       (async () => {
-        return (
+        const result = (
           await this._listDatabases({ readPreference: 'primaryPreferred' })
         ).databases.map((db) => db.name);
+
+        this._instanceState.messageBus.emit('mongosh:load-databases-complete');
+        return result;
       })(),
       (async () => {
         // See the comment in _getCollectionNamesForCompletion/database.ts

packages/shell-api/src/shell-api.spec.ts

@@ -916,6 +916,7 @@ describe('ShellApi', function () {
             ['maxTimeMS', null],
             ['enableTelemetry', false],
             ['editor', null],
+            ['disableSchemaSampling', false],
           ] as any);
 
           expect(shellResult.printable).to.deep.equal(expectedResult);

packages/shell-api/src/shell-instance-state.ts

@@ -483,17 +483,23 @@ export class ShellInstanceState {
       ): Promise<JSONSchema> => {
         const mongo = this.getMongoByConnectionId(connectionId);
         let docs: Document[] = [];
-        try {
-          docs = await mongo
-            ._getDb(databaseName)
-            .getCollection(collectionName)
-            ._getSampleDocsForCompletion();
-        } catch (err: any) {
-          if (
-            err?.code !== ShellApiErrors.NotConnected &&
-            err?.codeName !== 'Unauthorized'
-          ) {
-            throw err;
+        if (
+          (await this.evaluationListener.getConfig?.(
+            'disableSchemaSampling'
+          )) !== true
+        ) {
+          try {
+            docs = await mongo
+              ._getDb(databaseName)
+              .getCollection(collectionName)
+              ._getSampleDocsForCompletion();
+          } catch (err: any) {
+            if (
+              err?.code !== ShellApiErrors.NotConnected &&
+              err?.codeName !== 'Unauthorized'
+            ) {
+              throw err;
+            }
           }
         }
 
@@ -566,11 +572,13 @@ export class ShellInstanceState {
         try {
           const collectionNames =
             await this.currentDb._getCollectionNamesForCompletion();
-          return collectionNames.filter(
+          const result = collectionNames.filter(
             (name) =>
               name.toLowerCase().startsWith(collName.toLowerCase()) &&
               !CONTROL_CHAR_REGEXP.test(name)
           );
+          this.messageBus.emit('mongosh:load-collections-complete');
+          return result;
         } catch (err: any) {
           if (
             err?.code === ShellApiErrors.NotConnected ||
@@ -585,11 +593,13 @@ export class ShellInstanceState {
         try {
           const dbNames =
             await this.currentDb._mongo._getDatabaseNamesForCompletion();
-          return dbNames.filter(
+          const result = dbNames.filter(
             (name) =>
               name.toLowerCase().startsWith(dbName.toLowerCase()) &&
               !CONTROL_CHAR_REGEXP.test(name)
           );
+          this.messageBus.emit('mongosh:load-databases-complete');
+          return result;
         } catch (err: any) {
           if (
             err?.code === ShellApiErrors.NotConnected ||

packages/types/src/index.ts

@@ -335,10 +335,22 @@ export interface MongoshBusEventsMap extends ConnectEventMap {
    * Signals the completion of the autocomplete suggestion providers.
    * _ONLY AVAILABLE FOR TESTING._
    */
-  'mongosh:autocompletion-complete': (
-    replResults: string[],
-    mongoshResults: string[]
-  ) => void;
+  'mongosh:autocompletion-complete': () => void;
+  /**
+   * Signals the completion of the autocomplete helper.
+   * _ONLY AVAILABLE FOR TESTING._
+   */
+  'mongosh:load-databases-complete': () => void;
+  /**
+   * Signals the completion of the autocomplete helper.
+   * _ONLY AVAILABLE FOR TESTING._
+   */
+  'mongosh:load-collections-complete': () => void;
+  /**
+   * Signals the completion of the autocomplete helper.
+   * _ONLY AVAILABLE FOR TESTING._
+   */
+  'mongosh:load-sample-docs-complete': () => void;
   /**
    * Signals the completion of the asynchronous interrupt handler in MongoshRepl. Not fired for interrupts of _synchronous_ code.
    * _ONLY AVAILABLE FOR TESTING._
@@ -428,6 +440,7 @@ export class ShellUserConfig {
   enableTelemetry = false;
   editor: string | null = null;
   logLocation: string | undefined;
+  disableSchemaSampling = false;
 }
 
 export class ShellUserConfigValidator {
@@ -447,6 +460,7 @@ export class ShellUserConfigValidator {
           return `${key} must be null or a positive integer`;
         }
         return null;
+      case 'disableSchemaSampling':
       case 'enableTelemetry':
         if (typeof value !== 'boolean') {
           return `${key} must be a boolean`;