feat(kernel): support federated token providers

CHANGELOG.md

@@ -1,6 +1,7 @@
 # Release History
 
 ## Unreleased
+- Support `WithFederatedTokenProvider*` on the kernel backend by resolving its token once for kernel-side federation and forwarding the optional SP-wide client ID
 - Improve telemetry error reporting: driver failures are now categorized by cause instead of reported as a generic error (databricks/databricks-sql-go#414, #415, #417, #419, #424)
 
 ## v1.14.0 (2026-07-13)

CONNECTION_PARAMETERS.md

@@ -50,10 +50,13 @@ allowlist is tracked in PECOBLR-4153.
 | Personal access token (PAT) | `token:<t>@…`, or `accessToken=` / `authType=Pat` | `WithAccessToken` | ✅ | ✅ |
 | OAuth machine-to-machine (M2M) | `clientID=`+`clientSecret=` / `authType=OauthM2M` | `WithClientCredentials` | ✅ | ✅ |
 | OAuth user-to-machine (U2M) | `authType=OauthU2M` | `WithAuthenticator` (u2m) | ✅ | ✅ |
-| Custom token provider / external / static / federated | — | `WithTokenProvider`, `WithExternalToken`, `WithStaticToken`, `WithFederatedTokenProvider*` | ✅ | ❌ |
+| Custom / external / static token provider | — | `WithTokenProvider`, `WithExternalToken`, `WithStaticToken` | ✅ | ❌ |
+| Federated token provider | — | `WithFederatedTokenProvider*` | ✅ | ✅ |
 
 Notes for the SEA/kernel backend:
 
+- The kernel snapshots one `WithFederatedTokenProvider*` token during setup;
+  `AndClientID` also forwards the SP-wide client ID. Expired tokens require a new connection.
 - Custom OAuth **M2M scopes** are rejected on the kernel path (the kernel applies its
   own default scopes). Default scopes work on both.
 - **U2M** is interactive: on a cache miss, connecting launches the browser and a

KERNEL_REV

@@ -1 +1 @@
-9ac3f3d3f3e804d52d8890e890d9f8a8a617ec93
+eff8950428f4e6cc9975c663ec919f334962f7d0

README.md

@@ -269,7 +269,8 @@ groups. Telemetry parameters are covered under [Telemetry](#telemetry).
 | Personal access token (PAT) | `token:<t>@…`, or `accessToken=` / `authType=Pat` | `WithAccessToken` | Both |
 | OAuth machine-to-machine (M2M) | `clientID=`+`clientSecret=` / `authType=OauthM2M` | `WithClientCredentials` | Both |
 | OAuth user-to-machine (U2M) | `authType=OauthU2M` | `WithAuthenticator` (u2m) | Both |
-| Custom token provider / external / static / federated | — | `WithTokenProvider`, `WithExternalToken`, `WithStaticToken`, `WithFederatedTokenProvider*` | Thrift only |
+| Custom / external / static token provider | — | `WithTokenProvider`, `WithExternalToken`, `WithStaticToken` | Thrift only |
+| Federated token provider | — | `WithFederatedTokenProvider*` | Both |
 
 **PAT** (default): supply `token:<pat>@…` in the DSN, or `WithAccessToken`.
 
@@ -290,12 +291,14 @@ groups. Telemetry parameters are covered under [Telemetry](#telemetry).
 
 Notes for the SEA/kernel backend:
 
+- The kernel snapshots one `WithFederatedTokenProvider*` token during setup;
+  `AndClientID` also forwards the SP-wide client ID. Expired tokens require a new connection.
 - Custom OAuth **M2M scopes** are rejected on the kernel path (the kernel applies its
   own default scopes). Default scopes work on both.
 - **U2M** is interactive: on a cache miss, connecting launches the browser and a
   connect-context **deadline is not honored** during the login window. U2M scopes are at
   parity with Thrift. Use PAT or M2M for headless/deadline-bound connects.
-- Custom token-provider / external / static / federated authenticators are **Thrift
+- Custom token-provider / external / static authenticators are **Thrift
   only**.
 - OAuth token caching/refresh is owned by the kernel on the kernel path (no driver
   config).

connector.go

@@ -38,6 +38,13 @@ type interactiveU2MAuthenticator interface {
 	U2MClientID() string
 }
 
+// federatedTokenAuthenticator preserves the base provider for the kernel.
+type federatedTokenAuthenticator struct {
+	auth.Authenticator
+	provider tokenprovider.TokenProvider
+	clientID string
+}
+
 // Connect returns a connection to the Databricks database from a connection pool.
 func (c *connector) Connect(ctx context.Context) (driver.Conn, error) {
 	defer debuglog.Track(ctx, "connector.Connect", "host=%s", c.cfg.Host)()
@@ -567,7 +574,10 @@ func WithFederatedTokenProvider(baseProvider tokenprovider.TokenProvider) ConnOp
 		if baseProvider != nil {
 			// Wrap with federation provider that auto-detects need for token exchange
 			federationProvider := tokenprovider.NewFederationProvider(baseProvider, c.Host)
-			c.Authenticator = tokenprovider.NewAuthenticator(federationProvider)
+			c.Authenticator = &federatedTokenAuthenticator{
+				Authenticator: tokenprovider.NewAuthenticator(federationProvider),
+				provider:      baseProvider,
+			}
 		}
 	}
 }
@@ -578,7 +588,11 @@ func WithFederatedTokenProviderAndClientID(baseProvider tokenprovider.TokenProvi
 		if baseProvider != nil {
 			// Wrap with federation provider for SP-wide federation
 			federationProvider := tokenprovider.NewFederationProviderWithClientID(baseProvider, c.Host, clientID)
-			c.Authenticator = tokenprovider.NewAuthenticator(federationProvider)
+			c.Authenticator = &federatedTokenAuthenticator{
+				Authenticator: tokenprovider.NewAuthenticator(federationProvider),
+				provider:      baseProvider,
+				clientID:      clientID,
+			}
 		}
 	}
 }

connector_test.go

@@ -2,17 +2,59 @@ package dbsql
 
 import (
 	"net/http"
+	"net/http/httptest"
 	"testing"
 	"time"
 
 	"github.com/databricks/databricks-sql-go/auth/pat"
+	"github.com/databricks/databricks-sql-go/auth/tokenprovider"
 	"github.com/databricks/databricks-sql-go/internal/client"
 	"github.com/databricks/databricks-sql-go/internal/config"
+	"github.com/golang-jwt/jwt/v5"
 	"github.com/hashicorp/go-retryablehttp"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
+func TestFederatedTokenAuthenticatorPreservesThriftTokenExchange(t *testing.T) {
+	type exchangeRequest struct {
+		path         string
+		subjectToken string
+	}
+	exchangeRequests := make(chan exchangeRequest, 1)
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if err := r.ParseForm(); err != nil {
+			http.Error(w, err.Error(), http.StatusBadRequest)
+			return
+		}
+		exchangeRequests <- exchangeRequest{
+			path:         r.URL.Path,
+			subjectToken: r.FormValue("subject_token"),
+		}
+		w.Header().Set("Content-Type", "application/json")
+		_, _ = w.Write([]byte(`{"access_token":"exchanged-token","token_type":"Bearer","expires_in":3600}`))
+	}))
+	defer server.Close()
+
+	subjectToken, err := jwt.NewWithClaims(jwt.SigningMethodHS256, jwt.MapClaims{
+		"iss": "https://external.example.com/",
+	}).SignedString([]byte("test-key"))
+	require.NoError(t, err)
+
+	cfg := config.WithDefaults()
+	cfg.Host = server.URL
+	WithFederatedTokenProvider(tokenprovider.NewStaticTokenProvider(subjectToken))(cfg)
+
+	req, err := http.NewRequest(http.MethodGet, "https://example.com", nil)
+	require.NoError(t, err)
+	require.NoError(t, cfg.Authenticator.Authenticate(req))
+	require.Equal(t, "Bearer exchanged-token", req.Header.Get("Authorization"))
+
+	exchange := <-exchangeRequests
+	assert.Equal(t, "/oidc/v1/token", exchange.path)
+	assert.Equal(t, subjectToken, exchange.subjectToken)
+}
+
 func TestNewConnector(t *testing.T) {
 	t.Run("Connector initialized with functional options should have all options set", func(t *testing.T) {
 		host := "databricks-host"

doc.go

@@ -217,8 +217,9 @@ applied post-connect via USE CATALOG / USE SCHEMA); metric-view metadata
 (WithEnableMetricViewMetadata); the retry / backoff policy (WithRetries:
 RetryWaitMin / RetryWaitMax / RetryMax, including the disable form, forwarded to the
 kernel's HTTP retry config); and the TLS, proxy, and session-conf (query tags,
-statement timeout, time zone) options. Nothing is silently ignored: WithTimeout,
-token-provider / external / federated authenticators, and custom M2M OAuth scopes
+statement timeout, time zone) options. Federated token providers use one token snapshot
+during setup. Nothing is silently ignored: WithTimeout, token-provider / external /
+static authenticators, and custom M2M OAuth scopes
 (the kernel applies its own) are rejected at connect; staging (PUT/GET/REMOVE on a
 Unity Catalog volume) is rejected at execute. WithMaxRows is accepted but inert (the
 kernel manages fetching below the C ABI).

internal/backend/kernel/auth.go

@@ -23,16 +23,16 @@ const (
 // through to it by setAuth. resolveKernelAuth populates Scopes with the same
 // cloud-specific set the Thrift path requests (via oauth.GetScopes) so both
 // backends authorize identically; RedirectPort stays zero (no user option, kernel
-// default 8020) but is kept so kernel.Auth models the full set_auth_u2m surface —
+// default 8030) but is kept so kernel.Auth models the full set_auth_u2m surface —
 // a future WithOAuthRedirectPort becomes populating it, not re-plumbing the setter.
 // TestSetAuthByMode's "U2M full" case pins the marshalling of both.
 type Auth struct {
 	Mode         AuthMode
 	Token        string   // PAT
-	ClientID     string   // M2M + U2M (U2M: the cloud-inferred Go client id)
+	ClientID     string   // M2M + U2M; federated PAT uses the optional SP-wide client id
 	ClientSecret string   // M2M
 	Scopes       []string // U2M — Thrift-parity scopes from oauth.GetScopes; nil → kernel default
-	RedirectPort uint16   // U2M — no user option today; 0 → kernel default port (8020)
+	RedirectPort uint16   // U2M — no user option today; 0 → kernel default port (8030)
 }
 
 // M2MCredentialsProvider is implemented by the OAuth M2M authenticator to expose

internal/backend/kernel/backend.go

@@ -419,6 +419,15 @@ func (k *KernelBackend) setAuth(cfg *C.KernelSessionConfig) error {
 		}); err != nil {
 			return fmt.Errorf("kernel: set_auth_pat: %w", toConnError(err))
 		}
+		if k.cfg.Auth.ClientID != "" {
+			clientID := newCStr(k.cfg.Auth.ClientID)
+			defer clientID.free()
+			if err := call(func() C.KernelStatusCode {
+				return C.kernel_session_config_set_identity_federation_client_id(cfg, clientID.c)
+			}); err != nil {
+				return fmt.Errorf("kernel: set_identity_federation_client_id: %w", toConnError(err))
+			}
+		}
 	}
 	return nil
 }
@@ -476,7 +485,7 @@ func trySetProxy(cfg Config) error {
 // trySetRetry allocates a throwaway session config, applies the retry config from
 // cfg to it, and frees it — the analogous test seam to trySetProxy, so a tagged
 // test can exercise the real kernel_session_config_set_retry_config cgo setter
-// (the 4 knobs, plus the InvalidArgument rejections for a degenerate range) end to
+// (the 4 knobs, plus the InvalidArgument rejection for a zero minimum) end to
 // end. Not used in production.
 func trySetRetry(cfg Config) error {
 	var c *C.KernelSessionConfig

internal/backend/kernel/kernel_test.go

@@ -31,6 +31,7 @@ func TestSetAuthByMode(t *testing.T) {
 		auth Auth
 	}{
 		{"PAT", Auth{Mode: AuthPAT, Token: "dapi-x"}},
+		{"federated PAT", Auth{Mode: AuthPAT, Token: "subject-token", ClientID: "federation-client"}},
 		{"M2M", Auth{Mode: AuthM2M, ClientID: "cid", ClientSecret: "sec"}},
 		// "U2M full" populates Scopes/RedirectPort, which no production path sets today
 		// (resolveKernelAuth sources only the client id — see kernel.Auth docs). It is
@@ -104,8 +105,8 @@ func TestSetProxy(t *testing.T) {
 
 // TestSetRetry exercises the real kernel_session_config_set_retry_config cgo setter
 // via the trySetRetry seam: a valid range succeeds (incl. the disable form,
-// MaxRetries=0, and a non-zero overall budget), and a degenerate range (min=0 or
-// max<min) is rejected by the kernel as InvalidArgument. A no-op when Config.Retry
+// MaxRetries=0, and a non-zero overall budget). A zero minimum is rejected, but
+// max<min is corrected by the kernel. A no-op when Config.Retry
 // is nil. Proves the 4-arg marshalling and the C signature.
 func TestSetRetry(t *testing.T) {
 	cases := []struct {
@@ -117,9 +118,8 @@ func TestSetRetry(t *testing.T) {
 		{"disable (0 retries)", Config{Retry: &RetryConfig{MinWait: time.Second, MaxWait: 30 * time.Second, MaxRetries: 0}}, false},
 		{"with overall budget", Config{Retry: &RetryConfig{MinWait: time.Second, MaxWait: 30 * time.Second, MaxRetries: 4, OverallTimeout: 5 * time.Minute}}, false},
 		{"none (no-op)", Config{}, false},
-		// The kernel setter rejects a degenerate range: min==0 and max<min.
 		{"min zero rejected", Config{Retry: &RetryConfig{MinWait: 0, MaxWait: time.Second, MaxRetries: 3}}, true},
-		{"max below min rejected", Config{Retry: &RetryConfig{MinWait: 5 * time.Second, MaxWait: time.Second, MaxRetries: 3}}, true},
+		{"max below min corrected", Config{Retry: &RetryConfig{MinWait: 5 * time.Second, MaxWait: time.Second, MaxRetries: 3}}, false},
 	}
 	for _, c := range cases {
 		t.Run(c.name, func(t *testing.T) {

kernel_backend.go

@@ -15,12 +15,12 @@ import (
 // same config fields Thrift does and translates them to the kernel's flat
 // connection config, so the user-facing options are unchanged — only the routing
 // differs. The public API adds nothing beyond WithUseKernel.
-func newKernelBackend(_ context.Context, cfg *config.Config) (backend.Backend, error) {
+func newKernelBackend(ctx context.Context, cfg *config.Config) (backend.Backend, error) {
 	// Reject options the kernel path can't honor yet + resolve the auth form. The
 	// validation is pure Go and lives in kernel_config.go (untagged) so its tests —
 	// including the exhaustiveness guard against a dropped Config field — run in the
 	// default CGO_ENABLED=0 build. It returns kernel.Auth directly.
-	kauth, err := validateKernelConfig(cfg)
+	kauth, err := validateKernelConfigContext(ctx, cfg)
 	if err != nil {
 		return nil, err
 	}

kernel_config.go

@@ -1,6 +1,7 @@
 package dbsql
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"net/url"
@@ -37,6 +38,10 @@ import (
 // "kernel can't honor this option" case with errors.Is (e.g. to fall back to the
 // default backend) instead of matching on message text.
 func validateKernelConfig(cfg *config.Config) (kernel.Auth, error) {
+	return validateKernelConfigContext(context.Background(), cfg)
+}
+
+func validateKernelConfigContext(ctx context.Context, cfg *config.Config) (kernel.Auth, error) {
 	// Initial namespace (WithInitialNamespace) is forwarded, not rejected: the
 	// kernel C ABI has no catalog/schema setter, so KernelBackend.OpenSession
 	// selects it post-connect with USE CATALOG / USE SCHEMA. No per-backend handling
@@ -72,7 +77,7 @@ func validateKernelConfig(cfg *config.Config) (kernel.Auth, error) {
 	// the single source of truth. resolveKernelAuth rejects unsupported authenticators
 	// loudly so the failure names the cause instead of surfacing as an opaque
 	// Unauthenticated.
-	kauth, err := resolveKernelAuth(cfg)
+	kauth, err := resolveKernelAuthContext(ctx, cfg)
 	if err != nil {
 		return kernel.Auth{}, err
 	}
@@ -162,7 +167,7 @@ func buildKernelConfig(cfg *config.Config, kauth kernel.Auth) kernel.Config {
 // kernelRetryPlaceholderWaits are the backoff bounds substituted when the caller
 // gave no valid wait range but a definite attempt count to honor — the disable form
 // (RetryMax < 0), or WithRetries(n, 0, 0) where WithDefaults' waits were overwritten
-// to zero. The kernel setter validates the range (it rejects min == 0 / max < min),
+// to zero. The kernel setter rejects min == 0 and corrects max < min,
 // so a valid one must be passed even when the attempts make the backoff moot; any
 // positive min<=max works, and the kernel's own defaults (1s / 60s) are the natural
 // choice.
@@ -292,12 +297,26 @@ func resolveKernelProxy(cfg *config.Config, kc *kernel.Config) {
 // satisfy structurally:
 //   - implements M2MCredentialsProvider → M2M (client id + secret)
 //   - implements U2MCredentialsProvider → U2M (browser/PKCE; kernel-owned flow)
+//   - federated token provider          → PAT resolved once from the provider
 //   - PAT / nil / noop                  → PAT (from AccessToken or a *pat.PATAuth)
 //   - anything else                     → rejected loudly (token-provider / external
-//     / static / federated), so the failure names the cause instead of surfacing as
+//     / static), so the failure names the cause instead of surfacing as
 //     an opaque Unauthenticated.
 func resolveKernelAuth(cfg *config.Config) (kernel.Auth, error) {
+	return resolveKernelAuthContext(context.Background(), cfg)
+}
+
+func resolveKernelAuthContext(ctx context.Context, cfg *config.Config) (kernel.Auth, error) {
 	switch a := cfg.Authenticator.(type) {
+	case *federatedTokenAuthenticator:
+		token, err := a.provider.GetToken(ctx)
+		if err != nil {
+			return kernel.Auth{}, fmt.Errorf("databricks: failed to get a federated token for the kernel backend: %w", err)
+		}
+		if token == nil || token.AccessToken == "" {
+			return kernel.Auth{}, errors.New("databricks: the federated token provider returned an empty token")
+		}
+		return kernel.Auth{Mode: kernel.AuthPAT, Token: token.AccessToken, ClientID: a.clientID}, nil
 	case kernel.M2MCredentialsProvider:
 		// The kernel's set_auth_m2m takes no scopes and applies "all-apis" itself, so
 		// a custom scope set can't be forwarded — reject it instead of silently
@@ -317,7 +336,7 @@ func resolveKernelAuth(cfg *config.Config) (kernel.Auth, error) {
 		// kernel applied its own default set (all-apis + offline_access), which a
 		// workspace whose public client isn't granted all-apis rejects with
 		// access_denied. RedirectPort is still left zero (no user option; kernel
-		// default 8020). Passing nil to GetScopes yields the pure cloud-default set.
+		// default 8030). Passing nil to GetScopes yields the pure cloud-default set.
 		return kernel.Auth{Mode: kernel.AuthU2M, ClientID: a.U2MClientID(), Scopes: oauth.GetScopes(cfg.Host, nil)}, nil
 	case nil, *noop.NoopAuth, *pat.PATAuth:
 		// PAT (or no explicit authenticator). WithAccessToken sets both
@@ -348,7 +367,7 @@ func resolveKernelAuth(cfg *config.Config) (kernel.Auth, error) {
 		// kernel can't honor.)
 		return kernel.Auth{}, fmt.Errorf("databricks: this authenticator is %w; "+
 			"PAT (WithAccessToken) and OAuth M2M/U2M (WithClientCredentials / authType) are supported, but "+
-			"token-provider, external/static, and federated authenticators are not — "+
-			"use one of those or the default (Thrift) backend", dbsqlerr.ErrNotSupportedByKernel)
+			"custom token-provider and external/static authenticators are not — "+
+			"use PAT/OAuth or the default (Thrift) backend", dbsqlerr.ErrNotSupportedByKernel)
 	}
 }