Add support for SPOG host URLs

.yarnrc.yml

@@ -4,6 +4,8 @@ enableTelemetry: 0
 
 nodeLinker: node-modules
 
+npmRegistryServer: "https://npm-proxy.cloud.databricks.com/"
+
 plugins:
   - path: .yarn/plugins/@yarnpkg/plugin-workspace-tools.cjs
     spec: "@yarnpkg/plugin-workspace-tools"

packages/databricks-vscode-types/package.json

@@ -29,7 +29,7 @@
         "typescript": "^5.3.3"
     },
     "dependencies": {
-        "@databricks/sdk-experimental": "^0.16.0",
+        "@databricks/sdk-experimental": "^0.17.0",
         "databricks": "workspace:^"
     }
 }

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

@@ -143,6 +143,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;
     }
@@ -175,7 +181,8 @@ export abstract class AuthProvider {
                     host,
                     json.databricksPath ?? cli.cliPath,
                     cli,
-                    json.profile
+                    json.profile,
+                    json.workspaceId
                 );
 
             case "profile":
@@ -209,7 +216,8 @@ export abstract class AuthProvider {
                     host,
                     config.databricksCliPath ?? cli.cliPath,
                     cli,
-                    config.profile
+                    config.profile,
+                    config.workspaceId
                 );
 
             default:
@@ -222,6 +230,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);
@@ -249,10 +259,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() {
@@ -275,6 +290,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
@@ -320,7 +338,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);
     }
@@ -335,6 +354,7 @@ export class DatabricksCliAuthProvider extends AuthProvider {
             authType: this.authType,
             databricksPath: this.cliPath,
             ...(this.profile ? {profile: this.profile} : {}),
+            ...(this.workspaceId ? {workspaceId: this.workspaceId} : {}),
         };
     }
 
@@ -344,6 +364,12 @@ export class DatabricksCliAuthProvider extends AuthProvider {
             authType: "databricks-cli",
             databricksCliPath: this.cliPath,
             ...(this.profile ? {profile: this.profile} : {}),
+            ...(this.workspaceId
+                ? {
+                      workspaceId: this.workspaceId,
+                      experimentalIsUnifiedHost: true,
+                  }
+                : {}),
         });
     }
 
@@ -355,6 +381,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/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,10 @@ 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.16.0
+    "@databricks/sdk-experimental": ^0.17.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.16.0":
-  version: 0.16.0
-  resolution: "@databricks/sdk-experimental@npm:0.16.0"
+"@databricks/sdk-experimental@npm:^0.17.0":
+  version: 0.17.0
+  resolution: "@databricks/sdk-experimental@npm:0.17.0"
   dependencies:
     google-auth-library: ^10.5.0
     ini: ^6.0.0
     reflect-metadata: ^0.2.2
     semver: ^7.7.3
-  checksum: 621cbb34e260a902b46c67ae88a00184985679d10ee450dfaa297d4fb7baccca70d14af3fca2e66d333edd5d0304497fbb2219de231162e62b3293c5bd162933
+  checksum: 34de7d8708de12bf1fa44ae5ced7f8886a26d3ef9f072b47955c63dd7c3ae32db7c0b243074907b371026ee352d76c90d02faa9df0e95ad5d9e8aa22dc53d0b5
   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.16.0
+    "@databricks/sdk-experimental": ^0.17.0
     "@istanbuljs/nyc-config-typescript": ^1.0.2
     "@sinonjs/fake-timers": ^11.2.2
     "@types/bcryptjs": ^2.4.6