auth: extract ?o=/?a= from DATABRICKS_HOST env var (#5337)

## Why

Pasting a SPOG URL from the Databricks UI (e.g.
`https://acme.azuredatabricks.net/?o=12345`) into `DATABRICKS_HOST`
drops the workspace identifier before any API call goes out. The SDK
strips path and query from `Host` in `fixHostIfNeeded` without promoting
`?o=` to `WorkspaceID`, so the request goes to the SPOG hostname without
an `X-Databricks-Org-Id` header. The server can't route it and answers
with the login HTML page, which surfaces as:

```
$ DATABRICKS_HOST=https://acme.azuredatabricks.net/?o=12345 databricks bundle validate
Error: received HTML response instead of JSON
```

The bundle YAML `workspace.host` field is already normalized via
`NormalizeHostURL` in `bundle/config/workspace.go`, and `databricks api`
handles `?o=` per call (#5137). The env-var path was the remaining gap.

## Changes

Before: `DATABRICKS_HOST=https://acme.databricks.net/?o=12345` reached
the SDK with the query intact, the SDK dropped it, and `WorkspaceID`
stayed empty.

Now: `auth.NormalizeDatabricksHostEnv` runs once at the top of
`root.Execute`, before the SDK reads anything. It uses the existing
`ExtractHostQueryParams` helper to promote `?o=` / `?workspace_id=` to
`DATABRICKS_WORKSPACE_ID` and `?a=` / `?account_id=` to
`DATABRICKS_ACCOUNT_ID` (only when those env vars are unset), then
rewrites `DATABRICKS_HOST` without the query string.

A follow-up PR will push the same normalization into the SDK's
`fixHostIfNeeded` so Python/Java/JS SDK users and any direct Go-SDK
callers get the same fix without going through the CLI.

## Test plan

- [x] `go test ./libs/auth/` covers the new helper with table-driven
cases: `?o=` promotion, `?a=` + `?o=` together, existing
`DATABRICKS_WORKSPACE_ID` is preserved, hosts without query are
untouched, non-numeric `?o=` is dropped, unset `DATABRICKS_HOST` is a
no-op.
- [x] `go test ./cmd/root/` passes.
- [x] `./task checks`.
- [x] Manual end-to-end repro with a local SPOG-shaped test server:
before the fix the SDK sent no `X-Databricks-Org-Id` header and got HTML
back; after the fix the header is `258628866953061` and `bundle
validate` proceeds to the workspace API.

Author: simonfaltum

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

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)

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

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