feat(participant): export to a playground VSCODE-574 (#832)

* feat(participant): export to a playground VSCODE-574

* refactor: move logic to participant controller

* test: add export to playground tests

* refactor: rename

* refactor: rename more

* fix: return promise on cancelation progress

* test: increase waiting time of selection to fix ubuntu ci

* refactor: address pr comments

* refactor: clean up code around evaluate and cancelation tokens

* refactor: more clean up

* refactor: more cleanup

* test: use getMessageContent

.husky/pre-commit

@@ -1,6 +1,3 @@
-#!/usr/bin/env sh
-. "$(dirname -- "$0")/_/husky.sh"
-
 # List staged files only.
 fileList=$(git diff --diff-filter=AM --cached --name-only)
 

package.json

@@ -282,6 +282,10 @@
           "dark": "images/dark/play.svg"
         }
       },
+      {
+        "command": "mdb.exportCodeToPlayground",
+        "title": "Export Code to Playground"
+      },
       {
         "command": "mdb.exportToPython",
         "title": "MongoDB: Export To Python 3"
@@ -747,6 +751,17 @@
           "when": "mdb.isPlayground == true"
         }
       ],
+      "mdb.copilot": [
+        {
+          "command": "mdb.exportCodeToPlayground"
+        }
+      ],
+      "editor/context": [
+        {
+          "submenu": "mdb.copilot",
+          "group": "1_main@2"
+        }
+      ],
       "commandPalette": [
         {
           "command": "mdb.selectDatabaseWithParticipant",
@@ -948,6 +963,10 @@
           "command": "mdb.runPlayground",
           "when": "false"
         },
+        {
+          "command": "mdb.exportCodeToPlayground",
+          "when": "false"
+        },
         {
           "command": "mdb.createIndexFromTreeView",
           "when": "false"
@@ -994,6 +1013,12 @@
         }
       ]
     },
+    "submenus": [
+      {
+        "id": "mdb.copilot",
+        "label": "MongoDB Copilot Participant"
+      }
+    ],
     "keybindings": [
       {
         "command": "mdb.runSelectedPlaygroundBlocks",

src/commands/index.ts

@@ -15,6 +15,7 @@ enum EXTENSION_COMMANDS {
   MDB_RUN_SELECTED_PLAYGROUND_BLOCKS = 'mdb.runSelectedPlaygroundBlocks',
   MDB_RUN_ALL_PLAYGROUND_BLOCKS = 'mdb.runAllPlaygroundBlocks',
   MDB_RUN_ALL_OR_SELECTED_PLAYGROUND_BLOCKS = 'mdb.runPlayground',
+  MDB_EXPORT_CODE_TO_PLAYGROUND = 'mdb.exportCodeToPlayground',
 
   MDB_FIX_THIS_INVALID_INTERACTIVE_SYNTAX = 'mdb.fixThisInvalidInteractiveSyntax',
   MDB_FIX_ALL_INVALID_INTERACTIVE_SYNTAX = 'mdb.fixAllInvalidInteractiveSyntax',

src/editors/playgroundController.ts

@@ -445,13 +445,16 @@ export default class PlaygroundController {
     return this._createPlaygroundFileWithContent(content);
   }
 
-  async _evaluate({
-    codeToEvaluate,
-    filePath,
-  }: {
-    codeToEvaluate: string;
-    filePath?: string;
-  }): Promise<ShellEvaluateResult> {
+  async _evaluate(
+    {
+      codeToEvaluate,
+      filePath,
+    }: {
+      codeToEvaluate: string;
+      filePath?: string;
+    },
+    token: vscode.CancellationToken
+  ): Promise<ShellEvaluateResult> {
     const connectionId = this._connectionController.getActiveConnectionId();
 
     if (!connectionId) {
@@ -460,14 +463,17 @@ export default class PlaygroundController {
 
     this._statusView.showMessage('Getting results...');
 
-    let result: ShellEvaluateResult;
+    let result: ShellEvaluateResult = null;
     try {
       // Send a request to the language server to execute scripts from a playground.
-      result = await this._languageServerController.evaluate({
-        codeToEvaluate,
-        connectionId,
-        filePath,
-      });
+      result = await this._languageServerController.evaluate(
+        {
+          codeToEvaluate,
+          connectionId,
+          filePath,
+        },
+        token
+      );
     } catch (error) {
       const msg =
         'An internal error has occurred. The playground services have been restored. This can occur when the playground runner runs out of memory.';
@@ -504,37 +510,22 @@ export default class PlaygroundController {
       throw new Error(connectBeforeRunningMessage);
     }
 
-    try {
-      const progressResult = await vscode.window.withProgress(
-        {
-          location: ProgressLocation.Notification,
-          title: 'Running MongoDB playground...',
-          cancellable: true,
-        },
-        async (progress, token) => {
-          token.onCancellationRequested(() => {
-            // If a user clicked the cancel button terminate all playground scripts.
-            this._languageServerController.cancelAll();
-
-            return { result: undefined };
-          });
-
-          // Run all playground scripts.
-          const result: ShellEvaluateResult = await this._evaluate({
+    return await vscode.window.withProgress(
+      {
+        location: ProgressLocation.Notification,
+        title: 'Running MongoDB playground...',
+        cancellable: true,
+      },
+      (progress, token): Promise<ShellEvaluateResult> => {
+        return this._evaluate(
+          {
             codeToEvaluate,
             filePath,
-          });
-
-          return result;
-        }
-      );
-
-      return progressResult;
-    } catch (error) {
-      log.error('Evaluating playground with cancel modal failed', error);
-
-      return { result: undefined };
-    }
+          },
+          token
+        );
+      }
+    );
   }
 
   async _openInResultPane(result: PlaygroundResult): Promise<void> {
@@ -923,7 +914,7 @@ export default class PlaygroundController {
           language,
           num_stages: selectedText
             ? countAggregationStagesInString(selectedText)
-            : undefined,
+            : null,
           with_import_statements: importStatements,
           with_builders: builders,
           with_driver_syntax: driverSyntax,

src/language/languageServerController.ts

@@ -5,11 +5,7 @@ import type {
   LanguageClientOptions,
   ServerOptions,
 } from 'vscode-languageclient/node';
-import {
-  LanguageClient,
-  TransportKind,
-  CancellationTokenSource,
-} from 'vscode-languageclient/node';
+import { LanguageClient, TransportKind } from 'vscode-languageclient/node';
 import type { ExtensionContext } from 'vscode';
 import { workspace } from 'vscode';
 import util from 'util';
@@ -32,8 +28,6 @@ const log = createLogger('language server controller');
  */
 export default class LanguageServerController {
   _context: ExtensionContext;
-  _source?: CancellationTokenSource;
-  _isExecutingInProgress = false;
   _client: LanguageClient;
   _currentConnectionId: string | null = null;
   _currentConnectionString?: string;
@@ -182,32 +176,25 @@ export default class LanguageServerController {
   }
 
   async evaluate(
-    playgroundExecuteParameters: PlaygroundEvaluateParams
+    playgroundExecuteParameters: PlaygroundEvaluateParams,
+    token: vscode.CancellationToken
   ): Promise<ShellEvaluateResult> {
     log.info('Running a playground...', {
       connectionId: playgroundExecuteParameters.connectionId,
       filePath: playgroundExecuteParameters.filePath,
       inputLength: playgroundExecuteParameters.codeToEvaluate.length,
     });
-    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();
-
     // Send a request with a cancellation token
     // to the language server instance to execute scripts from a playground
     // and return results to the playground controller when ready.
     const res: ShellEvaluateResult = await this._client.sendRequest(
       ServerCommands.EXECUTE_CODE_FROM_PLAYGROUND,
       playgroundExecuteParameters,
-      this._source.token
+      token
     );
 
-    this._isExecutingInProgress = false;
-
     log.info('Evaluate response', {
       namespace: res?.result?.namespace,
       type: res?.result?.type,
@@ -272,18 +259,6 @@ export default class LanguageServerController {
     );
   }
 
-  cancelAll(): void {
-    log.info('Canceling a playground...');
-    // Send a request for cancellation. As a result
-    // the associated CancellationToken will be notified of the cancellation,
-    // the onCancellationRequested event will be fired,
-    // and IsCancellationRequested will return true.
-    if (this._isExecutingInProgress) {
-      this._source?.cancel();
-      this._isExecutingInProgress = false;
-    }
-  }
-
   async updateCurrentSessionFields(params): Promise<void> {
     await this._client.sendRequest(
       ServerCommands.UPDATE_CURRENT_SESSION_FIELDS,