Merge development to main (#421)

Co-authored-by: Lokesh Goel <113521973+lokesh-couchbase@users.noreply.github.com>
Co-authored-by: Lokesh Goel <lokesh.goel@couchbase.com>

Author: prajwal-pai77

package-lock.json

@@ -2,7 +2,7 @@
   "name": "vscode-couchbase",
   "displayName": "Couchbase",
   "description": "",
-  "version": "2.1.2",
+  "version": "2.1.4",
   "engines": {
     "vscode": "^1.63.1"
   },

package.json

@@ -104,6 +104,7 @@ export function activate(context: vscode.ExtensionContext) {
   const uriToCasMap = new Map<string, string>();
   const workbench = new QueryWorkbench();
   const searchWorkbench = new SearchWorkbench();
+  const queryKernel = new QueryKernel()
 
   let currentSearchIndexNode: SearchIndexNode;
   let currentSearchWorkbench: SearchWorkbench;
@@ -725,21 +726,33 @@ context.subscriptions.push(disposable);
     )
   );
 
+  // Query Context for workbench
   subscriptions.push(
     vscode.commands.registerCommand(Commands.queryContext, () => {
-      fetchQueryContext(workbench, context, globalStatusBarItem);
+      fetchQueryContext(workbench, queryKernel, context, globalStatusBarItem);
     })
   );
 
-  // subscription to make sure query context status bar is only visible on sqlpp files
+  // Subscription to handle changes in active workbench editor
   subscriptions.push(
     vscode.window.onDidChangeActiveTextEditor(async (editor) => {
-      await handleQueryContextStatusbar(editor, workbench, globalStatusBarItem);
+      const activeNotebookEditor = vscode.window.activeNotebookEditor;
+      await handleQueryContextStatusbar(editor, activeNotebookEditor, workbench, queryKernel, globalStatusBarItem);
     })
   );
-  // // Handle initial view of context status bar
+
+  // Subscription to handle changes in active notebook editor
+  subscriptions.push(
+    vscode.window.onDidChangeActiveNotebookEditor(async (notebookEditor) => {
+      const activeEditor = vscode.window.activeTextEditor;
+      await handleQueryContextStatusbar(activeEditor, notebookEditor, workbench, queryKernel, globalStatusBarItem);
+    })
+  );
+
+  // Handle initial view of context status bar
   const activeEditor = vscode.window.activeTextEditor;
-  handleQueryContextStatusbar(activeEditor, workbench, globalStatusBarItem);
+  const activeNotebookEditor = vscode.window.activeNotebookEditor;
+  handleQueryContextStatusbar(activeEditor, activeNotebookEditor, workbench, queryKernel, globalStatusBarItem);
 
   subscriptions.push(
     vscode.commands.registerCommand(
@@ -859,7 +872,7 @@ context.subscriptions.push(disposable);
     vscode.workspace.registerNotebookSerializer(
       Constants.notebookType, new QueryContentSerializer(), { transientOutputs: true }
     ),
-    new QueryKernel()
+    queryKernel
   );
 
   context.subscriptions.push(

src/extension.ts

@@ -1,12 +1,23 @@
 import * as vscode from 'vscode';
 import { QueryWorkbench } from '../workbench/queryWorkbench';
-import {  showQueryContextStatusbar } from '../util/queryContextUtils';
-export const handleQueryContextStatusbar = async (editor: vscode.TextEditor | undefined, workbench: QueryWorkbench, globalStatusBarItem: vscode.StatusBarItem) => {
-    if( editor && editor.document.languageId === "SQL++"){
-        // Case 1: Show Status bar
-        showQueryContextStatusbar(editor, workbench, globalStatusBarItem);
+import { showQueryContextStatusbar, showQueryContextStatusbarNotebook } from '../util/queryContextUtils';
+import { QueryKernel } from '../notebook/controller';
+
+export async function handleQueryContextStatusbar(
+    activeEditor: vscode.TextEditor | undefined,
+    activeNotebookEditor: vscode.NotebookEditor | undefined,
+    workbench: QueryWorkbench,
+    queryKernel: QueryKernel,
+    globalStatusBarItem: vscode.StatusBarItem
+) {
+    if (activeNotebookEditor) {
+
+        showQueryContextStatusbarNotebook(activeNotebookEditor, queryKernel, globalStatusBarItem)
+
+    } else if (activeEditor && activeEditor.document.languageId === "SQL++") {
+        showQueryContextStatusbar(activeEditor, workbench, globalStatusBarItem)
+
     } else {
-        // Case 2: Don't show status bar
         globalStatusBarItem.hide();
     }
-};
+}

src/handlers/handleQueryContextStatusbar.ts

@@ -16,7 +16,8 @@
 
 import * as vscode from 'vscode';
 import { getActiveConnection } from '../util/connections';
-import { CouchbaseError } from 'couchbase';
+import { CouchbaseError, QueryOptions } from 'couchbase';
+import { IQueryContext } from '../types/IQueryContext';
 
 /**
  * // This is a typescript class for creating and handling the execution of a new Couchbase Query Notebook,
@@ -33,6 +34,7 @@ export class QueryKernel {
 
     private _executionOrder = 0;
     private readonly _controller: vscode.NotebookController;
+    public notebookToContext: Map<string, IQueryContext>;
 
     // Constructor used to create new QueryKernel objects
     constructor() {
@@ -44,6 +46,7 @@ export class QueryKernel {
         this._controller.supportedLanguages = this._supportedLanguages;
         this._controller.supportsExecutionOrder = true;
         this._controller.executeHandler = this._executeAll.bind(this);
+        this.notebookToContext = new Map<string, IQueryContext>();
     }
 
     // Method to release any resources being held, currently it just disposes the notebook controller
@@ -61,21 +64,37 @@ export class QueryKernel {
         const execution = this._controller.createNotebookCellExecution(cell);
         execution.executionOrder = ++this._executionOrder;
         execution.start(Date.now());
+        let queryOptions: QueryOptions
+        const activeTextEditor = vscode.window.activeNotebookEditor;
+        if (activeTextEditor) {
+            const notebookUri = activeTextEditor.notebook.uri.toString();
 
+            // Retrieve the query context using the notebook URI
+            const queryContext = this.notebookToContext.get(notebookUri);
+            const queryContextString = queryContext
+                ? `${queryContext.bucketName}.${queryContext.scopeName}`
+                : '';
+
+            queryOptions = {
+                metrics: true,
+                queryContext: queryContextString
+            };
+    
         try {
             const activeConnection = getActiveConnection();
             if (!activeConnection) {
                 vscode.window.showErrorMessage("Connection Failure: Looks like you are not connected to cluster \nPlease check your cluster connection and try again", { modal: true });
                 throw Error('Connection Failed');
             };
             const result = await activeConnection.cluster?.query(
-                cell.document.getText()
+                cell.document.getText(), queryOptions
             );
             execution.replaceOutput([new vscode.NotebookCellOutput([
                 vscode.NotebookCellOutputItem.json(result?.rows)
             ])]);
 
             execution.end(true, Date.now());
+     
         } catch (err) {
             const errorArray = [];
             if (err instanceof CouchbaseError) {
@@ -100,4 +119,5 @@ export class QueryKernel {
         }
         execution.end(false, Date.now());
     }
+    }
 }

src/notebook/controller.ts

@@ -6,12 +6,14 @@ import { Bucket, BucketSettings } from "couchbase";
 import { QueryWorkbench } from "../../workbench/queryWorkbench";
 import {
     showQueryContextStatusbar,
+    showQueryContextStatusbarNotebook,
     showSearchContextStatusbar,
 } from "../../util/queryContextUtils";
 import { getActiveConnection } from "../../util/connections";
 import { SearchWorkbench } from "../../commands/fts/SearchWorkbench/searchWorkbench";
 import SearchIndexNode from "../../model/SearchIndexNode";
 import { Commands } from "../../commands/extensionCommands/commands";
+import { QueryKernel } from "../../notebook/controller";
 
 const fetchBucketNames = (
     bucketsSettings: BucketSettings[] | undefined,
@@ -33,6 +35,7 @@ const fetchBucketNames = (
 
 export async function fetchQueryContext(
     workbench: QueryWorkbench,
+    queryKernel: QueryKernel,
     context: vscode.ExtensionContext,
     globalStatusBarItem: any,
 ) {
@@ -47,9 +50,11 @@ export async function fetchQueryContext(
     try {
         // Fetch active editor
         const activeEditor = vscode.window.activeTextEditor;
-        if (!(activeEditor && activeEditor.document.languageId === "SQL++")) {
+        const activeNotebookEditor = vscode.window.activeNotebookEditor;
+        const isNotebook = !!activeNotebookEditor;
+        if (!activeNotebookEditor && !(activeEditor && activeEditor.document.languageId === "SQL++")) {
             vscode.window.showErrorMessage(
-                "Please ensure that the workbench is open/active",
+                "Please ensure that a SQL++ editor or a notebook is open and active"
             );
             return;
         }
@@ -96,14 +101,13 @@ export async function fetchQueryContext(
 
         const bucketNameSelected = selectedItem.label;
         if (bucketNameSelected === "Clear Context") {
-            workbench.editorToContext.delete(
-                activeEditor.document.uri.toString(),
-            );
-            showQueryContextStatusbar(
-                activeEditor,
-                workbench,
-                globalStatusBarItem,
-            );
+            if (isNotebook && activeNotebookEditor) {
+                queryKernel.notebookToContext.delete(activeNotebookEditor.notebook.uri.toString());
+                showQueryContextStatusbarNotebook(activeNotebookEditor, queryKernel, globalStatusBarItem);
+            } else if (activeEditor) {
+                workbench.editorToContext.delete(activeEditor.document.uri.toString());
+                showQueryContextStatusbar(activeEditor, workbench, globalStatusBarItem);
+            }
             return;
         }
         const scopes = await connection.cluster
@@ -128,11 +132,18 @@ export async function fetchQueryContext(
             vscode.window.showInformationMessage("No scope selected.");
             return;
         }
-        workbench.editorToContext.set(activeEditor.document.uri.toString(), {
+        const queryContext = {
             bucketName: bucketNameSelected,
             scopeName: scopeNameSelected.label,
-        });
-        showQueryContextStatusbar(activeEditor, workbench, globalStatusBarItem);
+        };
+        if (isNotebook && activeEditor) {
+            const notebookUri = activeNotebookEditor.notebook.uri.toString();
+            queryKernel.notebookToContext.set(notebookUri, queryContext);
+            showQueryContextStatusbarNotebook(activeNotebookEditor, queryKernel, globalStatusBarItem);
+        } else if (activeEditor) {
+            workbench.editorToContext.set(activeEditor.document.uri.toString(), queryContext);
+            showQueryContextStatusbar(activeEditor, workbench, globalStatusBarItem);
+        }
     } catch (err) {
         logger.error(`failed to open and set query context: ${err}`);
         logger.debug(err);

src/pages/queryContext/queryContext.ts

@@ -5,6 +5,7 @@ import { Constants } from './constants';
 import { Commands } from '../commands/extensionCommands/commands';
 import { SearchWorkbench } from '../commands/fts/SearchWorkbench/searchWorkbench';
 import SearchIndexNode from '../model/SearchIndexNode';
+import { QueryKernel } from '../notebook/controller';
 
 export const showQueryContextStatusbar = async (editor: vscode.TextEditor, workbench: QueryWorkbench, globalStatusBarItem: vscode.StatusBarItem) => {
     let queryContext = workbench.editorToContext.get(editor.document.uri.toString());
@@ -26,6 +27,27 @@ export const showQueryContextStatusbar = async (editor: vscode.TextEditor, workb
     globalStatusBarItem.show();
 };
 
+export const showQueryContextStatusbarNotebook = async (editor: vscode.NotebookEditor, queryKernel: QueryKernel, globalStatusBarItem: vscode.StatusBarItem) => {
+    const notebookUri = editor.notebook.uri.toString();
+    const queryContext = queryKernel.notebookToContext.get(notebookUri);
+    if (!queryContext) {
+        globalStatusBarItem.text = `$(group-by-ref-type) No Query Context Set`;
+    } else {
+        let bucketName = queryContext.bucketName;
+        if (bucketName.length > 15) {
+            bucketName = `${bucketName.substring(0, 13)}...`;
+        }
+        let scopeName = queryContext.scopeName;
+        if (scopeName.length > 15) {
+            scopeName = `${scopeName.substring(0, 13)}...`;
+        }
+        globalStatusBarItem.text = `$(group-by-ref-type) ${bucketName} > ${scopeName}`;
+    }
+    globalStatusBarItem.command = Commands.queryContext;
+    globalStatusBarItem.tooltip = "Query Context";
+    globalStatusBarItem.show();
+};
+
 
 export const showSearchContextStatusbar = async (editor: vscode.TextEditor, searchNode: SearchIndexNode, workbench: SearchWorkbench, globalStatusBarItem: vscode.StatusBarItem) => {
     let editorId = editor.document.uri.toString();