feature: CosmosDB execute stored procedure (#2235)

* feature: cosmos execute stored procedure * Support passing parameters * Don't add a learn more link for now * Show execute action above delete action * Use readonly document for execution result
microsoft · Jan 23, 2024 · 4f6f2d2 · 4f6f2d2
1 parent 3944a7c
commit 4f6f2d2
Show file tree

Hide file tree

Showing 3 changed files with 55 additions and 2 deletions.
diff --git a/package.json b/package.json
@@ -335,6 +335,11 @@
                 "command": "cosmosDB.executeAllMongoCommands",
                 "title": "Execute All MongoDB Commands"
             },
+            {
+                "category": "Core (SQL)",
+                "command": "cosmosDB.executeDocDBStoredProcedure",
+                "title": "Execute Stored Procedure..."
+            },
             {
                 "category": "MongoDB",
                 "command": "cosmosDB.executeMongoCommand",
@@ -723,6 +728,11 @@
                     "when": "view =~ /azure(ResourceGroups|Workspace|FocusView)/ && viewItem == cosmosDBStoredProcedure",
                     "group": "1@2"
                 },
+                {
+                    "command": "cosmosDB.executeDocDBStoredProcedure",
+                    "when": "view =~ /azure(ResourceGroups|Workspace|FocusView)/ && viewItem == cosmosDBStoredProcedure",
+                    "group": "1@1"
+                },
                 {
                     "command": "cosmosDB.deleteDocDBDatabase",
                     "when": "view =~ /azure(ResourceGroups|Workspace|FocusView)/ && viewItem == cosmosDBDocumentDatabase",

diff --git a/src/docdb/registerDocDBCommands.ts b/src/docdb/registerDocDBCommands.ts
@@ -8,6 +8,7 @@ import { ViewColumn, commands, languages } from "vscode";
 import { KeyValueStore } from "../KeyValueStore";
 import { doubleClickDebounceDelay, sqlFilter } from "../constants";
 import { ext } from "../extensionVariables";
+import { localize } from "../utils/localize";
 import * as vscodeUtil from "../utils/vscodeUtils";
 import { NoSqlCodeLensProvider, NoSqlQueryConnection, noSqlQueryConnectionKey } from "./NoSqlCodeLensProvider";
 import { getCosmosClient } from "./getCosmosClient";
@@ -72,6 +73,35 @@ export function registerDocDBCommands(): void {
         }
         await node.deleteTreeItem(context);
     });
+    registerCommandWithTreeNodeUnwrapping('cosmosDB.executeDocDBStoredProcedure', async (context: IActionContext, node?: DocDBStoredProcedureTreeItem) => {
+        const suppressCreateContext: ITreeItemPickerContext = context;
+        suppressCreateContext.suppressCreatePick = true;
+        if (!node) {
+            node = await pickDocDBAccount<DocDBStoredProcedureTreeItem>(context, DocDBStoredProcedureTreeItem.contextValue);
+        }
+
+        const partitionKey = await context.ui.showInputBox({
+            title: 'Partition Key',
+            // @todo: add a learnMoreLink
+        });
+
+        const paramString = await context.ui.showInputBox({
+            title: 'Parameters',
+            placeHolder: localize("executeCosmosStoredProcedureParameters", "empty or array of values e.g. [1, {key: value}]"),
+            // @todo: add a learnMoreLink
+        });
+
+        let parameters: (string | number | object)[] | undefined = undefined;
+        if (paramString !== "") {
+            try {
+                parameters = JSON.parse(paramString) as (string | number | object)[];
+            } catch (error) {
+                // Ignore parameters if they are invalid
+            }
+        }
+
+        await node.execute(context, partitionKey, parameters);
+    });
 }
 
 function setConnectedNoSqlContainer(node: DocDBCollectionTreeItem): void {

diff --git a/src/docdb/tree/DocDBStoredProcedureTreeItem.ts b/src/docdb/tree/DocDBStoredProcedureTreeItem.ts
@@ -4,10 +4,11 @@
  *--------------------------------------------------------------------------------------------*/
 
 import { Resource, StoredProcedureDefinition } from '@azure/cosmos';
-import { AzExtTreeItem, DialogResponses, IActionContext, TreeItemIconPath } from '@microsoft/vscode-azext-utils';
+import { AzExtTreeItem, DialogResponses, IActionContext, TreeItemIconPath, openReadOnlyJson, randomUtils } from '@microsoft/vscode-azext-utils';
 import * as vscode from "vscode";
 import { IEditableTreeItem } from '../../DatabasesFileSystem';
 import { ext } from '../../extensionVariables';
+import { localize } from '../../utils/localize';
 import { nonNullProp } from '../../utils/nonNull';
 import { DocDBStoredProceduresTreeItem } from './DocDBStoredProceduresTreeItem';
 import { IDocDBTreeRoot } from './IDocDBTreeRoot';
@@ -68,9 +69,21 @@ export class DocDBStoredProcedureTreeItem extends AzExtTreeItem implements IEdit
     }
 
     public async deleteTreeItemImpl(context: IActionContext): Promise<void> {
-        const message: string = `Are you sure you want to delete stored procedure '${this.label}'?`;
+        const message: string = localize("deleteCosmosStoredProcedure", `Are you sure you want to delete stored procedure '{0}'?`, this.label);
         await context.ui.showWarningMessage(message, { modal: true, stepName: 'deleteStoredProcedure' }, DialogResponses.deleteResponse);
         const client = this.root.getCosmosClient();
         await this.parent.getContainerClient(client).scripts.storedProcedure(this.id).delete();
     }
+
+    public async execute(context: IActionContext, partitionKey: string, parameters?: any[]): Promise<void> {
+        const client = this.root.getCosmosClient();
+        const result = await this.parent.getContainerClient(client).scripts.storedProcedure(this.id).execute(partitionKey, parameters);
+
+        try {
+            const resultFileName = `${this.label}-result`;
+            await openReadOnlyJson({ label: resultFileName, fullId: randomUtils.getRandomHexString() }, result);
+        } catch (error) {
+            await context.ui.showWarningMessage(`Unable to parse execution result`);
+        }
+    }
 }