Fix feature flags and U2M OAuth for SPOG support

Feature flags:
- Fix endpoint path: /api/2.0/feature-flags -> /api/2.0/connector-service/feature-flags/GOLANG/{version}
- Fix response parsing: map format -> array of {name, value} entries
- Add extraHeaders for SPOG routing (x-databricks-org-id)
- Extract ?o=<workspaceId> from httpPath in connector

U2M OAuth:
- Don't set ClientSecret for public apps (PKCE)
- Force AuthStyleInParams to prevent Basic auth with empty password
- Server rejects "Public app should not use a client secret" otherwise

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

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

Author: msrathore-db

auth/oauth/u2m/u2m.go

@@ -30,11 +30,14 @@ func GetConfig(ctx context.Context, hostName, clientID, clientSecret, callbackUR
 		RedirectURL: callbackURL,
 		Scopes:      scopes,
 	}
-	// Only set ClientSecret if non-empty. For U2M (public apps using PKCE),
-	// sending an empty client_secret causes the server to reject with
-	// "Public app should not use a client secret".
 	if clientSecret != "" {
 		config.ClientSecret = clientSecret
+	} else {
+		// For U2M (public apps using PKCE), force AuthStyleInParams to avoid
+		// sending Basic auth with empty password. AuthStyleInHeader sends
+		// "Authorization: Basic base64(clientID:)" which the server rejects
+		// with "Public app should not use a client secret".
+		config.Endpoint.AuthStyle = oauth2.AuthStyleInParams
 	}
 
 	return config, nil

connector.go

@@ -6,6 +6,7 @@ import (
 	"database/sql/driver"
 	"fmt"
 	"net/http"
+	"net/url"
 	"strings"
 	"time"
 
@@ -76,12 +77,16 @@ 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)
+
 	// Initialize telemetry: client config overlay decides; if unset, feature flags decide
 	conn.telemetry = telemetry.InitializeForConnection(
 		ctx,
 		c.cfg.Host,
 		c.client,
 		c.cfg.EnableTelemetry,
+		spogHeaders,
 	)
 	if conn.telemetry != nil {
 		log.Debug().Msg("telemetry initialized for connection")
@@ -107,6 +112,7 @@ func NewConnector(options ...ConnOption) (driver.Connector, error) {
 	// config with default options
 	cfg := config.WithDefaults()
 	cfg.DriverVersion = DriverVersion
+	telemetry.SetDriverVersion(DriverVersion)
 
 	for _, opt := range options {
 		opt(cfg)
@@ -117,6 +123,25 @@ 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.
+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 {
+		return nil
+	}
+	orgID := params.Get("o")
+	if orgID == "" {
+		return nil
+	}
+	return map[string]string{"x-databricks-org-id": orgID}
+}
+
 func withUserConfig(ucfg config.UserConfig) ConnOption {
 	return func(c *config.Config) {
 		c.UserConfig = ucfg

telemetry/config.go

@@ -102,7 +102,7 @@ func ParseTelemetryConfig(params map[string]string) *Config {
 //
 // Returns:
 //   - bool: true if telemetry should be enabled, false otherwise
-func isTelemetryEnabled(ctx context.Context, cfg *Config, host string, httpClient *http.Client) bool {
+func isTelemetryEnabled(ctx context.Context, cfg *Config, host string, httpClient *http.Client, extraHeaders map[string]string) bool {
 	// Priority 1: Client explicitly set (overrides server)
 	if cfg.EnableTelemetry.IsSet() {
 		val, _ := cfg.EnableTelemetry.Get()
@@ -111,7 +111,7 @@ func isTelemetryEnabled(ctx context.Context, cfg *Config, host string, httpClien
 
 	// Priority 2: Check server-side feature flag
 	flagCache := getFeatureFlagCache()
-	serverEnabled, err := flagCache.isTelemetryEnabled(ctx, host, httpClient)
+	serverEnabled, err := flagCache.isTelemetryEnabled(ctx, host, httpClient, extraHeaders)
 	if err != nil {
 		// Priority 3: Fail-safe default (disabled)
 		return false

telemetry/config_test.go

@@ -2,7 +2,7 @@ package telemetry
 
 import (
 	"context"
-	"encoding/json"
+
 	"net/http"
 	"net/http/httptest"
 	"testing"
@@ -206,12 +206,8 @@ func TestIsTelemetryEnabled_ClientOverrideEnabled(t *testing.T) {
 	// Setup: Create a server that returns disabled
 	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		// Server says disabled, but client override should win
-		resp := map[string]interface{}{
-			"flags": map[string]bool{
-				"databricks.partnerplatform.clientConfigsFeatureFlags.enableTelemetryForGoDriver": false,
-			},
-		}
-		_ = json.NewEncoder(w).Encode(resp)
+		w.Header().Set("Content-Type", "application/json")
+		_, _ = w.Write([]byte(`{"flags": [{"name": "databricks.partnerplatform.clientConfigsFeatureFlags.enableTelemetryForGoDriver", "value": "false"}]}`))
 	}))
 	defer server.Close()
 
@@ -228,7 +224,7 @@ func TestIsTelemetryEnabled_ClientOverrideEnabled(t *testing.T) {
 	defer flagCache.releaseContext(server.URL)
 
 	// Client override should bypass server check
-	result := isTelemetryEnabled(ctx, cfg, server.URL, httpClient)
+	result := isTelemetryEnabled(ctx, cfg, server.URL, httpClient, nil)
 
 	if !result {
 		t.Error("Expected telemetry to be enabled when client explicitly sets enableTelemetry=true, got disabled")
@@ -240,12 +236,8 @@ func TestIsTelemetryEnabled_ClientOverrideDisabled(t *testing.T) {
 	// Setup: Create a server that returns enabled
 	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		// Server says enabled, but client override should win
-		resp := map[string]interface{}{
-			"flags": map[string]bool{
-				"databricks.partnerplatform.clientConfigsFeatureFlags.enableTelemetryForGoDriver": true,
-			},
-		}
-		_ = json.NewEncoder(w).Encode(resp)
+		w.Header().Set("Content-Type", "application/json")
+		_, _ = w.Write([]byte(`{"flags": [{"name": "databricks.partnerplatform.clientConfigsFeatureFlags.enableTelemetryForGoDriver", "value": "true"}]}`))
 	}))
 	defer server.Close()
 
@@ -261,7 +253,7 @@ func TestIsTelemetryEnabled_ClientOverrideDisabled(t *testing.T) {
 	flagCache.getOrCreateContext(server.URL)
 	defer flagCache.releaseContext(server.URL)
 
-	result := isTelemetryEnabled(ctx, cfg, server.URL, httpClient)
+	result := isTelemetryEnabled(ctx, cfg, server.URL, httpClient, nil)
 
 	if result {
 		t.Error("Expected telemetry to be disabled when client explicitly sets enableTelemetry=false, got enabled")
@@ -272,12 +264,8 @@ func TestIsTelemetryEnabled_ClientOverrideDisabled(t *testing.T) {
 func TestIsTelemetryEnabled_ServerEnabled(t *testing.T) {
 	// Setup: Create a server that returns enabled
 	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		resp := map[string]interface{}{
-			"flags": map[string]bool{
-				"databricks.partnerplatform.clientConfigsFeatureFlags.enableTelemetryForGoDriver": true,
-			},
-		}
-		_ = json.NewEncoder(w).Encode(resp)
+		w.Header().Set("Content-Type", "application/json")
+		_, _ = w.Write([]byte(`{"flags": [{"name": "databricks.partnerplatform.clientConfigsFeatureFlags.enableTelemetryForGoDriver", "value": "true"}]}`))
 	}))
 	defer server.Close()
 
@@ -293,7 +281,7 @@ func TestIsTelemetryEnabled_ServerEnabled(t *testing.T) {
 	flagCache.getOrCreateContext(server.URL)
 	defer flagCache.releaseContext(server.URL)
 
-	result := isTelemetryEnabled(ctx, cfg, server.URL, httpClient)
+	result := isTelemetryEnabled(ctx, cfg, server.URL, httpClient, nil)
 
 	if !result {
 		t.Error("Expected telemetry to be enabled when server flag is true, got disabled")
@@ -304,12 +292,8 @@ func TestIsTelemetryEnabled_ServerEnabled(t *testing.T) {
 func TestIsTelemetryEnabled_ServerDisabled(t *testing.T) {
 	// Setup: Create a server that returns disabled
 	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		resp := map[string]interface{}{
-			"flags": map[string]bool{
-				"databricks.partnerplatform.clientConfigsFeatureFlags.enableTelemetryForGoDriver": false,
-			},
-		}
-		_ = json.NewEncoder(w).Encode(resp)
+		w.Header().Set("Content-Type", "application/json")
+		_, _ = w.Write([]byte(`{"flags": [{"name": "databricks.partnerplatform.clientConfigsFeatureFlags.enableTelemetryForGoDriver", "value": "false"}]}`))
 	}))
 	defer server.Close()
 
@@ -325,7 +309,7 @@ func TestIsTelemetryEnabled_ServerDisabled(t *testing.T) {
 	flagCache.getOrCreateContext(server.URL)
 	defer flagCache.releaseContext(server.URL)
 
-	result := isTelemetryEnabled(ctx, cfg, server.URL, httpClient)
+	result := isTelemetryEnabled(ctx, cfg, server.URL, httpClient, nil)
 
 	if result {
 		t.Error("Expected telemetry to be disabled when server flag is false, got enabled")
@@ -340,7 +324,7 @@ func TestIsTelemetryEnabled_FailSafeDefault(t *testing.T) {
 	httpClient := &http.Client{Timeout: 5 * time.Second}
 
 	// No server available, should default to disabled (fail-safe)
-	result := isTelemetryEnabled(ctx, cfg, "nonexistent-host", httpClient)
+	result := isTelemetryEnabled(ctx, cfg, "nonexistent-host", httpClient, nil)
 
 	if result {
 		t.Error("Expected telemetry to be disabled when server unavailable (fail-safe), got enabled")
@@ -367,7 +351,7 @@ func TestIsTelemetryEnabled_ServerError(t *testing.T) {
 	flagCache.getOrCreateContext(server.URL)
 	defer flagCache.releaseContext(server.URL)
 
-	result := isTelemetryEnabled(ctx, cfg, server.URL, httpClient)
+	result := isTelemetryEnabled(ctx, cfg, server.URL, httpClient, nil)
 
 	// On error, should default to disabled (fail-safe)
 	if result {
@@ -390,7 +374,7 @@ func TestIsTelemetryEnabled_ServerUnreachable(t *testing.T) {
 	flagCache.getOrCreateContext(unreachableHost)
 	defer flagCache.releaseContext(unreachableHost)
 
-	result := isTelemetryEnabled(ctx, cfg, unreachableHost, httpClient)
+	result := isTelemetryEnabled(ctx, cfg, unreachableHost, httpClient, nil)
 
 	// On error, should default to disabled (fail-safe)
 	if result {
@@ -418,7 +402,7 @@ func TestIsTelemetryEnabled_ClientOverridesServerError(t *testing.T) {
 	flagCache.getOrCreateContext(server.URL)
 	defer flagCache.releaseContext(server.URL)
 
-	result := isTelemetryEnabled(ctx, cfg, server.URL, httpClient)
+	result := isTelemetryEnabled(ctx, cfg, server.URL, httpClient, nil)
 
 	// Client override should work even when server errors
 	if !result {

telemetry/driver_integration.go

@@ -16,6 +16,7 @@ import (
 //   - host: Databricks host
 //   - httpClient: HTTP client for making requests
 //   - enableTelemetry: Client config overlay (unset = check server flag, true/false = override server)
+//   - extraHeaders: Additional HTTP headers for SPOG routing (e.g. x-databricks-org-id)
 //
 // Returns:
 //   - *Interceptor: Telemetry interceptor if enabled, nil otherwise
@@ -24,13 +25,14 @@ func InitializeForConnection(
 	host string,
 	httpClient *http.Client,
 	enableTelemetry config.ConfigValue[bool],
+	extraHeaders map[string]string,
 ) *Interceptor {
 	// Create telemetry config and apply client overlay
 	cfg := DefaultConfig()
 	cfg.EnableTelemetry = enableTelemetry
 
 	// Check if telemetry should be enabled
-	if !isTelemetryEnabled(ctx, cfg, host, httpClient) {
+	if !isTelemetryEnabled(ctx, cfg, host, httpClient, extraHeaders) {
 		return nil
 	}
 

telemetry/featureflag.go

@@ -24,6 +24,14 @@ const (
 	// flagEnableNewFeature = "databricks.partnerplatform.clientConfigsFeatureFlags.enableNewFeatureForGoDriver"
 )
 
+// driverVersion is set during initialization from the connector package.
+var driverVersion = "unknown"
+
+// SetDriverVersion sets the driver version used in feature flag endpoint paths.
+func SetDriverVersion(version string) {
+	driverVersion = version
+}
+
 // featureFlagCache manages feature flag state per host with reference counting.
 // This prevents rate limiting by caching feature flag responses.
 type featureFlagCache struct {
@@ -90,7 +98,7 @@ func (c *featureFlagCache) releaseContext(host string) {
 // getFeatureFlag retrieves a specific feature flag value for the host.
 // This is the generic method that handles caching and fetching for any flag.
 // Uses cached value if available and not expired.
-func (c *featureFlagCache) getFeatureFlag(ctx context.Context, host string, httpClient *http.Client, flagName string) (bool, error) {
+func (c *featureFlagCache) getFeatureFlag(ctx context.Context, host string, httpClient *http.Client, flagName string, extraHeaders map[string]string) (bool, error) {
 	c.mu.RLock()
 	flagCtx, exists := c.contexts[host]
 	c.mu.RUnlock()
@@ -111,7 +119,7 @@ func (c *featureFlagCache) getFeatureFlag(ctx context.Context, host string, http
 
 		// If we just created the context, make the initial blocking fetch
 		if !exists {
-			flags, err := fetchFeatureFlags(ctx, host, httpClient)
+			flags, err := fetchFeatureFlags(ctx, host, httpClient, extraHeaders)
 
 			flagCtx.mu.Lock()
 			flagCtx.fetching = false
@@ -155,7 +163,7 @@ func (c *featureFlagCache) getFeatureFlag(ctx context.Context, host string, http
 	flagCtx.mu.RUnlock()
 
 	// Fetch fresh values for all flags
-	flags, err := fetchFeatureFlags(ctx, host, httpClient)
+	flags, err := fetchFeatureFlags(ctx, host, httpClient, extraHeaders)
 
 	// Update cache (with proper locking)
 	flagCtx.mu.Lock()
@@ -184,8 +192,8 @@ func (c *featureFlagCache) getFeatureFlag(ctx context.Context, host string, http
 
 // 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, httpClient *http.Client) (bool, error) {
-	return c.getFeatureFlag(ctx, host, httpClient, flagEnableTelemetry)
+func (c *featureFlagCache) isTelemetryEnabled(ctx context.Context, host string, httpClient *http.Client, extraHeaders map[string]string) (bool, error) {
+	return c.getFeatureFlag(ctx, host, httpClient, flagEnableTelemetry, extraHeaders)
 }
 
 // isExpired returns true if the cache has expired.
@@ -203,35 +211,46 @@ func getAllFeatureFlags() []string {
 	}
 }
 
-// fetchFeatureFlags fetches multiple feature flag values from Databricks in a single request.
+// featureFlagEntry represents a single flag from the connector-service response.
+type featureFlagEntry struct {
+	Name  string `json:"name"`
+	Value string `json:"value"`
+}
+
+// featureFlagResponse represents the response from the connector-service endpoint.
+type featureFlagResponse struct {
+	Flags      []featureFlagEntry `json:"flags"`
+	TTLSeconds int                `json:"ttl_seconds,omitempty"`
+}
+
+// fetchFeatureFlags fetches feature flag values from the connector-service endpoint.
+// Endpoint: GET /api/2.0/connector-service/feature-flags/{CLIENT_TYPE}/{VERSION}
 // Returns a map of flag names to their boolean values.
-func fetchFeatureFlags(ctx context.Context, host string, httpClient *http.Client) (map[string]bool, error) {
+func fetchFeatureFlags(ctx context.Context, host string, httpClient *http.Client, extraHeaders map[string]string) (map[string]bool, error) {
 	// Add timeout to context if it doesn't have a deadline
 	if _, hasDeadline := ctx.Deadline(); !hasDeadline {
 		var cancel context.CancelFunc
 		ctx, cancel = context.WithTimeout(ctx, featureFlagHTTPTimeout)
 		defer cancel()
 	}
 
-	// Construct endpoint URL, adding https:// if not already present
+	// Construct endpoint URL using the connector-service path
 	var endpoint string
 	if strings.HasPrefix(host, "http://") || strings.HasPrefix(host, "https://") {
-		endpoint = fmt.Sprintf("%s/api/2.0/feature-flags", host)
+		endpoint = fmt.Sprintf("%s/api/2.0/connector-service/feature-flags/GOLANG/%s", host, driverVersion)
 	} else {
-		endpoint = fmt.Sprintf("https://%s/api/2.0/feature-flags", host)
+		endpoint = fmt.Sprintf("https://%s/api/2.0/connector-service/feature-flags/GOLANG/%s", host, driverVersion)
 	}
 
 	req, err := http.NewRequestWithContext(ctx, "GET", endpoint, nil)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create feature flag request: %w", err)
 	}
 
-	// Add query parameter with comma-separated list of feature flags
-	// This fetches all flags in a single request for efficiency
-	allFlags := getAllFeatureFlags()
-	q := req.URL.Query()
-	q.Add("flags", strings.Join(allFlags, ","))
-	req.URL.RawQuery = q.Encode()
+	// Add 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 {
@@ -245,18 +264,16 @@ func fetchFeatureFlags(ctx context.Context, host string, httpClient *http.Client
 		return nil, fmt.Errorf("feature flag check failed: %d", resp.StatusCode)
 	}
 
-	var result struct {
-		Flags map[string]bool `json:"flags"`
-	}
+	var result featureFlagResponse
 	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
 		return nil, fmt.Errorf("failed to decode feature flag response: %w", err)
 	}
 
-	// Return the full map of flags
-	// Flags not present in the response will have false value when accessed
-	if result.Flags == nil {
-		return make(map[string]bool), nil
+	// Convert array of {name, value} entries to a map of name -> bool
+	flags := make(map[string]bool, len(result.Flags))
+	for _, f := range result.Flags {
+		flags[f.Name] = strings.EqualFold(f.Value, "true")
 	}
 
-	return result.Flags, nil
+	return flags, nil
 }