Allow attaching a cosmos db account with tenantId.

l10n/bundle.l10n.json

@@ -80,6 +80,7 @@
   "Connection String": "Connection String",
   "Connection string cannot be empty": "Connection string cannot be empty",
   "Connection string must be of the form \"AccountEndpoint=…;AccountKey=…\"": "Connection string must be of the form \"AccountEndpoint=…;AccountKey=…\"",
+  "Connection string must include \"AccountEndpoint=…\" and either \"AccountKey=…\" or \"TenantId=…\"": "Connection string must include \"AccountEndpoint=…\" and either \"AccountKey=…\" or \"TenantId=…\"",
   "Connection timed out while verifying account credentials.": "Connection timed out while verifying account credentials.",
   "Connection timed out. Please verify the connection string and that the host is reachable.": "Connection timed out. Please verify the connection string and that the host is reachable.",
   "Container {0} not found": "Container {0} not found",

src/commands/newConnection/CosmosDBConnectionStringStep.ts

@@ -12,7 +12,7 @@ export class CosmosDBConnectionStringStep extends AzureWizardPromptStep<NewConne
     public async prompt(context: NewConnectionWizardContext): Promise<void> {
         context.connectionString = (
             await context.ui.showInputBox({
-                placeHolder: 'AccountEndpoint=…;AccountKey=…',
+                placeHolder: 'AccountEndpoint=…;AccountKey=…;TenantId=…',
                 prompt: l10n.t('Enter the connection string for your database account'),
                 validateInput: (connectionString?: string) => this.validateInput(connectionString),
                 asyncValidationTask: (connectionString: string) => this.validateConnectionString(connectionString),
@@ -45,7 +45,9 @@ export class CosmosDBConnectionStringStep extends AzureWizardPromptStep<NewConne
             if (error instanceof Error) {
                 return error.message;
             } else {
-                return l10n.t('Connection string must be of the form "AccountEndpoint=…;AccountKey=…"');
+                return l10n.t(
+                    'Connection string must include "AccountEndpoint=…" and either "AccountKey=…" or "TenantId=…"',
+                );
             }
         }
 

src/commands/newConnection/CosmosDBExecuteStep.ts

@@ -33,7 +33,11 @@ export class CosmosDBExecuteStep extends AzureWizardExecuteStep<NewConnectionWiz
                 const storageItem: StorageItem = {
                     id: parsedCS.accountId,
                     name: label,
-                    properties: { isEmulator: false, api },
+                    properties: {
+                        isEmulator: false,
+                        api,
+                        ...(parsedCS.tenantId && { tenantId: parsedCS.tenantId }),
+                    },
                     secrets: [connectionString],
                 };
 

src/cosmosdb/cosmosDBConnectionStrings.test.ts

@@ -9,6 +9,14 @@ import { parseCosmosDBConnectionString } from './cosmosDBConnectionStrings';
 describe('cosmosDBConnectionStrings', () => {
     // Testing different ordering, different use of ';', different casing, etc.
     describe('Without database name', () => {
+        it('Connection string with AccountEndpoint only', () => {
+            const parsedCS = parseCosmosDBConnectionString('AccountEndpoint=https://abcdef.documents.azure.com:443/');
+            expect(parsedCS.documentEndpoint).toEqual('https://abcdef.documents.azure.com:443/');
+            expect(parsedCS.masterKey).toEqual(undefined);
+            expect(parsedCS.databaseName).toEqual(undefined);
+            expect(parsedCS.tenantId).toEqual(undefined);
+        });
+
         it('Connection string with empty database', () => {
             const parsedCS = parseCosmosDBConnectionString(
                 'AccountEndpoint=https://abcdef.documents.azure.com:443/;AccountKey=abcdef==',
@@ -18,6 +26,16 @@ describe('cosmosDBConnectionStrings', () => {
             expect(parsedCS.databaseName).toEqual(undefined);
         });
 
+        it('Connection string with TenantId', () => {
+            const parsedCS = parseCosmosDBConnectionString(
+                'AccountEndpoint=https://abcdef.documents.azure.com:443/;TenantId=00000000-0000-0000-0000-000000000000',
+            );
+            expect(parsedCS.documentEndpoint).toEqual('https://abcdef.documents.azure.com:443/');
+            expect(parsedCS.masterKey).toEqual(undefined);
+            expect(parsedCS.databaseName).toEqual(undefined);
+            expect(parsedCS.tenantId).toEqual('00000000-0000-0000-0000-000000000000');
+        });
+
         it('Connection string with leading empty string in AccountEndpoint', () => {
             const parsedCS = parseCosmosDBConnectionString(
                 '    AccountEndpoint=https://abcdef.documents.azure.com:443/;AccountKey=abcdef==',

src/cosmosdb/cosmosDBConnectionStrings.ts

@@ -10,14 +10,15 @@ export function parseCosmosDBConnectionString(connectionString: string): ParsedC
     const endpoint = getPropertyFromConnectionString(connectionString, 'AccountEndpoint');
     const masterKey = getPropertyFromConnectionString(connectionString, 'AccountKey');
     const databaseName = getPropertyFromConnectionString(connectionString, 'Database');
+    const tenantId = getPropertyFromConnectionString(connectionString, 'TenantId');
 
     if (!endpoint) {
         throw new Error(l10n.t('Invalid Cosmos DB connection string.'));
     }
 
     const endpointUrl = new URL(endpoint);
 
-    return new ParsedCosmosDBConnectionString(connectionString, endpointUrl, masterKey, databaseName);
+    return new ParsedCosmosDBConnectionString(connectionString, endpointUrl, masterKey, databaseName, tenantId);
 }
 
 function getPropertyFromConnectionString(connectionString: string, property: string): string | undefined {
@@ -32,12 +33,14 @@ export class ParsedCosmosDBConnectionString extends ParsedConnectionString {
 
     public readonly documentEndpoint: string;
     public readonly masterKey: string | undefined;
+    public readonly tenantId: string | undefined;
 
     constructor(
         connectionString: string,
         endpoint: URL,
         masterKey: string | undefined,
         databaseName: string | undefined,
+        tenantId: string | undefined,
     ) {
         super(connectionString, databaseName);
 
@@ -48,6 +51,7 @@ export class ParsedCosmosDBConnectionString extends ParsedConnectionString {
         // since URL.toString() does not include the port if it is the default (80 or 443)
         this.documentEndpoint = `${endpoint.protocol}//${this.hostName}:${this.port}${endpoint.pathname}${endpoint.search}`;
         this.masterKey = masterKey;
+        this.tenantId = tenantId;
     }
 
     public get accountName(): string {

src/tree/cosmosdb/AccountInfo.ts

@@ -61,7 +61,7 @@ async function getAccountInfoForConnectionString(
         documentEndpoint: connectionString.documentEndpoint,
         isEmulator,
         masterKey: connectionString.masterKey,
-        tenantId: undefined,
+        tenantId: connectionString.tenantId,
     });
     const isServerless = false;
 
@@ -126,12 +126,13 @@ async function getAccountInfoForAttached(account: CosmosDBAttachedAccountModel):
     const isEmulator = account.isEmulator;
     const parsedCS = parseCosmosDBConnectionString(account.connectionString);
     const documentEndpoint = parsedCS.documentEndpoint;
+    const tenantId = account.tenantId || parsedCS.tenantId;
     const credentials = await getCosmosDBCredentials({
         accountName: name,
         documentEndpoint,
         isEmulator,
         masterKey: parsedCS.masterKey,
-        tenantId: undefined,
+        tenantId,
     });
     const isServerless = false;
 

src/tree/workspace-view/cosmosdb/CosmosDBAttachedAccountModel.ts

@@ -9,4 +9,5 @@ export type CosmosDBAttachedAccountModel = {
     storageId: string;
     isEmulator: boolean;
     name: string;
+    tenantId?: string;
 };

src/tree/workspace-view/cosmosdb/CosmosDBWorkspaceItem.ts

@@ -58,13 +58,16 @@ export class CosmosDBWorkspaceItem implements TreeElement, TreeElementWithContex
                     const api: API = nonNullValue(properties?.api, 'api') as API;
                     const isEmulator: boolean = !!nonNullValue(properties?.isEmulator, 'isEmulator');
                     const connectionString: string = nonNullValue(secrets?.[0], 'connectionString');
+                    const tenantId: string | undefined =
+                        typeof properties?.tenantId === 'string' ? (properties.tenantId as string) : undefined;
                     const experience = getExperienceFromApi(api);
                     const accountModel: CosmosDBAttachedAccountModel = {
                         id: `${this.id}/${id}`, // To enable TreeView.reveal, we need to have a unique nested id
                         storageId: id,
                         name,
                         connectionString,
                         isEmulator,
+                        tenantId,
                     };
 
                     if (experience?.api === API.Cassandra) {

src/vscodeUriHandler.ts

@@ -168,13 +168,15 @@ async function handleConnectionStringRequest(
         // Create storage item for the connection
         if (parsedConnection.api === API.Core) {
             const isEmulator = getIsEmulatorConnection(parsedConnection.connectionString);
+            const tenantId = params.tenantId || parsedConnection.connectionString.tenantId;
             const fullId = await createAttachedForConnection(
                 parsedConnection.connectionString.accountId,
                 parsedConnection.connectionString.accountName,
                 parsedConnection.api,
                 params.connectionString,
                 isEmulator,
                 parsedConnection.connectionString.port,
+                tenantId,
             );
             ext.cosmosDBWorkspaceBranchDataProvider.refresh();
             await revealAttachedInWorkspaceExplorer(fullId, params.database, params.container);
@@ -304,6 +306,7 @@ async function createAttachedForConnection(
     connectionString: string,
     isEmulator: boolean,
     emulatorPort?: string,
+    tenantId?: string,
     disableEmulatorSecurity?: boolean,
 ): Promise<string> {
     const rootId = `${WorkspaceResourceType.AttachedAccounts}`;
@@ -333,7 +336,12 @@ async function createAttachedForConnection(
         const storageItem: StorageItem = {
             id,
             name,
-            properties: { isEmulator, api, ...(disableEmulatorSecurity && { disableEmulatorSecurity }) },
+            properties: {
+                isEmulator,
+                api,
+                ...(tenantId && { tenantId }),
+                ...(disableEmulatorSecurity && { disableEmulatorSecurity }),
+            },
             secrets: [connectionString],
         };
 
@@ -495,7 +503,7 @@ function parseConnectionString(
 
     // All other connection strings are treated as Core API
     const parsedCS = parseCosmosDBConnectionString(connectionString);
-    [parsedCS.masterKey, parsedCS.databaseName]
+    [parsedCS.masterKey, parsedCS.databaseName, parsedCS.tenantId]
         .filter((value): value is string => Boolean(value))
         .forEach((value) => context.valuesToMask.push(value));
     return {
@@ -520,6 +528,7 @@ function extractParams(query: string): {
     subscriptionId?: string;
     resourceGroup?: string;
     connectionString?: string;
+    tenantId?: string;
     database?: string;
     container?: string;
 } {
@@ -529,6 +538,7 @@ function extractParams(query: string): {
         subscriptionId: queryParams.get('subscriptionId') ?? undefined,
         resourceGroup: queryParams.get('resourceGroup') ?? undefined,
         connectionString: queryParams.get('cs') ?? undefined,
+        tenantId: queryParams.get('tenantId') ?? undefined,
         database: queryParams.get('database') ?? undefined,
         container: queryParams.get('container') ?? undefined,
     };
@@ -548,6 +558,7 @@ interface UriParams {
     subscriptionId?: string | undefined;
     resourceGroup?: string | undefined;
     connectionString?: string | undefined;
+    tenantId?: string | undefined;
     database?: string | undefined;
     container?: string | undefined;
 }
@@ -574,6 +585,7 @@ function extractAndValidateParams(context: IActionContext, query: string): UriPa
             case 'connectionString':
             case 'database':
             case 'container':
+            case 'tenantId':
                 if (value !== undefined) {
                     context.valuesToMask.push(value);
                 }