diff --git a/packages/databricks-vscode-types/package.json b/packages/databricks-vscode-types/package.json index f50c1ba0b..489c1d875 100644 --- a/packages/databricks-vscode-types/package.json +++ b/packages/databricks-vscode-types/package.json @@ -29,7 +29,7 @@ "typescript": "^5.3.3" }, "dependencies": { - "@databricks/sdk-experimental": "^0.17.0", + "@databricks/sdk-experimental": "^0.18.0", "databricks": "workspace:^" } } diff --git a/packages/databricks-vscode/package.json b/packages/databricks-vscode/package.json index ba3309e33..44ac79ec5 100644 --- a/packages/databricks-vscode/package.json +++ b/packages/databricks-vscode/package.json @@ -1198,7 +1198,7 @@ }, "dependencies": { "@databricks/databricks-vscode-types": "workspace:^", - "@databricks/sdk-experimental": "^0.17.0", + "@databricks/sdk-experimental": "^0.18.0", "@types/lodash": "^4.14.202", "@types/shell-quote": "^1.7.5", "@vscode/debugadapter": "^1.64.0", diff --git a/packages/databricks-vscode/src/cli/CliWrapper.ts b/packages/databricks-vscode/src/cli/CliWrapper.ts index 925bf2552..776c79dd0 100644 --- a/packages/databricks-vscode/src/cli/CliWrapper.ts +++ b/packages/databricks-vscode/src/cli/CliWrapper.ts @@ -97,6 +97,7 @@ export interface ConfigEntry { name: string; host?: URL; accountId?: string; + workspaceId?: string; cloud: Cloud; authType: string; valid: boolean; @@ -414,6 +415,7 @@ export class CliWrapper { name: profile.name, host: UrlUtils.normalizeHost(profile.host), accountId: profile.account_id, + workspaceId: profile.workspace_id, cloud: profile.cloud, authType: profile.auth_type, valid: profile.valid, diff --git a/packages/databricks-vscode/src/cluster/ClusterManager.test.ts b/packages/databricks-vscode/src/cluster/ClusterManager.test.ts index 20b94284e..0169a7cfa 100644 --- a/packages/databricks-vscode/src/cluster/ClusterManager.test.ts +++ b/packages/databricks-vscode/src/cluster/ClusterManager.test.ts @@ -2,6 +2,7 @@ import {describe} from "mocha"; import {ClusterManager} from "./ClusterManager"; import { ApiClient, + Config, compute, Time, TimeUnits, @@ -28,6 +29,9 @@ describe(__filename, async () => { beforeEach(async () => { ({testClusterDetails} = await ClusterFixtures.getMockTestCluster()); mockedClient = mock(ApiClient); + const mockedConfig = mock(Config); + when(mockedConfig.ensureResolved()).thenResolve(); + when(mockedClient.config).thenReturn(instance(mockedConfig)); when( mockedClient.request( objectContaining({ diff --git a/packages/databricks-vscode/src/configuration/LoginWizard.ts b/packages/databricks-vscode/src/configuration/LoginWizard.ts index f986df4a5..db8d18dc2 100644 --- a/packages/databricks-vscode/src/configuration/LoginWizard.ts +++ b/packages/databricks-vscode/src/configuration/LoginWizard.ts @@ -481,7 +481,8 @@ async function validateDatabricksHost( if ( !url.hostname.match( /(\.databricks\.azure\.us|\.databricks\.azure\.cn|\.azuredatabricks\.net|\.gcp\.databricks\.com|\.cloud\.databricks\.com|\.dev\.databricks\.com)$/ - ) + ) && + !UrlUtils.isSpogHost(url) ) { return { message: @@ -503,7 +504,7 @@ function authMethodsForHostname(host: URL): Array { return ["databricks-cli", "google-id", "pat"]; } - if (UrlUtils.isAwsHost(host)) { + if (UrlUtils.isAwsHost(host) || UrlUtils.isSpogHost(host)) { return ["databricks-cli", "pat"]; } diff --git a/packages/databricks-vscode/src/configuration/auth/AuthProvider.ts b/packages/databricks-vscode/src/configuration/auth/AuthProvider.ts index 694d04927..9804a610b 100644 --- a/packages/databricks-vscode/src/configuration/auth/AuthProvider.ts +++ b/packages/databricks-vscode/src/configuration/auth/AuthProvider.ts @@ -134,6 +134,12 @@ export abstract class AuthProvider { if (config.databricksCliPath === undefined) { config.databricksCliPath = this._cli.cliPath; } + // Profiles with workspace_id target a specific workspace on a unified (SPOG) + // host. The SDK only sends X-Databricks-Org-Id when experimentalIsUnifiedHost + // is true, so we set it here after the profile has been loaded. + if (config.workspaceId) { + config.experimentalIsUnifiedHost = true; + } return config; } @@ -166,7 +172,8 @@ export abstract class AuthProvider { host, json.databricksPath ?? cli.cliPath, cli, - json.profile + json.profile, + json.workspaceId ); case "profile": @@ -200,7 +207,8 @@ export abstract class AuthProvider { host, config.databricksCliPath ?? cli.cliPath, cli, - config.profile + config.profile, + config.workspaceId ); default: @@ -213,6 +221,8 @@ export abstract class AuthProvider { } export class ProfileAuthProvider extends AuthProvider { + private _workspaceId?: string; + static async from(profile: string, cli: CliWrapper, checked = false) { const host = await ProfileAuthProvider.getSdkConfig(profile).getHost(); return new ProfileAuthProvider(host, profile, cli, checked); @@ -240,10 +250,15 @@ export class ProfileAuthProvider extends AuthProvider { } toEnv(): Record { - return { + const env: Record = { DATABRICKS_HOST: this.host.toString(), DATABRICKS_CONFIG_PROFILE: this.profile, }; + if (this._workspaceId) { + env["DATABRICKS_WORKSPACE_ID"] = this._workspaceId; + env["DATABRICKS_EXPERIMENTAL_IS_UNIFIED_HOST"] = "true"; + } + return env; } toIni() { @@ -266,6 +281,9 @@ export class ProfileAuthProvider extends AuthProvider { while (cancellationToken?.isCancellationRequested !== true) { try { const sdkConfig = await this.getSdkConfig(); + // Cache workspace_id so toEnv() can include SPOG routing vars + // for bundle commands that use the Go CLI directly. + this._workspaceId = sdkConfig.workspaceId; const authProvider = AuthProvider.fromSdkConfig( sdkConfig, this.cli @@ -311,7 +329,8 @@ export class DatabricksCliAuthProvider extends AuthProvider { host: URL, readonly cliPath: string, cli: CliWrapper, - readonly profile?: string + readonly profile?: string, + readonly workspaceId?: string ) { super(host, "databricks-cli", cli); } @@ -326,6 +345,7 @@ export class DatabricksCliAuthProvider extends AuthProvider { authType: this.authType, databricksPath: this.cliPath, ...(this.profile ? {profile: this.profile} : {}), + ...(this.workspaceId ? {workspaceId: this.workspaceId} : {}), }; } @@ -335,6 +355,12 @@ export class DatabricksCliAuthProvider extends AuthProvider { authType: "databricks-cli", databricksCliPath: this.cliPath, ...(this.profile ? {profile: this.profile} : {}), + ...(this.workspaceId + ? { + workspaceId: this.workspaceId, + experimentalIsUnifiedHost: true, + } + : {}), }); } @@ -346,6 +372,10 @@ export class DatabricksCliAuthProvider extends AuthProvider { if (this.profile) { env["DATABRICKS_CONFIG_PROFILE"] = this.profile; } + if (this.workspaceId) { + env["DATABRICKS_WORKSPACE_ID"] = this.workspaceId; + env["DATABRICKS_EXPERIMENTAL_IS_UNIFIED_HOST"] = "true"; + } return env; } diff --git a/packages/databricks-vscode/src/configuration/auth/DatabricksCliCheck.ts b/packages/databricks-vscode/src/configuration/auth/DatabricksCliCheck.ts index 1ea338d24..507506766 100644 --- a/packages/databricks-vscode/src/configuration/auth/DatabricksCliCheck.ts +++ b/packages/databricks-vscode/src/configuration/auth/DatabricksCliCheck.ts @@ -85,6 +85,12 @@ export class DatabricksCliCheck implements Disposable { authType: "databricks-cli", databricksCliPath: this.authProvider.cliPath, profile: this.authProvider.profile, + ...(this.authProvider.workspaceId + ? { + workspaceId: this.authProvider.workspaceId, + experimentalIsUnifiedHost: true, + } + : {}), }, { product: "databricks-vscode", diff --git a/packages/databricks-vscode/src/sdk-extensions/test/ClusterFixtures.ts b/packages/databricks-vscode/src/sdk-extensions/test/ClusterFixtures.ts index c57aaa33d..05b994815 100644 --- a/packages/databricks-vscode/src/sdk-extensions/test/ClusterFixtures.ts +++ b/packages/databricks-vscode/src/sdk-extensions/test/ClusterFixtures.ts @@ -9,7 +9,7 @@ import { anything, objectContaining, } from "ts-mockito"; -import {compute, ApiClient} from "@databricks/sdk-experimental"; +import {compute, ApiClient, Config} from "@databricks/sdk-experimental"; const testClusterDetails: compute.ClusterDetails = { cluster_id: "testClusterId", @@ -18,6 +18,9 @@ const testClusterDetails: compute.ClusterDetails = { export async function getMockTestCluster() { const mockedClient = mock(ApiClient); + const mockedConfig = mock(Config); + when(mockedConfig.ensureResolved()).thenResolve(); + when(mockedClient.config).thenReturn(instance(mockedConfig)); when( mockedClient.request( objectContaining({ diff --git a/packages/databricks-vscode/src/utils/envVarGenerators.ts b/packages/databricks-vscode/src/utils/envVarGenerators.ts index 9d5b8934d..da5201921 100644 --- a/packages/databricks-vscode/src/utils/envVarGenerators.ts +++ b/packages/databricks-vscode/src/utils/envVarGenerators.ts @@ -59,11 +59,21 @@ export function getAuthEnvVars(connectionManager: ConnectionManager) { return; } + // For SPOG (unified host) connections the Go CLI SDK must know the + // workspace_id so it can add the X-Databricks-Org-Id routing header. + const workspaceId = connectionManager.apiClient?.config?.workspaceId; + /* eslint-disable @typescript-eslint/naming-convention */ return { DATABRICKS_HOST: host, DATABRICKS_AUTH_TYPE: "metadata-service", DATABRICKS_METADATA_SERVICE_URL: connectionManager.metadataServiceUrl, + ...(workspaceId + ? { + DATABRICKS_WORKSPACE_ID: workspaceId, + DATABRICKS_EXPERIMENTAL_IS_UNIFIED_HOST: "true", + } + : {}), }; /* eslint-enable @typescript-eslint/naming-convention */ } diff --git a/packages/databricks-vscode/src/utils/urlUtils.test.ts b/packages/databricks-vscode/src/utils/urlUtils.test.ts index 901dacb4a..9074cd224 100644 --- a/packages/databricks-vscode/src/utils/urlUtils.test.ts +++ b/packages/databricks-vscode/src/utils/urlUtils.test.ts @@ -4,6 +4,7 @@ import { isAwsHost, isAzureHost, isGcpHost, + isSpogHost, normalizeHost, } from "./urlUtils"; @@ -41,4 +42,34 @@ describe(__filename, () => { }); }); }); + + it("should strip query params from host", () => { + const url = + "https://dbc-123456789012345.cloud.databricks.com/?o=789&other=foo"; + const normalized = normalizeHost(url); + assert.strictEqual(normalized.search, ""); + }); + + it("should identify SPOG hosts by *.databricks.com hostname", () => { + assert.ok(isSpogHost(new URL("https://db-deco-test.databricks.com"))); + assert.ok(isSpogHost(new URL("https://demo-spog.databricks.com"))); + }); + + it("should not classify standard cloud hosts as SPOG", () => { + assert.ok( + !isSpogHost( + new URL("https://dbc-123456789012345.cloud.databricks.com") + ) + ); + assert.ok( + !isSpogHost( + new URL("https://dbc-123456789012345.gcp.databricks.com") + ) + ); + assert.ok( + !isSpogHost( + new URL("https://dbc-123456789012345.dev.databricks.com") + ) + ); + }); }); diff --git a/packages/databricks-vscode/src/utils/urlUtils.ts b/packages/databricks-vscode/src/utils/urlUtils.ts index 03607a888..40ccf0bb6 100644 --- a/packages/databricks-vscode/src/utils/urlUtils.ts +++ b/packages/databricks-vscode/src/utils/urlUtils.ts @@ -47,3 +47,14 @@ export function isAwsHost(url: URL): boolean { /(\.cloud\.databricks\.com|\.dev\.databricks\.com)$/ ); } + +export function isSpogHost(url: URL): boolean { + // SPOG hosts are *.databricks.com but not the standard cloud-specific subdomains + // already classified as AWS (*.cloud.databricks.com, *.dev.databricks.com) + // or GCP (*.gcp.databricks.com). + return ( + !!url.hostname.match(/\.databricks\.com$/) && + !isAwsHost(url) && + !isGcpHost(url) + ); +} diff --git a/yarn.lock b/yarn.lock index 98e9e4cf8..bdb90dd18 100644 --- a/yarn.lock +++ b/yarn.lock @@ -343,7 +343,7 @@ __metadata: version: 0.0.0-use.local resolution: "@databricks/databricks-vscode-types@workspace:packages/databricks-vscode-types" dependencies: - "@databricks/sdk-experimental": ^0.17.0 + "@databricks/sdk-experimental": ^0.18.0 "@types/vscode": 1.86.0 databricks: "workspace:^" eslint: ^8.55.0 @@ -368,15 +368,15 @@ __metadata: languageName: unknown linkType: soft -"@databricks/sdk-experimental@npm:^0.17.0": - version: 0.17.0 - resolution: "@databricks/sdk-experimental@npm:0.17.0" +"@databricks/sdk-experimental@npm:^0.18.0": + version: 0.18.0 + resolution: "@databricks/sdk-experimental@npm:0.18.0" dependencies: google-auth-library: ^10.5.0 ini: ^6.0.0 reflect-metadata: ^0.2.2 semver: ^7.7.3 - checksum: 34de7d8708de12bf1fa44ae5ced7f8886a26d3ef9f072b47955c63dd7c3ae32db7c0b243074907b371026ee352d76c90d02faa9df0e95ad5d9e8aa22dc53d0b5 + checksum: eefb552284eaa5577baaa173f1612611833bd805cbd822f11373223d74b2e2b37d2a6dbe7bdf522312339bcecf8384a2510c5a74b51612664ddca7a4b51fa503 languageName: node linkType: hard @@ -3745,7 +3745,7 @@ __metadata: resolution: "databricks@workspace:packages/databricks-vscode" dependencies: "@databricks/databricks-vscode-types": "workspace:^" - "@databricks/sdk-experimental": ^0.17.0 + "@databricks/sdk-experimental": ^0.18.0 "@istanbuljs/nyc-config-typescript": ^1.0.2 "@sinonjs/fake-timers": ^11.2.2 "@types/bcryptjs": ^2.4.6