Add missing multi-tenant support.

src/commands/openNoSqlQueryEditor/openNoSqlQueryEditor.ts

@@ -5,7 +5,7 @@
 
 import { type IActionContext } from '@microsoft/vscode-azext-utils';
 import { AzExtResourceType } from '@microsoft/vscode-azureresources-api';
-import { getCosmosKeyCredential } from '../../docdb/getCosmosClient';
+import { getCosmosAuthCredential, getCosmosKeyCredential } from '../../docdb/getCosmosClient';
 import { type NoSqlQueryConnection } from '../../docdb/NoSqlCodeLensProvider';
 import { QueryEditorTab } from '../../panels/QueryEditorTab';
 import { type DocumentDBContainerResourceItem } from '../../tree/docdb/DocumentDBContainerResourceItem';
@@ -31,12 +31,14 @@ export async function openNoSqlQueryEditor(
 
     const accountInfo = node.model.accountInfo;
     const keyCred = getCosmosKeyCredential(accountInfo.credentials);
+    const tenantId = getCosmosAuthCredential(accountInfo.credentials)?.tenantId;
     const connection: NoSqlQueryConnection = {
         databaseId: node.model.database.id,
         containerId: node.model.container.id,
         endpoint: accountInfo.endpoint,
         masterKey: keyCred?.key,
         isEmulator: accountInfo.isEmulator,
+        tenantId: tenantId,
     };
 
     QueryEditorTab.render(connection);

src/docdb/NoSqlCodeLensProvider.ts

@@ -23,6 +23,7 @@ export type NoSqlQueryConnection = {
     endpoint: string;
     masterKey?: string;
     isEmulator: boolean;
+    tenantId: string | undefined;
 };
 
 export const noSqlQueryConnectionKey = 'NO_SQL_QUERY_CONNECTION_KEY.v1';

src/docdb/commands/executeNoSqlQuery.ts

@@ -37,13 +37,13 @@ export async function executeNoSqlQuery(
         );
     } else {
         // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
-        const { databaseId, containerId, endpoint, masterKey, isEmulator } =
+        const { databaseId, containerId, endpoint, masterKey, isEmulator, tenantId } =
             connectedCollection as NoSqlQueryConnection;
         const credentials: CosmosDBCredential[] = [];
         if (masterKey !== undefined) {
             credentials.push({ type: 'key', key: masterKey });
         }
-        credentials.push({ type: 'auth' });
+        credentials.push({ type: 'auth', tenantId: tenantId });
         const client = getCosmosClient(endpoint, credentials, isEmulator);
         const options = { populateQueryMetrics };
         const response = await client

src/docdb/commands/getNoSqlQueryPlan.ts

@@ -32,13 +32,13 @@ export async function getNoSqlQueryPlan(
         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 } =
+        const { databaseId, containerId, endpoint, masterKey, isEmulator, tenantId } =
             connectedCollection as NoSqlQueryConnection;
         const credentials: CosmosDBCredential[] = [];
         if (masterKey !== undefined) {
             credentials.push({ type: 'key', key: masterKey });
         }
-        credentials.push({ type: 'auth' });
+        credentials.push({ type: 'auth', tenantId: tenantId });
         const client = getCosmosClient(endpoint, credentials, isEmulator);
         const response = await client.database(databaseId).container(containerId).getQueryPlan(queryText);
         await vscodeUtil.showNewFile(

src/docdb/getCosmosClient.ts

@@ -20,6 +20,7 @@ export type CosmosDBKeyCredential = {
 
 export type CosmosDBAuthCredential = {
     type: 'auth';
+    tenantId: string | undefined;
 };
 
 export type CosmosDBCredential = CosmosDBKeyCredential | CosmosDBAuthCredential;
@@ -36,7 +37,7 @@ export function getCosmosClientByConnection(
     connection: NoSqlQueryConnection,
     options?: Partial<CosmosClientOptions>,
 ): CosmosClient {
-    const { endpoint, masterKey, isEmulator } = connection;
+    const { endpoint, masterKey, isEmulator, tenantId } = connection;
 
     const vscodeStrictSSL: boolean | undefined = vscode.workspace
         .getConfiguration()
@@ -59,7 +60,7 @@ export function getCosmosClientByConnection(
     } else {
         commonProperties.aadCredentials = {
             getToken: async (scopes, _options) => {
-                const session = await getSessionFromVSCode(scopes, undefined, { createIfNone: true });
+                const session = await getSessionFromVSCode(scopes, tenantId, { createIfNone: true });
                 return {
                     token: session?.accessToken ?? '',
                     expiresOnTimestamp: 0,
@@ -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,

src/docdb/session/DocumentSession.ts

@@ -40,12 +40,12 @@ export class DocumentSession {
     private isDisposed = false;
 
     constructor(connection: NoSqlQueryConnection, channel: Channel) {
-        const { databaseId, containerId, endpoint, masterKey, isEmulator } = connection;
+        const { databaseId, containerId, endpoint, masterKey, isEmulator, tenantId } = connection;
         const credentials: CosmosDBCredential[] = [];
         if (masterKey !== undefined) {
             credentials.push({ type: 'key', key: masterKey });
         }
-        credentials.push({ type: 'auth' });
+        credentials.push({ type: 'auth', tenantId: tenantId });
 
         this.id = uuid();
         this.channel = channel;

src/docdb/utils/NoSqlQueryConnection.ts

@@ -8,7 +8,7 @@ import { AzExtResourceType } from '@microsoft/vscode-azureresources-api';
 import { type DocumentDBContainerResourceItem } from '../../tree/docdb/DocumentDBContainerResourceItem';
 import { type DocumentDBItemsResourceItem } from '../../tree/docdb/DocumentDBItemsResourceItem';
 import { pickAppResource } from '../../utils/pickItem/pickAppResource';
-import { getCosmosKeyCredential } from '../getCosmosClient';
+import { getCosmosAuthCredential, getCosmosKeyCredential } from '../getCosmosClient';
 import { type NoSqlQueryConnection } from '../NoSqlCodeLensProvider';
 
 export function createNoSqlQueryConnection(
@@ -18,13 +18,15 @@ export function createNoSqlQueryConnection(
     const databaseId = node.model.database.id;
     const containerId = node.model.container.id;
     const keyCred = getCosmosKeyCredential(accountInfo.credentials);
+    const tenantId = getCosmosAuthCredential(accountInfo.credentials)?.tenantId;
 
     return {
         databaseId: databaseId,
         containerId: containerId,
         endpoint: accountInfo.endpoint,
         masterKey: keyCred?.key,
         isEmulator: accountInfo.isEmulator,
+        tenantId: tenantId,
     };
 }
 

src/docdb/utils/azureSessionHelper.ts

@@ -7,14 +7,20 @@
 import { getSessionFromVSCode } from '@microsoft/vscode-azext-azureauth/out/src/getSessionFromVSCode';
 import type * as vscode from 'vscode';
 
-export async function getSignedInPrincipalIdForAccountEndpoint(accountEndpoint: string): Promise<string | undefined> {
-    const session = await getSessionForDatabaseAccount(accountEndpoint);
+export async function getSignedInPrincipalIdForAccountEndpoint(
+    accountEndpoint: string,
+    tenantId: string | undefined,
+): Promise<string | undefined> {
+    const session = await getSessionForDatabaseAccount(accountEndpoint, tenantId);
     const principalId = session?.account.id.split('/')[1] ?? session?.account.id;
     return principalId;
 }
 
-async function getSessionForDatabaseAccount(endpoint: string): Promise<vscode.AuthenticationSession | undefined> {
+async function getSessionForDatabaseAccount(
+    endpoint: string,
+    tenantId: string | undefined,
+): Promise<vscode.AuthenticationSession | undefined> {
     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 });
 }

src/tree/SubscriptionTreeItem.ts

@@ -244,8 +244,9 @@ export class SubscriptionTreeItem extends SubscriptionTreeItemBase {
     //                     }
     //                 }
     //
+    //                const tenantId = parent.subscription.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,
     //                 );

src/tree/docdb/AccountInfo.ts

@@ -63,7 +63,8 @@ async function getAccountInfoForGeneric(account: CosmosAccountModel): Promise<Ac
     }
 
     const databaseAccount = await client.databaseAccounts.get(resourceGroup, name);
-    const credentials = await getCredentialsForGeneric(name, resourceGroup, client, databaseAccount);
+    const tenantId = account?.subscription?.tenantId;
+    const credentials = await getCredentialsForGeneric(name, resourceGroup, tenantId, client, databaseAccount);
     const documentEndpoint = nonNullProp(databaseAccount, 'documentEndpoint', `of the database account ${id}`);
     const isServerless = databaseAccount?.capabilities
         ? databaseAccount.capabilities.some((cap) => cap.name === SERVERLESS_CAPABILITY_NAME)
@@ -82,6 +83,7 @@ async function getAccountInfoForGeneric(account: CosmosAccountModel): Promise<Ac
 async function getCredentialsForGeneric(
     name: string,
     resourceGroup: string,
+    tenantId: string,
     client: CosmosDBManagementClient,
     databaseAccount: DatabaseAccountGetResults,
 ): Promise<CosmosDBCredential[]> {
@@ -133,9 +135,9 @@ async function getCredentialsForGeneric(
             }
         }
 
-        // OAuth is always enabled for Cosmos DB and will be used as a fallback if key auth is unavailable
-        const authCred = { type: 'auth' };
-        return [keyCred, authCred].filter((cred): cred is CosmosDBCredential => cred !== undefined);
+        // OAuth is always enabled for Cosmos DB and will be used as a fall back if key auth is unavailable
+        const authCred = { type: 'auth', tenantId: tenantId };
+        return [keyCred, authCred].filter((cred) => cred !== undefined) as CosmosDBCredential[];
     });
 
     return result ?? [];
@@ -217,8 +219,9 @@ async function getCredentialsForAttached(account: CosmosDBAttachedAccountModel):
         }
 
         // OAuth is always enabled for Cosmos DB and will be used as a fallback if key auth is unavailable
-        const authCred = { type: 'auth' };
-        return [keyCred, authCred].filter((cred): cred is CosmosDBCredential => cred !== undefined);
+        // TODO: we need to preserve the tenantId in the connection string, otherwise we can't use OAuth for foreign tenants
+        const authCred = { type: 'auth', tenantId: undefined };
+        return [keyCred, authCred].filter((cred) => cred !== undefined) as CosmosDBCredential[];
     });
 
     return result ?? [];

src/tree/docdb/DocumentDBAccountAttachedResourceItem.ts

@@ -7,7 +7,7 @@ import { type CosmosClient, type DatabaseDefinition, type Resource } from '@azur
 import { type TreeItem } from 'vscode';
 import { type Experience } from '../../AzureDBExperiences';
 import { getThemeAgnosticIconPath } from '../../constants';
-import { getCosmosClient } from '../../docdb/getCosmosClient';
+import { getCosmosAuthCredential, getCosmosClient } from '../../docdb/getCosmosClient';
 import { getSignedInPrincipalIdForAccountEndpoint } from '../../docdb/utils/azureSessionHelper';
 import { isRbacException, showRbacPermissionError } from '../../docdb/utils/rbacUtils';
 import { rejectOnTimeout } from '../../utils/timeout';
@@ -65,8 +65,9 @@ export abstract class DocumentDBAccountAttachedResourceItem extends CosmosDBAcco
         } catch (e) {
             if (e instanceof Error && isRbacException(e) && !this.hasShownRbacNotification) {
                 this.hasShownRbacNotification = true;
-
-                const principalId = (await getSignedInPrincipalIdForAccountEndpoint(accountInfo.endpoint)) ?? '';
+                const tenantId = getCosmosAuthCredential(accountInfo.credentials)?.tenantId;
+                const principalId =
+                    (await getSignedInPrincipalIdForAccountEndpoint(accountInfo.endpoint, tenantId)) ?? '';
                 void showRbacPermissionError(this.id, principalId);
             }
             throw e; // rethrowing tells the resources extension to show the exception message in the tree

src/tree/docdb/DocumentDBAccountResourceItem.ts

@@ -7,7 +7,7 @@ import { type CosmosClient, type DatabaseDefinition, type Resource } from '@azur
 import { type TreeItem } from 'vscode';
 import { type Experience } from '../../AzureDBExperiences';
 import { getThemeAgnosticIconPath } from '../../constants';
-import { getCosmosClient } from '../../docdb/getCosmosClient';
+import { getCosmosAuthCredential, getCosmosClient } from '../../docdb/getCosmosClient';
 import { getSignedInPrincipalIdForAccountEndpoint } from '../../docdb/utils/azureSessionHelper';
 import { ensureRbacPermissionV2, isRbacException, showRbacPermissionError } from '../../docdb/utils/rbacUtils';
 import { type CosmosAccountModel } from '../CosmosAccountModel';
@@ -65,7 +65,9 @@ export abstract class DocumentDBAccountResourceItem extends CosmosDBAccountResou
             if (e instanceof Error && isRbacException(e) && !this.hasShownRbacNotification) {
                 this.hasShownRbacNotification = true;
 
-                const principalId = (await getSignedInPrincipalIdForAccountEndpoint(accountInfo.endpoint)) ?? '';
+                const tenantId = getCosmosAuthCredential(accountInfo.credentials)?.tenantId;
+                const principalId =
+                    (await getSignedInPrincipalIdForAccountEndpoint(accountInfo.endpoint, tenantId)) ?? '';
                 // check if the principal ID matches the one that is signed in,
                 // otherwise this might be a security problem, hence show the error message
                 if (