auth: extract ?o=/?a= from DATABRICKS_HOST env var

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