From e30d820d6f84fc62fb6a41c9b9f335d4a180b41e Mon Sep 17 00:00:00 2001 From: samikshya-chand_data Date: Sat, 23 May 2026 20:37:09 +0530 Subject: [PATCH] Add SPOG support: x-databricks-org-id header for telemetry + feature-flag MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit SPOG (Single Panel of Glass) replaces workspace-specific hostnames with account-level vanity URLs. When httpPath carries `?o=`, endpoints that don't include the workspace in their URL path (telemetry, feature flags) need the workspace conveyed via the `x-databricks-org-id` header instead. Changes: - Parse `?o=` out of httpPath in DBSQLClient.connect() and stash the org-id as `x-databricks-org-id` on a new `ClientConfig.customHeaders` field. A user-supplied `customHeaders` entry (case-insensitive) takes precedence. - DatabricksTelemetryExporter spreads `config.customHeaders` into the outgoing POST headers. Auth headers still win on collision. - FeatureFlagCache does the same for the feature-flag GET. Not applicable to this driver (vs JDBC port in databricks/databricks-jdbc#1316): - httpPath property parser fix — Node.js passes `options.path` through unmodified. - Warehouse ID regex fix for SEA — driver uses Thrift only. - DBFS Volume header injection — driver exposes no Volume API. OAuth/OIDC token requests deliberately do NOT receive customHeaders. Co-authored-by: Isaac --- lib/DBSQLClient.ts | 26 +++++++++ lib/contracts/IClientContext.ts | 11 ++++ lib/contracts/IDBSQLClient.ts | 11 ++++ lib/telemetry/DatabricksTelemetryExporter.ts | 1 + lib/telemetry/FeatureFlagCache.ts | 1 + tests/unit/DBSQLClient.test.ts | 55 +++++++++++++++++++ .../DatabricksTelemetryExporter.test.ts | 40 ++++++++++++++ tests/unit/telemetry/FeatureFlagCache.test.ts | 39 +++++++++++++ 8 files changed, 184 insertions(+) diff --git a/lib/DBSQLClient.ts b/lib/DBSQLClient.ts index 340efd0c..25d738b1 100644 --- a/lib/DBSQLClient.ts +++ b/lib/DBSQLClient.ts @@ -330,6 +330,27 @@ export default class DBSQLClient extends EventEmitter implements IDBSQLClient, I return match ? match[1] : undefined; } + /** + * Build the customHeaders map applied to telemetry POSTs and feature-flag + * GETs (SPOG / Single Panel of Glass support). When `httpPath` carries + * `?o=` — account-level vanity URL routing — endpoints that + * don't include the workspace in their path need the workspace conveyed via + * the `x-databricks-org-id` header instead. A user-supplied value in + * `options.customHeaders` (case-insensitively keyed) wins over the parsed + * value. + */ + private buildCustomHeaders(options: ConnectionOptions): Record | undefined { + const merged: Record = { ...(options.customHeaders ?? {}) }; + const hasOrgIdAlready = Object.keys(merged).some((k) => k.toLowerCase() === 'x-databricks-org-id'); + if (!hasOrgIdAlready) { + const orgId = this.extractWorkspaceId(); + if (orgId) { + merged['x-databricks-org-id'] = orgId; + } + } + return Object.keys(merged).length > 0 ? merged : undefined; + } + /** * Build driver configuration for telemetry reporting. * @returns DriverConfiguration object with current driver settings @@ -561,6 +582,11 @@ export default class DBSQLClient extends EventEmitter implements IDBSQLClient, I this.config.userAgentEntry = options.userAgentEntry; } + // SPOG: parse `?o=` out of httpPath and stash it as + // `x-databricks-org-id` for the telemetry + feature-flag clients, which + // hit endpoints that don't carry the workspace in their URL path. + this.config.customHeaders = this.buildCustomHeaders(options); + this.authProvider = this.createAuthProvider(options, authProvider); this.connectionProvider = this.createConnectionProvider(options); diff --git a/lib/contracts/IClientContext.ts b/lib/contracts/IClientContext.ts index 955b9c52..43a47745 100644 --- a/lib/contracts/IClientContext.ts +++ b/lib/contracts/IClientContext.ts @@ -52,6 +52,17 @@ export interface ClientConfig { */ telemetryFlushOnExit?: boolean; userAgentEntry?: string; + + /** + * Extra HTTP headers attached to driver-owned out-of-band requests + * (telemetry, feature flags). Populated by `DBSQLClient.connect()` from + * `ConnectionOptions.customHeaders` plus an `x-databricks-org-id` header + * derived from the `?o=` query parameter on `httpPath` when present, to + * support SPOG (Single Panel of Glass) account-level routing on endpoints + * that don't carry `?o=` in their URL path. NOT applied to Thrift or + * OAuth/OIDC requests. + */ + customHeaders?: Record; } export default interface IClientContext { diff --git a/lib/contracts/IDBSQLClient.ts b/lib/contracts/IDBSQLClient.ts index 08592219..b75a8075 100644 --- a/lib/contracts/IDBSQLClient.ts +++ b/lib/contracts/IDBSQLClient.ts @@ -55,6 +55,17 @@ export type ConnectionOptions = { proxy?: ProxyOptions; enableMetricViewMetadata?: boolean; + /** + * Extra HTTP headers attached to driver-owned out-of-band requests + * (telemetry POSTs and feature-flag GETs). Not applied to the primary + * Thrift transport or to OAuth/OIDC token requests. + * + * When `path` contains `?o=` (SPOG account-level routing), + * the driver automatically injects an `x-databricks-org-id` header unless + * one is already present in this map. + */ + customHeaders?: Record; + /** * Whether the driver emits telemetry events (connection / statement / * cloud-fetch / error). Defaults to `true`. diff --git a/lib/telemetry/DatabricksTelemetryExporter.ts b/lib/telemetry/DatabricksTelemetryExporter.ts index 65fe8b64..bfa0f1b1 100644 --- a/lib/telemetry/DatabricksTelemetryExporter.ts +++ b/lib/telemetry/DatabricksTelemetryExporter.ts @@ -249,6 +249,7 @@ export default class DatabricksTelemetryExporter { let headers: Record = { 'Content-Type': 'application/json', 'User-Agent': userAgent, + ...(config.customHeaders ?? {}), }; if (authenticatedExport) { diff --git a/lib/telemetry/FeatureFlagCache.ts b/lib/telemetry/FeatureFlagCache.ts index 47fc9e8d..e060bd93 100644 --- a/lib/telemetry/FeatureFlagCache.ts +++ b/lib/telemetry/FeatureFlagCache.ts @@ -140,6 +140,7 @@ export default class FeatureFlagCache { const headers: Record = { 'Content-Type': 'application/json', 'User-Agent': this.userAgent, + ...(this.context.getConfig().customHeaders ?? {}), ...(await this.getAuthHeaders()), }; diff --git a/tests/unit/DBSQLClient.test.ts b/tests/unit/DBSQLClient.test.ts index 184e25a3..5e2b99fe 100644 --- a/tests/unit/DBSQLClient.test.ts +++ b/tests/unit/DBSQLClient.test.ts @@ -103,6 +103,18 @@ describe('DBSQLClient.connect', () => { logSpy.restore(); }); + + it('populates config.customHeaders with org-id parsed from ?o= (SPOG)', async () => { + const client = new DBSQLClient(); + await client.connect({ ...connectOptions, path: '/sql/1.0/warehouses/abc?o=12345678901234' }); + expect(client.getConfig().customHeaders).to.deep.equal({ 'x-databricks-org-id': '12345678901234' }); + }); + + it('leaves config.customHeaders undefined when path has no ?o= and none supplied', async () => { + const client = new DBSQLClient(); + await client.connect({ ...connectOptions, path: '/sql/1.0/warehouses/abc' }); + expect(client.getConfig().customHeaders).to.be.undefined; + }); }); describe('DBSQLClient.openSession', () => { @@ -785,6 +797,49 @@ describe('DBSQLClient telemetry paths', () => { }); }); + describe('buildCustomHeaders (SPOG)', () => { + it('injects x-databricks-org-id from ?o= in httpPath', () => { + const client = new DBSQLClient(); + (client as any).httpPath = '/sql/1.0/warehouses/abc?o=12345678901234'; + const headers = (client as any).buildCustomHeaders({ path: '/sql/1.0/warehouses/abc?o=12345678901234' }); + expect(headers).to.deep.equal({ 'x-databricks-org-id': '12345678901234' }); + }); + + it('returns undefined when no ?o= and no user-supplied customHeaders', () => { + const client = new DBSQLClient(); + (client as any).httpPath = '/sql/1.0/warehouses/abc'; + const headers = (client as any).buildCustomHeaders({ path: '/sql/1.0/warehouses/abc' }); + expect(headers).to.be.undefined; + }); + + it('preserves user-supplied customHeaders alongside parsed org-id', () => { + const client = new DBSQLClient(); + (client as any).httpPath = '/sql/1.0/warehouses/abc?o=42'; + const headers = (client as any).buildCustomHeaders({ + path: '/sql/1.0/warehouses/abc?o=42', + customHeaders: { 'x-trace-id': 'tid-001' }, + }); + expect(headers).to.deep.equal({ 'x-trace-id': 'tid-001', 'x-databricks-org-id': '42' }); + }); + + it('user-supplied x-databricks-org-id wins over ?o= parsed value (case-insensitive)', () => { + const client = new DBSQLClient(); + (client as any).httpPath = '/sql/1.0/warehouses/abc?o=42'; + const headers = (client as any).buildCustomHeaders({ + path: '/sql/1.0/warehouses/abc?o=42', + customHeaders: { 'X-Databricks-Org-Id': '999' }, + }); + expect(headers).to.deep.equal({ 'X-Databricks-Org-Id': '999' }); + }); + + it('does not inject org-id when ?o= value is non-numeric', () => { + const client = new DBSQLClient(); + (client as any).httpPath = '/sql/1.0/warehouses/abc?o=tenant_xyz'; + const headers = (client as any).buildCustomHeaders({ path: '/sql/1.0/warehouses/abc?o=tenant_xyz' }); + expect(headers).to.be.undefined; + }); + }); + describe('telemetry refcount on reconnect', () => { it('releases the prior refcount when connect() is called twice', async () => { const client = new DBSQLClient(); diff --git a/tests/unit/telemetry/DatabricksTelemetryExporter.test.ts b/tests/unit/telemetry/DatabricksTelemetryExporter.test.ts index fb347bf6..c3fd099d 100644 --- a/tests/unit/telemetry/DatabricksTelemetryExporter.test.ts +++ b/tests/unit/telemetry/DatabricksTelemetryExporter.test.ts @@ -120,6 +120,46 @@ describe('DatabricksTelemetryExporter', () => { } expect(threw).to.be.false; }); + + it('should attach config.customHeaders to the POST (SPOG)', async () => { + const context = new ClientContextStub({ + customHeaders: { 'x-databricks-org-id': '12345678901234' }, + } as any); + const registry = new CircuitBreakerRegistry(context); + const exporter = new DatabricksTelemetryExporter(context, 'host.example.com', registry, fakeAuthProvider); + const sendRequestStub = sinon.stub(exporter as any, 'sendRequest').returns(makeOkResponse()); + + await exporter.export([makeMetric()]); + + const init = sendRequestStub.firstCall.args[1] as { headers: Record }; + expect(init.headers['x-databricks-org-id']).to.equal('12345678901234'); + }); + + it('auth headers win over customHeaders on key collision', async () => { + const context = new ClientContextStub({ + customHeaders: { Authorization: 'Bearer not-the-real-token' }, + } as any); + const registry = new CircuitBreakerRegistry(context); + const exporter = new DatabricksTelemetryExporter(context, 'host.example.com', registry, fakeAuthProvider); + const sendRequestStub = sinon.stub(exporter as any, 'sendRequest').returns(makeOkResponse()); + + await exporter.export([makeMetric()]); + + const init = sendRequestStub.firstCall.args[1] as { headers: Record }; + expect(init.headers.Authorization).to.equal('Bearer test-token'); + }); + + it('does not attach customHeaders when none are configured', async () => { + const context = new ClientContextStub(); + const registry = new CircuitBreakerRegistry(context); + const exporter = new DatabricksTelemetryExporter(context, 'host.example.com', registry, fakeAuthProvider); + const sendRequestStub = sinon.stub(exporter as any, 'sendRequest').returns(makeOkResponse()); + + await exporter.export([makeMetric()]); + + const init = sendRequestStub.firstCall.args[1] as { headers: Record }; + expect(init.headers).to.not.have.property('x-databricks-org-id'); + }); }); describe('export() - retry logic', () => { diff --git a/tests/unit/telemetry/FeatureFlagCache.test.ts b/tests/unit/telemetry/FeatureFlagCache.test.ts index ed7bc79c..c12ca5b1 100644 --- a/tests/unit/telemetry/FeatureFlagCache.test.ts +++ b/tests/unit/telemetry/FeatureFlagCache.test.ts @@ -317,4 +317,43 @@ describe('FeatureFlagCache', () => { fetchStub.restore(); }); }); + + describe('customHeaders propagation (SPOG)', () => { + function makeJsonResponse(body: unknown) { + return Promise.resolve({ + ok: true, + status: 200, + statusText: 'OK', + json: () => Promise.resolve(body), + text: () => Promise.resolve(''), + }); + } + + it('attaches config.customHeaders to the feature-flag GET', async () => { + const context = new ClientContextStub({ + customHeaders: { 'x-databricks-org-id': '12345678901234' }, + } as any); + const cache = new FeatureFlagCache(context); + const stub = sinon.stub(cache as any, 'fetchWithRetry').returns(makeJsonResponse({ flags: [] })); + + await (cache as any).fetchFeatureFlag('host.example.com'); + + expect(stub.calledOnce).to.be.true; + const init = stub.firstCall.args[1] as { headers: Record }; + expect(init.headers['x-databricks-org-id']).to.equal('12345678901234'); + stub.restore(); + }); + + it('does not set x-databricks-org-id when customHeaders is empty', async () => { + const context = new ClientContextStub(); + const cache = new FeatureFlagCache(context); + const stub = sinon.stub(cache as any, 'fetchWithRetry').returns(makeJsonResponse({ flags: [] })); + + await (cache as any).fetchFeatureFlag('host.example.com'); + + const init = stub.firstCall.args[1] as { headers: Record }; + expect(init.headers).to.not.have.property('x-databricks-org-id'); + stub.restore(); + }); + }); });