From 780a5fbf9ecd0ed3db4402d7c8c809b4a807f9bb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mike=20Kr=C3=BCger?= Date: Fri, 16 Jan 2026 14:30:03 +0100 Subject: [PATCH] Allow attaching a cosmos db account with tenantId. Fixes #2620 --- l10n/bundle.l10n.json | 1 + .../CosmosDBConnectionStringStep.ts | 6 ++++-- .../newConnection/CosmosDBExecuteStep.ts | 6 +++++- src/cosmosdb/cosmosDBConnectionStrings.test.ts | 18 ++++++++++++++++++ src/cosmosdb/cosmosDBConnectionStrings.ts | 6 +++++- src/tree/cosmosdb/AccountInfo.ts | 5 +++-- .../cosmosdb/CosmosDBAttachedAccountModel.ts | 1 + .../cosmosdb/CosmosDBWorkspaceItem.ts | 3 +++ src/vscodeUriHandler.ts | 16 ++++++++++++++-- 9 files changed, 54 insertions(+), 8 deletions(-) diff --git a/l10n/bundle.l10n.json b/l10n/bundle.l10n.json index 1f7d378ff..cfe9592b6 100644 --- a/l10n/bundle.l10n.json +++ b/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=…\"", "Container {0} not found": "Container {0} not found", "Container name cannot be longer than 255 characters": "Container name cannot be longer than 255 characters", "Container name cannot contain the characters '\\', '/', '#', '?'": "Container name cannot contain the characters '\\', '/', '#', '?'", diff --git a/src/commands/newConnection/CosmosDBConnectionStringStep.ts b/src/commands/newConnection/CosmosDBConnectionStringStep.ts index 79b88ea96..387caac7d 100644 --- a/src/commands/newConnection/CosmosDBConnectionStringStep.ts +++ b/src/commands/newConnection/CosmosDBConnectionStringStep.ts @@ -12,7 +12,7 @@ export class CosmosDBConnectionStringStep extends AzureWizardPromptStep { 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 { // 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==', diff --git a/src/cosmosdb/cosmosDBConnectionStrings.ts b/src/cosmosdb/cosmosDBConnectionStrings.ts index 10a01d592..b7569b4c8 100644 --- a/src/cosmosdb/cosmosDBConnectionStrings.ts +++ b/src/cosmosdb/cosmosDBConnectionStrings.ts @@ -10,6 +10,7 @@ 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.')); @@ -17,7 +18,7 @@ export function parseCosmosDBConnectionString(connectionString: string): ParsedC 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 { diff --git a/src/tree/cosmosdb/AccountInfo.ts b/src/tree/cosmosdb/AccountInfo.ts index 03aeb8dcf..2902aa7c9 100644 --- a/src/tree/cosmosdb/AccountInfo.ts +++ b/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; diff --git a/src/tree/workspace-view/cosmosdb/CosmosDBAttachedAccountModel.ts b/src/tree/workspace-view/cosmosdb/CosmosDBAttachedAccountModel.ts index 563f025e7..1cc9b6584 100644 --- a/src/tree/workspace-view/cosmosdb/CosmosDBAttachedAccountModel.ts +++ b/src/tree/workspace-view/cosmosdb/CosmosDBAttachedAccountModel.ts @@ -9,4 +9,5 @@ export type CosmosDBAttachedAccountModel = { storageId: string; isEmulator: boolean; name: string; + tenantId?: string; }; diff --git a/src/tree/workspace-view/cosmosdb/CosmosDBWorkspaceItem.ts b/src/tree/workspace-view/cosmosdb/CosmosDBWorkspaceItem.ts index 2b29e211f..b0c5d4c08 100644 --- a/src/tree/workspace-view/cosmosdb/CosmosDBWorkspaceItem.ts +++ b/src/tree/workspace-view/cosmosdb/CosmosDBWorkspaceItem.ts @@ -58,6 +58,8 @@ 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 @@ -65,6 +67,7 @@ export class CosmosDBWorkspaceItem implements TreeElement, TreeElementWithContex name, connectionString, isEmulator, + tenantId, }; if (experience?.api === API.Cassandra) { diff --git a/src/vscodeUriHandler.ts b/src/vscodeUriHandler.ts index d7e922d47..79e8bc614 100644 --- a/src/vscodeUriHandler.ts +++ b/src/vscodeUriHandler.ts @@ -168,6 +168,7 @@ 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, @@ -175,6 +176,7 @@ async function handleConnectionStringRequest( 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 { 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); }