diff --git a/src/docdb/getCosmosClient.ts b/src/docdb/getCosmosClient.ts index 711904fa..4a20baa2 100644 --- a/src/docdb/getCosmosClient.ts +++ b/src/docdb/getCosmosClient.ts @@ -20,6 +20,7 @@ export type CosmosDBKeyCredential = { export type CosmosDBAuthCredential = { type: 'auth'; + tenantId: string; }; export type CosmosDBCredential = CosmosDBKeyCredential | CosmosDBAuthCredential; @@ -106,7 +107,7 @@ export function getCosmosClient( ...commonProperties, aadCredentials: { getToken: async (scopes, _options) => { - const session = await getSessionFromVSCode(scopes, undefined, { createIfNone: true }); + const session = await getSessionFromVSCode(scopes, authCred.tenantId, { createIfNone: true }); return { token: session?.accessToken ?? '', expiresOnTimestamp: 0, diff --git a/src/docdb/tree/DocDBAccountTreeItemBase.ts b/src/docdb/tree/DocDBAccountTreeItemBase.ts index 3c80426c..ec45d171 100644 --- a/src/docdb/tree/DocDBAccountTreeItemBase.ts +++ b/src/docdb/tree/DocDBAccountTreeItemBase.ts @@ -26,7 +26,12 @@ import { deleteCosmosDBAccount } from '../../commands/deleteDatabaseAccount/dele import { getThemeAgnosticIconPath, SERVERLESS_CAPABILITY_NAME } from '../../constants'; import { nonNullProp } from '../../utils/nonNull'; import { rejectOnTimeout } from '../../utils/timeout'; -import { getCosmosClient, getCosmosKeyCredential, type CosmosDBCredential } from '../getCosmosClient'; +import { + getCosmosAuthCredential, + getCosmosClient, + getCosmosKeyCredential, + type CosmosDBCredential, +} from '../getCosmosClient'; import { getSignedInPrincipalIdForAccountEndpoint } from '../utils/azureSessionHelper'; import { ensureRbacPermission, isRbacException, showRbacPermissionError } from '../utils/rbacUtils'; import { DocDBTreeItemBase } from './DocDBTreeItemBase'; @@ -131,8 +136,9 @@ export abstract class DocDBAccountTreeItemBase extends DocDBTreeItemBase { - const session = await getSessionForDatabaseAccount(accountEndpoint); +export async function getSignedInPrincipalIdForAccountEndpoint( + accountEndpoint: string, + tenantId: string | undefined, +): Promise { + const session = await getSessionForDatabaseAccount(accountEndpoint, tenantId); const principalId = session?.account.id.split('/')[1] ?? session?.account.id; return principalId; } -async function getSessionForDatabaseAccount(endpoint: string): Promise { +async function getSessionForDatabaseAccount( + endpoint: string, + tenantId: string | undefined, +): Promise { const endpointUrl = new URL(endpoint); const scrope = `${endpointUrl.origin}${endpointUrl.pathname}.default`; - return await getSessionFromVSCode(scrope, undefined, { createIfNone: false }); + return await getSessionFromVSCode(scrope, tenantId, { createIfNone: false }); } diff --git a/src/tree/SubscriptionTreeItem.ts b/src/tree/SubscriptionTreeItem.ts index 9d70c58b..ebc17a57 100644 --- a/src/tree/SubscriptionTreeItem.ts +++ b/src/tree/SubscriptionTreeItem.ts @@ -252,8 +252,9 @@ export class SubscriptionTreeItem extends SubscriptionTreeItemBase { } } + const tenantId = parent.subscription.tenantId ?? databaseAccount.identity?.tenantId; // OAuth is always enabled for Cosmos DB and will be used as a fall back if key auth is unavailable - const authCred = { type: 'auth' }; + const authCred = { type: 'auth', tenantId: tenantId }; const credentials = [keyCred, authCred].filter( (cred): cred is CosmosDBCredential => cred !== undefined, );