feat: Add Cosmos DB Shell (NoSQL) support to Cloud Shell (#2549)

* Add Cosmos DB Shell (NoSQL) support to Cloud Shell

Adds Cosmos DB (NoSQL) support to the Cloud Shell experience and wires up authentication for RBAC / local-auth-disabled accounts:

- Enable the Open Cosmos DB Shell button behind the enableCloudShell feature flag.
- Bootstrap .NET SDK 10 in Cloud Shell before installing the CosmosDBShell tool.
- Force gateway connection mode and add --verbose to surface connection errors.
- Select Entra ID vs account-key auth via isCloudShellEntraAuthEnabled (covers dataplane RBAC and disableLocalAuth accounts).
- Acquire a data-plane-scoped Entra token on demand when no cached aadToken exists, and omit --connect-tenant so a missing token falls through to DefaultAzureCredential (Cloud Shell az session) instead of interactive/device-code auth.

* Fix popup_window_error in CosmosDB Cloud Shell Entra auth

Acquiring an Entra token on demand via acquireMsalTokenForAccount triggered a browser loginPopup (even in silent mode when no MSAL account was cached), which fails with popup_window_error inside the hosted Cloud Shell context. getKey now returns the cached userContext.aadToken when present, or an empty string otherwise. With no token exported, the CosmosDBShell tool falls through to DefaultAzureCredential, which uses the Cloud Shell's signed-in az session.

* Force --connect-azure-cli whenever no credential env var is exported

In Azure Cloud Shell, DefaultAzureCredential tries ManagedIdentityCredential first, which cannot mint a token for the *.documents.azure.com audience (AudienceNotSupported). Force AzureCliCredential via --connect-azure-cli whenever no key/token env var is exported, covering both the Entra-no-token and key-auth-empty-key cases. Also auto-update the CosmosDBShell tool so cached installs pick up the new flag.

* Silently mint a Cosmos data-plane token for the Cloud Shell Entra path

A disableLocalAuth Cosmos account whose Data Explorer session is still in key mode has no cached userContext.aadToken, and the ephemeral Cloud Shell can obtain a Cosmos token from neither its managed identity (AudienceNotSupported for *.documents.azure.com) nor its az session (not logged in). Mint a Cosmos-scoped token in the browser and export it via COSMOSDB_SHELL_TOKEN. The acquisition is guarded on an existing cached MSAL account so it can never trigger an interactive popup, and any failure returns an empty string so the tool falls back to --connect-azure-cli.

* Only accept Y/N at the Cloud Shell consent prompt

askConfirmation now ignores any key that is not Y or N instead of treating every non-Y key as a decline, so an accidental keypress no longer aborts the Cloud Shell consent flow. Also add diagnostic warnings on the Cloud Shell token path so the browser console reveals whether the silent Cosmos token mint was skipped (no cached MSAL account) or returned empty.

* Use device-code auth when Cloud Shell has no token

* Pass the Data Explorer credential to the Cosmos DB shell instead of signing in interactively

Azure Cloud Shell cannot authenticate to Cosmos DB on its own: its managed identity
is rejected with AudienceNotSupported for the *.documents.azure.com audience, its az
session is not signed in, and neither the interactive browser nor the device-code flow
is usable from the embedded terminal.

Resolve the credential in Data Explorer and hand it to the shell out-of-band:
- getCosmosDBShellCredential() returns an Entra ID data-plane token (cached aadToken,
  or one minted silently only when an MSAL account already exists so it can never
  trigger a popup), falling back to the account master key unless local auth is
  disabled. It reports which kind it resolved so the correct env var is exported.
- CosmosDBShellHandler exports COSMOSDB_SHELL_TOKEN or COSMOSDB_SHELL_ACCOUNT_KEY and
  drops every credential flag, so the tool always lands on a terminal, non-interactive
  step of its credential chain.
- When nothing can be resolved, print actionable guidance instead of launching the
  tool with no credential.

* Fix TS7011 in ShellTypeFactory tests by typing the mocked getAllAccounts

The empty array literal in the getMsalInstance mock had no contextual type, so
tsc inferred an implicit any[] return under noImplicitAny. Route every mock
through a typed mockMsalAccounts helper.

* Reuse DE's cached credential for Cosmos DB Shell; add read-only key fallback

* Deliver Cosmos DB Shell key as a full connection string; harden key resolution

* Surface the specific reason a Cosmos DB Shell credential could not be resolved

* Deliver Cosmos DB Shell credential inline (export+connect on one line), mirroring Mongo handler

* Log which Cosmos DB Shell credential kind was resolved for debugging

* Increase ARM timeout for Cloud Shell provisioning calls to avoid spurious abort errors

* Gate Cosmos DB Shell button behind enableCosmosDBShell feature flag

* Fix TS2741: add missing enableCosmosDBShell to Features test fixture
This commit is contained in:
Mike Krüger
2026-08-04 10:04:27 +02:00
committed by GitHub
parent 5b3e066b23
commit 43fe8a002b
14 changed files with 835 additions and 18 deletions
+1
View File
@@ -406,6 +406,7 @@ export enum TerminalKind {
Cassandra = 2, Cassandra = 2,
Postgres = 3, Postgres = 3,
VCoreMongo = 4, VCoreMongo = 4,
CosmosDB = 5,
} }
export interface DataExplorerInputsFrame { export interface DataExplorerInputsFrame {
+4
View File
@@ -666,6 +666,10 @@ export default class Explorer {
title = "Mongo Shell"; title = "Mongo Shell";
break; break;
case ViewModels.TerminalKind.CosmosDB:
title = "Cosmos DB Shell";
break;
default: default:
throw new Error("Terminal kind: ${kind} not supported"); throw new Error("Terminal kind: ${kind} not supported");
} }
@@ -17,11 +17,11 @@ import SynapseIcon from "../../../../images/synapse-link.svg";
import VSCodeIcon from "../../../../images/vscode.svg"; import VSCodeIcon from "../../../../images/vscode.svg";
import { AuthType } from "../../../AuthType"; import { AuthType } from "../../../AuthType";
import * as Constants from "../../../Common/Constants"; import * as Constants from "../../../Common/Constants";
import { Platform, configContext } from "../../../ConfigContext"; import { configContext, Platform } from "../../../ConfigContext";
import * as ViewModels from "../../../Contracts/ViewModels"; import * as ViewModels from "../../../Contracts/ViewModels";
import { import {
userContext,
isVCoreMongoNativeAuthDisabled, isVCoreMongoNativeAuthDisabled,
userContext,
VCoreMongoNativeAuthDisabledMessage, VCoreMongoNativeAuthDisabledMessage,
VCoreMongoNativeAuthLearnMoreUrl, VCoreMongoNativeAuthLearnMoreUrl,
} from "../../../UserContext"; } from "../../../UserContext";
@@ -81,6 +81,15 @@ export function createStaticCommandBarButtons(
} }
} }
if (
userContext.apiType === "SQL" &&
userContext.features.enableCloudShell &&
userContext.features.enableCosmosDBShell
) {
addDivider();
buttons.push(createOpenTerminalButtonByKind(container, ViewModels.TerminalKind.CosmosDB));
}
if (!selectedNodeState.isDatabaseNodeOrNoneSelected()) { if (!selectedNodeState.isDatabaseNodeOrNoneSelected()) {
const isQuerySupported = userContext.apiType === "SQL" || userContext.apiType === "Gremlin"; const isQuerySupported = userContext.apiType === "SQL" || userContext.apiType === "Gremlin";
@@ -498,6 +507,8 @@ function createOpenTerminalButtonByKind(
return "PSQL"; return "PSQL";
case ViewModels.TerminalKind.VCoreMongo: case ViewModels.TerminalKind.VCoreMongo:
return "MongoDB (DocumentDB)"; return "MongoDB (DocumentDB)";
case ViewModels.TerminalKind.CosmosDB:
return "Cosmos DB";
default: default:
return ""; return "";
} }
@@ -507,7 +518,9 @@ function createOpenTerminalButtonByKind(
"This feature is not yet available in your account's region. View supported regions here: https://aka.ms/cosmos-enable-notebooks."; "This feature is not yet available in your account's region. View supported regions here: https://aka.ms/cosmos-enable-notebooks.";
const isNativeAuthDisabled = terminalKind === ViewModels.TerminalKind.VCoreMongo && isVCoreMongoNativeAuthDisabled(); const isNativeAuthDisabled = terminalKind === ViewModels.TerminalKind.VCoreMongo && isVCoreMongoNativeAuthDisabled();
const disableButton = const disableButton =
(!useNotebook.getState().isNotebooksEnabledForAccount && !useNotebook.getState().isNotebookEnabled) || (!useNotebook.getState().isNotebooksEnabledForAccount &&
!useNotebook.getState().isNotebookEnabled &&
!userContext.features.enableCloudShell) ||
isNativeAuthDisabled; isNativeAuthDisabled;
return { return {
iconSrc: HostedTerminalIcon, iconSrc: HostedTerminalIcon,
@@ -127,6 +127,7 @@ describe("CloudShellClient", () => {
vnetSettings: {}, vnetSettings: {},
}, },
}, },
timeoutMs: 30000,
}); });
expect(result).toEqual(mockResponse); expect(result).toEqual(mockResponse);
}); });
@@ -154,6 +155,7 @@ describe("CloudShellClient", () => {
vnetSettings: mockVNetSettings, vnetSettings: mockVNetSettings,
}, },
}, },
timeoutMs: 30000,
}); });
}); });
@@ -212,6 +214,7 @@ describe("CloudShellClient", () => {
path: "/subscriptions/sub-id/providers/Microsoft.CloudShell/register", path: "/subscriptions/sub-id/providers/Microsoft.CloudShell/register",
method: "POST", method: "POST",
apiVersion: "2022-12-01", apiVersion: "2022-12-01",
timeoutMs: 30000,
}); });
expect(result).toEqual(mockResponse); expect(result).toEqual(mockResponse);
}); });
@@ -227,6 +230,7 @@ describe("CloudShellClient", () => {
path: "/subscriptions/sub-id/providers/Microsoft.CloudShell/register", path: "/subscriptions/sub-id/providers/Microsoft.CloudShell/register",
method: "POST", method: "POST",
apiVersion: "2022-12-01", apiVersion: "2022-12-01",
timeoutMs: 30000,
}); });
}); });
}); });
@@ -251,6 +255,7 @@ describe("CloudShellClient", () => {
osType: OsType.Linux, osType: OsType.Linux,
}, },
}, },
timeoutMs: 30000,
}); });
expect(result).toEqual(mockResponse); expect(result).toEqual(mockResponse);
}); });
@@ -274,6 +279,7 @@ describe("CloudShellClient", () => {
osType: OsType.Linux, osType: OsType.Linux,
}, },
}, },
timeoutMs: 30000,
}); });
}); });
}); });
@@ -14,6 +14,13 @@ import {
} from "../Models/DataModels"; } from "../Models/DataModels";
import { getLocale } from "../Utils/CommonUtils"; import { getLocale } from "../Utils/CommonUtils";
// armRequest defaults to a 5s timeout with no retry for PUT/POST. Cloud Shell provisioning
// (registering the provider, applying settings, and provisioning the console itself) can
// legitimately take much longer than that on first use or when the console region differs
// from the user's usual region, so a much more generous timeout is used for these calls to
// avoid a spurious AbortError ("User aborted query.") aborting the whole session start.
const CLOUDSHELL_ARM_TIMEOUT_MS = 30000;
export const getUserSettings = async (): Promise<CloudShellSettings> => { export const getUserSettings = async (): Promise<CloudShellSettings> => {
return await armRequest<CloudShellSettings>({ return await armRequest<CloudShellSettings>({
host: configContext.ARM_ENDPOINT, host: configContext.ARM_ENDPOINT,
@@ -51,6 +58,7 @@ export const putEphemeralUserSettings = async (
method: "PUT", method: "PUT",
apiVersion: "2023-02-01-preview", apiVersion: "2023-02-01-preview",
body: ephemeralSettings, body: ephemeralSettings,
timeoutMs: CLOUDSHELL_ARM_TIMEOUT_MS,
}); });
}; };
@@ -69,6 +77,7 @@ export const registerCloudShellProvider = async (subscriptionId: string) => {
path: `/subscriptions/${subscriptionId}/providers/Microsoft.CloudShell/register`, path: `/subscriptions/${subscriptionId}/providers/Microsoft.CloudShell/register`,
method: "POST", method: "POST",
apiVersion: "2022-12-01", apiVersion: "2022-12-01",
timeoutMs: CLOUDSHELL_ARM_TIMEOUT_MS,
}); });
}; };
@@ -88,6 +97,7 @@ export const provisionConsole = async (consoleLocation: string): Promise<Provisi
"x-ms-console-preferred-location": consoleLocation, "x-ms-console-preferred-location": consoleLocation,
}, },
body: data, body: data,
timeoutMs: CLOUDSHELL_ARM_TIMEOUT_MS,
}); });
}; };
@@ -0,0 +1,122 @@
import { CosmosDBShellHandler } from "./CosmosDBShellHandler";
// Mock dependencies
jest.mock("../../../../UserContext", () => ({
userContext: {
databaseAccount: {
properties: {
documentEndpoint: "https://test-account.documents.azure.com:443/",
},
},
},
}));
describe("CosmosDBShellHandler", () => {
const mockKey = "testKey";
const endpoint = "https://test-account.documents.azure.com:443/";
const tokenConnectionCommand = `export COSMOSDB_SHELL_TOKEN='aadToken123'; cosmosdbshell --connect '${endpoint}' --connect-mode gateway --verbose`;
const keyConnectionCommand = `export COSMOSDB_SHELL_ACCOUNT_KEY='${mockKey}'; cosmosdbshell --connect '${endpoint}' --connect-mode gateway --verbose`;
let cosmosDBShellHandler: CosmosDBShellHandler;
beforeEach(() => {
cosmosDBShellHandler = new CosmosDBShellHandler({ kind: "key", value: mockKey });
jest.clearAllMocks();
});
afterEach(() => {
jest.clearAllMocks();
});
afterAll(() => {
jest.resetAllMocks();
jest.restoreAllMocks();
jest.resetModules();
});
describe("Positive Tests", () => {
it("should return correct shell name", () => {
expect(cosmosDBShellHandler.getShellName()).toBe("Cosmos DB");
});
it("should install the CosmosDBShell tool in setup commands", () => {
const commands = cosmosDBShellHandler.getSetUpCommands();
expect(Array.isArray(commands)).toBe(true);
expect(commands.some((c) => c.includes("dotnet tool install --global CosmosDBShell --prerelease"))).toBe(true);
expect(commands.some((c) => c.includes("dotnet tool update --global CosmosDBShell --prerelease"))).toBe(true);
expect(commands.some((c) => c.includes("$HOME/.dotnet/tools"))).toBe(true);
});
it("should bootstrap a .NET SDK 10 when the tool and SDK are missing", () => {
const commands = cosmosDBShellHandler.getSetUpCommands();
expect(commands.some((c) => c.includes("dotnet-install.sh") && c.includes("--channel 10.0"))).toBe(true);
expect(commands.some((c) => c === "export DOTNET_ROOT=$HOME/.dotnet")).toBe(true);
});
it("should not export any credential env var in setup commands for a key credential", () => {
const commands = cosmosDBShellHandler.getSetUpCommands();
expect(commands.some((c) => c.includes("COSMOSDB_SHELL_"))).toBe(false);
});
it("should not export any credential env var in setup commands for a token credential", () => {
const handler = new CosmosDBShellHandler({ kind: "token", value: "aadToken123" });
const commands = handler.getSetUpCommands();
expect(commands.some((c) => c.includes("COSMOSDB_SHELL_"))).toBe(false);
});
it("should export the account key and connect to the bare endpoint on the same command line", () => {
expect(cosmosDBShellHandler.getConnectionCommand()).toBe(keyConnectionCommand);
});
it("should export the token and connect to the bare endpoint on the same command line", () => {
const handler = new CosmosDBShellHandler({ kind: "token", value: "aadToken123" });
expect(handler.getConnectionCommand()).toBe(tokenConnectionCommand);
});
it("should never pass an interactive or ambient credential flag", () => {
const handler = new CosmosDBShellHandler({ kind: "token", value: "aadToken123" });
const connectionCommand = handler.getConnectionCommand();
expect(connectionCommand).not.toContain("--connect-tenant");
expect(connectionCommand).not.toContain("--connect-hint");
expect(connectionCommand).not.toContain("--connect-authority-host");
expect(connectionCommand).not.toContain("--connect-azure-cli");
expect(connectionCommand).not.toContain("--connect-managed-identity");
expect(connectionCommand).not.toContain("--connect-vscode-credential");
});
it("should return empty array for terminal suppressed data", () => {
expect(cosmosDBShellHandler.getTerminalSuppressedData()).toEqual([]);
});
});
describe("Negative Tests", () => {
it("should not export a credential env var when no credential was resolved", () => {
const handler = new CosmosDBShellHandler(undefined);
const commands = handler.getSetUpCommands();
expect(commands.some((c) => c.includes("COSMOSDB_SHELL_"))).toBe(false);
});
it("should not launch the shell when no credential was resolved", () => {
const handler = new CosmosDBShellHandler(undefined);
const connectionCommand = handler.getConnectionCommand();
expect(connectionCommand).not.toContain("cosmosdbshell --connect");
expect(connectionCommand).toContain("Unable to acquire a Cosmos DB credential");
expect(connectionCommand).toContain("Login for Entra ID");
});
it("should echo the specific reason when one is provided", () => {
const handler = new CosmosDBShellHandler(undefined, "listing the account keys failed (Forbidden)");
const connectionCommand = handler.getConnectionCommand();
expect(connectionCommand).toContain("Unable to acquire a Cosmos DB credential");
expect(connectionCommand).toContain("listing the account keys failed (Forbidden)");
});
});
});
@@ -0,0 +1,129 @@
import { userContext } from "../../../../UserContext";
import { AbstractShellHandler } from "./AbstractShellHandler";
/**
* A credential resolved by Data Explorer and handed to the Cosmos DB Shell.
*
* Both kinds are delivered the same way the Mongo shell delivers its credential: exported
* as an environment variable on the same command line as the `cosmosdbshell` invocation
* (`export VAR='...'; cosmosdbshell --connect <endpoint>`), immediately before the tool
* reads it, rather than as a separate setup step or embedded in a connection string.
*
* - `token` is an Entra ID (AAD) bearer token scoped to the account's data plane, exported
* as COSMOSDB_SHELL_TOKEN.
* - `key` is an account master (or read-only) key, exported as COSMOSDB_SHELL_ACCOUNT_KEY.
*
* The kind is tracked explicitly so the correct environment variable is always used.
*/
export interface CosmosDBShellCredential {
kind: "token" | "key";
value: string;
}
/**
* Shell handler for the Azure Cosmos DB Shell (https://github.com/Azure/CosmosDBShell),
* a .NET global tool that targets the Cosmos DB NoSQL (SQL Core) API.
*/
export class CosmosDBShellHandler extends AbstractShellHandler {
private _endpoint: string | undefined;
constructor(
private credential: CosmosDBShellCredential | undefined,
private unavailableReason?: string,
) {
super();
this._endpoint = userContext?.databaseAccount?.properties?.documentEndpoint;
}
public getShellName(): string {
return "Cosmos DB";
}
/**
* Setup commands for the Cosmos DB Shell:
*
* 1. Put the private .NET install dir and global tools dir on PATH (and set
* DOTNET_ROOT) so both the `dotnet` host and installed tool resolve.
* 2. Bootstrap a private .NET SDK 10 into $HOME/.dotnet when neither the
* cosmosdbshell tool nor a suitable SDK is already available. The
* CosmosDBShell global tool targets net10.0, so `dotnet tool install`
* requires the .NET SDK 10.0+, which the Azure Cloud Shell host does not
* ship by default.
* 3. Install the CosmosDBShell global tool if it is not already available, or
* update it to the latest release when it is (so an older cached install in
* the persistent Cloud Shell $HOME picks up newer connect options).
* 4. Persist the PATH/DOTNET_ROOT changes for future sessions.
*
* The credential itself is not exported here: it travels on the same command line as
* the `cosmosdbshell` invocation (see getConnectionCommand), mirroring how the Mongo
* shell handler builds its connection command.
*
* Installation steps run conditionally only if cosmosdbshell is not already
* present in the environment.
*/
public getSetUpCommands(): string[] {
return [
"export DOTNET_ROOT=$HOME/.dotnet",
"export PATH=$HOME/.dotnet:$HOME/.dotnet/tools:$PATH",
"if ! command -v cosmosdbshell &> /dev/null; then echo '⚠️ cosmosdbshell not found. Installing .NET SDK 10 and CosmosDBShell...'; fi",
"if ! command -v cosmosdbshell &> /dev/null && ! dotnet --list-sdks 2>/dev/null | grep -q '^10\\.'; then curl -sSL https://dot.net/v1/dotnet-install.sh | bash -s -- --channel 10.0 --install-dir $HOME/.dotnet; fi",
"if ! command -v cosmosdbshell &> /dev/null; then dotnet tool install --global CosmosDBShell --prerelease; else dotnet tool update --global CosmosDBShell --prerelease; fi",
"grep -qxF 'export DOTNET_ROOT=$HOME/.dotnet' ~/.bashrc || echo 'export DOTNET_ROOT=$HOME/.dotnet' >> ~/.bashrc",
"grep -qxF 'export PATH=$HOME/.dotnet:$HOME/.dotnet/tools:$PATH' ~/.bashrc || echo 'export PATH=$HOME/.dotnet:$HOME/.dotnet/tools:$PATH' >> ~/.bashrc",
];
}
private _getKeyConnectionCommand(key: string): string {
// Export the key immediately before invoking the tool, on the same command line, so it
// never appears as a --connect flag value or lands in a separate setup step. The tool
// reads the account key from COSMOSDB_SHELL_ACCOUNT_KEY and connects to the bare endpoint.
return `export COSMOSDB_SHELL_ACCOUNT_KEY='${key}'; cosmosdbshell --connect '${this._endpoint}' --connect-mode gateway --verbose`;
}
private _getTokenConnectionCommand(token: string): string {
// Same pattern for the Entra ID token: exported right before the invocation so it never
// appears in argv/ps, then the tool reads it from COSMOSDB_SHELL_TOKEN.
return `export COSMOSDB_SHELL_TOKEN='${token}'; cosmosdbshell --connect '${this._endpoint}' --connect-mode gateway --verbose`;
}
public getConnectionCommand(): string {
if (!this._endpoint) {
return `echo '${this.getShellName()} endpoint not found.'`;
}
if (!this.credential) {
// Never let the tool continue without a credential. With no account key and no
// COSMOSDB_SHELL_TOKEN exported, its credential chain would fall through to
// DefaultAzureCredential, which in Azure Cloud Shell tries the managed identity
// first and fails with "AudienceNotSupported" (the Cloud Shell MSI cannot mint a
// token for the *.documents.azure.com audience). The interactive browser and
// device-code flows are not usable from this embedded terminal either, so fail
// fast with actionable guidance instead.
//
// There is no way to inspect the browser dev tools console from inside this remote
// terminal, so the specific reason (when known) is echoed here too.
const reasonSuffix = this.unavailableReason ? `${this.unavailableReason}` : "";
return `echo 'Unable to acquire a ${this.getShellName()} credential${reasonSuffix}. Use "Login for Entra ID" in the Data Explorer toolbar, verify you have Cosmos DB data-plane RBAC access to this account, then reopen the shell.'`;
}
// Force gateway (HTTPS/443) connection mode. The shell otherwise defaults to
// direct (TCP) mode for real accounts, and Azure Cloud Shell blocks the
// direct-mode TCP ports, causing the connection to fail. `--verbose` surfaces
// the full exception details when a connection attempt fails.
//
// Surface which credential kind Explorer resolved in the browser console. A wrong
// guess here (e.g. a token when the account actually needs a key) otherwise fails
// silently inside the remote shell with no client-side signal to debug against.
console.warn(`CloudShell: connecting to Cosmos DB shell with a "${this.credential.kind}" credential`, {
endpoint: this._endpoint,
});
return this.credential.kind === "key"
? this._getKeyConnectionCommand(this.credential.value)
: this._getTokenConnectionCommand(this.credential.value);
}
public getTerminalSuppressedData(): string[] {
return [];
}
}
@@ -1,19 +1,23 @@
import { TerminalKind } from "../../../../Contracts/ViewModels"; import { TerminalKind } from "../../../../Contracts/ViewModels";
import { userContext } from "../../../../UserContext"; import { userContext } from "../../../../UserContext";
import { listKeys } from "../../../../Utils/arm/generatedClients/cosmos/databaseAccounts"; import { getReadOnlyKeys, listKeys } from "../../../../Utils/arm/generatedClients/cosmos/databaseAccounts";
import { acquireMsalTokenForAccount, getMsalInstance } from "../../../../Utils/AuthorizationUtils";
import { CassandraShellHandler } from "./CassandraShellHandler"; import { CassandraShellHandler } from "./CassandraShellHandler";
import { CosmosDBShellHandler } from "./CosmosDBShellHandler";
import { MongoShellHandler } from "./MongoShellHandler"; import { MongoShellHandler } from "./MongoShellHandler";
import { PostgresShellHandler } from "./PostgresShellHandler"; import { PostgresShellHandler } from "./PostgresShellHandler";
import { getHandler, getKey } from "./ShellTypeFactory"; import type { CosmosDBShellCredentialDiagnostics } from "./ShellTypeFactory";
import { getCosmosDBShellCredential, getHandler, getKey } from "./ShellTypeFactory";
import { VCoreMongoShellHandler } from "./VCoreMongoShellHandler"; import { VCoreMongoShellHandler } from "./VCoreMongoShellHandler";
interface UserContextType { interface UserContextType {
databaseAccount: { name: string }; databaseAccount: { name: string; properties?: { disableLocalAuth?: boolean } };
subscriptionId: string; subscriptionId: string;
resourceGroup: string; resourceGroup: string;
features: { enableAadDataPlane: boolean }; features: { enableAadDataPlane: boolean };
dataPlaneRbacEnabled: boolean; dataPlaneRbacEnabled: boolean;
aadToken?: string; aadToken?: string;
masterKey?: string;
apiType?: string; apiType?: string;
} }
@@ -30,13 +34,29 @@ jest.mock("../../../../UserContext", () => ({
jest.mock("../../../../Utils/arm/generatedClients/cosmos/databaseAccounts", () => ({ jest.mock("../../../../Utils/arm/generatedClients/cosmos/databaseAccounts", () => ({
listKeys: jest.fn(), listKeys: jest.fn(),
getReadOnlyKeys: jest.fn(),
})); }));
jest.mock("../../../../Utils/AuthorizationUtils", () => {
const actual = jest.requireActual("../../../../Utils/AuthorizationUtils");
return {
...actual,
getMsalInstance: jest.fn(),
acquireMsalTokenForAccount: jest.fn(),
};
});
describe("ShellTypeHandlerFactory", () => { describe("ShellTypeHandlerFactory", () => {
const mockKey = "testKey"; const mockKey = "testKey";
const mockMsalAccounts = (accounts: { username: string }[]): void => {
(getMsalInstance as jest.Mock).mockResolvedValue({ getAllAccounts: (): { username: string }[] => accounts });
};
beforeEach(() => { beforeEach(() => {
(listKeys as jest.Mock).mockResolvedValue({ primaryMasterKey: mockKey }); (listKeys as jest.Mock).mockResolvedValue({ primaryMasterKey: mockKey });
mockMsalAccounts([]);
(acquireMsalTokenForAccount as jest.Mock).mockResolvedValue("");
}); });
afterEach(() => { afterEach(() => {
@@ -69,7 +89,7 @@ describe("ShellTypeHandlerFactory", () => {
type DatabaseAccountType = { name: string }; type DatabaseAccountType = { name: string };
(userContext.databaseAccount as DatabaseAccountType).name = ""; (userContext.databaseAccount as DatabaseAccountType).name = "";
const key = await getKey(); const key = await getKey(false);
expect(key).toBe(""); expect(key).toBe("");
expect(listKeys).not.toHaveBeenCalled(); expect(listKeys).not.toHaveBeenCalled();
@@ -80,7 +100,7 @@ describe("ShellTypeHandlerFactory", () => {
it("should return empty string when listKeys returns null", async () => { it("should return empty string when listKeys returns null", async () => {
(listKeys as jest.Mock).mockResolvedValue(null); (listKeys as jest.Mock).mockResolvedValue(null);
const key = await getKey(); const key = await getKey(false);
expect(key).toBe(""); expect(key).toBe("");
}); });
@@ -89,7 +109,7 @@ describe("ShellTypeHandlerFactory", () => {
/* no primaryMasterKey */ /* no primaryMasterKey */
}); });
const key = await getKey(); const key = await getKey(false);
expect(key).toBe(""); expect(key).toBe("");
}); });
}); });
@@ -116,12 +136,57 @@ describe("ShellTypeHandlerFactory", () => {
expect(handler).toBeInstanceOf(CassandraShellHandler); expect(handler).toBeInstanceOf(CassandraShellHandler);
}); });
it("should return CosmosDBShellHandler with key for CosmosDB terminal kind", async () => {
const handler = await getHandler(TerminalKind.CosmosDB);
expect(handler).toBeInstanceOf(CosmosDBShellHandler);
});
it("should get key successfully when database name exists", async () => { it("should get key successfully when database name exists", async () => {
const key = await getKey(); const key = await getKey(false);
expect(key).toBe(mockKey); expect(key).toBe(mockKey);
expect(listKeys).toHaveBeenCalledWith("testSubId", "testResourceGroup", "testDbName"); expect(listKeys).toHaveBeenCalledWith("testSubId", "testResourceGroup", "testDbName");
}); });
it("should return the aadToken without listing keys when Entra ID auth is requested", async () => {
(userContext as UserContextType).aadToken = "aadToken123";
const key = await getKey(true);
expect(key).toBe("aadToken123");
expect(listKeys).not.toHaveBeenCalled();
});
it("should return an empty string when Entra ID auth is requested but no cached token or account exists", async () => {
(userContext as UserContextType).aadToken = undefined;
mockMsalAccounts([]);
const key = await getKey(true);
expect(key).toBe("");
expect(acquireMsalTokenForAccount).not.toHaveBeenCalled();
expect(listKeys).not.toHaveBeenCalled();
});
it("should silently mint a Cosmos token when Entra ID auth is requested and an MSAL account is cached", async () => {
(userContext as UserContextType).aadToken = undefined;
mockMsalAccounts([{ username: "user@contoso.com" }]);
(acquireMsalTokenForAccount as jest.Mock).mockResolvedValue("mintedToken123");
const key = await getKey(true);
expect(key).toBe("mintedToken123");
expect(acquireMsalTokenForAccount).toHaveBeenCalled();
expect(listKeys).not.toHaveBeenCalled();
});
it("should return an empty string when the silent token acquisition fails", async () => {
(userContext as UserContextType).aadToken = undefined;
mockMsalAccounts([{ username: "user@contoso.com" }]);
(acquireMsalTokenForAccount as jest.Mock).mockRejectedValue(new Error("interaction_required"));
jest.spyOn(console, "error").mockImplementation(() => undefined);
const key = await getKey(true);
expect(key).toBe("");
expect(listKeys).not.toHaveBeenCalled();
});
it("should return MongoShellHandler with primaryMasterKey for TerminalKind.Mongo when RBAC is disabled", async () => { it("should return MongoShellHandler with primaryMasterKey for TerminalKind.Mongo when RBAC is disabled", async () => {
(listKeys as jest.Mock).mockResolvedValue({ primaryMasterKey: "primaryKey123" }); (listKeys as jest.Mock).mockResolvedValue({ primaryMasterKey: "primaryKey123" });
(userContext as UserContextType).features.enableAadDataPlane = false; (userContext as UserContextType).features.enableAadDataPlane = false;
@@ -150,4 +215,148 @@ describe("ShellTypeHandlerFactory", () => {
); );
}); });
}); });
describe("getCosmosDBShellCredential", () => {
beforeEach(() => {
(userContext as UserContextType).aadToken = undefined;
(userContext as UserContextType).masterKey = undefined;
(userContext as UserContextType).apiType = "SQL";
(userContext as UserContextType).features.enableAadDataPlane = false;
(userContext as UserContextType).dataPlaneRbacEnabled = false;
(userContext as UserContextType).databaseAccount.properties = { disableLocalAuth: false };
});
it("should reuse the master key Data Explorer already resolved without calling ARM", async () => {
(userContext as UserContextType).masterKey = "cachedMasterKey";
const credential = await getCosmosDBShellCredential();
expect(credential).toEqual({ kind: "key", value: "cachedMasterKey" });
expect(listKeys).not.toHaveBeenCalled();
expect(getReadOnlyKeys).not.toHaveBeenCalled();
});
it("should fall back to the read-only keys when the caller lacks list-keys permission", async () => {
const authorizationFailed = Object.assign(new Error("AuthorizationFailed"), { code: "AuthorizationFailed" });
(listKeys as jest.Mock).mockRejectedValue(authorizationFailed);
(getReadOnlyKeys as jest.Mock).mockResolvedValue({ primaryReadonlyMasterKey: "readOnlyKey" });
const credential = await getCosmosDBShellCredential();
expect(credential).toEqual({ kind: "key", value: "readOnlyKey" });
expect(getReadOnlyKeys).toHaveBeenCalledWith("testSubId", "testResourceGroup", "testDbName");
});
it("should return undefined when both the read-write and read-only key fetches fail", async () => {
const authorizationFailed = Object.assign(new Error("AuthorizationFailed"), { code: "AuthorizationFailed" });
(listKeys as jest.Mock).mockRejectedValue(authorizationFailed);
(getReadOnlyKeys as jest.Mock).mockRejectedValue(new Error("Forbidden"));
jest.spyOn(console, "error").mockImplementation(() => undefined);
const credential = await getCosmosDBShellCredential();
expect(credential).toBeUndefined();
});
it("should return the account key when Entra ID auth is not enabled", async () => {
const credential = await getCosmosDBShellCredential();
expect(credential).toEqual({ kind: "key", value: mockKey });
expect(listKeys).toHaveBeenCalledWith("testSubId", "testResourceGroup", "testDbName");
});
it("should return the cached aadToken when Entra ID auth is enabled", async () => {
(userContext as UserContextType).dataPlaneRbacEnabled = true;
(userContext as UserContextType).aadToken = "aadToken123";
const credential = await getCosmosDBShellCredential();
expect(credential).toEqual({ kind: "token", value: "aadToken123" });
expect(listKeys).not.toHaveBeenCalled();
});
it("should return a silently minted token when an MSAL account is cached", async () => {
(userContext as UserContextType).dataPlaneRbacEnabled = true;
mockMsalAccounts([{ username: "user@contoso.com" }]);
(acquireMsalTokenForAccount as jest.Mock).mockResolvedValue("mintedToken123");
const credential = await getCosmosDBShellCredential();
expect(credential).toEqual({ kind: "token", value: "mintedToken123" });
expect(listKeys).not.toHaveBeenCalled();
});
it("should fall back to the account key when no token could be resolved", async () => {
(userContext as UserContextType).dataPlaneRbacEnabled = true;
const credential = await getCosmosDBShellCredential();
expect(acquireMsalTokenForAccount).not.toHaveBeenCalled();
expect(credential).toEqual({ kind: "key", value: mockKey });
});
it("should fall back to the account key when the silent token acquisition throws", async () => {
(userContext as UserContextType).dataPlaneRbacEnabled = true;
mockMsalAccounts([{ username: "user@contoso.com" }]);
(acquireMsalTokenForAccount as jest.Mock).mockRejectedValue(new Error("interaction_required"));
jest.spyOn(console, "error").mockImplementation(() => undefined);
const credential = await getCosmosDBShellCredential();
expect(credential).toEqual({ kind: "key", value: mockKey });
});
it("should not fall back to the account key when local auth is disabled", async () => {
(userContext as UserContextType).databaseAccount.properties = { disableLocalAuth: true };
jest.spyOn(console, "warn").mockImplementation(() => undefined);
const credential = await getCosmosDBShellCredential();
expect(credential).toBeUndefined();
expect(listKeys).not.toHaveBeenCalled();
});
it("should return undefined when listing the account keys fails", async () => {
(listKeys as jest.Mock).mockRejectedValue(new Error("Forbidden"));
jest.spyOn(console, "error").mockImplementation(() => undefined);
const credential = await getCosmosDBShellCredential();
expect(credential).toBeUndefined();
});
it("should return undefined when the database name is missing", async () => {
(userContext as UserContextType).databaseAccount.name = "";
const credential = await getCosmosDBShellCredential();
expect(credential).toBeUndefined();
expect(listKeys).not.toHaveBeenCalled();
(userContext as UserContextType).databaseAccount.name = "testDbName";
});
it("should populate a specific reason when both key fetches fail, for the shell to echo", async () => {
(listKeys as jest.Mock).mockRejectedValue(new Error("Forbidden"));
(getReadOnlyKeys as jest.Mock).mockRejectedValue(new Error("Forbidden"));
jest.spyOn(console, "error").mockImplementation(() => undefined);
const diagnostics: CosmosDBShellCredentialDiagnostics = {};
const credential = await getCosmosDBShellCredential(diagnostics);
expect(credential).toBeUndefined();
expect(diagnostics.reason).toContain("Forbidden");
});
it("should populate a reason when local auth is disabled and no token was resolved", async () => {
(userContext as UserContextType).databaseAccount.properties = { disableLocalAuth: true };
jest.spyOn(console, "warn").mockImplementation(() => undefined);
const diagnostics: CosmosDBShellCredentialDiagnostics = {};
const credential = await getCosmosDBShellCredential(diagnostics);
expect(credential).toBeUndefined();
expect(diagnostics.reason).toContain("local (key) auth is disabled");
});
});
}); });
@@ -1,9 +1,15 @@
import { TerminalKind } from "../../../../Contracts/ViewModels"; import { TerminalKind } from "../../../../Contracts/ViewModels";
import { userContext } from "../../../../UserContext"; import { userContext } from "../../../../UserContext";
import { listKeys } from "../../../../Utils/arm/generatedClients/cosmos/databaseAccounts"; import { getReadOnlyKeys, listKeys } from "../../../../Utils/arm/generatedClients/cosmos/databaseAccounts";
import { isDataplaneRbacEnabledForProxyApi } from "../../../../Utils/AuthorizationUtils"; import {
acquireMsalTokenForAccount,
getMsalInstance,
isCloudShellEntraAuthEnabled,
isDataplaneRbacEnabledForProxyApi,
} from "../../../../Utils/AuthorizationUtils";
import { AbstractShellHandler } from "./AbstractShellHandler"; import { AbstractShellHandler } from "./AbstractShellHandler";
import { CassandraShellHandler } from "./CassandraShellHandler"; import { CassandraShellHandler } from "./CassandraShellHandler";
import { CosmosDBShellCredential, CosmosDBShellHandler } from "./CosmosDBShellHandler";
import { MongoShellHandler } from "./MongoShellHandler"; import { MongoShellHandler } from "./MongoShellHandler";
import { PostgresShellHandler } from "./PostgresShellHandler"; import { PostgresShellHandler } from "./PostgresShellHandler";
import { VCoreMongoShellHandler } from "./VCoreMongoShellHandler"; import { VCoreMongoShellHandler } from "./VCoreMongoShellHandler";
@@ -16,23 +22,244 @@ export async function getHandler(shellType: TerminalKind): Promise<AbstractShell
case TerminalKind.Postgres: case TerminalKind.Postgres:
return new PostgresShellHandler(); return new PostgresShellHandler();
case TerminalKind.Mongo: case TerminalKind.Mongo:
return new MongoShellHandler(await getKey()); return new MongoShellHandler(await getKey(isDataplaneRbacEnabledForProxyApi(userContext)));
case TerminalKind.VCoreMongo: case TerminalKind.VCoreMongo:
return new VCoreMongoShellHandler(); return new VCoreMongoShellHandler();
case TerminalKind.Cassandra: case TerminalKind.Cassandra:
return new CassandraShellHandler(await getKey()); return new CassandraShellHandler(await getKey(isDataplaneRbacEnabledForProxyApi(userContext)));
case TerminalKind.CosmosDB: {
const diagnostics: CosmosDBShellCredentialDiagnostics = {};
const credential = await getCosmosDBShellCredential(diagnostics);
return new CosmosDBShellHandler(credential, diagnostics.reason);
}
default: default:
throw new Error(`Unsupported shell type: ${shellType}`); throw new Error(`Unsupported shell type: ${shellType}`);
} }
} }
export async function getKey(): Promise<string> { /**
* Populated by {@link getCosmosDBShellCredential} (and its helpers) with a human-readable
* explanation when no credential could be resolved, so the shell can echo a specific cause
* instead of a generic message. There is no way to inspect the browser dev tools console
* from inside the remote Cloud Shell terminal, so this is surfaced directly in the shell.
*/
export interface CosmosDBShellCredentialDiagnostics {
reason?: string;
}
/**
* Resolves the credential Data Explorer already holds for the current account and hands
* it to the Cosmos DB Shell.
*
* The Cloud Shell cannot authenticate to Cosmos DB on its own: its managed identity is
* rejected with "AudienceNotSupported" for the `*.documents.azure.com` audience, its `az`
* session is not signed in, and the interactive browser/device-code flows are not usable
* from the embedded terminal. The credential therefore has to come from Data Explorer.
*
* Resolution order:
* 1. Entra ID data-plane token — the cached `userContext.aadToken`, or a silently minted
* one. Silent acquisition is only attempted when an MSAL account is already cached so
* it can never trigger an interactive popup.
* 2. Account master key — the key Data Explorer already resolved (`userContext.masterKey`),
* or, if that is not populated, one fetched via ARM (`listKeys`, falling back to
* `getReadOnlyKeys` on any failure for read-only callers). The key is handed to the
* handler, which exports it as COSMOSDB_SHELL_ACCOUNT_KEY on the same command line as
* the `cosmosdbshell` invocation. Skipped when the account has local auth disabled
* (keys do not exist).
*
* Returns `undefined` when nothing could be resolved, which makes the handler surface
* actionable guidance instead of letting the tool attempt its own sign-in.
*/
export async function getCosmosDBShellCredential(
diagnostics?: CosmosDBShellCredentialDiagnostics,
): Promise<CosmosDBShellCredential | undefined> {
const dbName = userContext.databaseAccount?.name;
if (!dbName) {
if (diagnostics) {
diagnostics.reason = "no database account name is available";
}
return undefined;
}
if (isCloudShellEntraAuthEnabled(userContext)) {
const token = await resolveCosmosDataPlaneToken();
if (token) {
return { kind: "token", value: token };
}
}
if (userContext.databaseAccount?.properties?.disableLocalAuth) {
const reason =
"local (key) auth is disabled on this account and no Entra ID token could be resolved; sign in via " +
'"Login for Entra ID" in the toolbar and reopen the shell';
console.warn(`CloudShell: ${reason}.`);
if (diagnostics) {
diagnostics.reason = reason;
}
return undefined;
}
const key = await resolveAccountKey(dbName, diagnostics);
if (!key) {
if (!diagnostics?.reason) {
const reason =
"no account key could be resolved. Data Explorer may be authenticating through the portal proxy " +
"(per-request tokens), which cannot be reused by the shell. Ensure you have either data-plane RBAC " +
'access (then use "Login for Entra ID") or permission to list the account keys';
console.warn(`CloudShell: ${reason}.`);
if (diagnostics) {
diagnostics.reason = reason;
}
}
return undefined;
}
return { kind: "key", value: key };
}
/**
* Returns a Cosmos data-plane token without ever prompting. Silent acquisition is guarded
* on an existing cached MSAL account because `acquireMsalTokenForAccount` falls back to an
* interactive popup when none exists, which cannot complete in the hosted Cloud Shell.
*/
async function resolveCosmosDataPlaneToken(): Promise<string> {
if (userContext.aadToken) {
return userContext.aadToken;
}
try {
const msalInstance = await getMsalInstance();
if (msalInstance.getAllAccounts().length === 0) {
return "";
}
return (await acquireMsalTokenForAccount(userContext.databaseAccount, true)) || "";
} catch (error) {
console.error("Failed to silently acquire a Cosmos data-plane token for the Cloud Shell", error);
return "";
}
}
/**
* Resolves the subscription id and resource group for the current account. These are
* normally populated on `userContext`, but in some hosting contexts they can be missing,
* so they are parsed from the account's ARM resource id as a fallback. The id has the form
* `/subscriptions/<sub>/resourceGroups/<rg>/providers/Microsoft.DocumentDB/databaseAccounts/<name>`.
*/
function resolveAccountArmScope(): { subscriptionId: string; resourceGroup: string } {
let subscriptionId = userContext.subscriptionId;
let resourceGroup = userContext.resourceGroup;
const accountId = userContext.databaseAccount?.id;
if ((!subscriptionId || !resourceGroup) && accountId) {
const match = accountId.match(/\/subscriptions\/([^/]+)\/resourceGroups\/([^/]+)\//i);
if (match) {
subscriptionId = subscriptionId || match[1];
resourceGroup = resourceGroup || match[2];
}
}
return { subscriptionId, resourceGroup };
}
async function resolveAccountKey(dbName: string, diagnostics?: CosmosDBShellCredentialDiagnostics): Promise<string> {
// Prefer the key Data Explorer already resolved for this account so the shell reuses the
// exact credential DE is connected with and avoids a redundant ARM round-trip.
if (userContext.masterKey) {
return userContext.masterKey;
}
// Otherwise fetch it via ARM, mirroring DE's own `fetchAndUpdateKeys`: try the read-write
// keys first, then fall back to the read-only keys. The fallback is attempted on any
// failure (not just "AuthorizationFailed") because a read-only caller can surface the
// missing-permission error in different shapes; without it such a user gets no credential
// and the shell cannot connect.
const { subscriptionId, resourceGroup } = resolveAccountArmScope();
if (!subscriptionId || !resourceGroup) {
const reason = "the subscription id or resource group for this account could not be determined";
console.error(`CloudShell: cannot list Cosmos DB account keys because ${reason}.`);
if (diagnostics) {
diagnostics.reason = reason;
}
return "";
}
let listKeysErrorMessage: string | undefined;
try {
const keys = await listKeys(subscriptionId, resourceGroup, dbName);
if (keys?.primaryMasterKey) {
return keys.primaryMasterKey;
}
listKeysErrorMessage = "the response did not include a usable key";
} catch (error) {
listKeysErrorMessage = error instanceof Error ? error.message : String(error);
console.error("Failed to list read-write account keys for the Cloud Shell; trying read-only keys", error);
}
try {
const readOnlyKeys = await getReadOnlyKeys(subscriptionId, resourceGroup, dbName);
if (readOnlyKeys?.primaryReadonlyMasterKey) {
return readOnlyKeys.primaryReadonlyMasterKey;
}
if (diagnostics) {
diagnostics.reason =
`listing account keys failed (${listKeysErrorMessage}) and the read-only keys response did not ` +
"include a usable key either";
}
return "";
} catch (readOnlyError) {
console.error("Failed to list read-only account keys for the Cloud Shell", readOnlyError);
if (diagnostics) {
const readOnlyErrorMessage = readOnlyError instanceof Error ? readOnlyError.message : String(readOnlyError);
diagnostics.reason =
`listing the account keys failed (${listKeysErrorMessage}) and listing the read-only keys also failed ` +
`(${readOnlyErrorMessage}) — this is usually a missing listKeys/listReadOnlyKeys RBAC permission`;
}
return "";
}
}
/**
* Resolves the credential to inject into the Cloud Shell for a Mongo or Cassandra
* connection.
*
* @param useEntraIdAuth When true, returns an Entra ID (AAD) bearer token scoped to
* the account's data plane; otherwise returns the account master key. The caller
* decides which credential applies so it stays in sync with the connection command
* the matching handler builds.
*
* The Cosmos DB (NoSQL) shell does not use this function — see
* {@link getCosmosDBShellCredential}, which also reports which kind of credential it
* resolved so the correct environment variable is exported.
*/
export async function getKey(useEntraIdAuth: boolean): Promise<string> {
const dbName = userContext.databaseAccount.name; const dbName = userContext.databaseAccount.name;
if (!dbName) { if (!dbName) {
return ""; return "";
} }
if (isDataplaneRbacEnabledForProxyApi(userContext)) { if (useEntraIdAuth) {
return userContext.aadToken || ""; if (userContext.aadToken) {
return userContext.aadToken;
}
try {
const msalInstance = await getMsalInstance();
if (msalInstance.getAllAccounts().length === 0) {
// No cached account to silently acquire against; a token request here would
// require an interactive popup, which cannot complete in the Cloud Shell.
console.warn(
"CloudShell: no cached MSAL account available to mint a Cosmos data-plane token; " +
"the shell will use interactive device-code authentication.",
);
return "";
}
const token = (await acquireMsalTokenForAccount(userContext.databaseAccount, true)) || "";
if (!token) {
console.warn("CloudShell: silent Cosmos data-plane token acquisition returned an empty token.");
}
return token;
} catch (error) {
console.error("Failed to silently acquire a Cosmos data-plane token for the Cloud Shell", error);
return "";
}
} }
const keys = await listKeys(userContext.subscriptionId, userContext.resourceGroup, dbName); const keys = await listKeys(userContext.subscriptionId, userContext.resourceGroup, dbName);
@@ -0,0 +1,67 @@
import { Terminal } from "xterm";
import { askConfirmation } from "./CommonUtils";
describe("askConfirmation", () => {
const createTerminalMock = () => {
let keyHandler: ((event: { key: string }) => void) | undefined;
const dispose = jest.fn();
const terminal = {
writeln: jest.fn(),
focus: jest.fn(),
onKey: jest.fn((handler: (event: { key: string }) => void) => {
keyHandler = handler;
return { dispose };
}),
} as unknown as Terminal;
return {
terminal,
dispose,
pressKey: (key: string) => keyHandler?.({ key }),
};
};
it("resolves true when the user presses Y", async () => {
const { terminal, dispose, pressKey } = createTerminalMock();
const promise = askConfirmation(terminal, "Proceed?");
pressKey("Y");
await expect(promise).resolves.toBe(true);
expect(dispose).toHaveBeenCalledTimes(1);
});
it("resolves true when the user presses lowercase y", async () => {
const { terminal, pressKey } = createTerminalMock();
const promise = askConfirmation(terminal, "Proceed?");
pressKey("y");
await expect(promise).resolves.toBe(true);
});
it("resolves false when the user presses N", async () => {
const { terminal, dispose, pressKey } = createTerminalMock();
const promise = askConfirmation(terminal, "Proceed?");
pressKey("N");
await expect(promise).resolves.toBe(false);
expect(dispose).toHaveBeenCalledTimes(1);
});
it("ignores keys other than Y or N and keeps listening until a valid answer", async () => {
const { terminal, dispose, pressKey } = createTerminalMock();
const promise = askConfirmation(terminal, "Proceed?");
pressKey("a");
pressKey("1");
pressKey("\r");
expect(dispose).not.toHaveBeenCalled();
pressKey("y");
await expect(promise).resolves.toBe(true);
expect(dispose).toHaveBeenCalledTimes(1);
});
});
@@ -24,9 +24,15 @@ export const askConfirmation = async (terminal: Terminal, question: string): Pro
terminal.focus(); terminal.focus();
return new Promise<boolean>((resolve) => { return new Promise<boolean>((resolve) => {
const keyListener = terminal.onKey(({ key }: { key: string }) => { const keyListener = terminal.onKey(({ key }: { key: string }) => {
const normalizedKey = key.toLowerCase();
// Only "y" or "n" are accepted. Any other key is ignored so an accidental
// keypress does not abort the flow; keep listening until a valid answer.
if (normalizedKey !== "y" && normalizedKey !== "n") {
return;
}
keyListener.dispose(); keyListener.dispose();
terminal.writeln(key); terminal.writeln(key);
return resolve(key.toLowerCase() === "y"); return resolve(normalizedKey === "y");
}); });
}); });
}; };
@@ -46,6 +52,8 @@ export const getShellNameForDisplay = (terminalKind: TerminalKind): string => {
case TerminalKind.Mongo: case TerminalKind.Mongo:
case TerminalKind.VCoreMongo: case TerminalKind.VCoreMongo:
return "MongoDB"; return "MongoDB";
case TerminalKind.CosmosDB:
return "Cosmos DB";
default: default:
return ""; return "";
} }
+2
View File
@@ -35,6 +35,7 @@ export type Features = {
readonly disableConnectionStringLogin: boolean; readonly disableConnectionStringLogin: boolean;
readonly enableContainerCopy: boolean; readonly enableContainerCopy: boolean;
readonly enableCloudShell: boolean; readonly enableCloudShell: boolean;
readonly enableCosmosDBShell: boolean;
readonly enableRestoreContainer: boolean; // only for Fabric readonly enableRestoreContainer: boolean; // only for Fabric
readonly mongoDisableNativeAuth: boolean; readonly mongoDisableNativeAuth: boolean;
@@ -104,6 +105,7 @@ export function extractFeatures(given = new URLSearchParams(window.location.sear
enableContainerCopy: "true" === get("enablecontainercopy"), enableContainerCopy: "true" === get("enablecontainercopy"),
enableRestoreContainer: "true" === get("enablerestorecontainer"), enableRestoreContainer: "true" === get("enablerestorecontainer"),
enableCloudShell: true, enableCloudShell: true,
enableCosmosDBShell: "true" === get("enablecosmosdbshell"),
mongoDisableNativeAuth: "true" === get("mongodisablenativeauth"), mongoDisableNativeAuth: "true" === get("mongodisablenativeauth"),
}; };
} }
+1
View File
@@ -17,6 +17,7 @@ describe("AuthorizationUtils", () => {
updateUserContext({ updateUserContext({
features: { features: {
enableContainerCopy: false, enableContainerCopy: false,
enableCosmosDBShell: false,
enableAadDataPlane: enabled, enableAadDataPlane: enabled,
canExceedMaximumValue: false, canExceedMaximumValue: false,
cosmosdb: false, cosmosdb: false,
+18
View File
@@ -236,3 +236,21 @@ export function useDataplaneRbacAuthorization(userContext: UserContext): boolean
export function isDataplaneRbacEnabledForProxyApi(userContext: UserContext): boolean { export function isDataplaneRbacEnabledForProxyApi(userContext: UserContext): boolean {
return useDataplaneRbacAuthorization(userContext) && hasProxyServer(userContext.apiType); return useDataplaneRbacAuthorization(userContext) && hasProxyServer(userContext.apiType);
} }
/**
* Determines whether Cloud Shell connections should authenticate with an Entra ID
* token (COSMOSDB_SHELL_TOKEN) instead of the account master key.
*
* Returns true when data-plane RBAC authorization is in effect, or when the account
* has local (key) auth disabled — in the latter case a master key would be rejected
* by the service, so we must fall back to Entra ID even if the RBAC toggle is off.
*
* Unlike {@link isDataplaneRbacEnabledForProxyApi}, this does not require a proxy
* server, so it works for the SQL/CosmosDB API which connects directly to the
* data plane.
*/
export function isCloudShellEntraAuthEnabled(userContext: UserContext): boolean {
return (
useDataplaneRbacAuthorization(userContext) || userContext.databaseAccount?.properties?.disableLocalAuth === true
);
}