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

@@ -46,10 +46,13 @@ Any parameter not listed below (e.g. `ansi_mode`) is passed through as a
 | 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,
+			}
 		}
 	}
 }

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)
 	}
 }

kernel_config_test.go

@@ -10,14 +10,15 @@ import (
 
 	"github.com/databricks/databricks-sql-go/auth/oauth"
 	"github.com/databricks/databricks-sql-go/auth/pat"
+	"github.com/databricks/databricks-sql-go/auth/tokenprovider"
 	dbsqlerr "github.com/databricks/databricks-sql-go/errors"
 	"github.com/databricks/databricks-sql-go/internal/backend/kernel"
 	"github.com/databricks/databricks-sql-go/internal/client"
 	"github.com/databricks/databricks-sql-go/internal/config"
 )
 
 // nonPATAuth stands in for any non-PAT, non-OAuth authenticator (token-provider /
-// external / federated) — the kernel backend must reject it. It implements neither
+// external / static) — the kernel backend must reject it. It implements neither
 // auth.M2MCredentialsProvider nor auth.U2MCredentialsProvider.
 type nonPATAuth struct{}
 
@@ -191,6 +192,53 @@ func TestValidateKernelConfig(t *testing.T) {
 		}
 	})
 
+	t.Run("federated provider supplies PAT auth", func(t *testing.T) {
+		providerErr := errors.New("provider failure")
+		for _, tc := range []struct {
+			name, token, clientID string
+			providerErr           error
+		}{
+			{"account-wide", "subject-token", "", nil},
+			{"SP-wide", "subject-token", "federation-client", nil},
+			{"provider error", "", "", providerErr},
+			{"empty token", "", "", nil},
+		} {
+			t.Run(tc.name, func(t *testing.T) {
+				c := baseKernelConfig()
+				c.AccessToken = ""
+				calls := 0
+				provider := tokenprovider.NewExternalTokenProvider(func() (string, error) {
+					calls++
+					return tc.token, tc.providerErr
+				})
+				if tc.clientID == "" {
+					WithFederatedTokenProvider(provider)(c)
+				} else {
+					WithFederatedTokenProviderAndClientID(provider, tc.clientID)(c)
+				}
+				a, err := validateKernelConfig(c)
+				if tc.token == "" {
+					if err == nil || tc.providerErr != nil && !errors.Is(err, tc.providerErr) {
+						t.Fatalf("error = %v, want provider error %v", err, tc.providerErr)
+					}
+					return
+				}
+				if err != nil {
+					t.Fatalf("federated provider should validate, got %v", err)
+				}
+				if a.Mode != kernel.AuthPAT || a.Token != tc.token || a.ClientID != tc.clientID {
+					t.Errorf("auth = %+v, want PAT token=%q clientID=%q", a, tc.token, tc.clientID)
+				}
+				if mech, flow := kernelAuthMech(c); mech != "PAT" || flow != "" {
+					t.Errorf("kernelAuthMech = (%q, %q), want (PAT, empty)", mech, flow)
+				}
+				if calls != 1 {
+					t.Errorf("connection-config telemetry classification resolved the provider: calls = %d, want 1", calls)
+				}
+			})
+		}
+	})
+
 	t.Run("last-applied auth wins: M2M then PAT resolves to PAT", func(t *testing.T) {
 		// Regression for the auth-mode divergence: cfg.Authenticator is the single
 		// source of truth, so setting an M2M authenticator and then a PAT (a later
@@ -227,7 +275,7 @@ func TestValidateKernelConfig(t *testing.T) {
 		c.Authenticator = nonPATAuth{}
 		_, err := validateKernelConfig(c)
 		if err == nil {
-			t.Fatal("expected an error for a token-provider/external/federated authenticator")
+			t.Fatal("expected an error for a custom token-provider/external/static authenticator")
 		}
 		// An unsupported authenticator is a "kernel can't honor this" rejection, so it
 		// must wrap ErrNotSupportedByKernel like every other unsupported option — the

kernel_telemetry.go

@@ -96,6 +96,10 @@ func kernelAuthMech(cfg *config.Config) (mech, flow string) {
 		authFlowClientCreds = "CLIENT_CREDENTIALS" //nolint:gosec // G101: telemetry auth_flow enum value, not a credential
 		authFlowBrowser     = "BROWSER_BASED_AUTHENTICATION"
 	)
+	// Avoid a second provider snapshot when classifying federation telemetry.
+	if _, ok := cfg.Authenticator.(*federatedTokenAuthenticator); ok {
+		return authMechPAT, ""
+	}
 	ka, err := resolveKernelAuth(cfg)
 	if err != nil {
 		return "", ""