Add support for SPOG host URLs

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:^"
     }
 }

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",

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,

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({

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<AuthType> {
         return ["databricks-cli", "google-id", "pat"];
     }
 
-    if (UrlUtils.isAwsHost(host)) {
+    if (UrlUtils.isAwsHost(host) || UrlUtils.isSpogHost(host)) {
         return ["databricks-cli", "pat"];
     }
 

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<string, string> {
-        return {
+        const env: Record<string, string> = {
             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;
     }
 

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",

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({

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 */
 }

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")
+            )
+        );
+    });
 });

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)
+    );
+}

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