From bd2308eebab538637d34ce996b73f48d4194cabc Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Mon, 17 Aug 2026 22:08:24 +0000 Subject: [PATCH] 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 425c0c62..2ccfe57a 100644 --- a/README.md +++ b/README.md @@ -264,6 +264,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`. @@ -285,6 +286,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