From 33f2c35ea5eab8eef902d198ac85dd4e247b4b63 Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Mon, 17 Aug 2026 22:08:24 +0000 Subject: [PATCH 1/7] feat(kernel): forward identity federation client ID --- CHANGELOG.md | 1 + KERNEL_REV | 2 +- README.md | 3 ++ connector.go | 16 ++++++--- doc.go | 3 ++ internal/backend/kernel/auth.go | 4 +-- internal/backend/kernel/backend.go | 36 ++++++++++++++++++++- internal/backend/kernel/config.go | 3 ++ internal/backend/kernel/kernel_test.go | 32 ++++++++++++++---- internal/config/config.go | 33 +++++++++---------- kernel_config.go | 9 +++--- kernel_config_test.go | 23 +++++++++++++ kernel_experimental_test.go | 45 +++++++++++++++----------- 13 files changed, 156 insertions(+), 54 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 8acc5d8f..b64d7c27 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,7 @@ # Release History ## Unreleased +- Forward an optional SP-wide workload identity federation client ID through the kernel backend for PAT and OAuth authentication - 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) diff --git a/KERNEL_REV b/KERNEL_REV index a8ceb742..95cfce81 100644 --- a/KERNEL_REV +++ b/KERNEL_REV @@ -1 +1 @@ -9ac3f3d3f3e804d52d8890e890d9f8a8a617ec93 +eff8950428f4e6cc9975c663ec919f334962f7d0 diff --git a/README.md b/README.md index 702cbc18..b911bcbc 100644 --- a/README.md +++ b/README.md @@ -269,6 +269,7 @@ groups. Telemetry parameters are covered under [Telemetry](#telemetry). | Personal access token (PAT) | `token:@…`, 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 | +| SP-wide workload identity federation | — | `WithKernelIdentityFederationClientID(clientID)` | SEA only | | Custom token provider / external / static / federated | — | `WithTokenProvider`, `WithExternalToken`, `WithStaticToken`, `WithFederatedTokenProvider*` | Thrift only | **PAT** (default): supply `token:@…` in the DSN, or `WithAccessToken`. @@ -290,6 +291,8 @@ groups. Telemetry parameters are covered under [Telemetry](#telemetry). Notes for the SEA/kernel backend: +- `WithKernelIdentityFederationClientID` forwards a non-empty service-principal + client ID with PAT, OAuth M2M, or OAuth U2M to require SP-wide token exchange. - 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 diff --git a/connector.go b/connector.go index 040c3271..650bfdd7 100644 --- a/connector.go +++ b/connector.go @@ -53,10 +53,7 @@ func (c *connector) Connect(ctx context.Context) (driver.Conn, error) { be, err = newKernelBackend(ctx, c.cfg) } else { // The experimental WithKernel* options have no Thrift-path equivalent — reject - // them loudly rather than silently ignore, so a caller who sets one (a - // trusted-CA bundle, a hostname-verify skip, a proxy, a retry budget, or a - // CloudFetch chunk cap) and forgets WithUseKernel learns the option had no - // effect instead of connecting as if it were never set. Every WithKernel* + // them loudly rather than silently ignore. Every WithKernel* // option allocates KernelExperimental, so this one gate covers them all; the // message names the family rather than a stale subset that drifts as options // are added. @@ -602,6 +599,17 @@ func kernelExperimental(c *config.Config) *config.KernelExperimentalConfig { return c.KernelExperimental } +// WithKernelIdentityFederationClientID selects mandatory SP-wide workload +// identity federation for PAT, OAuth M2M, or OAuth U2M authentication. An empty +// client ID leaves BYOT / account-wide federation behavior unchanged. +// +// EXPERIMENTAL, kernel-only: the default (Thrift) backend rejects this at connect. +func WithKernelIdentityFederationClientID(clientID string) ConnOption { + return func(c *config.Config) { + kernelExperimental(c).IdentityFederationClientID = clientID + } +} + // WithKernelDecimalAsFloat makes the kernel path scan top-level DECIMAL columns to // a lossy float64 instead of the exact fixed-point string. The kernel still // receives native Arrow Decimal128; this only changes how the Go scanner diff --git a/doc.go b/doc.go index 266deff0..e31ea4d0 100644 --- a/doc.go +++ b/doc.go @@ -235,6 +235,9 @@ public client. Neither backend exposes a U2M-scopes option. Experimental kernel-only options (rejected by the default backend; the WithKernel* prefix marks them experimental): + - WithKernelIdentityFederationClientID(clientID) requires SP-wide workload identity + token exchange for PAT, OAuth M2M, or OAuth U2M. Empty preserves BYOT / account-wide + federation behavior. - WithKernelTrustedCerts(pem) adds a PEM CA bundle on top of the system roots (for a re-signing proxy or on-prem CA). Required because the kernel's TLS stack does not read SSL_CERT_FILE. diff --git a/internal/backend/kernel/auth.go b/internal/backend/kernel/auth.go index ddef6510..9b41f84e 100644 --- a/internal/backend/kernel/auth.go +++ b/internal/backend/kernel/auth.go @@ -23,7 +23,7 @@ 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 { @@ -32,7 +32,7 @@ type Auth struct { ClientID string // M2M + U2M (U2M: the cloud-inferred Go 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 diff --git a/internal/backend/kernel/backend.go b/internal/backend/kernel/backend.go index 500144f6..7a7d6431 100644 --- a/internal/backend/kernel/backend.go +++ b/internal/backend/kernel/backend.go @@ -173,6 +173,9 @@ func (k *KernelBackend) OpenSession(ctx context.Context) error { if err := k.setAuth(cfg); err != nil { return err } + if err := k.applyIdentityFederation(cfg); err != nil { + return err + } // User-Agent so query history attributes the kernel path to this driver. if k.cfg.UserAgent != "" { @@ -423,6 +426,22 @@ func (k *KernelBackend) setAuth(cfg *C.KernelSessionConfig) error { return nil } +// applyIdentityFederation forwards the optional SP-wide federation client ID. +// It is independent of the selected PAT, M2M, or U2M auth mode. +func (k *KernelBackend) applyIdentityFederation(cfg *C.KernelSessionConfig) error { + if k.cfg.IdentityFederationClientID == "" { + return nil + } + clientID := newCStr(k.cfg.IdentityFederationClientID) + 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 +} + // joinScopes renders U2M scopes as the comma-separated form the kernel U2M setter // expects. Empty (no scopes) yields "" so setAuth passes NULL and the kernel // applies its default scope set. @@ -444,6 +463,21 @@ func trySetAuth(auth Auth) error { return k.setAuth(cfg) } +// trySetIdentityFederation applies auth and the federation client ID to a +// throwaway config so tagged tests exercise the real C setters together. +func trySetIdentityFederation(cfg Config) error { + var c *C.KernelSessionConfig + if err := call(func() C.KernelStatusCode { return C.kernel_session_config_new(&c) }); err != nil { + return fmt.Errorf("config_new: %w", err) + } + defer C.kernel_session_config_free(c) + k := &KernelBackend{cfg: cfg} + if err := k.setAuth(c); err != nil { + return err + } + return k.applyIdentityFederation(c) +} + // trySetKernelTLS allocates a throwaway session config, applies the experimental // TLS knobs from cfg to it, and frees it — the analogous test seam to trySetAuth, // so a tagged test can exercise the real byte-buffer cgo setter (trusted certs) @@ -476,7 +510,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 diff --git a/internal/backend/kernel/config.go b/internal/backend/kernel/config.go index 0ae2c7c0..9cbb2004 100644 --- a/internal/backend/kernel/config.go +++ b/internal/backend/kernel/config.go @@ -16,6 +16,9 @@ type Config struct { HTTPPath string // e.g. /sql/1.0/warehouses/abc123 (carries ?o= org routing) WarehouseID string // bare warehouse id; preferred over HTTPPath when set Auth Auth // PAT / OAuth M2M / OAuth U2M + // IdentityFederationClientID selects mandatory SP-wide workload identity + // federation. Empty preserves BYOT / account-wide behavior. + IdentityFederationClientID string // UserAgent is forwarded as the User-Agent header so the kernel path is // attributed to this driver (not the kernel's built-in UA). Empty leaves it unset. diff --git a/internal/backend/kernel/kernel_test.go b/internal/backend/kernel/kernel_test.go index 9a175c7d..f4974a31 100644 --- a/internal/backend/kernel/kernel_test.go +++ b/internal/backend/kernel/kernel_test.go @@ -52,6 +52,27 @@ func TestSetAuthByMode(t *testing.T) { } } +func TestSetIdentityFederationClientID(t *testing.T) { + cases := []struct { + name string + auth Auth + id string + }{ + {"PAT", Auth{Mode: AuthPAT, Token: "dapi-x"}, "federation-client"}, + {"M2M", Auth{Mode: AuthM2M, ClientID: "cid", ClientSecret: "sec"}, "federation-client"}, + {"U2M", Auth{Mode: AuthU2M, ClientID: "u2m-cid"}, "federation-client"}, + {"empty omitted", Auth{Mode: AuthPAT, Token: "dapi-x"}, ""}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + cfg := Config{Auth: tc.auth, IdentityFederationClientID: tc.id} + if err := trySetIdentityFederation(cfg); err != nil { + t.Errorf("trySetIdentityFederation(%s) = %v, want nil", tc.name, err) + } + }) + } +} + // TestSetKernelTLS exercises the real cgo setters for the experimental kernel-only // TLS knobs (the byte-buffer trusted-CA bundle + the hostname-skip bool) via the // trySetKernelTLS seam — proving the (*C.uint8_t, C.size_t) marshalling and the C @@ -89,7 +110,7 @@ func TestSetProxy(t *testing.T) { }{ {"url only", Config{ProxyURL: "http://proxy:3128"}}, {"url + credentials", Config{ProxyURL: "http://proxy:3128", ProxyUsername: "u", ProxyPassword: "p"}}, - {"url + bypass", Config{ProxyURL: "http://proxy:3128", ProxyBypassHosts: "localhost,*.internal"}}, //nolint:gosec // G101: test literals, not real credentials + {"url + bypass", Config{ProxyURL: "http://proxy:3128", ProxyBypassHosts: "localhost,*.internal"}}, //nolint:gosec // G101: test literals, not real credentials {"all fields", Config{ProxyURL: "http://proxy:3128", ProxyUsername: "u", ProxyPassword: "p", ProxyBypassHosts: "localhost,*.internal"}}, //nolint:gosec // G101: test literals, not real credentials {"none (no-op)", Config{}}, } @@ -104,9 +125,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 Date: Thu, 20 Aug 2026 03:35:54 +0000 Subject: [PATCH 2/7] refactor(kernel): reuse federated token provider --- CHANGELOG.md | 2 +- CONNECTION_PARAMETERS.md | 8 ++- README.md | 14 ++--- connector.go | 36 ++++++++----- doc.go | 16 +++--- internal/backend/kernel/auth.go | 13 ++--- internal/backend/kernel/backend.go | 4 +- internal/backend/kernel/config.go | 4 +- internal/backend/kernel/kernel_test.go | 2 - internal/config/config.go | 33 ++++++------ kernel_backend.go | 4 +- kernel_config.go | 49 +++++++++++++----- kernel_config_test.go | 71 ++++++++++++++++++++------ kernel_experimental_test.go | 45 +++++++--------- kernel_telemetry.go | 9 +++- 15 files changed, 192 insertions(+), 118 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index b64d7c27..f1702586 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,7 +1,7 @@ # Release History ## Unreleased -- Forward an optional SP-wide workload identity federation client ID through the kernel backend for PAT and OAuth authentication +- 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) diff --git a/CONNECTION_PARAMETERS.md b/CONNECTION_PARAMETERS.md index 2b13f8e4..894e647c 100644 --- a/CONNECTION_PARAMETERS.md +++ b/CONNECTION_PARAMETERS.md @@ -46,10 +46,16 @@ Any parameter not listed below (e.g. `ansi_mode`) is passed through as a | Personal access token (PAT) | `token:@…`, 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*` | ✅ | ❌ | +| Federated token provider | — | `WithFederatedTokenProvider*` | ✅ | ✅ | +| Custom token provider / external / static | — | `WithTokenProvider`, `WithExternalToken`, `WithStaticToken` | ✅ | ❌ | Notes for the SEA/kernel backend: +- `WithFederatedTokenProvider*` resolves the provider once per connection and + supplies its token as PAT auth for kernel-side federation. The `AndClientID` + form also forwards the SP-wide federation client ID. The kernel cannot refresh + the provider token after the connection is created, so an expired token requires + 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 diff --git a/README.md b/README.md index b911bcbc..b8015c14 100644 --- a/README.md +++ b/README.md @@ -269,8 +269,8 @@ groups. Telemetry parameters are covered under [Telemetry](#telemetry). | Personal access token (PAT) | `token:@…`, 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 | -| SP-wide workload identity federation | — | `WithKernelIdentityFederationClientID(clientID)` | SEA only | -| Custom token provider / external / static / federated | — | `WithTokenProvider`, `WithExternalToken`, `WithStaticToken`, `WithFederatedTokenProvider*` | Thrift only | +| Federated token provider | — | `WithFederatedTokenProvider*` | Both | +| Custom token provider / external / static | — | `WithTokenProvider`, `WithExternalToken`, `WithStaticToken` | Thrift only | **PAT** (default): supply `token:@…` in the DSN, or `WithAccessToken`. @@ -291,15 +291,17 @@ groups. Telemetry parameters are covered under [Telemetry](#telemetry). Notes for the SEA/kernel backend: -- `WithKernelIdentityFederationClientID` forwards a non-empty service-principal - client ID with PAT, OAuth M2M, or OAuth U2M to require SP-wide token exchange. +- `WithFederatedTokenProvider*` resolves the provider once per connection and + supplies its token as PAT auth for kernel-side federation. The `AndClientID` + form also forwards the SP-wide federation client ID. The kernel cannot refresh + the provider token after the connection is created, so an expired token requires + 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 - only**. +- 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). diff --git a/connector.go b/connector.go index 650bfdd7..5ead5d09 100644 --- a/connector.go +++ b/connector.go @@ -38,6 +38,15 @@ type interactiveU2MAuthenticator interface { U2MClientID() string } +// federatedTokenAuthenticator keeps the existing federation option as the single +// auth source for both backends. Thrift uses the embedded FederationProvider; +// the kernel resolves the base provider once and performs federation itself. +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)() @@ -53,7 +62,10 @@ func (c *connector) Connect(ctx context.Context) (driver.Conn, error) { be, err = newKernelBackend(ctx, c.cfg) } else { // The experimental WithKernel* options have no Thrift-path equivalent — reject - // them loudly rather than silently ignore. Every WithKernel* + // them loudly rather than silently ignore, so a caller who sets one (a + // trusted-CA bundle, a hostname-verify skip, a proxy, a retry budget, or a + // CloudFetch chunk cap) and forgets WithUseKernel learns the option had no + // effect instead of connecting as if it were never set. Every WithKernel* // option allocates KernelExperimental, so this one gate covers them all; the // message names the family rather than a stale subset that drifts as options // are added. @@ -564,7 +576,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, + } } } } @@ -575,7 +590,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, + } } } } @@ -599,17 +618,6 @@ func kernelExperimental(c *config.Config) *config.KernelExperimentalConfig { return c.KernelExperimental } -// WithKernelIdentityFederationClientID selects mandatory SP-wide workload -// identity federation for PAT, OAuth M2M, or OAuth U2M authentication. An empty -// client ID leaves BYOT / account-wide federation behavior unchanged. -// -// EXPERIMENTAL, kernel-only: the default (Thrift) backend rejects this at connect. -func WithKernelIdentityFederationClientID(clientID string) ConnOption { - return func(c *config.Config) { - kernelExperimental(c).IdentityFederationClientID = clientID - } -} - // WithKernelDecimalAsFloat makes the kernel path scan top-level DECIMAL columns to // a lossy float64 instead of the exact fixed-point string. The kernel still // receives native Arrow Decimal128; this only changes how the Go scanner diff --git a/doc.go b/doc.go index e31ea4d0..6c78659a 100644 --- a/doc.go +++ b/doc.go @@ -217,11 +217,14 @@ 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 -(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). +statement timeout, time zone) options. Federated token providers are resolved once +per connection and passed as PAT auth for kernel-side federation; the SP-wide client +ID is also forwarded when present. The kernel cannot refresh that provider token after the +connection is created. Nothing is silently ignored: WithTimeout, custom +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). OAuth U2M is interactive: on a cache miss, connecting launches the system browser and blocks until login completes or the kernel's ~120s callback timeout expires. Because @@ -235,9 +238,6 @@ public client. Neither backend exposes a U2M-scopes option. Experimental kernel-only options (rejected by the default backend; the WithKernel* prefix marks them experimental): - - WithKernelIdentityFederationClientID(clientID) requires SP-wide workload identity - token exchange for PAT, OAuth M2M, or OAuth U2M. Empty preserves BYOT / account-wide - federation behavior. - WithKernelTrustedCerts(pem) adds a PEM CA bundle on top of the system roots (for a re-signing proxy or on-prem CA). Required because the kernel's TLS stack does not read SSL_CERT_FILE. diff --git a/internal/backend/kernel/auth.go b/internal/backend/kernel/auth.go index 9b41f84e..97181c94 100644 --- a/internal/backend/kernel/auth.go +++ b/internal/backend/kernel/auth.go @@ -27,12 +27,13 @@ const ( // 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) - 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 (8030) + Mode AuthMode + Token string // PAT + ClientID string // M2M + U2M (U2M: the cloud-inferred Go 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 (8030) + FederationClientID string // PAT from WithFederatedTokenProviderAndClientID; empty → account-wide federation } // M2MCredentialsProvider is implemented by the OAuth M2M authenticator to expose diff --git a/internal/backend/kernel/backend.go b/internal/backend/kernel/backend.go index 7a7d6431..753c7e51 100644 --- a/internal/backend/kernel/backend.go +++ b/internal/backend/kernel/backend.go @@ -426,8 +426,8 @@ func (k *KernelBackend) setAuth(cfg *C.KernelSessionConfig) error { return nil } -// applyIdentityFederation forwards the optional SP-wide federation client ID. -// It is independent of the selected PAT, M2M, or U2M auth mode. +// applyIdentityFederation forwards the optional SP-wide client ID paired with a +// token resolved from WithFederatedTokenProviderAndClientID. func (k *KernelBackend) applyIdentityFederation(cfg *C.KernelSessionConfig) error { if k.cfg.IdentityFederationClientID == "" { return nil diff --git a/internal/backend/kernel/config.go b/internal/backend/kernel/config.go index 9cbb2004..32f60149 100644 --- a/internal/backend/kernel/config.go +++ b/internal/backend/kernel/config.go @@ -16,8 +16,8 @@ type Config struct { HTTPPath string // e.g. /sql/1.0/warehouses/abc123 (carries ?o= org routing) WarehouseID string // bare warehouse id; preferred over HTTPPath when set Auth Auth // PAT / OAuth M2M / OAuth U2M - // IdentityFederationClientID selects mandatory SP-wide workload identity - // federation. Empty preserves BYOT / account-wide behavior. + // IdentityFederationClientID comes from the federated token provider's + // AndClientID form. Empty preserves BYOT / account-wide behavior. IdentityFederationClientID string // UserAgent is forwarded as the User-Agent header so the kernel path is diff --git a/internal/backend/kernel/kernel_test.go b/internal/backend/kernel/kernel_test.go index f4974a31..117aaee4 100644 --- a/internal/backend/kernel/kernel_test.go +++ b/internal/backend/kernel/kernel_test.go @@ -59,8 +59,6 @@ func TestSetIdentityFederationClientID(t *testing.T) { id string }{ {"PAT", Auth{Mode: AuthPAT, Token: "dapi-x"}, "federation-client"}, - {"M2M", Auth{Mode: AuthM2M, ClientID: "cid", ClientSecret: "sec"}, "federation-client"}, - {"U2M", Auth{Mode: AuthU2M, ClientID: "u2m-cid"}, "federation-client"}, {"empty omitted", Auth{Mode: AuthPAT, Token: "dapi-x"}, ""}, } for _, tc := range cases { diff --git a/internal/config/config.go b/internal/config/config.go index 5ae31f04..6434eb15 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -49,9 +49,15 @@ type Config struct { ThriftProtocolVersion cli_service.TProtocolVersion ThriftDebugClientProtocol bool - // KernelExperimental carries kernel-only options with no Thrift equivalent. - // Keeping them off UserConfig avoids expanding the stable DSN surface. The - // Thrift backend rejects a non-nil value; the kernel backend forwards it. + // KernelExperimental carries experimental, kernel-backend-only options that + // have no equivalent on the default (Thrift) path — currently the richer TLS + // surface (a trusted-CA bundle and an independent hostname-skip) the kernel + // exposes over its C ABI. It lives here on Config, NOT on UserConfig, so it + // stays off the stable exported/DSN surface (the same treatment TLSConfig and + // ArrowConfig get). nil means no experimental option was set. The Thrift + // backend rejects a non-nil value loudly; the kernel backend forwards it to + // the kernel C ABI. Mirrors Node's non-exported InternalConnectionOptions / + // Python's underscore-prefixed kwargs. KernelExperimental *KernelExperimentalConfig } @@ -61,10 +67,6 @@ type Config struct { // the exhaustiveness guard TestKernelExperimentalFieldsClassified asserts this so // a newly-added field can't slip through unclassified. type KernelExperimentalConfig struct { - // IdentityFederationClientID selects mandatory SP-wide workload identity - // federation for PAT, OAuth M2M, or OAuth U2M authentication. - IdentityFederationClientID string - // TLSTrustedCertsPEM is a PEM CA bundle added to the kernel's trust store on // top of the system roots (maps to kernel_session_config_set_tls_trusted_certs). // Needed because the kernel's rustls stack ignores SSL_CERT_FILE, so a custom @@ -119,15 +121,14 @@ func (k *KernelExperimentalConfig) DeepCopy() *KernelExperimentalConfig { return nil } cp := &KernelExperimentalConfig{ - IdentityFederationClientID: k.IdentityFederationClientID, - TLSSkipHostnameVerify: k.TLSSkipHostnameVerify, - ProxyURL: k.ProxyURL, - ProxyUsername: k.ProxyUsername, - ProxyPassword: k.ProxyPassword, - ProxyBypassHosts: k.ProxyBypassHosts, - RetryOverallTimeout: k.RetryOverallTimeout, - MaxChunksInMemory: k.MaxChunksInMemory, - DecimalAsFloat: k.DecimalAsFloat, + TLSSkipHostnameVerify: k.TLSSkipHostnameVerify, + ProxyURL: k.ProxyURL, + ProxyUsername: k.ProxyUsername, + ProxyPassword: k.ProxyPassword, + ProxyBypassHosts: k.ProxyBypassHosts, + RetryOverallTimeout: k.RetryOverallTimeout, + MaxChunksInMemory: k.MaxChunksInMemory, + DecimalAsFloat: k.DecimalAsFloat, } if k.TLSTrustedCertsPEM != nil { cp.TLSTrustedCertsPEM = append([]byte(nil), k.TLSTrustedCertsPEM...) diff --git a/kernel_backend.go b/kernel_backend.go index 35800d64..e7fad0df 100644 --- a/kernel_backend.go +++ b/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 } diff --git a/kernel_config.go b/kernel_config.go index b4609afd..51d59b04 100644 --- a/kernel_config.go +++ b/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 := resolveKernelAuth(ctx, cfg) if err != nil { return kernel.Auth{}, err } @@ -113,11 +118,12 @@ func validateKernelConfig(cfg *config.Config) (kernel.Auth, error) { // with newKernelBackend's kernel.Config assembly. func buildKernelConfig(cfg *config.Config, kauth kernel.Auth) kernel.Config { kc := kernel.Config{ - Host: cfg.Host, - HTTPPath: cfg.HTTPPath, - WarehouseID: cfg.WarehouseID, - Auth: kauth, - Location: cfg.Location, + Host: cfg.Host, + HTTPPath: cfg.HTTPPath, + WarehouseID: cfg.WarehouseID, + Auth: kauth, + IdentityFederationClientID: kauth.FederationClientID, + Location: cfg.Location, // Same UA the Thrift path sends, so query history attributes both alike. UserAgent: client.BuildUserAgent(cfg), // Initial namespace: no kernel config setter, so the kernel backend applies @@ -134,10 +140,11 @@ func buildKernelConfig(cfg *config.Config, kauth kernel.Auth) kernel.Config { if cfg.TLSConfig != nil && cfg.TLSConfig.InsecureSkipVerify { kc.TLSSkipVerify = true } - // Experimental kernel-only knobs have no Thrift-path equivalent and are - // forwarded to the kernel backend here. + // Experimental kernel-only TLS knobs (WithKernelTrustedCerts / + // WithKernelSkipHostnameVerify), if any. These have no Thrift-path equivalent + // (the connector rejects them on that path) and are forwarded verbatim to the + // kernel C ABI in OpenSession. if ke := cfg.KernelExperimental; ke != nil { - kc.IdentityFederationClientID = ke.IdentityFederationClientID kc.TLSTrustedCertsPEM = ke.TLSTrustedCertsPEM kc.TLSSkipHostnameVerify = ke.TLSSkipHostnameVerify // Kernel-only CloudFetch in-memory-chunk knob (WithKernelMaxChunksInMemory). @@ -291,12 +298,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 +// - anything else → rejected loudly (custom token-provider / +// external / static), so the failure names the cause instead of surfacing as // an opaque Unauthenticated. -func resolveKernelAuth(cfg *config.Config) (kernel.Auth, error) { +func resolveKernelAuth(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, + FederationClientID: 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 @@ -347,7 +368,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) } } diff --git a/kernel_config_test.go b/kernel_config_test.go index 04b38683..ca6373bb 100644 --- a/kernel_config_test.go +++ b/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 +// nonPATAuth stands in for any non-PAT, non-OAuth authenticator (custom token +// provider / external / static) — the kernel backend must reject it. It implements neither // auth.M2MCredentialsProvider nor auth.U2MCredentialsProvider. type nonPATAuth struct{} @@ -191,6 +192,52 @@ func TestValidateKernelConfig(t *testing.T) { } }) + t.Run("federated provider resolves once to PAT with SP-wide client ID", func(t *testing.T) { + c := baseKernelConfig() + c.AccessToken = "" + calls := 0 + baseProvider := tokenprovider.NewExternalTokenProvider(func() (string, error) { + calls++ + return "federated-token", nil + }) + WithFederatedTokenProviderAndClientID( + baseProvider, + "federation-client", + )(c) + if got := c.Authenticator.(*federatedTokenAuthenticator).provider; got != baseProvider { + t.Fatal("kernel federation should resolve the base provider, not the Thrift exchange wrapper") + } + a, err := validateKernelConfig(c) + if err != nil { + t.Fatalf("federated provider should validate, got %v", err) + } + if a.Mode != kernel.AuthPAT || a.Token != "federated-token" || a.FederationClientID != "federation-client" { + t.Errorf("auth = %+v, want PAT token=federated-token FederationClientID=federation-client", a) + } + if calls != 1 { + t.Errorf("provider calls = %d, want 1", calls) + } + if mech, flow := kernelAuthMech(c); mech != "PAT" || flow != "" { + t.Errorf("kernelAuthMech = (%q, %q), want (PAT, empty)", mech, flow) + } + if calls != 1 { + t.Errorf("telemetry resolved the provider again: calls = %d, want 1", calls) + } + }) + + t.Run("account-wide federated provider omits client ID", func(t *testing.T) { + c := baseKernelConfig() + c.AccessToken = "" + WithFederatedTokenProvider(tokenprovider.NewStaticTokenProvider("federated-token"))(c) + a, err := validateKernelConfig(c) + if err != nil { + t.Fatalf("federated provider should validate, got %v", err) + } + if a.Mode != kernel.AuthPAT || a.Token != "federated-token" || a.FederationClientID != "" { + t.Errorf("auth = %+v, want PAT token=federated-token and no FederationClientID", a) + } + }) + 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 +274,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 @@ -403,19 +450,11 @@ func TestKernelConfigFieldsClassified(t *testing.T) { // TestKernelExperimentalFieldsClassified only asserts the disposition map, not the // runtime copy). These run in the default CGO_ENABLED=0 build. func TestBuildKernelConfig(t *testing.T) { - t.Run("identity federation client ID forwarded for every auth mode", func(t *testing.T) { - for _, auth := range []kernel.Auth{ - {Mode: kernel.AuthPAT, Token: "dapi-x"}, - {Mode: kernel.AuthM2M, ClientID: "cid", ClientSecret: "secret"}, - {Mode: kernel.AuthU2M, ClientID: "u2m-cid"}, - } { - c := baseKernelConfig() - c.KernelExperimental = &config.KernelExperimentalConfig{IdentityFederationClientID: "federation-client"} - kc := buildKernelConfig(c, auth) - if kc.IdentityFederationClientID != "federation-client" { - t.Errorf("auth mode %v: IdentityFederationClientID = %q, want %q", - auth.Mode, kc.IdentityFederationClientID, "federation-client") - } + t.Run("federated auth client ID forwarded", func(t *testing.T) { + auth := kernel.Auth{Mode: kernel.AuthPAT, Token: "federated-token", FederationClientID: "federation-client"} + kc := buildKernelConfig(baseKernelConfig(), auth) + if kc.IdentityFederationClientID != "federation-client" { + t.Errorf("IdentityFederationClientID = %q, want %q", kc.IdentityFederationClientID, "federation-client") } }) diff --git a/kernel_experimental_test.go b/kernel_experimental_test.go index ac9b032a..de9ef508 100644 --- a/kernel_experimental_test.go +++ b/kernel_experimental_test.go @@ -27,16 +27,15 @@ import ( // deliberate decision and a setter in KernelBackend.OpenSession so it can't be // silently dropped. var kernelExperimentalFieldDisposition = map[string]string{ - "IdentityFederationClientID": "forwarded", // set_identity_federation_client_id - "TLSTrustedCertsPEM": "forwarded", // set_tls_trusted_certs - "TLSSkipHostnameVerify": "forwarded", // set_tls_skip_hostname_verification - "ProxyURL": "forwarded", // set_proxy (url) - "ProxyUsername": "forwarded", // set_proxy (username) - "ProxyPassword": "forwarded", // set_proxy (password) - "ProxyBypassHosts": "forwarded", // set_proxy (bypass_hosts) - "RetryOverallTimeout": "forwarded", // set_retry_config (overall_timeout_ms, 4th knob) - "MaxChunksInMemory": "forwarded", // set_session_conf (cloudfetch_max_chunks_in_memory, client-only) - "DecimalAsFloat": "forwarded", // kernel.Config.DecimalAsFloat → kernelOp → arrowscan (client-side scan choice) + "TLSTrustedCertsPEM": "forwarded", // set_tls_trusted_certs + "TLSSkipHostnameVerify": "forwarded", // set_tls_skip_hostname_verification + "ProxyURL": "forwarded", // set_proxy (url) + "ProxyUsername": "forwarded", // set_proxy (username) + "ProxyPassword": "forwarded", // set_proxy (password) + "ProxyBypassHosts": "forwarded", // set_proxy (bypass_hosts) + "RetryOverallTimeout": "forwarded", // set_retry_config (overall_timeout_ms, 4th knob) + "MaxChunksInMemory": "forwarded", // set_session_conf (cloudfetch_max_chunks_in_memory, client-only) + "DecimalAsFloat": "forwarded", // kernel.Config.DecimalAsFloat → kernelOp → arrowscan (client-side scan choice) } func TestKernelExperimentalFieldsClassified(t *testing.T) { @@ -64,7 +63,7 @@ func TestKernelExperimentalFieldsClassified(t *testing.T) { // branch is what rejects it. We assert the option→config wiring here (a non-nil // KernelExperimental after applying a WithKernel* option is the signal the Thrift // branch keys off). -func TestWithKernelOptionsSetExperimental(t *testing.T) { +func TestWithKernelTLSOptionsSetExperimental(t *testing.T) { cases := []struct { name string opt ConnOption @@ -73,9 +72,6 @@ func TestWithKernelOptionsSetExperimental(t *testing.T) { {"trusted certs", WithKernelTrustedCerts([]byte("ca")), func(k *config.KernelExperimentalConfig) bool { return string(k.TLSTrustedCertsPEM) == "ca" }}, - {"identity federation", WithKernelIdentityFederationClientID("federation-client"), func(k *config.KernelExperimentalConfig) bool { - return k.IdentityFederationClientID == "federation-client" - }}, {"skip hostname", WithKernelSkipHostnameVerify(), func(k *config.KernelExperimentalConfig) bool { return k.TLSSkipHostnameVerify }}, @@ -120,7 +116,6 @@ func TestWithKernelOptionsRejectedOnThriftPath(t *testing.T) { opt ConnOption }{ {"trusted certs", WithKernelTrustedCerts([]byte("ca"))}, - {"identity federation", WithKernelIdentityFederationClientID("federation-client")}, {"skip hostname", WithKernelSkipHostnameVerify()}, {"proxy", WithKernelProxy(KernelProxy{URL: "http://proxy:3128"})}, {"retry overall timeout", WithKernelRetryOverallTimeout(5 * time.Minute)}, @@ -185,23 +180,19 @@ func TestWithKernelTrustedCertsCopiesPEM(t *testing.T) { // mutation reach another. func TestKernelExperimentalDeepCopy(t *testing.T) { orig := &config.KernelExperimentalConfig{ - IdentityFederationClientID: "federation-client", - TLSTrustedCertsPEM: []byte("ca-bundle"), - TLSSkipHostnameVerify: true, - ProxyURL: "http://proxy:3128", - ProxyUsername: "u", - ProxyPassword: "p", - ProxyBypassHosts: "*.internal", - RetryOverallTimeout: 5 * time.Minute, - MaxChunksInMemory: 4, + TLSTrustedCertsPEM: []byte("ca-bundle"), + TLSSkipHostnameVerify: true, + ProxyURL: "http://proxy:3128", + ProxyUsername: "u", + ProxyPassword: "p", + ProxyBypassHosts: "*.internal", + RetryOverallTimeout: 5 * time.Minute, + MaxChunksInMemory: 4, } cp := orig.DeepCopy() if cp == nil || string(cp.TLSTrustedCertsPEM) != "ca-bundle" || !cp.TLSSkipHostnameVerify { t.Fatalf("DeepCopy lost data: %+v", cp) } - if cp.IdentityFederationClientID != "federation-client" { - t.Errorf("DeepCopy lost IdentityFederationClientID: %q", cp.IdentityFederationClientID) - } if cp.ProxyURL != "http://proxy:3128" || cp.ProxyUsername != "u" || cp.ProxyPassword != "p" || cp.ProxyBypassHosts != "*.internal" { t.Errorf("DeepCopy lost proxy fields: %+v", cp) diff --git a/kernel_telemetry.go b/kernel_telemetry.go index 2639f51d..55bdd1fe 100644 --- a/kernel_telemetry.go +++ b/kernel_telemetry.go @@ -1,6 +1,8 @@ package dbsql import ( + "context" + "github.com/databricks/databricks-sql-go/internal/backend/kernel" "github.com/databricks/databricks-sql-go/internal/config" "github.com/databricks/databricks-sql-go/telemetry" @@ -96,7 +98,12 @@ 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" ) - ka, err := resolveKernelAuth(cfg) + // Resolving a federated provider obtains a token. Telemetry must not trigger a + // second provider call after connection setup; the kernel consumes it as PAT. + if _, ok := cfg.Authenticator.(*federatedTokenAuthenticator); ok { + return authMechPAT, "" + } + ka, err := resolveKernelAuth(context.Background(), cfg) if err != nil { return "", "" } From e9c46c62da543b9dfd062b6e1583ef64bea5e980 Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Thu, 20 Aug 2026 03:48:57 +0000 Subject: [PATCH 3/7] refactor(kernel): minimize federation plumbing --- CONNECTION_PARAMETERS.md | 13 ++-- README.md | 13 ++-- connector.go | 5 +- doc.go | 16 ++--- internal/backend/kernel/auth.go | 13 ++-- internal/backend/kernel/backend.go | 43 +++--------- internal/backend/kernel/config.go | 3 - internal/backend/kernel/kernel_test.go | 22 +------ kernel_config.go | 17 +++-- kernel_config_test.go | 90 ++++++++++---------------- kernel_telemetry.go | 4 +- 11 files changed, 77 insertions(+), 162 deletions(-) diff --git a/CONNECTION_PARAMETERS.md b/CONNECTION_PARAMETERS.md index 894e647c..2ebdda60 100644 --- a/CONNECTION_PARAMETERS.md +++ b/CONNECTION_PARAMETERS.md @@ -46,23 +46,18 @@ Any parameter not listed below (e.g. `ansi_mode`) is passed through as a | Personal access token (PAT) | `token:@…`, or `accessToken=` / `authType=Pat` | `WithAccessToken` | ✅ | ✅ | | OAuth machine-to-machine (M2M) | `clientID=`+`clientSecret=` / `authType=OauthM2M` | `WithClientCredentials` | ✅ | ✅ | | OAuth user-to-machine (U2M) | `authType=OauthU2M` | `WithAuthenticator` (u2m) | ✅ | ✅ | -| Federated token provider | — | `WithFederatedTokenProvider*` | ✅ | ✅ | -| Custom token provider / external / static | — | `WithTokenProvider`, `WithExternalToken`, `WithStaticToken` | ✅ | ❌ | +| Custom / external / static / federated token provider | — | `WithTokenProvider`, `WithExternalToken`, `WithStaticToken`, `WithFederatedTokenProvider*` | ✅ | ✅ federated only | Notes for the SEA/kernel backend: -- `WithFederatedTokenProvider*` resolves the provider once per connection and - supplies its token as PAT auth for kernel-side federation. The `AndClientID` - form also forwards the SP-wide federation client ID. The kernel cannot refresh - the provider token after the connection is created, so an expired token requires - a new connection. +- The kernel snapshots one `WithFederatedTokenProvider*` token during setup; see + the README authentication notes for the refresh limitation. - 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. -- OAuth token caching/refresh is owned by the kernel on the kernel path (no driver - config). +- OAuth M2M/U2M token caching and refresh are owned by the kernel (no driver config). ## Query execution diff --git a/README.md b/README.md index b8015c14..72f1e56c 100644 --- a/README.md +++ b/README.md @@ -269,8 +269,7 @@ groups. Telemetry parameters are covered under [Telemetry](#telemetry). | Personal access token (PAT) | `token:@…`, 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 | -| Federated token provider | — | `WithFederatedTokenProvider*` | Both | -| Custom token provider / external / static | — | `WithTokenProvider`, `WithExternalToken`, `WithStaticToken` | Thrift only | +| Custom / external / static / federated token provider | — | `WithTokenProvider`, `WithExternalToken`, `WithStaticToken`, `WithFederatedTokenProvider*` | Thrift only; federated: Both | **PAT** (default): supply `token:@…` in the DSN, or `WithAccessToken`. @@ -291,19 +290,15 @@ groups. Telemetry parameters are covered under [Telemetry](#telemetry). Notes for the SEA/kernel backend: -- `WithFederatedTokenProvider*` resolves the provider once per connection and - supplies its token as PAT auth for kernel-side federation. The `AndClientID` - form also forwards the SP-wide federation client ID. The kernel cannot refresh - the provider token after the connection is created, so an expired token requires - a new connection. +- 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 authenticators are **Thrift only**. -- OAuth token caching/refresh is owned by the kernel on the kernel path (no driver - config). +- OAuth M2M/U2M token caching and refresh are owned by the kernel (no driver config). ## Cloud Fetch diff --git a/connector.go b/connector.go index 5ead5d09..a92c46b6 100644 --- a/connector.go +++ b/connector.go @@ -38,9 +38,8 @@ type interactiveU2MAuthenticator interface { U2MClientID() string } -// federatedTokenAuthenticator keeps the existing federation option as the single -// auth source for both backends. Thrift uses the embedded FederationProvider; -// the kernel resolves the base provider once and performs federation itself. +// federatedTokenAuthenticator lets Thrift use FederationProvider while the +// kernel snapshots the base provider's token. type federatedTokenAuthenticator struct { auth.Authenticator provider tokenprovider.TokenProvider diff --git a/doc.go b/doc.go index 6c78659a..abb4a7e5 100644 --- a/doc.go +++ b/doc.go @@ -216,15 +216,13 @@ context cancellation during execute; the initial namespace (WithInitialNamespace 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. Federated token providers are resolved once -per connection and passed as PAT auth for kernel-side federation; the SP-wide client -ID is also forwarded when present. The kernel cannot refresh that provider token after the -connection is created. Nothing is silently ignored: WithTimeout, custom -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). +kernel's HTTP retry config); the TLS, proxy, and session-conf (query tags, +statement timeout, time zone) options; and federated token providers (one token +snapshot during setup). Nothing is silently ignored: WithTimeout, custom 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). OAuth U2M is interactive: on a cache miss, connecting launches the system browser and blocks until login completes or the kernel's ~120s callback timeout expires. Because diff --git a/internal/backend/kernel/auth.go b/internal/backend/kernel/auth.go index 97181c94..50e62ad4 100644 --- a/internal/backend/kernel/auth.go +++ b/internal/backend/kernel/auth.go @@ -27,13 +27,12 @@ const ( // 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) - 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 (8030) - FederationClientID string // PAT from WithFederatedTokenProviderAndClientID; empty → account-wide federation + Mode AuthMode + Token string // PAT + 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 (8030) } // M2MCredentialsProvider is implemented by the OAuth M2M authenticator to expose diff --git a/internal/backend/kernel/backend.go b/internal/backend/kernel/backend.go index 753c7e51..95bbfeb3 100644 --- a/internal/backend/kernel/backend.go +++ b/internal/backend/kernel/backend.go @@ -173,9 +173,6 @@ func (k *KernelBackend) OpenSession(ctx context.Context) error { if err := k.setAuth(cfg); err != nil { return err } - if err := k.applyIdentityFederation(cfg); err != nil { - return err - } // User-Agent so query history attributes the kernel path to this driver. if k.cfg.UserAgent != "" { @@ -422,22 +419,15 @@ func (k *KernelBackend) setAuth(cfg *C.KernelSessionConfig) error { }); err != nil { return fmt.Errorf("kernel: set_auth_pat: %w", toConnError(err)) } - } - return nil -} - -// applyIdentityFederation forwards the optional SP-wide client ID paired with a -// token resolved from WithFederatedTokenProviderAndClientID. -func (k *KernelBackend) applyIdentityFederation(cfg *C.KernelSessionConfig) error { - if k.cfg.IdentityFederationClientID == "" { - return nil - } - clientID := newCStr(k.cfg.IdentityFederationClientID) - 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)) + 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 } @@ -463,21 +453,6 @@ func trySetAuth(auth Auth) error { return k.setAuth(cfg) } -// trySetIdentityFederation applies auth and the federation client ID to a -// throwaway config so tagged tests exercise the real C setters together. -func trySetIdentityFederation(cfg Config) error { - var c *C.KernelSessionConfig - if err := call(func() C.KernelStatusCode { return C.kernel_session_config_new(&c) }); err != nil { - return fmt.Errorf("config_new: %w", err) - } - defer C.kernel_session_config_free(c) - k := &KernelBackend{cfg: cfg} - if err := k.setAuth(c); err != nil { - return err - } - return k.applyIdentityFederation(c) -} - // trySetKernelTLS allocates a throwaway session config, applies the experimental // TLS knobs from cfg to it, and frees it — the analogous test seam to trySetAuth, // so a tagged test can exercise the real byte-buffer cgo setter (trusted certs) diff --git a/internal/backend/kernel/config.go b/internal/backend/kernel/config.go index 32f60149..0ae2c7c0 100644 --- a/internal/backend/kernel/config.go +++ b/internal/backend/kernel/config.go @@ -16,9 +16,6 @@ type Config struct { HTTPPath string // e.g. /sql/1.0/warehouses/abc123 (carries ?o= org routing) WarehouseID string // bare warehouse id; preferred over HTTPPath when set Auth Auth // PAT / OAuth M2M / OAuth U2M - // IdentityFederationClientID comes from the federated token provider's - // AndClientID form. Empty preserves BYOT / account-wide behavior. - IdentityFederationClientID string // UserAgent is forwarded as the User-Agent header so the kernel path is // attributed to this driver (not the kernel's built-in UA). Empty leaves it unset. diff --git a/internal/backend/kernel/kernel_test.go b/internal/backend/kernel/kernel_test.go index 117aaee4..266efb20 100644 --- a/internal/backend/kernel/kernel_test.go +++ b/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 @@ -52,25 +53,6 @@ func TestSetAuthByMode(t *testing.T) { } } -func TestSetIdentityFederationClientID(t *testing.T) { - cases := []struct { - name string - auth Auth - id string - }{ - {"PAT", Auth{Mode: AuthPAT, Token: "dapi-x"}, "federation-client"}, - {"empty omitted", Auth{Mode: AuthPAT, Token: "dapi-x"}, ""}, - } - for _, tc := range cases { - t.Run(tc.name, func(t *testing.T) { - cfg := Config{Auth: tc.auth, IdentityFederationClientID: tc.id} - if err := trySetIdentityFederation(cfg); err != nil { - t.Errorf("trySetIdentityFederation(%s) = %v, want nil", tc.name, err) - } - }) - } -} - // TestSetKernelTLS exercises the real cgo setters for the experimental kernel-only // TLS knobs (the byte-buffer trusted-CA bundle + the hostname-skip bool) via the // trySetKernelTLS seam — proving the (*C.uint8_t, C.size_t) marshalling and the C @@ -108,7 +90,7 @@ func TestSetProxy(t *testing.T) { }{ {"url only", Config{ProxyURL: "http://proxy:3128"}}, {"url + credentials", Config{ProxyURL: "http://proxy:3128", ProxyUsername: "u", ProxyPassword: "p"}}, - {"url + bypass", Config{ProxyURL: "http://proxy:3128", ProxyBypassHosts: "localhost,*.internal"}}, //nolint:gosec // G101: test literals, not real credentials + {"url + bypass", Config{ProxyURL: "http://proxy:3128", ProxyBypassHosts: "localhost,*.internal"}}, //nolint:gosec // G101: test literals, not real credentials {"all fields", Config{ProxyURL: "http://proxy:3128", ProxyUsername: "u", ProxyPassword: "p", ProxyBypassHosts: "localhost,*.internal"}}, //nolint:gosec // G101: test literals, not real credentials {"none (no-op)", Config{}}, } diff --git a/kernel_config.go b/kernel_config.go index 51d59b04..7ce4ab3e 100644 --- a/kernel_config.go +++ b/kernel_config.go @@ -118,12 +118,11 @@ func validateKernelConfigContext(ctx context.Context, cfg *config.Config) (kerne // with newKernelBackend's kernel.Config assembly. func buildKernelConfig(cfg *config.Config, kauth kernel.Auth) kernel.Config { kc := kernel.Config{ - Host: cfg.Host, - HTTPPath: cfg.HTTPPath, - WarehouseID: cfg.WarehouseID, - Auth: kauth, - IdentityFederationClientID: kauth.FederationClientID, - Location: cfg.Location, + Host: cfg.Host, + HTTPPath: cfg.HTTPPath, + WarehouseID: cfg.WarehouseID, + Auth: kauth, + Location: cfg.Location, // Same UA the Thrift path sends, so query history attributes both alike. UserAgent: client.BuildUserAgent(cfg), // Initial namespace: no kernel config setter, so the kernel backend applies @@ -314,9 +313,9 @@ func resolveKernelAuth(ctx context.Context, cfg *config.Config) (kernel.Auth, er return kernel.Auth{}, errors.New("databricks: the federated token provider returned an empty token") } return kernel.Auth{ - Mode: kernel.AuthPAT, - Token: token.AccessToken, - FederationClientID: a.clientID, + 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 diff --git a/kernel_config_test.go b/kernel_config_test.go index ca6373bb..e109f023 100644 --- a/kernel_config_test.go +++ b/kernel_config_test.go @@ -192,49 +192,40 @@ func TestValidateKernelConfig(t *testing.T) { } }) - t.Run("federated provider resolves once to PAT with SP-wide client ID", func(t *testing.T) { - c := baseKernelConfig() - c.AccessToken = "" - calls := 0 - baseProvider := tokenprovider.NewExternalTokenProvider(func() (string, error) { - calls++ - return "federated-token", nil - }) - WithFederatedTokenProviderAndClientID( - baseProvider, - "federation-client", - )(c) - if got := c.Authenticator.(*federatedTokenAuthenticator).provider; got != baseProvider { - t.Fatal("kernel federation should resolve the base provider, not the Thrift exchange wrapper") - } - a, err := validateKernelConfig(c) - if err != nil { - t.Fatalf("federated provider should validate, got %v", err) - } - if a.Mode != kernel.AuthPAT || a.Token != "federated-token" || a.FederationClientID != "federation-client" { - t.Errorf("auth = %+v, want PAT token=federated-token FederationClientID=federation-client", a) - } - if calls != 1 { - t.Errorf("provider calls = %d, want 1", calls) - } - if mech, flow := kernelAuthMech(c); mech != "PAT" || flow != "" { - t.Errorf("kernelAuthMech = (%q, %q), want (PAT, empty)", mech, flow) - } - if calls != 1 { - t.Errorf("telemetry resolved the provider again: calls = %d, want 1", calls) - } - }) - - t.Run("account-wide federated provider omits client ID", func(t *testing.T) { - c := baseKernelConfig() - c.AccessToken = "" - WithFederatedTokenProvider(tokenprovider.NewStaticTokenProvider("federated-token"))(c) - a, err := validateKernelConfig(c) - if err != nil { - t.Fatalf("federated provider should validate, got %v", err) + t.Run("federated provider supplies PAT auth", func(t *testing.T) { + cases := []struct { + name string + option func(tokenprovider.TokenProvider) ConnOption + clientID string + }{ + {"account-wide", WithFederatedTokenProvider, ""}, + {"SP-wide", func(p tokenprovider.TokenProvider) ConnOption { + return WithFederatedTokenProviderAndClientID(p, "federation-client") + }, "federation-client"}, } - if a.Mode != kernel.AuthPAT || a.Token != "federated-token" || a.FederationClientID != "" { - t.Errorf("auth = %+v, want PAT token=federated-token and no FederationClientID", a) + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + c := baseKernelConfig() + c.AccessToken = "" + calls := 0 + tc.option(tokenprovider.NewExternalTokenProvider(func() (string, error) { + calls++ + return "subject-token", nil + }))(c) + a, err := validateKernelConfig(c) + if err != nil { + t.Fatalf("federated provider should validate, got %v", err) + } + if a.Mode != kernel.AuthPAT || a.Token != "subject-token" || a.ClientID != tc.clientID { + t.Errorf("auth = %+v, want PAT token=subject-token clientID=%q", a, 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) + } + }) } }) @@ -450,21 +441,6 @@ func TestKernelConfigFieldsClassified(t *testing.T) { // TestKernelExperimentalFieldsClassified only asserts the disposition map, not the // runtime copy). These run in the default CGO_ENABLED=0 build. func TestBuildKernelConfig(t *testing.T) { - t.Run("federated auth client ID forwarded", func(t *testing.T) { - auth := kernel.Auth{Mode: kernel.AuthPAT, Token: "federated-token", FederationClientID: "federation-client"} - kc := buildKernelConfig(baseKernelConfig(), auth) - if kc.IdentityFederationClientID != "federation-client" { - t.Errorf("IdentityFederationClientID = %q, want %q", kc.IdentityFederationClientID, "federation-client") - } - }) - - t.Run("identity federation client ID omitted when unset", func(t *testing.T) { - kc := buildKernelConfig(baseKernelConfig(), kernel.Auth{Mode: kernel.AuthPAT, Token: "dapi-x"}) - if kc.IdentityFederationClientID != "" { - t.Errorf("IdentityFederationClientID = %q, want empty", kc.IdentityFederationClientID) - } - }) - t.Run("experimental TLS fields forwarded", func(t *testing.T) { c := baseKernelConfig() c.KernelExperimental = &config.KernelExperimentalConfig{ diff --git a/kernel_telemetry.go b/kernel_telemetry.go index 55bdd1fe..737ff08a 100644 --- a/kernel_telemetry.go +++ b/kernel_telemetry.go @@ -98,8 +98,8 @@ 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" ) - // Resolving a federated provider obtains a token. Telemetry must not trigger a - // second provider call after connection setup; the kernel consumes it as PAT. + // Classify federation for connection-config telemetry without taking another + // provider snapshot; the kernel consumes its snapshot as PAT. if _, ok := cfg.Authenticator.(*federatedTokenAuthenticator); ok { return authMechPAT, "" } From e222ebc3dfd53b05bd19164bc74d3b11eeb4b87b Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Thu, 20 Aug 2026 03:56:01 +0000 Subject: [PATCH 4/7] refactor(kernel): trim federation diff --- CONNECTION_PARAMETERS.md | 8 +++---- README.md | 9 +++++--- connector.go | 3 +-- doc.go | 14 ++++++------ internal/backend/kernel/kernel_test.go | 6 ++--- kernel_config.go | 18 +++++++-------- kernel_config_test.go | 31 +++++++++++--------------- kernel_telemetry.go | 7 ++---- 8 files changed, 45 insertions(+), 51 deletions(-) diff --git a/CONNECTION_PARAMETERS.md b/CONNECTION_PARAMETERS.md index 2ebdda60..d3281259 100644 --- a/CONNECTION_PARAMETERS.md +++ b/CONNECTION_PARAMETERS.md @@ -46,18 +46,18 @@ Any parameter not listed below (e.g. `ansi_mode`) is passed through as a | Personal access token (PAT) | `token:@…`, 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 / external / static / federated token provider | — | `WithTokenProvider`, `WithExternalToken`, `WithStaticToken`, `WithFederatedTokenProvider*` | ✅ | ✅ federated only | +| 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; see - the README authentication notes for the refresh limitation. - 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. -- OAuth M2M/U2M token caching and refresh are owned by the kernel (no driver config). +- OAuth token caching/refresh is owned by the kernel on the kernel path (no driver + config). ## Query execution diff --git a/README.md b/README.md index 72f1e56c..33676483 100644 --- a/README.md +++ b/README.md @@ -269,7 +269,8 @@ groups. Telemetry parameters are covered under [Telemetry](#telemetry). | Personal access token (PAT) | `token:@…`, 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 / external / static / federated token provider | — | `WithTokenProvider`, `WithExternalToken`, `WithStaticToken`, `WithFederatedTokenProvider*` | Thrift only; federated: Both | +| Custom / external / static token provider | — | `WithTokenProvider`, `WithExternalToken`, `WithStaticToken` | Thrift only | +| Federated token provider | — | `WithFederatedTokenProvider*` | Both | **PAT** (default): supply `token:@…` in the DSN, or `WithAccessToken`. @@ -297,8 +298,10 @@ Notes for the SEA/kernel backend: - **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 authenticators are **Thrift only**. -- OAuth M2M/U2M token caching and refresh are owned by the kernel (no driver config). +- 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). ## Cloud Fetch diff --git a/connector.go b/connector.go index a92c46b6..3e226bc3 100644 --- a/connector.go +++ b/connector.go @@ -38,8 +38,7 @@ type interactiveU2MAuthenticator interface { U2MClientID() string } -// federatedTokenAuthenticator lets Thrift use FederationProvider while the -// kernel snapshots the base provider's token. +// federatedTokenAuthenticator preserves the base provider for the kernel. type federatedTokenAuthenticator struct { auth.Authenticator provider tokenprovider.TokenProvider diff --git a/doc.go b/doc.go index abb4a7e5..18673400 100644 --- a/doc.go +++ b/doc.go @@ -216,13 +216,13 @@ context cancellation during execute; the initial namespace (WithInitialNamespace 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); the TLS, proxy, and session-conf (query tags, -statement timeout, time zone) options; and federated token providers (one token -snapshot during setup). Nothing is silently ignored: WithTimeout, custom 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). +kernel's HTTP retry config); and the TLS, proxy, and session-conf (query tags, +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). OAuth U2M is interactive: on a cache miss, connecting launches the system browser and blocks until login completes or the kernel's ~120s callback timeout expires. Because diff --git a/internal/backend/kernel/kernel_test.go b/internal/backend/kernel/kernel_test.go index 266efb20..04e061e2 100644 --- a/internal/backend/kernel/kernel_test.go +++ b/internal/backend/kernel/kernel_test.go @@ -105,8 +105,9 @@ 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, a non-zero overall budget, and max Date: Thu, 20 Aug 2026 04:08:21 +0000 Subject: [PATCH 5/7] test(kernel): cover federation provider errors --- kernel_config.go | 2 +- kernel_config_test.go | 29 ++++++++++++++++++++++------- 2 files changed, 23 insertions(+), 8 deletions(-) diff --git a/kernel_config.go b/kernel_config.go index a0755afb..3fa09e71 100644 --- a/kernel_config.go +++ b/kernel_config.go @@ -167,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. diff --git a/kernel_config_test.go b/kernel_config_test.go index 11eb3e49..6671fa24 100644 --- a/kernel_config_test.go +++ b/kernel_config_test.go @@ -193,26 +193,41 @@ func TestValidateKernelConfig(t *testing.T) { }) t.Run("federated provider supplies PAT auth", func(t *testing.T) { - for _, clientID := range []string{"", "federation-client"} { - t.Run("clientID="+clientID, 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 "subject-token", nil + return tc.token, tc.providerErr }) - if clientID == "" { + if tc.clientID == "" { WithFederatedTokenProvider(provider)(c) } else { - WithFederatedTokenProviderAndClientID(provider, clientID)(c) + 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 != "subject-token" || a.ClientID != clientID { - t.Errorf("auth = %+v, want PAT token=subject-token clientID=%q", a, clientID) + 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) From 0dbbf2b76581357593217ae2c59b21b4af7e159e Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Thu, 20 Aug 2026 04:16:25 +0000 Subject: [PATCH 6/7] docs: note kernel federation token snapshot --- CONNECTION_PARAMETERS.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CONNECTION_PARAMETERS.md b/CONNECTION_PARAMETERS.md index d3281259..9dab24a6 100644 --- a/CONNECTION_PARAMETERS.md +++ b/CONNECTION_PARAMETERS.md @@ -51,6 +51,8 @@ Any parameter not listed below (e.g. `ansi_mode`) is passed through as a 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 From 14f6d8090f6c91cf531f095c43c73015c4db3c22 Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Fri, 21 Aug 2026 22:20:32 +0000 Subject: [PATCH 7/7] test: preserve thrift federation delegation --- connector_test.go | 42 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 42 insertions(+) diff --git a/connector_test.go b/connector_test.go index 9eff225b..eb2267da 100644 --- a/connector_test.go +++ b/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"