Skip to content

Commit 0acc023

Browse files
lerouxblerouxb
authored
feat: allow opting out of schema sampling with config disableSchemaSampling=true MONGOSH-2370 (#2503)
* allow opting out of schema sampling with config disableSchemaSampling=false * copilot * move the config * test fixes * better way of waiting for completion * remove excess promises * also emit the events for old autocomplete * not needed * emit the event later * wait more? * add a tick * reorder the tests * print debug info * includePrerelease * remove logging
1 parent b23127a commit 0acc023

File tree

11 files changed

+134
-52
lines changed

11 files changed

+134
-52
lines changed

‎packages/autocomplete/src/index.ts

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -293,7 +293,9 @@ export async function completer(
293293
// from https://github.com/mongodb-js/devtools-shared/commit/e4a5b00a83b19a76bdf380799a421511230168db
294294
function satisfiesVersion(v1: string, v2: string): boolean {
295295
const isGTECheck = /^\d+?\.\d+?\.\d+?$/.test(v2);
296-
return semver.satisfies(v1, isGTECheck ? `>=${v2}` : v2);
296+
return semver.satisfies(v1, isGTECheck ? `>=${v2}` : v2, {
297+
includePrerelease: true,
298+
});
297299
}
298300

299301
function isAcceptable(
@@ -329,6 +331,7 @@ function isAcceptable(
329331
: connectionInfo.is_atlas || connectionInfo.is_local_atlas
330332
? entry.env.includes(ATLAS)
331333
: entry.env.includes(ON_PREM));
334+
332335
return isAcceptableVersion && isAcceptableEnvironment;
333336
}
334337

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

Lines changed: 35 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -309,6 +309,7 @@ describe('CliRepl', function () {
309309
'maxTimeMS',
310310
'enableTelemetry',
311311
'editor',
312+
'disableSchemaSampling',
312313
'snippetIndexSourceURLs',
313314
'snippetRegistryURL',
314315
'snippetAutoload',
@@ -2478,10 +2479,6 @@ describe('CliRepl', function () {
24782479
// be listed
24792480
wantWatch = true;
24802481
wantShardDistribution = true;
2481-
2482-
// TODO: we need MQL support in mongodb-ts-autocomplete in order for it
2483-
// to autocomplete collection field names
2484-
wantQueryOperators = false;
24852482
}
24862483
}
24872484

@@ -2492,8 +2489,11 @@ describe('CliRepl', function () {
24922489
await tick();
24932490
input.write('\u0009');
24942491
await waitCompletion(cliRepl.bus);
2492+
await tick();
24952493
};
24962494

2495+
let docsLoadedPromise: Promise<void>;
2496+
24972497
beforeEach(async function () {
24982498
if (testServer === null) {
24992499
cliReplOptions.shellCliOptions = { nodb: true };
@@ -2508,9 +2508,16 @@ describe('CliRepl', function () {
25082508
if (!(testServer as any)?._opts.args?.includes('--auth')) {
25092509
// make sure there are some collections we can autocomplete on below
25102510
input.write('db.coll.insertOne({})\n');
2511-
input.write('db.movies.insertOne({})\n');
2511+
await waitEval(cliRepl.bus);
2512+
input.write('db.movies.insertOne({ year: 1 })\n');
25122513
await waitEval(cliRepl.bus);
25132514
}
2515+
2516+
docsLoadedPromise = new Promise<void>((resolve) => {
2517+
cliRepl.bus.once('mongosh:load-sample-docs-complete', () => {
2518+
resolve();
2519+
});
2520+
});
25142521
});
25152522

25162523
afterEach(async function () {
@@ -2521,6 +2528,29 @@ describe('CliRepl', function () {
25212528
await cliRepl.mongoshRepl.close();
25222529
});
25232530

2531+
it(`${
2532+
wantQueryOperators ? 'completes' : 'does not complete'
2533+
} query operators`, async function () {
2534+
input.write('db.movies.find({year: {$g');
2535+
await tabCompletion();
2536+
2537+
if (wantQueryOperators) {
2538+
if (process.env.USE_NEW_AUTOCOMPLETE) {
2539+
// wait for the documents to finish loading to be sure that the next
2540+
// tabCompletion() call will work
2541+
await docsLoadedPromise;
2542+
}
2543+
}
2544+
2545+
await tabCompletion();
2546+
2547+
if (wantQueryOperators) {
2548+
expect(output).to.include('db.movies.find({year: {$gte');
2549+
} else {
2550+
expect(output).to.not.include('db.movies.find({year: {$gte');
2551+
}
2552+
});
2553+
25242554
it(`${
25252555
wantWatch ? 'completes' : 'does not complete'
25262556
} the watch method`, async function () {
@@ -2644,19 +2674,6 @@ describe('CliRepl', function () {
26442674
expect(output).to.include('use admin');
26452675
});
26462676

2647-
it(`${
2648-
wantQueryOperators ? 'completes' : 'does not complete'
2649-
} query operators`, async function () {
2650-
input.write('db.movies.find({year: {$g');
2651-
await tabCompletion();
2652-
await tabCompletion();
2653-
if (wantQueryOperators) {
2654-
expect(output).to.include('db.movies.find({year: {$gte');
2655-
} else {
2656-
expect(output).to.not.include('db.movies.find({year: {$gte');
2657-
}
2658-
});
2659-
26602677
it('completes properties of shell API result types', async function () {
26612678
if (!hasCollectionNames) return this.skip();
26622679

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

Lines changed: 36 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,8 @@ describe('MongoshNodeRepl', function () {
5555
let config: Record<string, any>;
5656
const tmpdir = useTmpdir();
5757

58+
let docsLoadedPromise: Promise<void>;
59+
5860
beforeEach(function () {
5961
input = new PassThrough();
6062
outputStream = new PassThrough();
@@ -116,6 +118,12 @@ describe('MongoshNodeRepl', function () {
116118
ioProvider: ioProvider,
117119
};
118120
mongoshRepl = new MongoshNodeRepl(mongoshReplOptions);
121+
122+
docsLoadedPromise = new Promise<void>((resolve) => {
123+
mongoshRepl.bus.once('mongosh:load-sample-docs-complete', () => {
124+
resolve();
125+
});
126+
});
119127
});
120128

121129
let originalEnvVars: Record<string, string | undefined>;
@@ -379,6 +387,7 @@ describe('MongoshNodeRepl', function () {
379387
});
380388
const initialized = await mongoshRepl.initialize(serviceProvider);
381389
await mongoshRepl.startRepl(initialized);
390+
await mongoshRepl.setConfig('disableSchemaSampling', false);
382391
});
383392

384393
it('provides an editor action', async function () {
@@ -481,14 +490,37 @@ describe('MongoshNodeRepl', function () {
481490
await tick();
482491
expect(output, output).to.include('db.coll.updateOne');
483492
});
484-
// this will eventually be supported in the new autocomplete
485-
it.skip('autocompletes collection schema fields', async function () {
486-
input.write('db.coll.find({');
487-
await tabtab();
493+
it('autocompletes collection schema fields', async function () {
494+
if (!process.env.USE_NEW_AUTOCOMPLETE) {
495+
// auto-completing collection field names only supported by new autocomplete
496+
return this.skip();
497+
}
498+
input.write('db.coll.find({fo');
499+
await tab();
500+
await docsLoadedPromise;
501+
await tab();
488502
await waitCompletion(bus);
489503
await tick();
490504
expect(output, output).to.include('db.coll.find({foo');
491505
});
506+
507+
it('does not autocomplete collection schema fields if disableSchemaSampling=true', async function () {
508+
if (!process.env.USE_NEW_AUTOCOMPLETE) {
509+
// auto-completing collection field names only supported by new autocomplete
510+
return this.skip();
511+
}
512+
await mongoshRepl.setConfig('disableSchemaSampling', true);
513+
try {
514+
input.write('db.coll.find({fo');
515+
await tab();
516+
await tab();
517+
await waitCompletion(bus);
518+
await tick();
519+
expect(output, output).to.not.include('db.coll.find({foo');
520+
} finally {
521+
await mongoshRepl.setConfig('disableSchemaSampling', false);
522+
}
523+
});
492524
it('autocompletes shell-api methods (once)', async function () {
493525
input.write('vers');
494526
await tabtab();

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

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -469,11 +469,6 @@ class MongoshNodeRepl implements EvaluationListener {
469469
}
470470
})(),
471471
]);
472-
this.bus.emit(
473-
'mongosh:autocompletion-complete',
474-
replResults,
475-
mongoshResults
476-
); // For testing.
477472

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

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

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -71,10 +71,7 @@ async function waitEval(bus: MongoshBus) {
7171
}
7272

7373
async function waitCompletion(bus: MongoshBus) {
74-
await Promise.race([
75-
waitBus(bus, 'mongosh:autocompletion-complete'),
76-
new Promise((resolve) => setTimeout(resolve, 10_000)?.unref?.()),
77-
]);
74+
await waitBus(bus, 'mongosh:autocompletion-complete');
7875
await tick();
7976
}
8077

‎packages/shell-api/src/collection.ts

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2590,7 +2590,11 @@ export class Collection<
25902590
async _getSampleDocsForCompletion(): Promise<Document[]> {
25912591
return await Promise.race([
25922592
(async () => {
2593-
return await this._getSampleDocs();
2593+
const result = await this._getSampleDocs();
2594+
this._mongo._instanceState.messageBus.emit(
2595+
'mongosh:load-sample-docs-complete'
2596+
);
2597+
return result;
25942598
})(),
25952599
(async () => {
25962600
// 200ms should be a good compromise between giving the server a chance

‎packages/shell-api/src/database.ts

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -295,9 +295,13 @@ export class Database<
295295
async _getCollectionNamesForCompletion(): Promise<string[]> {
296296
return await Promise.race([
297297
(async () => {
298-
return await this._getCollectionNames({
298+
const result = await this._getCollectionNames({
299299
readPreference: 'primaryPreferred',
300300
});
301+
this._mongo._instanceState.messageBus.emit(
302+
'mongosh:load-collections-complete'
303+
);
304+
return result;
301305
})(),
302306
(async () => {
303307
// 200ms should be a good compromise between giving the server a chance

‎packages/shell-api/src/mongo.ts

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -362,9 +362,12 @@ export default class Mongo<
362362
async _getDatabaseNamesForCompletion(): Promise<string[]> {
363363
return await Promise.race([
364364
(async () => {
365-
return (
365+
const result = (
366366
await this._listDatabases({ readPreference: 'primaryPreferred' })
367367
).databases.map((db) => db.name);
368+
369+
this._instanceState.messageBus.emit('mongosh:load-databases-complete');
370+
return result;
368371
})(),
369372
(async () => {
370373
// See the comment in _getCollectionNamesForCompletion/database.ts

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -916,6 +916,7 @@ describe('ShellApi', function () {
916916
['maxTimeMS', null],
917917
['enableTelemetry', false],
918918
['editor', null],
919+
['disableSchemaSampling', false],
919920
] as any);
920921

921922
expect(shellResult.printable).to.deep.equal(expectedResult);

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

Lines changed: 23 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -483,17 +483,23 @@ export class ShellInstanceState {
483483
): Promise<JSONSchema> => {
484484
const mongo = this.getMongoByConnectionId(connectionId);
485485
let docs: Document[] = [];
486-
try {
487-
docs = await mongo
488-
._getDb(databaseName)
489-
.getCollection(collectionName)
490-
._getSampleDocsForCompletion();
491-
} catch (err: any) {
492-
if (
493-
err?.code !== ShellApiErrors.NotConnected &&
494-
err?.codeName !== 'Unauthorized'
495-
) {
496-
throw err;
486+
if (
487+
(await this.evaluationListener.getConfig?.(
488+
'disableSchemaSampling'
489+
)) !== true
490+
) {
491+
try {
492+
docs = await mongo
493+
._getDb(databaseName)
494+
.getCollection(collectionName)
495+
._getSampleDocsForCompletion();
496+
} catch (err: any) {
497+
if (
498+
err?.code !== ShellApiErrors.NotConnected &&
499+
err?.codeName !== 'Unauthorized'
500+
) {
501+
throw err;
502+
}
497503
}
498504
}
499505

@@ -566,11 +572,13 @@ export class ShellInstanceState {
566572
try {
567573
const collectionNames =
568574
await this.currentDb._getCollectionNamesForCompletion();
569-
return collectionNames.filter(
575+
const result = collectionNames.filter(
570576
(name) =>
571577
name.toLowerCase().startsWith(collName.toLowerCase()) &&
572578
!CONTROL_CHAR_REGEXP.test(name)
573579
);
580+
this.messageBus.emit('mongosh:load-collections-complete');
581+
return result;
574582
} catch (err: any) {
575583
if (
576584
err?.code === ShellApiErrors.NotConnected ||
@@ -585,11 +593,13 @@ export class ShellInstanceState {
585593
try {
586594
const dbNames =
587595
await this.currentDb._mongo._getDatabaseNamesForCompletion();
588-
return dbNames.filter(
596+
const result = dbNames.filter(
589597
(name) =>
590598
name.toLowerCase().startsWith(dbName.toLowerCase()) &&
591599
!CONTROL_CHAR_REGEXP.test(name)
592600
);
601+
this.messageBus.emit('mongosh:load-databases-complete');
602+
return result;
593603
} catch (err: any) {
594604
if (
595605
err?.code === ShellApiErrors.NotConnected ||

0 commit comments

Comments
 (0)