VSCODE-504: streams support in playground (#633)

* chore: autocomplete support for streams

* chore: print console logs continuously

* address PR feedbacks

Author: shaketbaby

src/editors/playgroundController.ts

@@ -1,11 +1,10 @@
 import * as vscode from 'vscode';
 import path from 'path';
-import type { OutputChannel, TextEditor } from 'vscode';
+import type { TextEditor } from 'vscode';
 import { ProgressLocation } from 'vscode';
 import vm from 'vm';
 import os from 'os';
 import transpiler from 'bson-transpilers';
-import util from 'util';
 
 import type ActiveConnectionCodeLensProvider from './activeConnectionCodeLensProvider';
 import type PlaygroundSelectedCodeActionProvider from './playgroundSelectedCodeActionProvider';
@@ -111,7 +110,6 @@ export default class PlaygroundController {
 
   _isPartialRun = false;
 
-  _outputChannel: OutputChannel;
   private _activeConnectionCodeLensProvider: ActiveConnectionCodeLensProvider;
   private _playgroundResultViewColumn?: vscode.ViewColumn;
   private _playgroundResultTextDocument?: vscode.TextDocument;
@@ -145,8 +143,6 @@ export default class PlaygroundController {
     this._telemetryService = telemetryService;
     this._statusView = statusView;
     this._playgroundResultViewProvider = playgroundResultViewProvider;
-    this._outputChannel =
-      vscode.window.createOutputChannel('Playground output');
     this._activeConnectionCodeLensProvider = activeConnectionCodeLensProvider;
     this._exportToLanguageCodeLensProvider = exportToLanguageCodeLensProvider;
     this._playgroundSelectedCodeActionProvider =
@@ -467,7 +463,7 @@ export default class PlaygroundController {
             // If a user clicked the cancel button terminate all playground scripts.
             this._languageServerController.cancelAll();
 
-            return { outputLines: undefined, result: undefined };
+            return { result: undefined };
           });
 
           // Run all playground scripts.
@@ -483,10 +479,7 @@ export default class PlaygroundController {
     } catch (error) {
       log.error('Evaluating playground with cancel modal failed', error);
 
-      return {
-        outputLines: undefined,
-        result: undefined,
-      };
+      return { result: undefined };
     }
   }
 
@@ -567,27 +560,10 @@ export default class PlaygroundController {
       }
     }
 
-    this._outputChannel.clear();
-
     const evaluateResponse: ShellEvaluateResult =
       await this._evaluateWithCancelModal();
 
-    if (evaluateResponse?.outputLines?.length) {
-      for (const line of evaluateResponse.outputLines) {
-        this._outputChannel.appendLine(
-          typeof line.content === 'string'
-            ? line.content
-            : util.inspect(line.content)
-        );
-      }
-
-      this._outputChannel.show(true);
-    }
-
-    if (
-      !evaluateResponse ||
-      (!evaluateResponse.outputLines && !evaluateResponse.result)
-    ) {
+    if (!evaluateResponse || !evaluateResponse.result) {
       return false;
     }
 

src/language/languageServerController.ts

@@ -12,6 +12,7 @@ import {
 } from 'vscode-languageclient/node';
 import type { ExtensionContext } from 'vscode';
 import { workspace } from 'vscode';
+import util from 'util';
 
 import { createLogger } from '../logging';
 import type {
@@ -38,6 +39,9 @@ export default class LanguageServerController {
   _currentConnectionString?: string;
   _currentConnectionOptions?: MongoClientOptions;
 
+  _consoleOutputChannel =
+    vscode.window.createOutputChannel('Playground output');
+
   constructor(context: ExtensionContext) {
     this._context = context;
 
@@ -151,6 +155,19 @@ export default class LanguageServerController {
         void vscode.window.showErrorMessage(messsage);
       }
     );
+
+    this._client.onNotification(
+      ServerCommands.SHOW_CONSOLE_OUTPUT,
+      (outputs) => {
+        for (const line of outputs) {
+          this._consoleOutputChannel.appendLine(
+            typeof line === 'string' ? line : util.inspect(line)
+          );
+        }
+
+        this._consoleOutputChannel.show(true);
+      }
+    );
   }
 
   deactivate(): Thenable<void> | undefined {
@@ -173,6 +190,8 @@ export default class LanguageServerController {
     });
     this._isExecutingInProgress = true;
 
+    this._consoleOutputChannel.clear();
+
     // Instantiate a new CancellationTokenSource object
     // that generates a cancellation token for each run of a playground.
     this._source = new CancellationTokenSource();

src/language/mongoDBService.ts

@@ -21,6 +21,7 @@ import parseSchema from 'mongodb-schema';
 import path from 'path';
 import { signatures } from '@mongosh/shell-api';
 import translator from '@mongosh/i18n';
+import { isAtlasStream } from 'mongodb-build-info';
 import { Worker as WorkerThreads } from 'worker_threads';
 
 import { ExportToLanguageMode } from '../types/playgroundType';
@@ -65,6 +66,7 @@ export default class MongoDBService {
   _currentConnectionOptions?: MongoClientOptions;
 
   _databaseCompletionItems: CompletionItem[] = [];
+  _streamProcessorCompletionItems: CompletionItem[] = [];
   _shellSymbolCompletionItems: { [symbol: string]: CompletionItem[] } = {};
   _globalSymbolCompletionItems: CompletionItem[] = [];
   _collections: { [database: string]: string[] } = {};
@@ -147,6 +149,7 @@ export default class MongoDBService {
         databases: true,
         collections: true,
         fields: true,
+        streamProcessors: true,
       });
       await this._closeCurrentConnection();
     }
@@ -174,6 +177,22 @@ export default class MongoDBService {
       );
     }
 
+    if (isAtlasStream(connectionString || '')) {
+      await this._getAndCacheStreamProcessors();
+    } else {
+      await this._getAndCacheDatabases();
+    }
+
+    this._connection.console.log(
+      `CliServiceProvider active connection has changed: { connectionId: ${connectionId} }`
+    );
+    return {
+      successfullyConnected: true,
+      connectionId,
+    };
+  }
+
+  async _getAndCacheDatabases() {
     try {
       // Get database names for the current connection.
       const databases = await this._getDatabases();
@@ -184,14 +203,17 @@ export default class MongoDBService {
         `LS get databases error: ${util.inspect(error)}`
       );
     }
+  }
 
-    this._connection.console.log(
-      `CliServiceProvider active connection has changed: { connectionId: ${connectionId} }`
-    );
-    return {
-      successfullyConnected: true,
-      connectionId,
-    };
+  async _getAndCacheStreamProcessors() {
+    try {
+      const processors = await this._getStreamProcessors();
+      this._cacheStreamProcessorCompletionItems(processors);
+    } catch (error) {
+      this._connection.console.error(
+        `LS get stream processors error: ${util.inspect(error)}`
+      );
+    }
   }
 
   /**
@@ -245,9 +267,16 @@ export default class MongoDBService {
           )
         );
 
-        worker?.on(
-          'message',
-          ({ error, data }: { data?: ShellEvaluateResult; error?: any }) => {
+        worker?.on('message', ({ name, payload }) => {
+          if (name === ServerCommands.SHOW_CONSOLE_OUTPUT) {
+            void this._connection.sendNotification(name, payload);
+          }
+
+          if (name === ServerCommands.CODE_EXECUTION_RESULT) {
+            const { error, data } = payload as {
+              data?: ShellEvaluateResult;
+              error?: any;
+            };
             if (error) {
               this._connection.console.error(
                 `WORKER error: ${util.inspect(error)}`
@@ -261,7 +290,7 @@ export default class MongoDBService {
               resolve(data);
             });
           }
-        );
+        });
 
         worker.postMessage({
           name: ServerCommands.EXECUTE_CODE_FROM_PLAYGROUND,
@@ -294,6 +323,24 @@ export default class MongoDBService {
     });
   }
 
+  /**
+   * Get stream processors names for the current connection.
+   */
+  async _getStreamProcessors(): Promise<Document[]> {
+    if (this._serviceProvider) {
+      try {
+        const cmd = { listStreamProcessors: 1 };
+        const result = await this._serviceProvider.runCommand('admin', cmd);
+        return result.streamProcessors ?? [];
+      } catch (error) {
+        this._connection.console.error(
+          `LS get stream processors error: ${error}`
+        );
+      }
+    }
+    return [];
+  }
+
   /**
    * Get database names for the current connection.
    */
@@ -377,7 +424,7 @@ export default class MongoDBService {
   }
 
   /**
-   * Return 'db' and 'use' completion items.
+   * Return 'db', 'sp' and 'use' completion items.
    */
   _cacheGlobalSymbolCompletionItems() {
     this._globalSymbolCompletionItems = [
@@ -386,6 +433,11 @@ export default class MongoDBService {
         kind: CompletionItemKind.Method,
         preselect: true,
       },
+      {
+        label: 'sp',
+        kind: CompletionItemKind.Method,
+        preselect: true,
+      },
       {
         label: 'use',
         kind: CompletionItemKind.Function,
@@ -783,6 +835,18 @@ export default class MongoDBService {
     }
   }
 
+  /**
+   * If the current node is 'sp.processor.<trigger>' or 'sp["processor"].<trigger>'.
+   */
+  _provideStreamProcessorSymbolCompletionItems(state: CompletionState) {
+    if (state.isStreamProcessorSymbol) {
+      this._connection.console.log(
+        'VISITOR found stream processor symbol completions'
+      );
+      return this._shellSymbolCompletionItems.StreamProcessor;
+    }
+  }
+
   /**
    * If the current node is 'db.collection.find().<trigger>'.
    */
@@ -895,6 +959,37 @@ export default class MongoDBService {
     }
   }
 
+  /**
+   * If the current node is 'sp.<trigger>'.
+   */
+  _provideSpSymbolCompletionItems(state: CompletionState) {
+    if (state.isSpSymbol) {
+      if (state.isStreamProcessorName) {
+        this._connection.console.log(
+          'VISITOR found sp symbol and stream processor name completions'
+        );
+        return this._shellSymbolCompletionItems.Streams.concat(
+          this._streamProcessorCompletionItems
+        );
+      }
+
+      this._connection.console.log('VISITOR found sp symbol completions');
+      return this._shellSymbolCompletionItems.Streams;
+    }
+  }
+
+  /**
+   * If the current node is 'sp.get(<trigger>)'.
+   */
+  _provideStreamProcessorNameCompletionItems(state: CompletionState) {
+    if (state.isStreamProcessorName) {
+      this._connection.console.log(
+        'VISITOR found stream processor name completions'
+      );
+      return this._streamProcessorCompletionItems;
+    }
+  }
+
   /**
    * If the current node can be used as a collection name
    * e.g. 'db.<trigger>.find()' or 'let a = db.<trigger>'.
@@ -965,6 +1060,7 @@ export default class MongoDBService {
       this._provideIdentifierObjectValueCompletionItems.bind(this, state),
       this._provideTextObjectValueCompletionItems.bind(this, state),
       this._provideCollectionSymbolCompletionItems.bind(this, state),
+      this._provideStreamProcessorSymbolCompletionItems.bind(this, state),
       this._provideFindCursorCompletionItems.bind(this, state),
       this._provideAggregationCursorCompletionItems.bind(this, state),
       this._provideGlobalSymbolCompletionItems.bind(this, state),
@@ -974,13 +1070,15 @@ export default class MongoDBService {
         currentLineText,
         position
       ),
+      this._provideSpSymbolCompletionItems.bind(this, state),
       this._provideCollectionNameCompletionItems.bind(
         this,
         state,
         currentLineText,
         position
       ),
       this._provideDbNameCompletionItems.bind(this, state),
+      this._provideStreamProcessorNameCompletionItems.bind(this, state),
     ];
 
     for (const func of completionOptions) {
@@ -1117,6 +1215,18 @@ export default class MongoDBService {
     this._collections[database] = collections.map((item) => item.name);
   }
 
+  _cacheStreamProcessorCompletionItems(processors: Document[]): void {
+    this._streamProcessorCompletionItems = processors.map(({ name }) => ({
+      kind: CompletionItemKind.Folder,
+      preselect: true,
+      label: name,
+    }));
+  }
+
+  clearCachedStreamProcessors(): void {
+    this._streamProcessorCompletionItems = [];
+  }
+
   clearCachedFields(): void {
     this._fields = {};
   }
@@ -1142,13 +1252,16 @@ export default class MongoDBService {
 
   clearCachedCompletions(clear: ClearCompletionsCache): void {
     if (clear.fields) {
-      this._fields = {};
+      this.clearCachedFields();
     }
     if (clear.databases) {
-      this._databaseCompletionItems = [];
+      this.clearCachedDatabases();
     }
     if (clear.collections) {
-      this._collections = {};
+      this.clearCachedCollections();
+    }
+    if (clear.streamProcessors) {
+      this.clearCachedStreamProcessors();
     }
   }
 }