Refactor SPOG header injection: use transport wrapper, drop telemetry API churn

Replace the per-function extraHeaders parameter threading (added during
the previous merge with main) with a minimal http.Client wrapper in the
top-level dbsql package.

Before (5 files, including 3 in telemetry/*):
  connector.go extracted ?o=<workspaceId>, then passed it as an
  ExtraHeaders field through telemetry.TelemetryInitOptions →
  isTelemetryEnabled → featureFlagCache.isTelemetryEnabled →
  fetchFeatureFlag, where it was applied to the outbound request.
  The telemetry push path (telemetry/exporter.go) was NOT covered.

After (2 files — connector.go and auth/oauth/u2m/u2m.go):
  connector.go extracts ?o= as before, but now wraps the driver's
  *http.Client with a headerInjectingTransport that sets the SPOG
  header on every outbound request through that client. Passes the
  wrapped client (not c.client) into TelemetryInitOptions.HTTPClient.

Advantages:
  - telemetry/*.go files revert to identical-to-main. No API churn.
  - Both feature-flag and telemetry-push paths automatically get the
    SPOG header (previously only feature-flag did).
  - Future HTTP paths that reuse the telemetry http.Client inherit
    SPOG routing for free.

Thrift is unaffected: it uses c.client directly (not the wrapper) and
routes via ?o= in the URL path. The transport wrapper is only applied
to the HTTP client handed to telemetry.

Signed-off-by: Madhavendra Rathore
Signed-off-by: Madhavendra Rathore <madhavendra.rathore@databricks.com>

Author: msrathore-db

connector.go

@@ -81,20 +81,26 @@ 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
-	spogHeaders := extractSpogHeaders(c.cfg.HTTPPath)
+	// 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,
 		RetryCount:      c.cfg.TelemetryRetryCount,
 		RetryDelay:      c.cfg.TelemetryRetryDelay,
-		ExtraHeaders:    spogHeaders,
 	})
 	if conn.telemetry != nil {
 		log.Debug().Msg("telemetry initialized for connection")
@@ -171,6 +177,51 @@ func extractSpogHeaders(httpPath string) map[string]string {
 	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

telemetry/config.go

@@ -126,12 +126,12 @@ func ParseTelemetryConfig(params map[string]string) *Config {
 //     (databricks.partnerplatform.clientConfigsFeatureFlags.enableTelemetryForGoDriver).
 //
 // In all other cases — explicit opt-out or server flag absent/unreachable — returns false.
-func isTelemetryEnabled(ctx context.Context, cfg *Config, host string, driverVersion string, httpClient *http.Client, extraHeaders map[string]string) bool {
+func isTelemetryEnabled(ctx context.Context, cfg *Config, host string, driverVersion string, httpClient *http.Client) bool {
 	if cfg.EnableTelemetry != nil {
 		return *cfg.EnableTelemetry
 	}
 
-	serverEnabled, err := getFeatureFlagCache().isTelemetryEnabled(ctx, host, driverVersion, httpClient, extraHeaders)
+	serverEnabled, err := getFeatureFlagCache().isTelemetryEnabled(ctx, host, driverVersion, httpClient)
 	if err != nil {
 		return false
 	}

telemetry/driver_integration.go

@@ -41,14 +41,6 @@ type TelemetryInitOptions struct {
 
 	// RetryDelay is the base delay between retries (0 = use default 100ms).
 	RetryDelay time.Duration
-
-	// ExtraHeaders are additional HTTP headers to attach to feature-flag
-	// check requests and telemetry-push requests. Primarily used to carry
-	// x-databricks-org-id for SPOG (Custom URL) workspace routing — see
-	// extractSpogHeaders in the top-level dbsql package.
-	//
-	// May be nil.
-	ExtraHeaders map[string]string
 }
 
 // InitializeForConnection initializes telemetry for a database connection.
@@ -86,7 +78,7 @@ func InitializeForConnection(ctx context.Context, opts TelemetryInitOptions) *In
 	flagCache.getOrCreateContext(opts.Host)
 
 	// Check if telemetry should be enabled
-	enabled := isTelemetryEnabled(ctx, cfg, opts.Host, opts.DriverVersion, opts.HTTPClient, opts.ExtraHeaders)
+	enabled := isTelemetryEnabled(ctx, cfg, opts.Host, opts.DriverVersion, opts.HTTPClient)
 	if !enabled {
 		flagCache.releaseContext(opts.Host)
 		return nil

telemetry/featureflag.go

@@ -86,7 +86,7 @@ func (c *featureFlagCache) releaseContext(host string) {
 
 // isTelemetryEnabled checks if telemetry is enabled for the host.
 // Uses cached value if available and not expired.
-func (c *featureFlagCache) isTelemetryEnabled(ctx context.Context, host string, driverVersion string, httpClient *http.Client, extraHeaders map[string]string) (bool, error) {
+func (c *featureFlagCache) isTelemetryEnabled(ctx context.Context, host string, driverVersion string, httpClient *http.Client) (bool, error) {
 	c.mu.RLock()
 	flagCtx, exists := c.contexts[host]
 	c.mu.RUnlock()
@@ -135,7 +135,7 @@ func (c *featureFlagCache) isTelemetryEnabled(ctx context.Context, host string,
 	flagCtx.mu.Unlock()
 
 	// Fetch fresh value (outside lock so other readers are not blocked).
-	enabled, err := fetchFeatureFlag(ctx, host, driverVersion, httpClient, extraHeaders)
+	enabled, err := fetchFeatureFlag(ctx, host, driverVersion, httpClient)
 
 	// Update cache.
 	flagCtx.mu.Lock()
@@ -166,7 +166,7 @@ func (c *featureFlagContext) isExpired() bool {
 }
 
 // fetchFeatureFlag fetches the feature flag value from Databricks.
-func fetchFeatureFlag(ctx context.Context, host string, driverVersion string, httpClient *http.Client, extraHeaders map[string]string) (bool, error) {
+func fetchFeatureFlag(ctx context.Context, host string, driverVersion string, httpClient *http.Client) (bool, error) {
 	// Add timeout to context if it doesn't have a deadline
 	if _, hasDeadline := ctx.Deadline(); !hasDeadline {
 		var cancel context.CancelFunc
@@ -183,11 +183,6 @@ func fetchFeatureFlag(ctx context.Context, host string, driverVersion string, ht
 		return false, fmt.Errorf("failed to create feature flag request: %w", err)
 	}
 
-	// Attach extra headers (e.g. x-databricks-org-id for SPOG routing).
-	for k, v := range extraHeaders {
-		req.Header.Set(k, v)
-	}
-
 	resp, err := httpClient.Do(req)
 	if err != nil {
 		return false, fmt.Errorf("failed to fetch feature flag: %w", err)