Add SPOG support: x-databricks-org-id header for telemetry + feature-flag

SPOG (Single Panel of Glass) replaces workspace-specific hostnames with
account-level vanity URLs. When httpPath carries `?o=<workspaceId>`,
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=<digits>` 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

Author: samikshya-db

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=<workspaceId>` — 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<string, string> | undefined {
+    const merged: Record<string, string> = { ...(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=<workspaceId>` 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);

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<string, string>;
 }
 
 export default interface IClientContext {

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=<workspaceId>` (SPOG account-level routing),
+   * the driver automatically injects an `x-databricks-org-id` header unless
+   * one is already present in this map.
+   */
+  customHeaders?: Record<string, string>;
+
   /**
    * Whether the driver emits telemetry events (connection / statement /
    * cloud-fetch / error). Defaults to `true`.

lib/telemetry/DatabricksTelemetryExporter.ts

@@ -249,6 +249,7 @@ export default class DatabricksTelemetryExporter {
     let headers: Record<string, string> = {
       'Content-Type': 'application/json',
       'User-Agent': userAgent,
+      ...(config.customHeaders ?? {}),
     };
 
     if (authenticatedExport) {

lib/telemetry/FeatureFlagCache.ts

@@ -140,6 +140,7 @@ export default class FeatureFlagCache {
       const headers: Record<string, string> = {
         'Content-Type': 'application/json',
         'User-Agent': this.userAgent,
+        ...(this.context.getConfig().customHeaders ?? {}),
         ...(await this.getAuthHeaders()),
       };
 

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();

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<string, string> };
+      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<string, string> };
+      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<string, string> };
+      expect(init.headers).to.not.have.property('x-databricks-org-id');
+    });
   });
 
   describe('export() - retry logic', () => {

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<string, string> };
+      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<string, string> };
+      expect(init.headers).to.not.have.property('x-databricks-org-id');
+      stub.restore();
+    });
+  });
 });