Skip to content

Commit 9261ac6

Browse files
msrathore-dbmsrathore-db
authored
Extract SPOG org-id from cluster httpPath for non-Thrift requests (#367)
## Summary - Fixes silent telemetry loss on SPOG (custom-URL) hosts when connecting to an all-purpose cluster via an httpPath like `sql/protocolv1/o/<workspace-id>/<cluster-id>`. - `extractSpogHeaders` now extracts the workspace ID from the `/o/<wsid>/` path segment as a fallback when `?o=<wsid>` is not present, and emits an `x-databricks-org-id` header so the wrapped `telemetryClient` / feature-flag client can route correctly on SPOG. - Priority order preserved: `?o=` in `httpPath` ▶ `/sql/protocolv1/o/<wsid>/` path segment. Caller-set request headers still win, including mixed-case `X-Databricks-Org-Id`. ## Why On a SPOG host the workspace identity has to be in either the URL or the `x-databricks-org-id` header for PoPP to route a request to the correct workspace. For all-purpose cluster Thrift this is free — the workspace ID is in the `/o/<wsid>/` segment of httpPath, so PoPP routes Thrift via `routing_reason=workspace-id` and the session opens fine without an explicit `?o=`. The telemetry and feature-flag transports built by `connector.Connect` wrap `c.client` with `withSpogHeaders` only when `extractSpogHeaders(c.cfg.HTTPPath)` returns a non-nil map. Before this PR that only happened when `?o=` was present, so cluster URLs without `?o=` produced no `x-databricks-org-id` header, PoPP fell back to default (account) routing on `/telemetry-ext`, and the responses were 303 redirects to `/login` — silently dropping telemetry on SPOG. ## What changes `connector.go` - New `clusterPathOrgIDPattern` regex (`(?:^|/)sql/protocolv1/o/(\d+)/[^/?]+`) compiled once at package init. - `extractSpogHeaders` checks `?o=` first (existing behavior), then falls back to the cluster path segment, then returns nil. The malformed-query-string path now also falls through to path inspection instead of returning early. Log messages indicate which source produced the workspace ID. - Updated the connection comment to describe both query-param and cluster-path extraction. - `regexp` added to imports. `connector_spog_test.go` - Four new table entries under `TestExtractSpogHeaders`: - Cluster path without `?o=` → header extracted from `/o/<wsid>/` - Cluster path with leading `/` → same - Cluster path with `?o=` → query-param value wins - Warehouse path without `?o=` → still nil (regression guard: the new regex must not match warehouse paths) - New transport regression test proving a mixed-case caller-set `X-Databricks-Org-Id` is not overwritten by the SPOG wrapper. ## Test plan - [x] `go test -run 'TestExtractSpogHeaders|TestHeaderInjectingTransport' -v .` — focused SPOG extraction and transport tests pass. - [x] `go test -short .` — root package suite passes. - [x] Behavior validated on the OSS JDBC equivalent fix against Prod SPOG (`peco.azuredatabricks.net`) all-purpose cluster: telemetry POST `/telemetry-ext` flipped from `HTTP 303 → /login` to `HTTP 200 OK` once the path-segment extraction populated `x-databricks-org-id`. Same shape of fix here. ## Out of scope - The corresponding OSS JDBC driver fix is opened as databricks/databricks-jdbc#1475. - The Python connector (`databricks/databricks-sql-python`) has a sibling upstream-head PR for the same fix: databricks/databricks-sql-python#817. - The Node.js connector (`databricks/databricks-sql-nodejs`) already extracts org ID from both query param and path segment (see `extractWorkspaceId` in `lib/DBSQLClient.ts`). This pull request and its description were written with assistance from Claude Code. --------- Signed-off-by: Madhavendra Rathore <madhavendra.rathore@databricks.com>
1 parent 9d711e7 commit 9261ac6

2 files changed

Lines changed: 141 additions & 39 deletions

File tree

‎connector.go‎

Lines changed: 60 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import (
77
"fmt"
88
"net/http"
99
"net/url"
10+
"regexp"
1011
"strings"
1112
"time"
1213

@@ -81,11 +82,11 @@ func (c *connector) Connect(ctx context.Context) (driver.Conn, error) {
8182
}
8283
log := logger.WithContext(conn.id, driverctx.CorrelationIdFromContext(ctx), "")
8384

84-
// Extract SPOG routing headers from ?o= in HTTPPath. When ?o=<workspaceId>
85-
// is present (Custom URL / SPOG hosts), wrap the HTTP client used for
86-
// telemetry + feature-flag calls with a transport that injects
87-
// x-databricks-org-id. Thrift routes via the URL so its own c.client
88-
// doesn't need wrapping.
85+
// Extract SPOG routing headers from HTTPPath. When the workspace ID is
86+
// available via ?o=<workspaceId> or a cluster /o/<workspaceId>/ path segment,
87+
// wrap the HTTP client used for telemetry + feature-flag calls with a
88+
// transport that injects x-databricks-org-id. Thrift routes via the URL so
89+
// its own c.client doesn't need wrapping.
8990
telemetryClient := c.client
9091
if spogHeaders := extractSpogHeaders(c.cfg.HTTPPath); len(spogHeaders) > 0 {
9192
telemetryClient = withSpogHeaders(c.client, spogHeaders)
@@ -136,44 +137,67 @@ func NewConnector(options ...ConnOption) (driver.Connector, error) {
136137
return &connector{cfg: cfg, client: client}, nil
137138
}
138139

139-
// extractSpogHeaders extracts ?o=<workspaceId> from httpPath and returns
140-
// an x-databricks-org-id header for SPOG routing.
140+
// clusterPathOrgIDPattern matches the workspace ID inside an all-purpose-compute
141+
// Thrift path of the form [/]sql/protocolv1/o/<workspace-id>/<cluster-id>[/...].
142+
var (
143+
orgIDPattern = regexp.MustCompile(`^[0-9]+$`)
144+
clusterPathOrgIDPattern = regexp.MustCompile(`^/?sql/protocolv1/o/([0-9]+)/[^/?]+`)
145+
)
146+
147+
// extractSpogHeaders inspects httpPath for the workspace ID and returns it as an
148+
// x-databricks-org-id header dict for SPOG routing.
141149
//
142-
// On SPOG (Custom URL) workspaces, httpPath is of the form
143-
// /sql/1.0/warehouses/<id>?o=<workspaceId>. The ?o= parameter keeps Thrift
144-
// requests routed to the correct workspace via the URL itself, but other
145-
// endpoints (telemetry, feature flags) run on separate hosts and need the
146-
// x-databricks-org-id header. This function extracts ?o= from httpPath once
147-
// and returns it so those paths can inject it as an HTTP header.
150+
// Two sources are checked, in priority order:
151+
// 1. ?o=<workspace-id> query parameter (warehouse paths on SPOG typically use
152+
// this form, e.g. /sql/1.0/warehouses/<id>?o=<workspace-id>).
153+
// 2. /sql/protocolv1/o/<workspace-id>/<cluster-id> path segment (all-purpose
154+
// cluster paths embed the workspace in the path itself).
148155
//
149-
// Returns nil if:
150-
// - httpPath has no query string ("?"), or
151-
// - the query string is malformed and can't be parsed, or
152-
// - the ?o= parameter is missing or empty.
156+
// Thrift requests are routed by the URL itself, but other endpoints
157+
// (telemetry, feature flags) run on separate paths that don't carry the
158+
// workspace ID — without this header, PoPP on SPOG hosts can't determine the
159+
// workspace and redirects the request to /login.
160+
//
161+
// Returns nil if no workspace ID can be determined.
153162
func extractSpogHeaders(httpPath string) map[string]string {
154-
if !strings.Contains(httpPath, "?") {
163+
if httpPath == "" {
155164
return nil
156165
}
157-
// Parse query string from httpPath
158-
parts := strings.SplitN(httpPath, "?", 2)
159-
params, err := url.ParseQuery(parts[1])
160-
if err != nil {
161-
logger.Debug().Msgf(
162-
"SPOG header extraction: malformed query string in httpPath, skipping org-id extraction: %s",
163-
err)
164-
return nil
166+
167+
// 1) ?o=<wsid> query parameter.
168+
if strings.Contains(httpPath, "?") {
169+
parts := strings.SplitN(httpPath, "?", 2)
170+
params, err := url.ParseQuery(parts[1])
171+
if err != nil {
172+
logger.Debug().Msgf(
173+
"SPOG header extraction: malformed query string in httpPath, falling back to path inspection: %s",
174+
err)
175+
} else if orgID := params.Get("o"); orgID != "" {
176+
if !orgIDPattern.MatchString(orgID) {
177+
logger.Debug().Msg(
178+
"SPOG header extraction: ignoring non-numeric ?o= value in httpPath, falling back to path inspection")
179+
} else {
180+
logger.Debug().Msgf(
181+
"SPOG header extraction: injecting x-databricks-org-id=%s (extracted from ?o= in httpPath)",
182+
orgID)
183+
return map[string]string{"x-databricks-org-id": orgID}
184+
}
185+
}
165186
}
166-
orgID := params.Get("o")
167-
if orgID == "" {
168-
logger.Debug().Msg(
169-
"SPOG header extraction: httpPath has query string but no ?o= param, " +
170-
"skipping x-databricks-org-id injection")
171-
return nil
187+
188+
// 2) /sql/protocolv1/o/<wsid>/<cluster> path segment.
189+
if match := clusterPathOrgIDPattern.FindStringSubmatch(httpPath); match != nil {
190+
orgID := match[1]
191+
logger.Debug().Msgf(
192+
"SPOG header extraction: injecting x-databricks-org-id=%s (extracted from cluster path segment)",
193+
orgID)
194+
return map[string]string{"x-databricks-org-id": orgID}
172195
}
173-
logger.Debug().Msgf(
174-
"SPOG header extraction: injecting x-databricks-org-id=%s (extracted from ?o= in httpPath)",
175-
orgID)
176-
return map[string]string{"x-databricks-org-id": orgID}
196+
197+
logger.Debug().Msg(
198+
"SPOG header extraction: no workspace ID found in httpPath, " +
199+
"skipping x-databricks-org-id injection")
200+
return nil
177201
}
178202

179203
// withSpogHeaders returns a new *http.Client that reuses the transport of the

‎connector_spog_test.go‎

Lines changed: 81 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -48,15 +48,70 @@ func TestExtractSpogHeaders(t *testing.T) {
4848
want: map[string]string{"x-databricks-org-id": "12345"},
4949
},
5050
{
51-
name: "first o= wins when duplicated",
52-
httpPath: "/sql/1.0/warehouses/abc?o=first&o=second",
53-
want: map[string]string{"x-databricks-org-id": "first"},
51+
name: "first numeric o= wins when duplicated",
52+
httpPath: "/sql/1.0/warehouses/abc?o=111&o=222",
53+
want: map[string]string{"x-databricks-org-id": "111"},
54+
},
55+
{
56+
name: "non-numeric o= value returns nil",
57+
httpPath: "/sql/1.0/warehouses/abc?o=abc123",
58+
want: nil,
59+
},
60+
{
61+
name: "control-character o= value returns nil",
62+
httpPath: "/sql/1.0/warehouses/abc?o=123%0D%0AX-Injected:%20yes",
63+
want: nil,
64+
},
65+
{
66+
name: "invalid o= falls back to valid cluster path segment",
67+
httpPath: "sql/protocolv1/o/6051921418418893/0528-220959-uzmcn1qt?o=abc123",
68+
want: map[string]string{"x-databricks-org-id": "6051921418418893"},
5469
},
5570
{
5671
name: "just ? with nothing after returns nil",
5772
httpPath: "/sql/1.0/warehouses/abc?",
5873
want: nil,
5974
},
75+
{
76+
// All-purpose cluster paths embed the workspace ID in /o/<wsid>/<cluster>.
77+
// Without ?o=, the driver must still extract it so non-Thrift endpoints
78+
// (telemetry, feature flags) get x-databricks-org-id on SPOG hosts.
79+
name: "cluster path without ?o= extracts org id from path segment",
80+
httpPath: "sql/protocolv1/o/6051921418418893/0528-220959-uzmcn1qt",
81+
want: map[string]string{"x-databricks-org-id": "6051921418418893"},
82+
},
83+
{
84+
name: "cluster path with leading slash also extracts",
85+
httpPath: "/sql/protocolv1/o/6051921418418893/0528-220959-uzmcn1qt",
86+
want: map[string]string{"x-databricks-org-id": "6051921418418893"},
87+
},
88+
{
89+
name: "?o= query param wins over cluster path segment",
90+
httpPath: "sql/protocolv1/o/111/0528-220959-uzmcn1qt?o=222",
91+
want: map[string]string{"x-databricks-org-id": "222"},
92+
},
93+
{
94+
name: "nested cluster path prefix returns nil",
95+
httpPath: "evil/sql/protocolv1/o/999/0528-220959-uzmcn1qt",
96+
want: nil,
97+
},
98+
{
99+
name: "incomplete cluster path returns nil",
100+
httpPath: "sql/protocolv1/o/999/",
101+
want: nil,
102+
},
103+
{
104+
name: "warehouse path containing cluster-looking suffix returns nil",
105+
httpPath: "/sql/1.0/warehouses/sql/protocolv1/o/999/cluster-id",
106+
want: nil,
107+
},
108+
{
109+
// Regression guard: the new cluster-path regex must not match
110+
// warehouse paths (which never embed the workspace ID).
111+
name: "warehouse path without ?o= still returns nil",
112+
httpPath: "/sql/1.0/warehouses/abc123",
113+
want: nil,
114+
},
60115
}
61116
for _, tc := range tests {
62117
t.Run(tc.name, func(t *testing.T) {
@@ -111,6 +166,29 @@ func TestHeaderInjectingTransport_DoesNotOverrideCallerSet(t *testing.T) {
111166
assert.Equal(t, "from-caller", gotHeader, "caller-set header must not be overridden")
112167
}
113168

169+
func TestHeaderInjectingTransport_DoesNotOverrideCallerSetMixedCase(t *testing.T) {
170+
var gotHeader string
171+
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
172+
gotHeader = r.Header.Get("x-databricks-org-id")
173+
w.WriteHeader(http.StatusOK)
174+
}))
175+
defer srv.Close()
176+
177+
client := withSpogHeaders(&http.Client{}, map[string]string{
178+
"x-databricks-org-id": "from-wrapper",
179+
})
180+
181+
req, err := http.NewRequest("GET", srv.URL, nil)
182+
require.NoError(t, err)
183+
req.Header.Set("X-Databricks-Org-Id", "from-caller")
184+
resp, err := client.Do(req)
185+
require.NoError(t, err)
186+
_, _ = io.Copy(io.Discard, resp.Body)
187+
_ = resp.Body.Close()
188+
189+
assert.Equal(t, "from-caller", gotHeader, "caller-set header must not be overridden")
190+
}
191+
114192
func TestHeaderInjectingTransport_PreservesOtherHeaders(t *testing.T) {
115193
var gotAuth, gotSpog, gotCustom string
116194
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {

0 commit comments

Comments
 (0)