Reorganize DocDB command registration (#2240)

src/docdb/commands/connectNoSqlContainer.ts

@@ -0,0 +1,28 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { IActionContext } from "@microsoft/vscode-azext-utils";
+import { KeyValueStore } from "../../KeyValueStore";
+import { ext } from "../../extensionVariables";
+import { NoSqlQueryConnection, noSqlQueryConnectionKey } from "../NoSqlCodeLensProvider";
+import { DocDBCollectionTreeItem } from "../tree/DocDBCollectionTreeItem";
+import { pickDocDBAccount } from "./pickDocDBAccount";
+
+export function setConnectedNoSqlContainer(node: DocDBCollectionTreeItem): void {
+    const noSqlQueryConnection: NoSqlQueryConnection = {
+        databaseId: node.parent.id,
+        containerId: node.id,
+        endpoint: node.root.endpoint,
+        masterKey: node.root.masterKey,
+        isEmulator: !!node.root.isEmulator
+    };
+    KeyValueStore.instance.set(noSqlQueryConnectionKey, noSqlQueryConnection);
+    ext.noSqlCodeLensProvider.updateCodeLens();
+}
+
+export async function connectNoSqlContainer(context: IActionContext): Promise<void> {
+    const node = await pickDocDBAccount<DocDBCollectionTreeItem>(context, DocDBCollectionTreeItem.contextValue);
+    setConnectedNoSqlContainer(node);
+}

src/docdb/commands/createDocDBCollection.ts

@@ -0,0 +1,15 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { IActionContext } from "@microsoft/vscode-azext-utils";
+import { DocDBDatabaseTreeItem } from "../tree/DocDBDatabaseTreeItem";
+import { pickDocDBAccount } from "./pickDocDBAccount";
+
+export async function createDocDBCollection(context: IActionContext, node?: DocDBDatabaseTreeItem): Promise<void> {
+    if (!node) {
+        node = await pickDocDBAccount<DocDBDatabaseTreeItem>(context, DocDBDatabaseTreeItem.contextValue);
+    }
+    await node.createChild(context);
+}

src/docdb/commands/createDocDBDatabase.ts

@@ -0,0 +1,18 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { IActionContext } from "@microsoft/vscode-azext-utils";
+import { DocDBAccountTreeItem } from "../tree/DocDBAccountTreeItem";
+import { DocDBDatabaseTreeItem } from "../tree/DocDBDatabaseTreeItem";
+import { pickDocDBAccount } from "./pickDocDBAccount";
+
+
+export async function createDocDBDatabase(context: IActionContext, node?: DocDBAccountTreeItem): Promise<void> {
+    if (!node) {
+        node = await pickDocDBAccount<DocDBAccountTreeItem>(context);
+    }
+    const databaseNode: DocDBDatabaseTreeItem = await node.createChild(context);
+    await databaseNode.createChild(context);
+}

src/docdb/commands/createDocDBDocument.ts

@@ -0,0 +1,18 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { IActionContext } from "@microsoft/vscode-azext-utils";
+import { commands } from "vscode";
+import { DocDBDocumentTreeItem } from "../tree/DocDBDocumentTreeItem";
+import { DocDBDocumentsTreeItem } from "../tree/DocDBDocumentsTreeItem";
+import { pickDocDBAccount } from "./pickDocDBAccount";
+
+export async function createDocDBDocument(context: IActionContext, node?: DocDBDocumentsTreeItem) {
+    if (!node) {
+        node = await pickDocDBAccount<DocDBDocumentsTreeItem>(context, DocDBDocumentsTreeItem.contextValue);
+    }
+    const documentNode = <DocDBDocumentTreeItem>await node.createChild(context);
+    await commands.executeCommand("cosmosDB.openDocument", documentNode);
+}

src/docdb/commands/createDocDBStoredProcedure.ts

@@ -0,0 +1,17 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { IActionContext } from "@microsoft/vscode-azext-utils";
+import { commands } from "vscode";
+import { DocDBStoredProceduresTreeItem } from "../tree/DocDBStoredProceduresTreeItem";
+import { pickDocDBAccount } from "./pickDocDBAccount";
+
+export async function createDocDBStoredProcedure(context: IActionContext, node?: DocDBStoredProceduresTreeItem) {
+    if (!node) {
+        node = await pickDocDBAccount<DocDBStoredProceduresTreeItem>(context, DocDBStoredProceduresTreeItem.contextValue);
+    }
+    const childNode = await node.createChild(context);
+    await commands.executeCommand("cosmosDB.openStoredProcedure", childNode);
+}

src/docdb/commands/createDocDBTrigger.ts

@@ -0,0 +1,17 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { IActionContext } from "@microsoft/vscode-azext-utils";
+import { commands } from "vscode";
+import { DocDBTriggersTreeItem } from "../tree/DocDBTriggersTreeItem";
+import { pickDocDBAccount } from "./pickDocDBAccount";
+
+export async function createDocDBTrigger(context: IActionContext, node?: DocDBTriggersTreeItem) {
+    if (!node) {
+        node = await pickDocDBAccount<DocDBTriggersTreeItem>(context, DocDBTriggersTreeItem.contextValue);
+    }
+    const childNode = await node.createChild(context);
+    await commands.executeCommand("cosmosDB.openTrigger", childNode);
+}

src/docdb/commands/deleteDocDBCollection.ts

@@ -0,0 +1,17 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { IActionContext, ITreeItemPickerContext } from "@microsoft/vscode-azext-utils";
+import { DocDBCollectionTreeItem } from "../tree/DocDBCollectionTreeItem";
+import { pickDocDBAccount } from "./pickDocDBAccount";
+
+export async function deleteDocDBCollection(context: IActionContext, node?: DocDBCollectionTreeItem): Promise<void> {
+    const suppressCreateContext: ITreeItemPickerContext = context;
+    suppressCreateContext.suppressCreatePick = true;
+    if (!node) {
+        node = await pickDocDBAccount<DocDBCollectionTreeItem>(context, DocDBCollectionTreeItem.contextValue);
+    }
+    await node.deleteTreeItem(context);
+}

src/docdb/commands/deleteDocDBDatabase.ts

@@ -0,0 +1,17 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { IActionContext, ITreeItemPickerContext } from "@microsoft/vscode-azext-utils";
+import { DocDBDatabaseTreeItem } from "../tree/DocDBDatabaseTreeItem";
+import { pickDocDBAccount } from "./pickDocDBAccount";
+
+export async function deleteDocDBDatabase(context: IActionContext, node?: DocDBDatabaseTreeItem): Promise<void> {
+    const suppressCreateContext: ITreeItemPickerContext = context;
+    suppressCreateContext.suppressCreatePick = true;
+    if (!node) {
+        node = await pickDocDBAccount<DocDBDatabaseTreeItem>(context, DocDBDatabaseTreeItem.contextValue);
+    }
+    await node.deleteTreeItem(context);
+}

src/docdb/commands/deleteDocDBDocument.ts

@@ -0,0 +1,17 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { IActionContext, ITreeItemPickerContext } from "@microsoft/vscode-azext-utils";
+import { DocDBDocumentTreeItem } from "../tree/DocDBDocumentTreeItem";
+import { pickDocDBAccount } from "./pickDocDBAccount";
+
+export async function deleteDocDBDocument(context: IActionContext, node?: DocDBDocumentTreeItem) {
+    const suppressCreateContext: ITreeItemPickerContext = context;
+    suppressCreateContext.suppressCreatePick = true;
+    if (!node) {
+        node = await pickDocDBAccount<DocDBDocumentTreeItem>(context, DocDBDocumentTreeItem.contextValue);
+    }
+    await node.deleteTreeItem(context);
+}

src/docdb/commands/deleteDocDBStoredProcedure.ts

@@ -0,0 +1,18 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { IActionContext, ITreeItemPickerContext } from "@microsoft/vscode-azext-utils";
+import { DocDBStoredProcedureTreeItem } from "../tree/DocDBStoredProcedureTreeItem";
+import { pickDocDBAccount } from "./pickDocDBAccount";
+
+export async function deleteDocDBStoredProcedure(context: IActionContext, node?: DocDBStoredProcedureTreeItem) {
+    const suppressCreateContext: ITreeItemPickerContext = context;
+    suppressCreateContext.suppressCreatePick = true;
+    if (!node) {
+        node = await pickDocDBAccount<DocDBStoredProcedureTreeItem>(context, DocDBStoredProcedureTreeItem.contextValue);
+
+    }
+    await node.deleteTreeItem(context);
+}

src/docdb/commands/deleteDocDBTrigger.ts

@@ -0,0 +1,18 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { IActionContext, ITreeItemPickerContext } from "@microsoft/vscode-azext-utils";
+import { DocDBTriggerTreeItem } from "../tree/DocDBTriggerTreeItem";
+import { pickDocDBAccount } from "./pickDocDBAccount";
+
+export async function deleteDocDBTrigger(context: IActionContext, node?: DocDBTriggerTreeItem) {
+    const suppressCreateContext: ITreeItemPickerContext = context;
+    suppressCreateContext.suppressCreatePick = true;
+    if (!node) {
+        node = await pickDocDBAccount<DocDBTriggerTreeItem>(context, DocDBTriggerTreeItem.contextValue);
+
+    }
+    await node.deleteTreeItem(context);
+}

src/docdb/commands/executeDocDBStoredProcedure.ts

@@ -0,0 +1,39 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { IActionContext, ITreeItemPickerContext } from "@microsoft/vscode-azext-utils";
+import { localize } from "../../utils/localize";
+import { DocDBStoredProcedureTreeItem } from "../tree/DocDBStoredProcedureTreeItem";
+import { pickDocDBAccount } from "./pickDocDBAccount";
+
+export async function executeDocDBStoredProcedure(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);
+}

src/docdb/commands/executeNoSqlQuery.ts

@@ -0,0 +1,37 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { IActionContext } from "@microsoft/vscode-azext-utils";
+import { ViewColumn } from "vscode";
+import { KeyValueStore } from "../../KeyValueStore";
+import * as vscodeUtil from "../../utils/vscodeUtils";
+import { NoSqlQueryConnection, noSqlQueryConnectionKey } from "../NoSqlCodeLensProvider";
+import { getCosmosClient } from "../getCosmosClient";
+
+export async function executeNoSqlQuery(_context: IActionContext, args: { queryText: string, populateQueryMetrics?: boolean }): Promise<void> {
+    if (!args) {
+        throw new Error("Unable to execute query due to missing args. Please connect to a Cosmos DB collection.");
+    }
+    const { queryText, populateQueryMetrics } = args;
+    const connectedCollection = KeyValueStore.instance.get(noSqlQueryConnectionKey);
+    if (!connectedCollection) {
+        throw new Error("Unable to execute query due to missing node data. Please connect to a Cosmos DB collection node.");
+    } else {
+        // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
+        const { databaseId, containerId, endpoint, masterKey, isEmulator } = connectedCollection as NoSqlQueryConnection;
+        const client = getCosmosClient(endpoint, masterKey, isEmulator);
+        const options = { populateQueryMetrics };
+        const response = await client.database(databaseId).container(containerId).items.query(queryText, options).fetchAll();
+        const resultDocumentTitle = `query results for ${containerId}`;
+        if (populateQueryMetrics === true) {
+            await vscodeUtil.showNewFile(JSON.stringify({
+                result: response.resources,
+                queryMetrics: response.queryMetrics
+            }, undefined, 2), resultDocumentTitle, ".json", ViewColumn.Beside);
+        } else {
+            await vscodeUtil.showNewFile(JSON.stringify(response.resources, undefined, 2), resultDocumentTitle, ".json", ViewColumn.Beside);
+        }
+    }
+}

src/docdb/commands/getNoSqlQueryPlan.ts

@@ -0,0 +1,28 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { IActionContext } from "@microsoft/vscode-azext-utils";
+import { ViewColumn } from "vscode";
+import { KeyValueStore } from "../../KeyValueStore";
+import * as vscodeUtil from "../../utils/vscodeUtils";
+import { NoSqlQueryConnection, noSqlQueryConnectionKey } from "../NoSqlCodeLensProvider";
+import { getCosmosClient } from "../getCosmosClient";
+
+export async function getNoSqlQueryPlan(_context: IActionContext, args: { queryText: string } | undefined): Promise<void> {
+    if (!args) {
+        throw new Error("Unable to get query plan due to missing args. Please connect to a Cosmos DB collection node.");
+    }
+    const queryText = args.queryText;
+    const connectedCollection = KeyValueStore.instance.get(noSqlQueryConnectionKey);
+    if (!connectedCollection) {
+        throw new Error("Unable to get query plan due to missing node data. Please connect to a Cosmos DB collection.");
+    } else {
+        // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
+        const { databaseId, containerId, endpoint, masterKey, isEmulator } = connectedCollection as NoSqlQueryConnection;
+        const client = getCosmosClient(endpoint, masterKey, isEmulator);
+        const response = await client.database(databaseId).container(containerId).getQueryPlan(queryText);
+        await vscodeUtil.showNewFile(JSON.stringify(response.result, undefined, 2), `query results for ${containerId}`, ".json", ViewColumn.Beside);
+    }
+}

src/docdb/commands/openStoredProcedure.ts

@@ -0,0 +1,16 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { IActionContext } from "@microsoft/vscode-azext-utils";
+import { ext } from "../../extensionVariables";
+import { DocDBStoredProcedureTreeItem } from "../tree/DocDBStoredProcedureTreeItem";
+import { pickDocDBAccount } from "./pickDocDBAccount";
+
+export async function openStoredProcedure(context: IActionContext, node?: DocDBStoredProcedureTreeItem) {
+    if (!node) {
+        node = await pickDocDBAccount<DocDBStoredProcedureTreeItem>(context, DocDBStoredProcedureTreeItem.contextValue);
+    }
+    await ext.fileSystem.showTextDocument(node);
+}

src/docdb/commands/openTrigger.ts

@@ -0,0 +1,16 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { IActionContext } from "@microsoft/vscode-azext-utils";
+import { ext } from "../../extensionVariables";
+import { DocDBTriggerTreeItem } from "../tree/DocDBTriggerTreeItem";
+import { pickDocDBAccount } from "./pickDocDBAccount";
+
+export async function openTrigger(context: IActionContext, node?: DocDBTriggerTreeItem) {
+    if (!node) {
+        node = await pickDocDBAccount<DocDBTriggerTreeItem>(context, DocDBTriggerTreeItem.contextValue);
+    }
+    await ext.fileSystem.showTextDocument(node);
+}

src/docdb/commands/pickDocDBAccount.ts

@@ -0,0 +1,20 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { AzExtTreeItem, IActionContext } from "@microsoft/vscode-azext-utils";
+import { sqlFilter } from "../../constants";
+import { ext } from "../../extensionVariables";
+
+export async function pickDocDBAccount<T extends AzExtTreeItem>(
+    context: IActionContext,
+    expectedContextValue?: string | RegExp | (string | RegExp)[]
+): Promise<T> {
+    return await ext.rgApi.pickAppResource<T>(context, {
+        filter: [
+            sqlFilter
+        ],
+        expectedChildContextValue: expectedContextValue
+    });
+}