Enable CosmosDB OAuth support (#2277)

* Enable CosmosDB OAuth support

* Fix connectionString construction
Deduplicate key credential extraction

* Fix lint

package.json

@@ -1069,6 +1069,11 @@
                     "type": "integer",
                     "default": 8081,
                     "description": "Port to use when connecting to a CosmosDB Emulator instance"
+                },
+                "azureDatabases.testCosmosAuth": {
+                    "type": "boolean",
+                    "description": "Whether to only use Auth credential for Cosmos DB resources",
+                    "default": true
                 }
             }
         }
@@ -1129,8 +1134,9 @@
         "@azure/arm-postgresql": "^6.0.0",
         "@azure/arm-postgresql-flexible": "^7.1.0",
         "@azure/cosmos": "^3.6.3",
-        "@microsoft/vscode-azext-azureutils": "^2.0.2",
-        "@microsoft/vscode-azext-utils": "^2.1.0",
+        "@microsoft/vscode-azext-azureauth": "^2.3.0",
+        "@microsoft/vscode-azext-azureutils": "^3.0.1",
+        "@microsoft/vscode-azext-utils": "^2.4.0",
         "antlr4ts": "^0.4.1-alpha.0",
         "bson": "^6.0.0",
         "fs-extra": "^8.0.0",

src/commands/api/DatabaseAccountTreeItemInternal.ts

@@ -5,6 +5,7 @@
 
 import { callWithTelemetryAndErrorHandling, IActionContext } from '@microsoft/vscode-azext-utils';
 import { API } from '../../AzureDBExperiences';
+import { getCosmosKeyCredential } from '../../docdb/getCosmosClient';
 import { DocDBAccountTreeItemBase } from '../../docdb/tree/DocDBAccountTreeItemBase';
 import { ext } from '../../extensionVariables';
 import { ParsedMongoConnectionString } from '../../mongo/mongoConnectionStrings';
@@ -57,10 +58,15 @@ export class DatabaseAccountTreeItemInternal implements DatabaseAccountTreeItem
 
     public get docDBData(): { masterKey: string; documentEndpoint: string; } | undefined {
         if (this._accountNode instanceof DocDBAccountTreeItemBase) {
-            return {
-                documentEndpoint: this._accountNode.root.endpoint,
-                masterKey: this._accountNode.root.masterKey
-            };
+            const keyCred = getCosmosKeyCredential(this._accountNode.root.credentials);
+            if (keyCred) {
+                return {
+                    documentEndpoint: this._accountNode.root.endpoint,
+                    masterKey: keyCred.key
+                };
+            } else {
+                return undefined;
+            }
         } else {
             return undefined;
         }

src/docdb/NoSqlCodeLensProvider.ts

@@ -21,7 +21,7 @@ export type NoSqlQueryConnection = {
     databaseId: string;
     containerId: string;
     endpoint: string;
-    masterKey: string;
+    masterKey?: string;
     isEmulator: boolean;
 };
 

src/docdb/commands/connectNoSqlContainer.ts

@@ -7,16 +7,19 @@ import { IActionContext } from "@microsoft/vscode-azext-utils";
 import { KeyValueStore } from "../../KeyValueStore";
 import { ext } from "../../extensionVariables";
 import { NoSqlQueryConnection, noSqlQueryConnectionKey } from "../NoSqlCodeLensProvider";
+import { getCosmosKeyCredential } from "../getCosmosClient";
 import { DocDBCollectionTreeItem } from "../tree/DocDBCollectionTreeItem";
 import { pickDocDBAccount } from "./pickDocDBAccount";
 
 export function setConnectedNoSqlContainer(node: DocDBCollectionTreeItem): void {
+    const root = node.root;
+    const keyCred = getCosmosKeyCredential(root.credentials);
     const noSqlQueryConnection: NoSqlQueryConnection = {
         databaseId: node.parent.id,
         containerId: node.id,
-        endpoint: node.root.endpoint,
-        masterKey: node.root.masterKey,
-        isEmulator: !!node.root.isEmulator
+        endpoint: root.endpoint,
+        masterKey: keyCred?.key,
+        isEmulator: !!root.isEmulator
     };
     KeyValueStore.instance.set(noSqlQueryConnectionKey, noSqlQueryConnection);
     ext.noSqlCodeLensProvider.updateCodeLens();

src/docdb/commands/executeNoSqlQuery.ts

@@ -10,7 +10,7 @@ import { KeyValueStore } from "../../KeyValueStore";
 import { localize } from "../../utils/localize";
 import * as vscodeUtil from "../../utils/vscodeUtils";
 import { NoSqlQueryConnection, noSqlQueryConnectionKey } from "../NoSqlCodeLensProvider";
-import { getCosmosClient } from "../getCosmosClient";
+import { CosmosDBCredential, getCosmosClient } from "../getCosmosClient";
 
 export async function executeNoSqlQuery(_context: IActionContext, args: { queryText: string, populateQueryMetrics?: boolean }): Promise<void> {
     let queryText: string;
@@ -33,7 +33,12 @@ export async function executeNoSqlQuery(_context: IActionContext, args: { queryT
     } 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 credentials: CosmosDBCredential[] = [];
+        if (masterKey !== undefined) {
+            credentials.push({ type: "key", key: masterKey });
+        }
+        credentials.push({ type: "auth" });
+        const client = getCosmosClient(endpoint, credentials, isEmulator);
         const options = { populateQueryMetrics };
         const response = await client.database(databaseId).container(containerId).items.query(queryText, options).fetchAll();
         const resultDocumentTitle = `query results for ${containerId}`;

src/docdb/commands/getNoSqlQueryPlan.ts

@@ -10,7 +10,7 @@ import { KeyValueStore } from "../../KeyValueStore";
 import { localize } from "../../utils/localize";
 import * as vscodeUtil from "../../utils/vscodeUtils";
 import { NoSqlQueryConnection, noSqlQueryConnectionKey } from "../NoSqlCodeLensProvider";
-import { getCosmosClient } from "../getCosmosClient";
+import { CosmosDBCredential, getCosmosClient } from "../getCosmosClient";
 
 export async function getNoSqlQueryPlan(_context: IActionContext, args: { queryText: string } | undefined): Promise<void> {
     let queryText: string;
@@ -30,7 +30,12 @@ export async function getNoSqlQueryPlan(_context: IActionContext, args: { queryT
     } 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 credentials: CosmosDBCredential[] = [];
+        if (masterKey !== undefined) {
+            credentials.push({ type: "key", key: masterKey });
+        }
+        credentials.push({ type: "auth" });
+        const client = getCosmosClient(endpoint, credentials, 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/getCosmosClient.ts

@@ -9,10 +9,66 @@ import * as https from "https";
 import * as vscode from 'vscode';
 import { ext } from "../extensionVariables";
 
-export function getCosmosClient(endpoint: string, key: string, isEmulator: boolean | undefined): CosmosClient {
+// eslint-disable-next-line import/no-internal-modules
+import { getSessionFromVSCode } from "@microsoft/vscode-azext-azureauth/out/src/getSessionFromVSCode";
+
+export type CosmosDBKeyCredential = {
+    type: "key";
+    key: string;
+};
+
+export type CosmosDBAuthCredential = {
+    type: "auth";
+};
+
+export type CosmosDBCredential = CosmosDBKeyCredential | CosmosDBAuthCredential;
+
+export function getCosmosKeyCredential(credentials: CosmosDBCredential[]): CosmosDBKeyCredential | undefined {
+    return credentials.filter((cred): cred is CosmosDBKeyCredential => cred.type === "key")[0];
+}
+
+export function getCosmosAuthCredential(credentials: CosmosDBCredential[]): CosmosDBAuthCredential | undefined {
+    return credentials.filter((cred): cred is CosmosDBAuthCredential => cred.type === "auth")[0];
+}
+
+export function getCosmosClient(
+    endpoint: string,
+    credentials: CosmosDBCredential[],
+    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 };
-    return new CosmosClient({ endpoint, key, userAgentSuffix: appendExtensionUserAgent(), agent: new https.Agent({ rejectUnauthorized: isEmulator ? !isEmulator : vscodeStrictSSL }), connectionPolicy: connectionPolicy });
 
+    const keyCred = getCosmosKeyCredential(credentials);
+    const authCred = getCosmosAuthCredential(credentials);
+
+    const commonProperties = {
+        endpoint,
+        userAgentSuffix: appendExtensionUserAgent(),
+        agent: new https.Agent({ rejectUnauthorized: isEmulator ? !isEmulator : vscodeStrictSSL }),
+        connectionPolicy
+    };
+    // @todo: Add telemetry to monitor usage of each credential type
+    if (keyCred) {
+        return new CosmosClient({
+            ...commonProperties,
+            key: keyCred.key
+        });
+    } else if (authCred) {
+        return new CosmosClient({
+            ...commonProperties,
+            aadCredentials: {
+                getToken: async (scopes, _options) => {
+                    const session = await getSessionFromVSCode(scopes, undefined, { createIfNone: true });
+                    return {
+                        token: session?.accessToken ?? "",
+                        expiresOnTimestamp: 0
+                    };
+                }
+            }
+        });
+    } else {
+        throw Error("No credential available to create CosmosClient.");
+    }
 }

src/docdb/tree/DocDBAccountTreeItemBase.ts

@@ -12,7 +12,7 @@ import { deleteCosmosDBAccount } from '../../commands/deleteDatabaseAccount/dele
 import { SERVERLESS_CAPABILITY_NAME, getThemeAgnosticIconPath } from '../../constants';
 import { nonNullProp } from '../../utils/nonNull';
 import { rejectOnTimeout } from '../../utils/timeout';
-import { getCosmosClient } from '../getCosmosClient';
+import { CosmosDBCredential, getCosmosClient, getCosmosKeyCredential } from '../getCosmosClient';
 import { DocDBTreeItemBase } from './DocDBTreeItemBase';
 
 /**
@@ -24,22 +24,38 @@ export abstract class DocDBAccountTreeItemBase extends DocDBTreeItemBase<Databas
     public readonly childTypeLabel: string = "Database";
 
 
-    constructor(parent: AzExtParentTreeItem, id: string, label: string, endpoint: string, masterKey: string, isEmulator: boolean | undefined, readonly databaseAccount?: DatabaseAccountGetResults) {
+    constructor(
+        parent: AzExtParentTreeItem,
+        id: string,
+        label: string,
+        endpoint: string,
+        credentials: CosmosDBCredential[],
+        isEmulator: boolean | undefined,
+        readonly databaseAccount?: DatabaseAccountGetResults
+    ) {
         super(parent);
         this.id = id;
         this.label = label;
         this.root = {
             endpoint,
-            masterKey,
+            credentials,
             isEmulator,
-            getCosmosClient: () => getCosmosClient(endpoint, masterKey, isEmulator)
+            getCosmosClient: () => getCosmosClient(endpoint, credentials, isEmulator)
         };
 
-        this.valuesToMask.push(id, endpoint, masterKey);
+        const keys = credentials
+            .map((cred) => cred.type === "key" ? cred.key : undefined)
+            .filter((value): value is string => value !== undefined);
+        this.valuesToMask.push(id, endpoint, ...keys);
     }
 
     public get connectionString(): string {
-        return `AccountEndpoint=${this.root.endpoint};AccountKey=${this.root.masterKey}`;
+        const firstKey = getCosmosKeyCredential(this.root.credentials);
+        if (firstKey) {
+            return `AccountEndpoint=${this.root.endpoint};AccountKey=${firstKey.key}`;
+        } else {
+            return `AccountEndpoint=${this.root.endpoint}`;
+        }
     }
 
     public get iconPath(): string | vscode.Uri | { light: string | vscode.Uri; dark: string | vscode.Uri } {

src/docdb/tree/DocDBDocumentTreeItem.ts

@@ -106,7 +106,7 @@ export class DocDBDocumentTreeItem extends AzExtTreeItem implements IEditableTre
         }
     }
 
-    private getPartitionKeyValue(): string | undefined {
+    private getPartitionKeyValue(): string | number | undefined {
         const partitionKey = this.parent.parent.partitionKey;
         if (!partitionKey) { //Fixed collections -> no partitionKeyValue
             return undefined;

src/docdb/tree/IDocDBTreeRoot.ts

@@ -5,10 +5,11 @@
  *--------------------------------------------------------------------------------------------*/
 
 import { CosmosClient } from "@azure/cosmos";
+import { CosmosDBCredential } from "../getCosmosClient";
 
 export interface IDocDBTreeRoot {
     endpoint: string;
-    masterKey: string;
+    credentials: CosmosDBCredential[];
     isEmulator: boolean | undefined;
     getCosmosClient(): CosmosClient;
 }

src/graph/tree/GraphAccountTreeItem.ts

@@ -6,9 +6,10 @@
 import { DatabaseAccountGetResults } from '@azure/arm-cosmosdb/src/models';
 import { DatabaseDefinition, Resource } from '@azure/cosmos';
 import { AzExtParentTreeItem } from '@microsoft/vscode-azext-utils';
+import { CosmosDBCredential } from '../../docdb/getCosmosClient';
 import { DocDBAccountTreeItemBase } from '../../docdb/tree/DocDBAccountTreeItemBase';
-import { DocDBStoredProceduresTreeItem } from '../../docdb/tree/DocDBStoredProceduresTreeItem';
 import { DocDBStoredProcedureTreeItem } from '../../docdb/tree/DocDBStoredProcedureTreeItem';
+import { DocDBStoredProceduresTreeItem } from '../../docdb/tree/DocDBStoredProceduresTreeItem';
 import { IGremlinEndpoint } from '../../vscode-cosmosdbgraph.api';
 import { GraphCollectionTreeItem } from './GraphCollectionTreeItem';
 import { GraphDatabaseTreeItem } from './GraphDatabaseTreeItem';
@@ -18,8 +19,17 @@ export class GraphAccountTreeItem extends DocDBAccountTreeItemBase {
     public static contextValue: string = "cosmosDBGraphAccount";
     public contextValue: string = GraphAccountTreeItem.contextValue;
 
-    constructor(parent: AzExtParentTreeItem, id: string, label: string, documentEndpoint: string, private _gremlinEndpoint: IGremlinEndpoint | undefined, masterKey: string, isEmulator: boolean | undefined, readonly databaseAccount?: DatabaseAccountGetResults) {
-        super(parent, id, label, documentEndpoint, masterKey, isEmulator, databaseAccount);
+    constructor(
+        parent: AzExtParentTreeItem,
+        id: string,
+        label: string,
+        documentEndpoint: string,
+        private _gremlinEndpoint: IGremlinEndpoint | undefined,
+        credentials: CosmosDBCredential[],
+        isEmulator: boolean | undefined,
+        readonly databaseAccount?: DatabaseAccountGetResults
+    ) {
+        super(parent, id, label, documentEndpoint, credentials, isEmulator, databaseAccount);
         this.valuesToMask.push(documentEndpoint);
         if (_gremlinEndpoint) {
             this.valuesToMask.push(_gremlinEndpoint.host);

src/tree/AttachedAccountsTreeItem.ts

@@ -10,6 +10,7 @@ import { API, getExperienceFromApi, getExperienceQuickPick, getExperienceQuickPi
 import { removeTreeItemFromCache } from '../commands/api/apiCache';
 import { emulatorPassword, isWindows } from '../constants';
 import { parseDocDBConnectionString } from '../docdb/docDBConnectionStrings';
+import { CosmosDBCredential } from '../docdb/getCosmosClient';
 import { DocDBAccountTreeItem } from '../docdb/tree/DocDBAccountTreeItem';
 import { DocDBAccountTreeItemBase } from '../docdb/tree/DocDBAccountTreeItemBase';
 import { ext } from '../extensionVariables';
@@ -332,15 +333,17 @@ export class AttachedAccountsTreeItem extends AzExtParentTreeItem {
             const parsedCS = parseDocDBConnectionString(connectionString);
 
             label = label || `${parsedCS.accountId} (${getExperienceFromApi(api).shortName})`;
+
+            const credentials: CosmosDBCredential[] = [{ type: "key", key: parsedCS.masterKey }];
             switch (api) {
                 case API.Table:
-                    treeItem = new TableAccountTreeItem(this, parsedCS.accountId, label, parsedCS.documentEndpoint, parsedCS.masterKey, isEmulator);
+                    treeItem = new TableAccountTreeItem(this, parsedCS.accountId, label, parsedCS.documentEndpoint, credentials, isEmulator);
                     break;
                 case API.Graph:
-                    treeItem = new GraphAccountTreeItem(this, parsedCS.accountId, label, parsedCS.documentEndpoint, undefined, parsedCS.masterKey, isEmulator);
+                    treeItem = new GraphAccountTreeItem(this, parsedCS.accountId, label, parsedCS.documentEndpoint, undefined, credentials, isEmulator);
                     break;
                 case API.Core:
-                    treeItem = new DocDBAccountTreeItem(this, parsedCS.accountId, label, parsedCS.documentEndpoint, parsedCS.masterKey, isEmulator);
+                    treeItem = new DocDBAccountTreeItem(this, parsedCS.accountId, label, parsedCS.documentEndpoint, credentials, isEmulator);
                     break;
                 default:
                     throw new Error(`Unexpected defaultExperience "${api}".`);