1+ import crypto from 'crypto' ;
12import completer from '@mongosh/autocomplete' ;
23import { MongoshInternalError , MongoshWarning } from '@mongosh/errors' ;
34import { changeHistory } from '@mongosh/history' ;
@@ -7,6 +8,7 @@ import type {
78} from '@mongosh/service-provider-core' ;
89import type {
910 EvaluationListener ,
11+ Mongo ,
1012 OnLoadResult ,
1113 ShellCliOptions ,
1214} from '@mongosh/shell-api' ;
@@ -50,6 +52,9 @@ import { installPasteSupport } from './repl-paste-support';
5052import util from 'util' ;
5153
5254import { MongoDBAutocompleter } from '@mongodb-js/mongodb-ts-autocomplete' ;
55+ import type { AutocompletionContext } from '@mongodb-js/mongodb-ts-autocomplete' ;
56+ import type { JSONSchema } from 'mongodb-schema' ;
57+ import { analyzeDocuments } from 'mongodb-schema' ;
5358
5459declare const __non_webpack_require__ : any ;
5560
@@ -433,6 +438,88 @@ class MongoshNodeRepl implements EvaluationListener {
433438 this . runtimeState ( ) . context . history = history ;
434439 }
435440
441+ // TODO: probably better to have this on instanceState
442+ public _getMongoByConnectionId (
443+ instanceState : ShellInstanceState ,
444+ connectionId : string
445+ ) : Mongo {
446+ for ( const mongo of instanceState . mongos ) {
447+ if ( connectionIdFromURI ( mongo . getURI ( ) ) === connectionId ) {
448+ return mongo ;
449+ }
450+ }
451+ throw new Error ( `mongo with connection id ${ connectionId } not found` ) ;
452+ }
453+
454+ public getAutocompletionContext (
455+ instanceState : ShellInstanceState
456+ ) : AutocompletionContext {
457+ return {
458+ currentDatabaseAndConnection : ( ) => {
459+ return {
460+ connectionId : connectionIdFromURI (
461+ instanceState . currentDb . getMongo ( ) . getURI ( )
462+ ) ,
463+ databaseName : instanceState . currentDb . getName ( ) ,
464+ } ;
465+ } ,
466+ databasesForConnection : async (
467+ connectionId : string
468+ ) : Promise < string [ ] > => {
469+ const mongo = this . _getMongoByConnectionId ( instanceState , connectionId ) ;
470+ try {
471+ const dbNames = await mongo . _getDatabaseNamesForCompletion ( ) ;
472+ return dbNames . filter (
473+ ( name : string ) => ! CONTROL_CHAR_REGEXP . test ( name )
474+ ) ;
475+ } catch ( err : any ) {
476+ // TODO: move this code to a method in the shell instance so we don't
477+ // have to hardcode the error code or export it.
478+ if ( err ?. code === 'SHAPI-10004' || err ?. codeName === 'Unauthorized' ) {
479+ return [ ] ;
480+ }
481+ throw err ;
482+ }
483+ } ,
484+ collectionsForDatabase : async (
485+ connectionId : string ,
486+ databaseName : string
487+ ) : Promise < string [ ] > => {
488+ const mongo = this . _getMongoByConnectionId ( instanceState , connectionId ) ;
489+ try {
490+ const collectionNames = await mongo
491+ . _getDb ( databaseName )
492+ . _getCollectionNamesForCompletion ( ) ;
493+ return collectionNames . filter (
494+ ( name : string ) => ! CONTROL_CHAR_REGEXP . test ( name )
495+ ) ;
496+ } catch ( err : any ) {
497+ // TODO: move this code to a method in the shell instance so we don't
498+ // have to hardcode the error code or export it.
499+ if ( err ?. code === 'SHAPI-10004' || err ?. codeName === 'Unauthorized' ) {
500+ return [ ] ;
501+ }
502+ throw err ;
503+ }
504+ } ,
505+ schemaInformationForCollection : async (
506+ connectionId : string ,
507+ databaseName : string ,
508+ collectionName : string
509+ ) : Promise < JSONSchema > => {
510+ const mongo = this . _getMongoByConnectionId ( instanceState , connectionId ) ;
511+ const docs = await mongo
512+ . _getDb ( databaseName )
513+ . getCollection ( collectionName )
514+ . _getSampleDocsForCompletion ( ) ;
515+ const schemaAccessor = await analyzeDocuments ( docs ) ;
516+
517+ const schema = await schemaAccessor . getMongoDBJsonSchema ( ) ;
518+ return schema ;
519+ } ,
520+ } ;
521+ }
522+
436523 private async finishInitializingNodeRepl ( ) : Promise < void > {
437524 const { repl, instanceState } = this . runtimeState ( ) ;
438525 if ( ! repl ) return ;
@@ -445,7 +532,8 @@ class MongoshNodeRepl implements EvaluationListener {
445532 line : string
446533 ) => Promise < [ string [ ] , string , 'exclusive' ] | [ string [ ] , string ] > ;
447534 if ( process . env . USE_NEW_AUTOCOMPLETE ) {
448- const autocompletionContext = instanceState . getAutocompletionContext ( ) ;
535+ const autocompletionContext =
536+ this . getAutocompletionContext ( instanceState ) ;
449537 newMongoshCompleter = new MongoDBAutocompleter ( {
450538 context : autocompletionContext ,
451539 } ) ;
@@ -1347,4 +1435,12 @@ async function enterAsynchronousExecutionForPrompt(): Promise<void> {
13471435 await new Promise ( setImmediate ) ;
13481436}
13491437
1438+ function connectionIdFromURI ( uri : string ) : string {
1439+ // turn the uri into something we can safely use as part of a "filename"
1440+ // inside autocomplete
1441+ const hash = crypto . createHash ( 'sha256' ) ;
1442+ hash . update ( uri ) ;
1443+ return hash . digest ( 'hex' ) ;
1444+ }
1445+
13501446export default MongoshNodeRepl ;
0 commit comments