Add SPOG (Custom URL) routing support via x-databricks-org-id header (#347)

## Summary

On SPOG (Custom URL / account-level) workspaces, `httpPath` has the form
`/sql/1.0/warehouses/<id>?o=<workspaceId>`. The `?o=` parameter routes
Thrift calls correctly via the URL, but other endpoints (telemetry push,
feature-flag check) run on separate hosts and need `x-databricks-org-id`
as an HTTP header to route to the right workspace. Without it, those
requests 404 or get misrouted on SPOG hosts.

## Change

All contained in `connector.go`:

1. `extractSpogHeaders(httpPath string) map[string]string` — parses the
`?o=` query param using `url.ParseQuery` (stdlib, not regex). Returns
`{"x-databricks-org-id": "<workspaceId>"}` or `nil`. Three DEBUG log
paths cover: malformed query, missing `?o=`, and successful extraction.

2. `headerInjectingTransport` — a lightweight `http.RoundTripper`
wrapper that clones the request per the contract and sets the provided
headers if not already set by the caller.

3. `withSpogHeaders(base *http.Client, headers map[string]string)
*http.Client` — returns a new client with the same settings but a
wrapped transport.

4. In `Connect()`, when `extractSpogHeaders` returns non-nil, the driver
passes a wrapped client into `TelemetryInitOptions.HTTPClient`. The
wrapped client is used for both the feature-flag check and the telemetry
push. The driver's own `c.client` is left alone, so Thrift routing
(which uses `?o=` in the URL) is unaffected.

## Why a transport wrapper instead of threading a parameter

An earlier revision of this PR threaded an `extraHeaders` parameter
through `telemetry.TelemetryInitOptions` → `isTelemetryEnabled` →
`featureFlagCache.isTelemetryEnabled` → `fetchFeatureFlag`. That
approach:

- Required API-surface changes in 3 telemetry files (`config.go`,
`featureflag.go`, `driver_integration.go`).
- Only covered the feature-flag check; `telemetry/exporter.go`
(telemetry push) still sent `Content-Type` as the only header — SPOG
routing would 404 at push time.

The RoundTripper wrapper:

- Keeps `telemetry/*` identical to `origin/main`. Zero API churn.
- Automatically applies to every outbound request using the wrapped
client — feature-flag check, telemetry push, and any future HTTP paths
that reuse it.
- Respects caller-set headers (if a request already has
`x-databricks-org-id` for some reason, the wrapper does not override).

## Endpoints covered

| Endpoint | Uses wrapped client? | Gets `x-databricks-org-id`? |
| :---- | :----: | :----: |
| Feature flags `/api/2.0/connector-service/feature-flags/GOLANG/{v}` |
Yes | ✅ |
| Telemetry push | Yes | ✅ |
| Thrift | No (uses `c.client` directly; already routes via URL) | —
(not needed — URL-routed) |
| OAuth token exchange | No (separate client, talks to
`login.microsoftonline.com`) | — (not needed) |
| CloudFetch / Volume operations | No (presigned URLs) | — (not needed)
|

## Verification

- `go build ./...` clean.
- Debug log output confirms extraction on SPOG URLs:
  ```
SPOG header extraction: injecting x-databricks-org-id=<id> (extracted
from ?o= in httpPath)
  ```

## Note on the earlier `auth/oauth/u2m/u2m.go` change

The two earliest commits on this branch (`23697e5`, `0ec7e06`) modified
`auth/oauth/u2m/u2m.go` to avoid sending an empty `client_secret` on the
PKCE public-app flow, documented as a fix for the server's `"Public app
should not use a client secret"` rejection. **That change was
empirically verified to not be needed** (see commit `3576c92` which
reverts it):

- **Prod Legacy** (`adb-6436897454825492.12.azuredatabricks.net`):
Unpatched `u2m.go` PASSES end-to-end. Server accepts the empty
`client_secret`.
- **Stg Legacy** (`adb-7064161269814046.2.staging.azuredatabricks.net`):
Unpatched `u2m.go` FAILS with `unexpected HTTP status 400 Bad Request`
during token exchange.

Since production tolerates the current behavior, customers aren't
impacted. Keeping this PR minimal; if staging-level strictness later
rolls out to prod, we can re-add the u2m fix then.

## Test plan

- [ ] Unit tests pass (`go test ./...`)
- [ ] Manually verified SPOG header injection appears in
`x-databricks-org-id` on feature-flag check and telemetry push against a
SPOG workspace (not gated by this PR, but prerequisite for SPOG
telemetry to work at all)
- [ ] No regressions on non-SPOG workspaces (wrapper is a no-op when
`extractSpogHeaders` returns nil)

---------

Signed-off-by: Madhavendra Rathore <madhavendra.rathore@databricks.com>
Signed-off-by: Madhavendra Rathore
Co-authored-by: Samikshya Chand <148681192+samikshya-db@users.noreply.github.com>

Author: msrathore-db

connector.go

@@ -6,6 +6,7 @@ import (
 	"database/sql/driver"
 	"fmt"
 	"net/http"
+	"net/url"
 	"strings"
 	"time"
 
@@ -80,11 +81,21 @@ func (c *connector) Connect(ctx context.Context) (driver.Conn, error) {
 	}
 	log := logger.WithContext(conn.id, driverctx.CorrelationIdFromContext(ctx), "")
 
+	// Extract SPOG routing headers from ?o= in HTTPPath. When ?o=<workspaceId>
+	// is present (Custom URL / SPOG hosts), wrap the HTTP client used for
+	// telemetry + feature-flag calls with a transport that injects
+	// x-databricks-org-id. Thrift routes via the URL so its own c.client
+	// doesn't need wrapping.
+	telemetryClient := c.client
+	if spogHeaders := extractSpogHeaders(c.cfg.HTTPPath); len(spogHeaders) > 0 {
+		telemetryClient = withSpogHeaders(c.client, spogHeaders)
+	}
+
 	// Initialize telemetry: client config overlay decides; if unset, feature flags decide
 	conn.telemetry = telemetry.InitializeForConnection(ctx, telemetry.TelemetryInitOptions{
 		Host:            c.cfg.Host,
 		DriverVersion:   c.cfg.DriverVersion,
-		HTTPClient:      c.client,
+		HTTPClient:      telemetryClient,
 		EnableTelemetry: c.cfg.EnableTelemetry,
 		BatchSize:       c.cfg.TelemetryBatchSize,
 		FlushInterval:   c.cfg.TelemetryFlushInterval,
@@ -126,6 +137,91 @@ func NewConnector(options ...ConnOption) (driver.Connector, error) {
 	return &connector{cfg: cfg, client: client}, nil
 }
 
+// extractSpogHeaders extracts ?o=<workspaceId> from httpPath and returns
+// an x-databricks-org-id header for SPOG routing.
+//
+// On SPOG (Custom URL) workspaces, httpPath is of the form
+// /sql/1.0/warehouses/<id>?o=<workspaceId>. The ?o= parameter keeps Thrift
+// requests routed to the correct workspace via the URL itself, but other
+// endpoints (telemetry, feature flags) run on separate hosts and need the
+// x-databricks-org-id header. This function extracts ?o= from httpPath once
+// and returns it so those paths can inject it as an HTTP header.
+//
+// Returns nil if:
+//   - httpPath has no query string ("?"), or
+//   - the query string is malformed and can't be parsed, or
+//   - the ?o= parameter is missing or empty.
+func extractSpogHeaders(httpPath string) map[string]string {
+	if !strings.Contains(httpPath, "?") {
+		return nil
+	}
+	// Parse query string from httpPath
+	parts := strings.SplitN(httpPath, "?", 2)
+	params, err := url.ParseQuery(parts[1])
+	if err != nil {
+		logger.Debug().Msgf(
+			"SPOG header extraction: malformed query string in httpPath, skipping org-id extraction: %s",
+			err)
+		return nil
+	}
+	orgID := params.Get("o")
+	if orgID == "" {
+		logger.Debug().Msg(
+			"SPOG header extraction: httpPath has query string but no ?o= param, " +
+				"skipping x-databricks-org-id injection")
+		return nil
+	}
+	logger.Debug().Msgf(
+		"SPOG header extraction: injecting x-databricks-org-id=%s (extracted from ?o= in httpPath)",
+		orgID)
+	return map[string]string{"x-databricks-org-id": orgID}
+}
+
+// withSpogHeaders returns a new *http.Client that reuses the transport of the
+// provided client, wrapped to inject the given SPOG headers on every outbound
+// request. The original client is left unchanged. If a request already has a
+// given header set (e.g., the caller set it explicitly), the wrapper does not
+// override it.
+//
+// This is how the driver gets x-databricks-org-id onto both the feature-flag
+// check and the telemetry push without touching the telemetry package's
+// signatures.
+func withSpogHeaders(base *http.Client, headers map[string]string) *http.Client {
+	baseTransport := base.Transport
+	if baseTransport == nil {
+		baseTransport = http.DefaultTransport
+	}
+	return &http.Client{
+		Transport: &headerInjectingTransport{
+			base:    baseTransport,
+			headers: headers,
+		},
+		CheckRedirect: base.CheckRedirect,
+		Jar:           base.Jar,
+		Timeout:       base.Timeout,
+	}
+}
+
+// headerInjectingTransport wraps an http.RoundTripper and sets a fixed set of
+// headers on every outbound request. Caller-supplied headers with the same
+// name are not overridden.
+type headerInjectingTransport struct {
+	base    http.RoundTripper
+	headers map[string]string
+}
+
+// RoundTrip implements http.RoundTripper.
+func (t *headerInjectingTransport) RoundTrip(req *http.Request) (*http.Response, error) {
+	// Clone per RoundTripper contract — must not mutate the caller's request.
+	req2 := req.Clone(req.Context())
+	for k, v := range t.headers {
+		if req2.Header.Get(k) == "" {
+			req2.Header.Set(k, v)
+		}
+	}
+	return t.base.RoundTrip(req2)
+}
+
 func withUserConfig(ucfg config.UserConfig) ConnOption {
 	return func(c *config.Config) {
 		c.UserConfig = ucfg

connector_spog_test.go

@@ -0,0 +1,164 @@
+package dbsql
+
+import (
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestExtractSpogHeaders(t *testing.T) {
+	tests := []struct {
+		name     string
+		httpPath string
+		want     map[string]string
+	}{
+		{
+			name:     "no query string returns nil",
+			httpPath: "/sql/1.0/warehouses/abc123",
+			want:     nil,
+		},
+		{
+			name:     "empty httpPath returns nil",
+			httpPath: "",
+			want:     nil,
+		},
+		{
+			name:     "query string with o= extracts org id",
+			httpPath: "/sql/1.0/warehouses/abc123?o=7064161269814046",
+			want:     map[string]string{"x-databricks-org-id": "7064161269814046"},
+		},
+		{
+			name:     "query string without o= returns nil",
+			httpPath: "/sql/1.0/warehouses/abc123?other=val",
+			want:     nil,
+		},
+		{
+			name:     "empty o= value returns nil",
+			httpPath: "/sql/1.0/warehouses/abc123?o=",
+			want:     nil,
+		},
+		{
+			name:     "o= among multiple params extracts correctly",
+			httpPath: "/sql/1.0/warehouses/abc?foo=1&o=12345&bar=2",
+			want:     map[string]string{"x-databricks-org-id": "12345"},
+		},
+		{
+			name:     "first o= wins when duplicated",
+			httpPath: "/sql/1.0/warehouses/abc?o=first&o=second",
+			want:     map[string]string{"x-databricks-org-id": "first"},
+		},
+		{
+			name:     "just ? with nothing after returns nil",
+			httpPath: "/sql/1.0/warehouses/abc?",
+			want:     nil,
+		},
+	}
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			got := extractSpogHeaders(tc.httpPath)
+			assert.Equal(t, tc.want, got)
+		})
+	}
+}
+
+func TestHeaderInjectingTransport_InjectsHeader(t *testing.T) {
+	var gotHeader string
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		gotHeader = r.Header.Get("x-databricks-org-id")
+		w.WriteHeader(http.StatusOK)
+	}))
+	defer srv.Close()
+
+	client := withSpogHeaders(&http.Client{}, map[string]string{
+		"x-databricks-org-id": "7064161269814046",
+	})
+
+	req, err := http.NewRequest("GET", srv.URL, nil)
+	require.NoError(t, err)
+	resp, err := client.Do(req)
+	require.NoError(t, err)
+	_, _ = io.Copy(io.Discard, resp.Body)
+	_ = resp.Body.Close()
+
+	assert.Equal(t, "7064161269814046", gotHeader, "SPOG header should be injected")
+}
+
+func TestHeaderInjectingTransport_DoesNotOverrideCallerSet(t *testing.T) {
+	var gotHeader string
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		gotHeader = r.Header.Get("x-databricks-org-id")
+		w.WriteHeader(http.StatusOK)
+	}))
+	defer srv.Close()
+
+	client := withSpogHeaders(&http.Client{}, map[string]string{
+		"x-databricks-org-id": "from-wrapper",
+	})
+
+	req, err := http.NewRequest("GET", srv.URL, nil)
+	require.NoError(t, err)
+	req.Header.Set("x-databricks-org-id", "from-caller")
+	resp, err := client.Do(req)
+	require.NoError(t, err)
+	_, _ = io.Copy(io.Discard, resp.Body)
+	_ = resp.Body.Close()
+
+	assert.Equal(t, "from-caller", gotHeader, "caller-set header must not be overridden")
+}
+
+func TestHeaderInjectingTransport_PreservesOtherHeaders(t *testing.T) {
+	var gotAuth, gotSpog, gotCustom string
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		gotAuth = r.Header.Get("Authorization")
+		gotSpog = r.Header.Get("x-databricks-org-id")
+		gotCustom = r.Header.Get("X-Custom")
+		w.WriteHeader(http.StatusOK)
+	}))
+	defer srv.Close()
+
+	client := withSpogHeaders(&http.Client{}, map[string]string{
+		"x-databricks-org-id": "abc",
+	})
+
+	req, err := http.NewRequest("GET", srv.URL, nil)
+	require.NoError(t, err)
+	req.Header.Set("Authorization", "Bearer xxx")
+	req.Header.Set("X-Custom", "hello")
+	resp, err := client.Do(req)
+	require.NoError(t, err)
+	_, _ = io.Copy(io.Discard, resp.Body)
+	_ = resp.Body.Close()
+
+	assert.Equal(t, "Bearer xxx", gotAuth)
+	assert.Equal(t, "hello", gotCustom)
+	assert.Equal(t, "abc", gotSpog)
+}
+
+func TestWithSpogHeaders_OriginalClientUntouched(t *testing.T) {
+	originalTransport := &countingTransport{}
+	original := &http.Client{Transport: originalTransport}
+
+	wrapped := withSpogHeaders(original, map[string]string{"x-databricks-org-id": "x"})
+
+	// Original client's transport should NOT be the wrapper type.
+	_, isWrapped := original.Transport.(*headerInjectingTransport)
+	assert.False(t, isWrapped, "original client's transport must not be mutated")
+
+	// Wrapped client MUST have the wrapper.
+	_, isWrapped = wrapped.Transport.(*headerInjectingTransport)
+	assert.True(t, isWrapped, "wrapped client must have headerInjectingTransport")
+}
+
+type countingTransport struct {
+	count int
+}
+
+func (c *countingTransport) RoundTrip(req *http.Request) (*http.Response, error) {
+	c.count++
+	return &http.Response{StatusCode: 200, Body: io.NopCloser(strings.NewReader(""))}, nil
+}