From ab33f0e43bd8833da92ce043ebf8bff62691f62a Mon Sep 17 00:00:00 2001 From: simon Date: Wed, 27 May 2026 12:56:51 +0200 Subject: [PATCH 1/6] auth: extract ?o=/?a= from DATABRICKS_HOST env var Pasting a SPOG URL from the Databricks UI into DATABRICKS_HOST drops the workspace identifier on the way into the SDK: the SDK strips query parameters in fixHostIfNeeded without promoting them to WorkspaceID. API calls then hit the SPOG without an X-Databricks-Org-Id header and the server answers with HTML, surfacing as "received HTML response instead of JSON". The bundle YAML host field is already normalized via NormalizeHostURL, but the env var bypassed that path. Run the existing ExtractHostQueryParams normalization on DATABRICKS_HOST at the top of root.Execute before any SDK code runs. Promote ?o= and ?a= to DATABRICKS_WORKSPACE_ID and DATABRICKS_ACCOUNT_ID when those env vars are unset, and rewrite DATABRICKS_HOST without the query string. Co-authored-by: Isaac --- NEXT_CHANGELOG.md | 1 + cmd/root/root.go | 6 ++ libs/auth/host_env.go | 39 +++++++++++++ libs/auth/host_env_test.go | 110 +++++++++++++++++++++++++++++++++++++ 4 files changed, 156 insertions(+) create mode 100644 libs/auth/host_env.go create mode 100644 libs/auth/host_env_test.go diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index 004cd4da4c0..ac2aa9ec6a6 100644 --- a/NEXT_CHANGELOG.md +++ b/NEXT_CHANGELOG.md @@ -5,6 +5,7 @@ ### Notable Changes ### CLI +* Recognize `?o=` and `?a=` in `DATABRICKS_HOST` so SPOG URLs pasted from the Databricks UI route correctly without manually splitting out `DATABRICKS_WORKSPACE_ID` / `DATABRICKS_ACCOUNT_ID`. ### Bundles * The error reported when a direct-only resource (catalogs, external locations, vector search endpoints) is used with the terraform engine now also suggests setting `bundle.engine: direct` in `databricks.yml`, in addition to the `DATABRICKS_BUNDLE_ENGINE` environment variable ([#5295](https://github.com/databricks/cli/pull/5295)). diff --git a/cmd/root/root.go b/cmd/root/root.go index 6b6de2a9baa..f9fd2ea6713 100644 --- a/cmd/root/root.go +++ b/cmd/root/root.go @@ -131,6 +131,12 @@ Stack Trace: %s`, version, r, string(trace)) }() + // Promote ?o=/?a= query parameters in DATABRICKS_HOST to + // DATABRICKS_WORKSPACE_ID/DATABRICKS_ACCOUNT_ID before the SDK reads the + // env var. Without this, SPOG URLs pasted from the Databricks UI lose + // their workspace routing identifier and API calls return HTML. + auth.NormalizeDatabricksHostEnv() + // Configure a telemetry logger and store it in the context. ctx = telemetry.WithNewLogger(ctx) diff --git a/libs/auth/host_env.go b/libs/auth/host_env.go new file mode 100644 index 00000000000..b828ac6cdf1 --- /dev/null +++ b/libs/auth/host_env.go @@ -0,0 +1,39 @@ +package auth + +import "os" + +// SPOG URLs from the Databricks UI carry the workspace ID as a ?o= query +// parameter and the account ID as ?a=, e.g. +// https://acme.databricks.net/?o=12345. The SDK strips path and query from +// Host in fixHostIfNeeded without extracting these IDs, so pasting such a +// URL into DATABRICKS_HOST drops the workspace identifier and API calls hit +// the SPOG without an X-Databricks-Org-Id header, which the server answers +// with HTML (a login page) instead of JSON. +const ( + envHost = "DATABRICKS_HOST" + envWorkspaceID = "DATABRICKS_WORKSPACE_ID" + envAccountID = "DATABRICKS_ACCOUNT_ID" +) + +// NormalizeDatabricksHostEnv extracts ?o=/?workspace_id= and ?a=/?account_id= +// from DATABRICKS_HOST and promotes them to DATABRICKS_WORKSPACE_ID and +// DATABRICKS_ACCOUNT_ID respectively, then rewrites DATABRICKS_HOST without +// the query string. Existing values of the destination env vars are never +// overwritten. Safe to call when DATABRICKS_HOST is unset or has no query. +func NormalizeDatabricksHostEnv() { + host, ok := os.LookupEnv(envHost) + if !ok || host == "" { + return + } + params := ExtractHostQueryParams(host) + if params.Host == host { + return + } + os.Setenv(envHost, params.Host) + if params.WorkspaceID != "" && os.Getenv(envWorkspaceID) == "" { + os.Setenv(envWorkspaceID, params.WorkspaceID) + } + if params.AccountID != "" && os.Getenv(envAccountID) == "" { + os.Setenv(envAccountID, params.AccountID) + } +} diff --git a/libs/auth/host_env_test.go b/libs/auth/host_env_test.go new file mode 100644 index 00000000000..5bdedef276e --- /dev/null +++ b/libs/auth/host_env_test.go @@ -0,0 +1,110 @@ +package auth + +import ( + "os" + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestNormalizeDatabricksHostEnv(t *testing.T) { + tests := []struct { + name string + host string + workspaceID string + accountID string + wantHost string + wantHostSet bool + wantWorkspaceID string + wantWorkspaceSet bool + wantAccountID string + wantAccountIDSet bool + preserveHostUnset bool + }{ + { + name: "spog url promotes workspace id", + host: "https://acme.databricks.net/?o=12345", + wantHost: "https://acme.databricks.net", + wantHostSet: true, + wantWorkspaceID: "12345", + wantWorkspaceSet: true, + }, + { + name: "spog url with account id", + host: "https://acme.databricks.net/?a=abc&o=12345", + wantHost: "https://acme.databricks.net", + wantHostSet: true, + wantWorkspaceID: "12345", + wantWorkspaceSet: true, + wantAccountID: "abc", + wantAccountIDSet: true, + }, + { + name: "host without query is left alone", + host: "https://acme.databricks.net", + wantHost: "https://acme.databricks.net", + wantHostSet: true, + }, + { + name: "existing workspace id is preserved", + host: "https://acme.databricks.net/?o=12345", + workspaceID: "99999", + wantHost: "https://acme.databricks.net", + wantHostSet: true, + wantWorkspaceID: "99999", + wantWorkspaceSet: true, + }, + { + name: "no host env unsets nothing", + preserveHostUnset: true, + }, + { + name: "non-numeric o is dropped, host trailing slash trimmed", + host: "https://acme.databricks.net/?o=notanumber", + wantHost: "https://acme.databricks.net", + wantHostSet: true, + wantWorkspaceID: "", + wantWorkspaceSet: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if tt.preserveHostUnset { + os.Unsetenv(envHost) + } else { + t.Setenv(envHost, tt.host) + } + if tt.workspaceID != "" { + t.Setenv(envWorkspaceID, tt.workspaceID) + } else { + os.Unsetenv(envWorkspaceID) + } + if tt.accountID != "" { + t.Setenv(envAccountID, tt.accountID) + } else { + os.Unsetenv(envAccountID) + } + + NormalizeDatabricksHostEnv() + + gotHost, hostSet := os.LookupEnv(envHost) + assert.Equal(t, tt.wantHostSet, hostSet) + if tt.wantHostSet { + assert.Equal(t, tt.wantHost, gotHost) + } + + gotWS, wsSet := os.LookupEnv(envWorkspaceID) + assert.Equal(t, tt.wantWorkspaceSet, wsSet) + if tt.wantWorkspaceSet { + assert.Equal(t, tt.wantWorkspaceID, gotWS) + } + + gotAcc, accSet := os.LookupEnv(envAccountID) + assert.Equal(t, tt.wantAccountIDSet, accSet) + if tt.wantAccountIDSet { + assert.Equal(t, tt.wantAccountID, gotAcc) + } + }) + } +} From 7f3cfa571279b32e8bfdd1507c718b43a44abe03 Mon Sep 17 00:00:00 2001 From: simon Date: Wed, 27 May 2026 13:26:50 +0200 Subject: [PATCH 2/6] Drop NEXT_CHANGELOG entry Co-authored-by: Isaac --- NEXT_CHANGELOG.md | 1 - 1 file changed, 1 deletion(-) diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index ac2aa9ec6a6..004cd4da4c0 100644 --- a/NEXT_CHANGELOG.md +++ b/NEXT_CHANGELOG.md @@ -5,7 +5,6 @@ ### Notable Changes ### CLI -* Recognize `?o=` and `?a=` in `DATABRICKS_HOST` so SPOG URLs pasted from the Databricks UI route correctly without manually splitting out `DATABRICKS_WORKSPACE_ID` / `DATABRICKS_ACCOUNT_ID`. ### Bundles * The error reported when a direct-only resource (catalogs, external locations, vector search endpoints) is used with the terraform engine now also suggests setting `bundle.engine: direct` in `databricks.yml`, in addition to the `DATABRICKS_BUNDLE_ENGINE` environment variable ([#5295](https://github.com/databricks/cli/pull/5295)). From a83c25e167b983e6861258c9b674720eb19d79f8 Mon Sep 17 00:00:00 2001 From: simon Date: Wed, 27 May 2026 13:35:29 +0200 Subject: [PATCH 3/6] Use os.LookupEnv instead of os.Getenv for lint Co-authored-by: Isaac --- libs/auth/host_env.go | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/libs/auth/host_env.go b/libs/auth/host_env.go index b828ac6cdf1..7827c783a85 100644 --- a/libs/auth/host_env.go +++ b/libs/auth/host_env.go @@ -30,10 +30,14 @@ func NormalizeDatabricksHostEnv() { return } os.Setenv(envHost, params.Host) - if params.WorkspaceID != "" && os.Getenv(envWorkspaceID) == "" { - os.Setenv(envWorkspaceID, params.WorkspaceID) + if params.WorkspaceID != "" { + if _, set := os.LookupEnv(envWorkspaceID); !set { + os.Setenv(envWorkspaceID, params.WorkspaceID) + } } - if params.AccountID != "" && os.Getenv(envAccountID) == "" { - os.Setenv(envAccountID, params.AccountID) + if params.AccountID != "" { + if _, set := os.LookupEnv(envAccountID); !set { + os.Setenv(envAccountID, params.AccountID) + } } } From 5084f3e7a045acd0039e61578b1ec172f99f194e Mon Sep 17 00:00:00 2001 From: simon Date: Wed, 27 May 2026 14:50:53 +0200 Subject: [PATCH 4/6] auth: add TODO pointing host_env normalizer at SDK #1699 Stopgap until the SDK fix lands and the CLI bumps to a version that includes it; then this normalizer can be deleted. Co-authored-by: Isaac --- libs/auth/host_env.go | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/libs/auth/host_env.go b/libs/auth/host_env.go index 7827c783a85..5eb2c4824e1 100644 --- a/libs/auth/host_env.go +++ b/libs/auth/host_env.go @@ -9,6 +9,10 @@ import "os" // URL into DATABRICKS_HOST drops the workspace identifier and API calls hit // the SPOG without an X-Databricks-Org-Id header, which the server answers // with HTML (a login page) instead of JSON. +// +// TODO: stopgap. The matching SDK fix is databricks/databricks-sdk-go#1699, +// which handles ?o=/?a= directly in fixHostIfNeeded. Delete this normalizer +// on the next SDK bump that includes that change. const ( envHost = "DATABRICKS_HOST" envWorkspaceID = "DATABRICKS_WORKSPACE_ID" From 9db4b3b40e5b3231f7fa3b602c1e414f1caf6d07 Mon Sep 17 00:00:00 2001 From: simon Date: Wed, 27 May 2026 15:31:18 +0200 Subject: [PATCH 5/6] auth: address review feedback on host_env normalizer - Treat an explicitly-empty DATABRICKS_WORKSPACE_ID/ACCOUNT_ID the same as unset, so a paste-empty env doesn't block ?o=/?a= promotion. - Add a TODO at the env-var-precedence branch so the silent-preserve behavior is easy to revisit when the SDK fix lands. - Use t.Setenv exclusively in the table-driven test (parallel-safe, auto-cleanup) and assert against os.Getenv directly. Co-authored-by: Isaac --- libs/auth/host_env.go | 8 ++- libs/auth/host_env_test.go | 110 ++++++++++++------------------------- 2 files changed, 40 insertions(+), 78 deletions(-) diff --git a/libs/auth/host_env.go b/libs/auth/host_env.go index 5eb2c4824e1..72be0a85b2a 100644 --- a/libs/auth/host_env.go +++ b/libs/auth/host_env.go @@ -34,13 +34,17 @@ func NormalizeDatabricksHostEnv() { return } os.Setenv(envHost, params.Host) + // TODO: existing env vars take precedence over the query params here; + // they are treated as a more explicit signal than a value embedded in + // the host URL. Revisit once the SDK fix lands (see top of file), + // erroring on a real conflict may be preferable to silent-preserve. if params.WorkspaceID != "" { - if _, set := os.LookupEnv(envWorkspaceID); !set { + if cur, _ := os.LookupEnv(envWorkspaceID); cur == "" { os.Setenv(envWorkspaceID, params.WorkspaceID) } } if params.AccountID != "" { - if _, set := os.LookupEnv(envAccountID); !set { + if cur, _ := os.LookupEnv(envAccountID); cur == "" { os.Setenv(envAccountID, params.AccountID) } } diff --git a/libs/auth/host_env_test.go b/libs/auth/host_env_test.go index 5bdedef276e..43155cb5a3f 100644 --- a/libs/auth/host_env_test.go +++ b/libs/auth/host_env_test.go @@ -9,102 +9,60 @@ import ( func TestNormalizeDatabricksHostEnv(t *testing.T) { tests := []struct { - name string - host string - workspaceID string - accountID string - wantHost string - wantHostSet bool - wantWorkspaceID string - wantWorkspaceSet bool - wantAccountID string - wantAccountIDSet bool - preserveHostUnset bool + name string + host string + workspaceID string + accountID string + wantHost string + wantWorkspaceID string + wantAccountID string }{ { - name: "spog url promotes workspace id", - host: "https://acme.databricks.net/?o=12345", - wantHost: "https://acme.databricks.net", - wantHostSet: true, - wantWorkspaceID: "12345", - wantWorkspaceSet: true, + name: "spog url promotes workspace id", + host: "https://acme.databricks.net/?o=12345", + wantHost: "https://acme.databricks.net", + wantWorkspaceID: "12345", }, { - name: "spog url with account id", - host: "https://acme.databricks.net/?a=abc&o=12345", - wantHost: "https://acme.databricks.net", - wantHostSet: true, - wantWorkspaceID: "12345", - wantWorkspaceSet: true, - wantAccountID: "abc", - wantAccountIDSet: true, + name: "spog url with account id", + host: "https://acme.databricks.net/?a=abc&o=12345", + wantHost: "https://acme.databricks.net", + wantWorkspaceID: "12345", + wantAccountID: "abc", }, { - name: "host without query is left alone", - host: "https://acme.databricks.net", - wantHost: "https://acme.databricks.net", - wantHostSet: true, + name: "host without query is left alone", + host: "https://acme.databricks.net", + wantHost: "https://acme.databricks.net", }, { - name: "existing workspace id is preserved", - host: "https://acme.databricks.net/?o=12345", - workspaceID: "99999", - wantHost: "https://acme.databricks.net", - wantHostSet: true, - wantWorkspaceID: "99999", - wantWorkspaceSet: true, + name: "existing workspace id is preserved", + host: "https://acme.databricks.net/?o=12345", + workspaceID: "99999", + wantHost: "https://acme.databricks.net", + wantWorkspaceID: "99999", }, { - name: "no host env unsets nothing", - preserveHostUnset: true, + name: "empty host is a no-op", }, { - name: "non-numeric o is dropped, host trailing slash trimmed", - host: "https://acme.databricks.net/?o=notanumber", - wantHost: "https://acme.databricks.net", - wantHostSet: true, - wantWorkspaceID: "", - wantWorkspaceSet: false, + name: "non-numeric o is dropped, host trailing slash trimmed", + host: "https://acme.databricks.net/?o=notanumber", + wantHost: "https://acme.databricks.net", }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - if tt.preserveHostUnset { - os.Unsetenv(envHost) - } else { - t.Setenv(envHost, tt.host) - } - if tt.workspaceID != "" { - t.Setenv(envWorkspaceID, tt.workspaceID) - } else { - os.Unsetenv(envWorkspaceID) - } - if tt.accountID != "" { - t.Setenv(envAccountID, tt.accountID) - } else { - os.Unsetenv(envAccountID) - } + t.Setenv(envHost, tt.host) + t.Setenv(envWorkspaceID, tt.workspaceID) + t.Setenv(envAccountID, tt.accountID) NormalizeDatabricksHostEnv() - gotHost, hostSet := os.LookupEnv(envHost) - assert.Equal(t, tt.wantHostSet, hostSet) - if tt.wantHostSet { - assert.Equal(t, tt.wantHost, gotHost) - } - - gotWS, wsSet := os.LookupEnv(envWorkspaceID) - assert.Equal(t, tt.wantWorkspaceSet, wsSet) - if tt.wantWorkspaceSet { - assert.Equal(t, tt.wantWorkspaceID, gotWS) - } - - gotAcc, accSet := os.LookupEnv(envAccountID) - assert.Equal(t, tt.wantAccountIDSet, accSet) - if tt.wantAccountIDSet { - assert.Equal(t, tt.wantAccountID, gotAcc) - } + assert.Equal(t, tt.wantHost, os.Getenv(envHost)) + assert.Equal(t, tt.wantWorkspaceID, os.Getenv(envWorkspaceID)) + assert.Equal(t, tt.wantAccountID, os.Getenv(envAccountID)) }) } } From 8a44c5e618da454788b7143b595bb9b53d766c4d Mon Sep 17 00:00:00 2001 From: simon Date: Wed, 27 May 2026 16:15:19 +0200 Subject: [PATCH 6/6] auth: scope host normalization to SDK Config, no env-var mutation Replaces NormalizeDatabricksHostEnv (which called os.Setenv on DATABRICKS_HOST/WORKSPACE_ID/ACCOUNT_ID at the top of root.Execute) with NormalizeDatabricksConfigFromEnv(ctx, cfg), which sets the SDK Config fields directly. Same end result (?o=/?a= promoted, host stripped), but the process environment is left untouched. Called from the three places that build a Config from env: MustAccountClient, MustWorkspaceClient, and cmd/api's makeCommand. Subprocesses (terraform) still get the normalized values via auth.Env(cfg) -> b.AuthEnv(), which derives env from cfg rather than reading os.Environ. Co-authored-by: Isaac --- cmd/api/api.go | 2 ++ cmd/root/auth.go | 2 ++ cmd/root/root.go | 6 ---- libs/auth/host_env.go | 61 +++++++++++++++++++------------------- libs/auth/host_env_test.go | 44 ++++++++++++++++----------- 5 files changed, 61 insertions(+), 54 deletions(-) diff --git a/cmd/api/api.go b/cmd/api/api.go index ab70ca8b753..a832f018c1a 100644 --- a/cmd/api/api.go +++ b/cmd/api/api.go @@ -94,6 +94,8 @@ func makeCommand(method string) *cobra.Command { cfg.Profile = databrickscfg.ResolveDefaultProfile(cmd.Context()) } + auth.NormalizeDatabricksConfigFromEnv(cmd.Context(), cfg) + api, err := client.New(cfg) if err != nil { return err diff --git a/cmd/root/auth.go b/cmd/root/auth.go index 8b9bfd0810b..d40bbe7d8c8 100644 --- a/cmd/root/auth.go +++ b/cmd/root/auth.go @@ -152,6 +152,7 @@ func MustAccountClient(cmd *cobra.Command, args []string) error { } ctx := cmd.Context() + auth.NormalizeDatabricksConfigFromEnv(ctx, cfg) ctx = cmdctx.SetConfigUsed(ctx, cfg) cmd.SetContext(ctx) @@ -250,6 +251,7 @@ func MustWorkspaceClient(cmd *cobra.Command, args []string) error { cfg.Profile = profile } + auth.NormalizeDatabricksConfigFromEnv(ctx, cfg) resolveDefaultProfile(ctx, cfg) _, isTargetFlagSet := targetFlagValue(cmd) diff --git a/cmd/root/root.go b/cmd/root/root.go index f9fd2ea6713..6b6de2a9baa 100644 --- a/cmd/root/root.go +++ b/cmd/root/root.go @@ -131,12 +131,6 @@ Stack Trace: %s`, version, r, string(trace)) }() - // Promote ?o=/?a= query parameters in DATABRICKS_HOST to - // DATABRICKS_WORKSPACE_ID/DATABRICKS_ACCOUNT_ID before the SDK reads the - // env var. Without this, SPOG URLs pasted from the Databricks UI lose - // their workspace routing identifier and API calls return HTML. - auth.NormalizeDatabricksHostEnv() - // Configure a telemetry logger and store it in the context. ctx = telemetry.WithNewLogger(ctx) diff --git a/libs/auth/host_env.go b/libs/auth/host_env.go index 72be0a85b2a..23326e0269f 100644 --- a/libs/auth/host_env.go +++ b/libs/auth/host_env.go @@ -1,31 +1,40 @@ package auth -import "os" +import ( + "context" + + "github.com/databricks/cli/libs/env" + sdkconfig "github.com/databricks/databricks-sdk-go/config" +) // SPOG URLs from the Databricks UI carry the workspace ID as a ?o= query // parameter and the account ID as ?a=, e.g. // https://acme.databricks.net/?o=12345. The SDK strips path and query from -// Host in fixHostIfNeeded without extracting these IDs, so pasting such a -// URL into DATABRICKS_HOST drops the workspace identifier and API calls hit +// Host in fixHostIfNeeded without extracting these IDs, so a DATABRICKS_HOST +// env var with such a URL drops the workspace identifier and API calls hit // the SPOG without an X-Databricks-Org-Id header, which the server answers // with HTML (a login page) instead of JSON. // // TODO: stopgap. The matching SDK fix is databricks/databricks-sdk-go#1699, -// which handles ?o=/?a= directly in fixHostIfNeeded. Delete this normalizer -// on the next SDK bump that includes that change. -const ( - envHost = "DATABRICKS_HOST" - envWorkspaceID = "DATABRICKS_WORKSPACE_ID" - envAccountID = "DATABRICKS_ACCOUNT_ID" -) +// which handles ?o=/?a= directly in fixHostIfNeeded. Delete this helper on +// the next SDK bump that includes that change. -// NormalizeDatabricksHostEnv extracts ?o=/?workspace_id= and ?a=/?account_id= -// from DATABRICKS_HOST and promotes them to DATABRICKS_WORKSPACE_ID and -// DATABRICKS_ACCOUNT_ID respectively, then rewrites DATABRICKS_HOST without -// the query string. Existing values of the destination env vars are never -// overwritten. Safe to call when DATABRICKS_HOST is unset or has no query. -func NormalizeDatabricksHostEnv() { - host, ok := os.LookupEnv(envHost) +// NormalizeDatabricksConfigFromEnv promotes ?o=/?workspace_id= and +// ?a=/?account_id= query parameters from the DATABRICKS_HOST env var into +// the matching fields on cfg, and sets cfg.Host to the stripped URL. It +// does not mutate process env, so the effect is scoped to the SDK config +// built from this cfg (and any subprocess env derived from it via +// auth.Env). +// +// Only fills in empty fields. If cfg.Host is already set, the query +// params aren't promoted at all (an explicit host takes priority). If a +// dedicated env var (DATABRICKS_WORKSPACE_ID, DATABRICKS_ACCOUNT_ID) is +// set, that more explicit signal wins over the query param. +func NormalizeDatabricksConfigFromEnv(ctx context.Context, cfg *sdkconfig.Config) { + if cfg.Host != "" { + return + } + host, ok := env.Lookup(ctx, "DATABRICKS_HOST") if !ok || host == "" { return } @@ -33,19 +42,11 @@ func NormalizeDatabricksHostEnv() { if params.Host == host { return } - os.Setenv(envHost, params.Host) - // TODO: existing env vars take precedence over the query params here; - // they are treated as a more explicit signal than a value embedded in - // the host URL. Revisit once the SDK fix lands (see top of file), - // erroring on a real conflict may be preferable to silent-preserve. - if params.WorkspaceID != "" { - if cur, _ := os.LookupEnv(envWorkspaceID); cur == "" { - os.Setenv(envWorkspaceID, params.WorkspaceID) - } + cfg.Host = params.Host + if cfg.WorkspaceID == "" && params.WorkspaceID != "" && env.Get(ctx, "DATABRICKS_WORKSPACE_ID") == "" { + cfg.WorkspaceID = params.WorkspaceID } - if params.AccountID != "" { - if cur, _ := os.LookupEnv(envAccountID); cur == "" { - os.Setenv(envAccountID, params.AccountID) - } + if cfg.AccountID == "" && params.AccountID != "" && env.Get(ctx, "DATABRICKS_ACCOUNT_ID") == "" { + cfg.AccountID = params.AccountID } } diff --git a/libs/auth/host_env_test.go b/libs/auth/host_env_test.go index 43155cb5a3f..152f4d22d47 100644 --- a/libs/auth/host_env_test.go +++ b/libs/auth/host_env_test.go @@ -1,18 +1,20 @@ package auth import ( - "os" "testing" + "github.com/databricks/cli/libs/env" + sdkconfig "github.com/databricks/databricks-sdk-go/config" "github.com/stretchr/testify/assert" ) -func TestNormalizeDatabricksHostEnv(t *testing.T) { +func TestNormalizeDatabricksConfigFromEnv(t *testing.T) { tests := []struct { name string host string - workspaceID string - accountID string + envWorkspaceID string + envAccountID string + cfgInHost string wantHost string wantWorkspaceID string wantAccountID string @@ -31,19 +33,24 @@ func TestNormalizeDatabricksHostEnv(t *testing.T) { wantAccountID: "abc", }, { - name: "host without query is left alone", - host: "https://acme.databricks.net", - wantHost: "https://acme.databricks.net", + name: "host without query is a no-op", + host: "https://acme.databricks.net", }, { - name: "existing workspace id is preserved", + name: "env workspace id wins over query param", host: "https://acme.databricks.net/?o=12345", - workspaceID: "99999", + envWorkspaceID: "99999", wantHost: "https://acme.databricks.net", - wantWorkspaceID: "99999", + wantWorkspaceID: "", + }, + { + name: "cfg host already set leaves env alone", + host: "https://other.databricks.net/?o=12345", + cfgInHost: "https://acme.databricks.net", + wantHost: "https://acme.databricks.net", }, { - name: "empty host is a no-op", + name: "no host env is a no-op", }, { name: "non-numeric o is dropped, host trailing slash trimmed", @@ -54,15 +61,16 @@ func TestNormalizeDatabricksHostEnv(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - t.Setenv(envHost, tt.host) - t.Setenv(envWorkspaceID, tt.workspaceID) - t.Setenv(envAccountID, tt.accountID) + ctx := env.Set(t.Context(), "DATABRICKS_HOST", tt.host) + ctx = env.Set(ctx, "DATABRICKS_WORKSPACE_ID", tt.envWorkspaceID) + ctx = env.Set(ctx, "DATABRICKS_ACCOUNT_ID", tt.envAccountID) - NormalizeDatabricksHostEnv() + cfg := &sdkconfig.Config{Host: tt.cfgInHost} + NormalizeDatabricksConfigFromEnv(ctx, cfg) - assert.Equal(t, tt.wantHost, os.Getenv(envHost)) - assert.Equal(t, tt.wantWorkspaceID, os.Getenv(envWorkspaceID)) - assert.Equal(t, tt.wantAccountID, os.Getenv(envAccountID)) + assert.Equal(t, tt.wantHost, cfg.Host) + assert.Equal(t, tt.wantWorkspaceID, cfg.WorkspaceID) + assert.Equal(t, tt.wantAccountID, cfg.AccountID) }) } }