allow opting out of schema sampling with config disableSchemaSampling=false

Author: lerouxb

packages/autocomplete/src/index.ts

@@ -411,7 +411,9 @@ function filterShellAPI(
 
 type AutocompleteShellInstanceState = {
   getAutocompleteParameters: () => AutocompleteParameters;
-  getAutocompletionContext: () => AutocompletionContext;
+  getAutocompletionContext: (options: {
+    disableSchemaSampling?: boolean;
+  }) => AutocompletionContext;
 };
 
 function transformAutocompleteResults(
@@ -429,15 +431,22 @@ export async function initNewAutocompleter(
   instanceState: Pick<
     AutocompleteShellInstanceState,
     'getAutocompletionContext'
-  >
+  >,
+  {
+    disableSchemaSampling = false,
+  }: {
+    disableSchemaSampling?: boolean;
+  }
 ): Promise<(text: string) => Promise<CompletionResults>> {
   // only import the autocompleter code the first time we need it to
   // hide the time it takes from the initial startup time
   const { MongoDBAutocompleter } = await import(
     '@mongodb-js/mongodb-ts-autocomplete'
   );
 
-  const autocompletionContext = instanceState.getAutocompletionContext();
+  const autocompletionContext = instanceState.getAutocompletionContext({
+    disableSchemaSampling,
+  });
   const mongoDBCompleter = new MongoDBAutocompleter({
     context: autocompletionContext,
   });

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

@@ -379,6 +379,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 +482,34 @@ 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({');
+        it('autocompletes collection schema fields', async function () {
+          if (!process.env.USE_NEW_AUTOCOMPLETE) {
+            // auto-completing collection field names only supported by new autocomplete
+            this.skip();
+          }
+          input.write('db.coll.find({fo');
           await tabtab();
           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
+            this.skip();
+          }
+          await mongoshRepl.setConfig('disableSchemaSampling', true);
+          try {
+            input.write('db.coll.find({fo');
+            await tabtab();
+            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

@@ -457,7 +457,11 @@ class MongoshNodeRepl implements EvaluationListener {
         (async () => {
           if (process.env.USE_NEW_AUTOCOMPLETE) {
             if (!newMongoshCompleter) {
-              newMongoshCompleter = await initNewAutocompleter(instanceState);
+              newMongoshCompleter = await initNewAutocompleter(instanceState, {
+                disableSchemaSampling: await this.getConfig(
+                  'disableSchemaSampling'
+                ),
+              });
             }
 
             return newMongoshCompleter(text);

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

@@ -415,7 +415,11 @@ export class ShellInstanceState {
     throw new Error(`mongo with connection id ${connectionId} not found`);
   }
 
-  public getAutocompletionContext(): AutocompletionContext {
+  public getAutocompletionContext({
+    disableSchemaSampling = false,
+  }: {
+    disableSchemaSampling?: boolean;
+  }): AutocompletionContext {
     return {
       currentDatabaseAndConnection: ():
         | {
@@ -483,17 +487,19 @@ 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 (!disableSchemaSampling) {
+          try {
+            docs = await mongo
+              ._getDb(databaseName)
+              .getCollection(collectionName)
+              ._getSampleDocsForCompletion();
+          } catch (err: any) {
+            if (
+              err?.code !== ShellApiErrors.NotConnected &&
+              err?.codeName !== 'Unauthorized'
+            ) {
+              throw err;
+            }
           }
         }
 

packages/types/src/index.ts

@@ -520,6 +520,7 @@ export class CliUserConfig extends SnippetShellUserConfig {
   logMaxFileCount = 100;
   logCompressionEnabled = false;
   logRetentionGB: number | undefined = undefined;
+  disableSchemaSampling = false;
 }
 
 export class CliUserConfigValidator extends SnippetShellUserConfigValidator {
@@ -557,6 +558,7 @@ export class CliUserConfigValidator extends SnippetShellUserConfigValidator {
       case 'forceDisableTelemetry':
       case 'showStackTraces':
       case 'logCompressionEnabled':
+      case 'disableSchemaSampling':
         if (typeof value !== 'boolean') {
           return `${key} must be a boolean`;
         }