diff --git a/CHANGELOG.md b/CHANGELOG.md index 8acc5d8f..f1702586 100644 --- a/CHANGELOG.md +++ b/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) diff --git a/CONNECTION_PARAMETERS.md b/CONNECTION_PARAMETERS.md index e50c6753..1f2d3566 100644 --- a/CONNECTION_PARAMETERS.md +++ b/CONNECTION_PARAMETERS.md @@ -50,10 +50,13 @@ allowlist is tracked in PECOBLR-4153. | 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*` | ✅ | ❌ | +| 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 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..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 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:@…` 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). diff --git a/connector.go b/connector.go index 040c3271..3e226bc3 100644 --- a/connector.go +++ b/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, + } } } } 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" diff --git a/doc.go b/doc.go index 266deff0..18673400 100644 --- a/doc.go +++ b/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). diff --git a/internal/backend/kernel/auth.go b/internal/backend/kernel/auth.go index ddef6510..50e62ad4 100644 --- a/internal/backend/kernel/auth.go +++ b/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 diff --git a/internal/backend/kernel/backend.go b/internal/backend/kernel/backend.go index 500144f6..95bbfeb3 100644 --- a/internal/backend/kernel/backend.go +++ b/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 diff --git a/internal/backend/kernel/kernel_test.go b/internal/backend/kernel/kernel_test.go index 9a175c7d..04e061e2 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 @@ -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