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/libs/auth/host_env.go b/libs/auth/host_env.go new file mode 100644 index 00000000000..23326e0269f --- /dev/null +++ b/libs/auth/host_env.go @@ -0,0 +1,52 @@ +package auth + +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 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 helper on +// the next SDK bump that includes that change. + +// 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 + } + params := ExtractHostQueryParams(host) + if params.Host == host { + return + } + cfg.Host = params.Host + if cfg.WorkspaceID == "" && params.WorkspaceID != "" && env.Get(ctx, "DATABRICKS_WORKSPACE_ID") == "" { + cfg.WorkspaceID = params.WorkspaceID + } + 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 new file mode 100644 index 00000000000..152f4d22d47 --- /dev/null +++ b/libs/auth/host_env_test.go @@ -0,0 +1,76 @@ +package auth + +import ( + "testing" + + "github.com/databricks/cli/libs/env" + sdkconfig "github.com/databricks/databricks-sdk-go/config" + "github.com/stretchr/testify/assert" +) + +func TestNormalizeDatabricksConfigFromEnv(t *testing.T) { + tests := []struct { + name string + host string + envWorkspaceID string + envAccountID string + cfgInHost 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", + wantWorkspaceID: "12345", + }, + { + 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 a no-op", + host: "https://acme.databricks.net", + }, + { + name: "env workspace id wins over query param", + host: "https://acme.databricks.net/?o=12345", + envWorkspaceID: "99999", + wantHost: "https://acme.databricks.net", + 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: "no host env 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", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + 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) + + cfg := &sdkconfig.Config{Host: tt.cfgInHost} + NormalizeDatabricksConfigFromEnv(ctx, cfg) + + assert.Equal(t, tt.wantHost, cfg.Host) + assert.Equal(t, tt.wantWorkspaceID, cfg.WorkspaceID) + assert.Equal(t, tt.wantAccountID, cfg.AccountID) + }) + } +}