NoSQL execute query (#2217)

JasonYeMSFT · web-flow · commit d404afe702ab · 2023-12-05T15:50:44.000-08:00
* Saved work

* Simple working query editor

* Test getQueryPlan

* Better code lens

* Show query result/plan side by side prototype

* Support query with metrics

* Add connect code lens

* Add package json declaration

* Rename new query title
diff --git a/extension.bundle.ts b/extension.bundle.ts
@@ -27,23 +27,23 @@ export * from './src/docdb/registerDocDBCommands';
 export { activateInternal, cosmosDBCopyConnectionString, createServer, deactivateInternal, deleteAccount } from './src/extension';
 export { ext } from './src/extensionVariables';
 export * from './src/graph/registerGraphCommands';
-export { connectToMongoClient, isCosmosEmulatorConnectionString } from './src/mongo/connectToMongoClient';
 export { MongoCommand } from './src/mongo/MongoCommand';
-export { addDatabaseToAccountConnectionString, encodeMongoConnectionString, getDatabaseNameFromConnectionString } from './src/mongo/mongoConnectionStrings';
 export { findCommandAtPosition, getAllCommandsFromText } from './src/mongo/MongoScrapbook';
 export { MongoShell } from './src/mongo/MongoShell';
+export { connectToMongoClient, isCosmosEmulatorConnectionString } from './src/mongo/connectToMongoClient';
+export { addDatabaseToAccountConnectionString, encodeMongoConnectionString, getDatabaseNameFromConnectionString } from './src/mongo/mongoConnectionStrings';
 export * from './src/mongo/registerMongoCommands';
 export { IDatabaseInfo } from './src/mongo/tree/MongoAccountTreeItem';
 export { addDatabaseToConnectionString } from './src/postgres/postgresConnectionStrings';
 export { AttachedAccountsTreeItem, MONGO_CONNECTION_EXPECTED } from './src/tree/AttachedAccountsTreeItem';
 export { AzureAccountTreeItemWithAttached } from './src/tree/AzureAccountTreeItemWithAttached';
 export * from './src/utils/azureClients';
+export { getPublicIpv4, isIpInRanges } from './src/utils/getIp';
 export { improveError } from './src/utils/improveError';
 export { randomUtils } from './src/utils/randomUtils';
 export { getGlobalSetting, updateGlobalSetting } from './src/utils/settingUtils';
 export { rejectOnTimeout, valueOnTimeout } from './src/utils/timeout';
-export { getDocumentTreeItemLabel, IDisposable } from './src/utils/vscodeUtils';
+export { IDisposable, getDocumentTreeItemLabel } from './src/utils/vscodeUtils';
 export { wrapError } from './src/utils/wrapError';
-export { isIpInRanges, getPublicIpv4 } from './src/utils/getIp'
 
 // NOTE: The auto-fix action "source.organizeImports" does weird things with this file, but there doesn't seem to be a way to disable it on a per-file basis so we'll just let it happen
diff --git a/package.json b/package.json
@@ -151,6 +151,16 @@
                 "extensions": [
                     ".psql"
                 ]
+            },
+            {
+                "id": "nosql",
+                "aliases": [
+                    "Cosmos NoSQL",
+                    "nosql"
+                ],
+                "extensions": [
+                    ".nosql"
+                ]
             }
         ],
         "grammars": [
@@ -214,6 +224,12 @@
                 "command": "cosmosDB.connectMongoDB",
                 "title": "Connect to Database..."
             },
+            {
+                "category": "Cosmos DB",
+                "command": "cosmosDB.connectNoSqlContainer",
+                "title": "Connect to NoSQL container",
+                "when": "false"
+            },
             {
                 "category": "Cosmos DB",
                 "command": "cosmosDB.copyConnectionString",
@@ -324,6 +340,18 @@
                 "command": "cosmosDB.executeMongoCommand",
                 "title": "Execute MongoDB Command"
             },
+            {
+                "category": "Cosmos DB",
+                "command": "cosmosDB.executeNoSqlQuery",
+                "title": "Execute NoSQL Query",
+                "when": "false"
+            },
+            {
+                "category": "Cosmos DB",
+                "command": "cosmosDB.getNoSqlQueryPlan",
+                "title": "Get NoSQL Query Plan",
+                "when": "false"
+            },
             {
                 "category": "Cosmos DB",
                 "command": "cosmosDB.importDocument",
@@ -360,6 +388,11 @@
                 "command": "cosmosDB.openStoredProcedure",
                 "title": "Open Stored Procedure"
             },
+            {
+                "category": "Core (SQL)",
+                "command": "cosmosDB.writeNoSqlQuery",
+                "title": "Create New NoSQL Query"
+            },
             {
                 "category": "PostgreSQL",
                 "command": "postgreSQL.configureFirewall",
@@ -575,6 +608,11 @@
                     "when": "view =~ /azure(ResourceGroups|Workspace|FocusView)/ && viewItem == cosmosDBDocumentsGroup",
                     "group": "1@1"
                 },
+                {
+                    "command": "cosmosDB.writeNoSqlQuery",
+                    "when": "view =~ /azure(ResourceGroups|Workspace|FocusView)/ && viewItem == cosmosDBDocumentCollection",
+                    "group": "1@2"
+                },
                 {
                     "command": "cosmosDB.createDocDBStoredProcedure",
                     "when": "view =~ /azure(ResourceGroups|Workspace|FocusView)/ && viewItem == cosmosDBStoredProceduresGroup",
diff --git a/src/KeyValueStore.ts b/src/KeyValueStore.ts
@@ -0,0 +1,35 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+/**
+ * A global key value store that associates tree nodes with open editors.
+ */
+export class KeyValueStore {
+    private _items: Map<string, object>;
+
+    private static _instance: KeyValueStore | undefined;
+
+    public static get instance(): KeyValueStore {
+        if (!this._instance) {
+            this._instance = new KeyValueStore();
+        }
+        return this._instance;
+    }
+
+    /**
+     * Prevent external instantiation.
+     */
+    private constructor() {
+        this._items = new Map<string, object>();
+    }
+
+    public get(key: string): object | undefined {
+        return this._items.get(key);
+    }
+
+    public set(key: string, value: object): void {
+        this._items.set(key, value);
+    }
+}
diff --git a/src/docdb/NoSqlCodeLensProvider.ts b/src/docdb/NoSqlCodeLensProvider.ts
@@ -0,0 +1,100 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { IActionContext, callWithTelemetryAndErrorHandling } from "@microsoft/vscode-azext-utils";
+import {
+    CancellationToken,
+    CodeLens,
+    CodeLensProvider,
+    Event,
+    EventEmitter,
+    Position,
+    ProviderResult,
+    Range,
+    TextDocument
+} from "vscode";
+import { KeyValueStore } from "../KeyValueStore";
+
+export type NoSqlQueryConnection = {
+    databaseId: string;
+    containerId: string;
+    endpoint: string;
+    masterKey: string;
+    isEmulator: boolean;
+};
+
+export const noSqlQueryConnectionKey = "NO_SQL_QUERY_CONNECTION_KEY.v1";
+
+export class NoSqlCodeLensProvider implements CodeLensProvider {
+    private _onDidChangeEmitter: EventEmitter<void> = new EventEmitter<void>();
+
+    public get onDidChangeCodeLenses(): Event<void> {
+        return this._onDidChangeEmitter.event;
+    }
+
+    public updateCodeLens(): void {
+        this._onDidChangeEmitter.fire();
+    }
+
+    public provideCodeLenses(document: TextDocument, _token: CancellationToken): ProviderResult<CodeLens[]> {
+        return callWithTelemetryAndErrorHandling("nosql.provideCodeLenses", (context: IActionContext) => {
+            context.telemetry.suppressIfSuccessful = true;
+            const text = document.getText();
+            const queryText = text;
+
+            // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment, @typescript-eslint/no-explicit-any
+            const connectedCollection: NoSqlQueryConnection | undefined = KeyValueStore.instance.get(noSqlQueryConnectionKey) as any;
+            let connectCodeLens: CodeLens;
+            if (!connectedCollection) {
+                connectCodeLens = new CodeLens(
+                    new Range(new Position(0, 0), new Position(0, 0)),
+                    {
+                        title: "Not connected",
+                        command: "cosmosDB.connectNoSqlContainer",
+                        arguments: []
+                    }
+                );
+            } else {
+                connectCodeLens = new CodeLens(
+                    new Range(new Position(0, 0), new Position(0, 0)),
+                    {
+                        title: `Connected to ${connectedCollection.databaseId}.${connectedCollection.containerId}`,
+                        command: "cosmosDB.connectNoSqlContainer",
+                        arguments: []
+                    }
+                );
+            }
+            const lenses: CodeLens[] = [
+                connectCodeLens,
+                new CodeLens(
+                    new Range(new Position(0, 0), new Position(0, 0)),
+                    {
+                        title: "Execute",
+                        command: "cosmosDB.executeNoSqlQuery",
+                        arguments: [{ queryText }]
+                    }
+                ),
+                new CodeLens(
+                    new Range(new Position(0, 0), new Position(0, 0)),
+                    {
+                        title: "Execute with Query Metrics",
+                        command: "cosmosDB.executeNoSqlQuery",
+                        arguments: [{ queryText, populateQueryMetrics: true }]
+                    }
+                ),
+                new CodeLens(
+                    new Range(new Position(0, 0), new Position(0, 0)),
+                    {
+                        title: "Get Query Plan",
+                        command: "cosmosDB.getNoSqlQueryPlan",
+                        arguments: [{ queryText }]
+                    }
+                )
+            ];
+
+            return lenses;
+        });
+    }
+}
diff --git a/src/docdb/getCosmosClient.ts b/src/docdb/getCosmosClient.ts
@@ -10,7 +10,6 @@ import * as vscode from 'vscode';
 import { ext } from "../extensionVariables";
 
 export function getCosmosClient(endpoint: string, key: string, isEmulator: boolean | undefined): CosmosClient {
-
     const vscodeStrictSSL: boolean | undefined = vscode.workspace.getConfiguration().get<boolean>(ext.settingsKeys.vsCode.proxyStrictSSL);
     const enableEndpointDiscovery: boolean | undefined = vscode.workspace.getConfiguration().get<boolean>(ext.settingsKeys.enableEndpointDiscovery);
     const connectionPolicy = { enableEndpointDiscovery: (enableEndpointDiscovery === undefined) ? true : enableEndpointDiscovery };
diff --git a/src/docdb/registerDocDBCommands.ts b/src/docdb/registerDocDBCommands.ts
@@ -3,20 +3,33 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
-import { AzExtTreeItem, IActionContext, ITreeItemPickerContext, registerCommandWithTreeNodeUnwrapping } from "@microsoft/vscode-azext-utils";
-import { commands } from "vscode";
+import { AzExtTreeItem, IActionContext, ITreeItemPickerContext, registerCommand, registerCommandWithTreeNodeUnwrapping } from "@microsoft/vscode-azext-utils";
+import { ViewColumn, commands, languages } from "vscode";
+import { KeyValueStore } from "../KeyValueStore";
 import { doubleClickDebounceDelay, sqlFilter } from "../constants";
 import { ext } from "../extensionVariables";
+import * as vscodeUtil from "../utils/vscodeUtils";
+import { NoSqlCodeLensProvider, NoSqlQueryConnection, noSqlQueryConnectionKey } from "./NoSqlCodeLensProvider";
+import { getCosmosClient } from "./getCosmosClient";
 import { DocDBAccountTreeItem } from "./tree/DocDBAccountTreeItem";
 import { DocDBCollectionTreeItem } from "./tree/DocDBCollectionTreeItem";
 import { DocDBDatabaseTreeItem } from "./tree/DocDBDatabaseTreeItem";
-import { DocDBDocumentsTreeItem } from "./tree/DocDBDocumentsTreeItem";
 import { DocDBDocumentTreeItem } from "./tree/DocDBDocumentTreeItem";
-import { DocDBStoredProceduresTreeItem } from "./tree/DocDBStoredProceduresTreeItem";
+import { DocDBDocumentsTreeItem } from "./tree/DocDBDocumentsTreeItem";
 import { DocDBStoredProcedureTreeItem } from "./tree/DocDBStoredProcedureTreeItem";
+import { DocDBStoredProceduresTreeItem } from "./tree/DocDBStoredProceduresTreeItem";
+
+const nosqlLanguageId = "nosql";
 
 export function registerDocDBCommands(): void {
+    ext.noSqlCodeLensProvider = new NoSqlCodeLensProvider();
+    ext.context.subscriptions.push(languages.registerCodeLensProvider(nosqlLanguageId, ext.noSqlCodeLensProvider));
+
+    registerCommand("cosmosDB.connectNoSqlContainer", connectNoSqlContainer);
+    registerCommand("cosmosDB.executeNoSqlQuery", executeNoSqlQuery);
+    registerCommand("cosmosDB.getNoSqlQueryPlan", getNoSqlQueryPlan);
     registerCommandWithTreeNodeUnwrapping('cosmosDB.createDocDBDatabase', createDocDBDatabase);
+    registerCommandWithTreeNodeUnwrapping('cosmosDB.writeNoSqlQuery', writeNoSqlQuery);
     registerCommandWithTreeNodeUnwrapping('cosmosDB.createDocDBCollection', createDocDBCollection);
     registerCommandWithTreeNodeUnwrapping('cosmosDB.createDocDBDocument', async (context: IActionContext, node?: DocDBDocumentsTreeItem) => {
         if (!node) {
@@ -61,6 +74,72 @@ export function registerDocDBCommands(): void {
     });
 }
 
+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();
+}
+
+async function writeNoSqlQuery(_context: IActionContext, node: DocDBCollectionTreeItem): Promise<void> {
+    setConnectedNoSqlContainer(node);
+    const sampleQuery = `SELECT * FROM ${node.id}`;
+    await vscodeUtil.showNewFile(sampleQuery, `query for ${node.label}`, ".nosql");
+}
+
+async function connectNoSqlContainer(context: IActionContext): Promise<void> {
+    const node = await pickDocDBAccount<DocDBCollectionTreeItem>(context, DocDBCollectionTreeItem.contextValue);
+    setConnectedNoSqlContainer(node);
+}
+
+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);
+        }
+    }
+}
+
+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);
+    }
+}
+
 export async function createDocDBDatabase(context: IActionContext, node?: DocDBAccountTreeItem): Promise<void> {
     if (!node) {
         node = await pickDocDBAccount<DocDBAccountTreeItem>(context);
diff --git a/src/docdb/tree/DocDBDocumentTreeItem.ts b/src/docdb/tree/DocDBDocumentTreeItem.ts
@@ -106,7 +106,7 @@ export class DocDBDocumentTreeItem extends AzExtTreeItem implements IEditableTre
         }
     }
 
-    private getPartitionKeyValue(): string | undefined | Object {
+    private getPartitionKeyValue(): string | undefined {
         const partitionKey = this.parent.parent.partitionKey;
         if (!partitionKey) { //Fixed collections -> no partitionKeyValue
             return undefined;
diff --git a/src/extensionVariables.ts b/src/extensionVariables.ts
@@ -7,6 +7,7 @@ import { AzExtTreeDataProvider, AzExtTreeItem, IAzExtOutputChannel } from "@micr
 import { AzureHostExtensionApi } from "@microsoft/vscode-azext-utils/hostapi";
 import { ExtensionContext, SecretStorage, TreeView } from "vscode";
 import { DatabasesFileSystem } from "./DatabasesFileSystem";
+import { NoSqlCodeLensProvider } from "./docdb/NoSqlCodeLensProvider";
 import { MongoDBLanguageClient } from "./mongo/languageClient";
 import { MongoCodeLensProvider } from "./mongo/services/MongoCodeLensProvider";
 import { MongoDatabaseTreeItem } from "./mongo/tree/MongoDatabaseTreeItem";
@@ -33,6 +34,7 @@ export namespace ext {
     export const prefix: string = 'azureDatabases';
     export let fileSystem: DatabasesFileSystem;
     export let mongoCodeLensProvider: MongoCodeLensProvider;
+    export let noSqlCodeLensProvider: NoSqlCodeLensProvider;
     export let mongoLanguageClient: MongoDBLanguageClient;
     export let rgApi: AzureHostExtensionApi;
 

Original file line number	Diff line number	Diff line change
`@@ -106,7 +106,7 @@ export class DocDBDocumentTreeItem extends AzExtTreeItem implements IEditableTre`
`106`	`106`	`}`
`107`	`107`	`}`
`108`	`108`
`109`		`- private getPartitionKeyValue(): string \| undefined \| Object {`
	`109`	`+ private getPartitionKeyValue(): string \| undefined {`
`110`	`110`	`const partitionKey = this.parent.parent.partitionKey;`
`111`	`111`	`if (!partitionKey) { //Fixed collections -> no partitionKeyValue`
`112`	`112`	`return undefined;`