diff --git a/auth/oauth/m2m/m2m.go b/auth/oauth/m2m/m2m.go index 76b5e513..89b53d92 100644 --- a/auth/oauth/m2m/m2m.go +++ b/auth/oauth/m2m/m2m.go @@ -5,6 +5,7 @@ import ( "fmt" "net/http" "sync" + "time" "github.com/databricks/databricks-sql-go/auth" "github.com/databricks/databricks-sql-go/auth/oauth" @@ -55,43 +56,93 @@ func (c *authClient) M2MScopes() []string { // Auth will start the OAuth Authorization Flow to authenticate the cli client // using the users credentials in the browser. Compatible with SSO. func (c *authClient) Authenticate(r *http.Request) error { - c.mx.Lock() - defer c.mx.Unlock() - if c.tokenSource != nil { - token, err := c.tokenSource.Token() - if err != nil { - return err - } - token.SetAuthHeader(r) - return nil - } - - config, err := GetConfig(context.Background(), c.hostName, c.clientID, c.clientSecret, c.scopes) + ts, first, err := c.getOrInitTokenSource() if err != nil { - return fmt.Errorf("unable to generate clientCredentials.Config: %w", err) + return err } - c.tokenSource = GetTokenSource(config) - token, err := c.tokenSource.Token() + // c.mx guards only the lazy init above. Token refreshes are still serialized + // inside the oauth2.ReuseTokenSource (it holds its own mutex for the whole + // exchange), so concurrent callers wait on an in-flight refresh. That wait is + // bounded by roughly 2 × tokenExchangeTimeout (see tokenExchangeTimeout). + token, err := ts.Token() if err != nil { - logger.Err(err).Msg("failed to get token") - + if first { + logger.Err(err).Msg("failed to get token") + } return err } - // Log via the driver's configurable logger (defaults to Warn) rather than the - // global zerolog logger, so this line honors SetLogLevel like every other. - logger.Debug().Msgf("databricks OAuth token fetched successfully") + if first { + // Log via the driver's configurable logger (defaults to Warn) rather than the + // global zerolog logger, so this line honors SetLogLevel like every other. + logger.Debug().Msgf("databricks OAuth token fetched successfully") + } token.SetAuthHeader(r) return nil +} + +// getOrInitTokenSource returns the cached token source, building it on first use. +// c.mx guards only this lazy initialization; first reports whether this call +// created the token source. +func (c *authClient) getOrInitTokenSource() (ts oauth2.TokenSource, first bool, err error) { + c.mx.Lock() + defer c.mx.Unlock() + if c.tokenSource != nil { + return c.tokenSource, false, nil + } + config, err := getConfig(context.Background(), c.hostName, c.clientID, c.clientSecret, c.scopes) + if err != nil { + return nil, false, fmt.Errorf("unable to generate clientCredentials.Config: %w", err) + } + + c.tokenSource = GetTokenSource(config) + return c.tokenSource, true, nil } +// tokenExchangeTimeout bounds each client-credentials token request so a hung +// token endpoint surfaces as an error instead of blocking every request on the +// connector (endpoint discovery is similarly capped by oauth.hostConfigTimeout). +// +// It is applied as http.Client.Timeout, so it limits each HTTP request, not the +// whole exchange. GetConfig leaves AuthStyle at AuthStyleAutoDetect, under which +// x/oauth2 retries a failed header-auth request (a timeout included) with +// params-auth until one style has succeeded for this token URL. The worst-case +// duration of a single Token() call against a hung endpoint is therefore about +// 2 × tokenExchangeTimeout. +// +// Known limitations: the value is fixed (it is a package variable only so unit +// tests can shorten it) and is not exposed as a ConnOption, so callers behind a +// slow IdP cannot tune it. The exchange uses a copy of *http.DefaultClient (as +// read when GetTokenSource is called) with only Timeout replaced, so a +// process-wide proxy/mTLS transport, CheckRedirect or Jar installed on +// http.DefaultClient still applies, matching the previous http.DefaultClient +// behaviour; any Timeout set on http.DefaultClient is superseded. Connector +// transport/TLS/proxy settings (WithTransport, WithSkipTLSHostVerify, …) do not +// apply, and the request context (r.Context()) cancellation is not propagated to +// the exchange. +var tokenExchangeTimeout = 30 * time.Second + +// GetTokenSource returns a reusable client-credentials token source. The token +// exchange uses a shallow copy of *http.DefaultClient (Transport, CheckRedirect +// and Jar are preserved) taken when GetTokenSource is called, with Timeout +// overridden by tokenExchangeTimeout. Later changes to http.DefaultClient are +// not picked up by an already-built token source. func GetTokenSource(config clientcredentials.Config) oauth2.TokenSource { - tokenSource := config.TokenSource(context.Background()) + client := *http.DefaultClient + client.Timeout = tokenExchangeTimeout + ctx := context.WithValue(context.Background(), oauth2.HTTPClient, &client) + tokenSource := config.TokenSource(ctx) return tokenSource } +// getConfig is the endpoint-resolution step Authenticate uses on its first call. +// It is a variable only so unit tests can point the token URL at a local server: +// oauth.GetEndpoint requires a Databricks hostname and a publicly trusted TLS +// certificate, which an httptest server cannot provide. +var getConfig = GetConfig + func GetConfig(ctx context.Context, issuerURL, clientID, clientSecret string, scopes []string) (clientcredentials.Config, error) { // Get the endpoint based on the host name endpoint, err := oauth.GetEndpoint(ctx, issuerURL) diff --git a/auth/oauth/m2m/m2m_test.go b/auth/oauth/m2m/m2m_test.go index 1ff8ce97..297bded7 100644 --- a/auth/oauth/m2m/m2m_test.go +++ b/auth/oauth/m2m/m2m_test.go @@ -1,11 +1,123 @@ package m2m import ( + "context" + "net/http" + "net/http/httptest" + "sync/atomic" "testing" + "time" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "golang.org/x/oauth2/clientcredentials" ) +// Issue #475 (part 4): the M2M token exchange must be bounded, like endpoint +// discovery already is (oauth.hostConfigTimeout). A hung token endpoint must +// make Authenticate return an error instead of blocking forever. +func TestM2MTokenExchangeIsBounded(t *testing.T) { + release := make(chan struct{}) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + select { + case <-release: + case <-r.Context().Done(): + } + })) + defer srv.Close() + defer close(release) + + orig := tokenExchangeTimeout + tokenExchangeTimeout = 200 * time.Millisecond + defer func() { tokenExchangeTimeout = orig }() + + c := NewAuthenticator("id", "secret", "unused.example.com").(*authClient) + c.tokenSource = GetTokenSource(clientcredentials.Config{ + ClientID: "id", + ClientSecret: "secret", + TokenURL: srv.URL + "/oidc/v1/token", + Scopes: c.scopes, + }) + + done := make(chan error, 1) + go func() { + req, _ := http.NewRequest(http.MethodGet, "https://example.com", nil) + done <- c.Authenticate(req) + }() + + select { + case err := <-done: + require.Error(t, err, "a hung token endpoint must surface as an error") + case <-time.After(5 * time.Second): + t.Fatal("Authenticate blocked on a hung token endpoint; the token exchange is unbounded") + } +} + +// The first-call branch of Authenticate (no cached token source) resolves the +// config and builds the token source under c.mx, then fetches a token. A hung +// token endpoint there must also surface as an error, and must not leave the lock +// held so later calls on the same connector can proceed. +func TestM2MFirstAuthenticateIsBounded(t *testing.T) { + release := make(chan struct{}) + var tokenHits int32 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + atomic.AddInt32(&tokenHits, 1) + select { + case <-release: + case <-r.Context().Done(): + } + })) + defer srv.Close() + defer close(release) + + origTimeout := tokenExchangeTimeout + tokenExchangeTimeout = 200 * time.Millisecond + defer func() { tokenExchangeTimeout = origTimeout }() + + // oauth.GetEndpoint needs a Databricks hostname and a publicly trusted TLS + // cert, so stub only the discovery step; GetTokenSource and Token() run for real. + origGetConfig := getConfig + var configCalls int32 + getConfig = func(_ context.Context, _, clientID, clientSecret string, scopes []string) (clientcredentials.Config, error) { + atomic.AddInt32(&configCalls, 1) + return clientcredentials.Config{ + ClientID: clientID, + ClientSecret: clientSecret, + TokenURL: srv.URL + "/oidc/v1/token", + Scopes: scopes, + }, nil + } + defer func() { getConfig = origGetConfig }() + + c := NewAuthenticator("id", "secret", "unused.example.com").(*authClient) + require.Nil(t, c.tokenSource, "test must exercise the first-call branch") + + authenticate := func() <-chan error { + done := make(chan error, 1) + go func() { + req, _ := http.NewRequest(http.MethodGet, "https://example.com", nil) + done <- c.Authenticate(req) + }() + return done + } + + // oauth2's AuthStyle auto-detection may issue more than one POST per exchange, + // so assert the endpoint was reached on each call rather than an exact count. + var prevHits int32 + for _, name := range []string{"first call", "call after a timed-out first call"} { + select { + case err := <-authenticate(): + require.Error(t, err, "%s: a hung token endpoint must surface as an error", name) + case <-time.After(5 * time.Second): + t.Fatalf("%s: Authenticate blocked on a hung token endpoint (unbounded exchange or c.mx never released)", name) + } + hits := atomic.LoadInt32(&tokenHits) + assert.Greater(t, hits, prevHits, "%s: token endpoint should have been called", name) + prevHits = hits + } + assert.EqualValues(t, 1, atomic.LoadInt32(&configCalls), "discovery should run only on the first call") +} + func TestM2MScopes(t *testing.T) { t.Run("default should be [all-apis]", func(t *testing.T) { auth := NewAuthenticator("id", "secret", "staging.cloud.company.com").(*authClient) diff --git a/connector.go b/connector.go index 4ef687e2..46d3ac02 100644 --- a/connector.go +++ b/connector.go @@ -534,6 +534,59 @@ func WithTransport(t http.RoundTripper) ConnOption { } } +// WithCloudFetchHTTPClient sets the *http.Client used to download CloudFetch +// result files from their presigned URLs. Use it to bound individual download +// attempts (e.g. via http.Client.Timeout) without replacing the Thrift +// transport. +// +// http.Client.Timeout is a per-attempt limit, not a hard cap per link: a +// timeout expiry is treated as a transient error and the download is retried +// up to retryMax times (see WithRetries) with backoff, so the worst-case time +// per link is roughly (retryMax+1) × Timeout plus backoff waits. The query +// context does NOT bound downloads: result fetching is deliberately detached +// from the QueryContext cancellation/deadline. To impose an overall bound, +// call Rows.Close() (which aborts in-flight downloads) or cap WithRetries +// together with Timeout. Timeout also covers reading +// the full response body, so size it for downloading an entire result file, +// not just receiving headers; a value too small for large files fails every +// attempt and then the query. +// +// When set, it +// takes precedence over the client derived from WithTransport, regardless of +// option order. Default is http.DefaultClient (no timeout). A nil client is +// ignored, so it never discards a client derived from WithTransport. +// +// Thrift backend only: on the kernel backend (WithUseKernel) CloudFetch runs +// inside the kernel's own HTTP stack, so this client — including its Timeout — +// is accepted but has no effect. +func WithCloudFetchHTTPClient(client *http.Client) ConnOption { + return func(c *config.Config) { + if client != nil { + c.HTTPClient = client + } + } +} + +// WithCloudFetchMinTimeToExpiry sets the safety margin applied to CloudFetch +// presigned links: a link whose expiry is less than d away is treated as +// expired. An expired link is terminal for the result set (the query must be +// re-run); links are not re-fetched. Default is 0. Negative values are +// treated as 0. Link expiry times have one-second resolution, so d is +// truncated to whole seconds: a sub-second value such as 500ms has no +// effect, and 1.9s behaves as 1s. +// +// Thrift backend only: on the kernel backend (WithUseKernel) CloudFetch link +// handling is internal to the kernel, so this option is accepted but has no +// effect. +func WithCloudFetchMinTimeToExpiry(d time.Duration) ConnOption { + return func(c *config.Config) { + if d < 0 { + d = 0 + } + c.MinTimeToExpiry = d + } +} + // WithCloudFetch sets up the use of cloud fetch for query execution. Default is true. func WithCloudFetch(useCloudFetch bool) ConnOption { return func(c *config.Config) { diff --git a/connector_test.go b/connector_test.go index 3b1382dd..0095866c 100644 --- a/connector_test.go +++ b/connector_test.go @@ -483,6 +483,68 @@ func TestNewConnector(t *testing.T) { }) } +// Issue #475: CloudFetch download client and link-expiry margin must be +// configurable without displacing the Thrift transport. +func TestCloudFetchConnOptions(t *testing.T) { + base := []ConnOption{ + WithServerHostname("databricks-host"), + WithAccessToken("token"), + WithHTTPPath("http-path"), + } + build := func(t *testing.T, opts ...ConnOption) *connector { + con, err := NewConnector(append(append([]ConnOption{}, base...), opts...)...) + require.NoError(t, err) + coni, ok := con.(*connector) + require.True(t, ok) + return coni + } + + t.Run("WithCloudFetchHTTPClient sets HTTPClient and leaves Transport nil", func(t *testing.T) { + c := &http.Client{Timeout: 42 * time.Second} + coni := build(t, WithCloudFetchHTTPClient(c)) + assert.Same(t, c, coni.cfg.HTTPClient) + assert.Equal(t, 42*time.Second, coni.cfg.HTTPClient.Timeout) + assert.Nil(t, coni.cfg.Transport) + }) + + t.Run("explicit CloudFetch client wins over WithTransport in either order", func(t *testing.T) { + tr := &http.Transport{MaxIdleConns: 3} + c := &http.Client{Timeout: 7 * time.Second} + + coni := build(t, WithTransport(tr), WithCloudFetchHTTPClient(c)) + assert.Same(t, c, coni.cfg.HTTPClient) + assert.Equal(t, tr, coni.cfg.Transport) + + coni = build(t, WithCloudFetchHTTPClient(c), WithTransport(tr)) + assert.Same(t, c, coni.cfg.HTTPClient) + assert.Equal(t, tr, coni.cfg.Transport) + }) + + t.Run("nil CloudFetch client is ignored and keeps the WithTransport-derived client", func(t *testing.T) { + tr := &http.Transport{MaxIdleConns: 3} + + for _, opts := range [][]ConnOption{ + {WithTransport(tr), WithCloudFetchHTTPClient(nil)}, + {WithCloudFetchHTTPClient(nil), WithTransport(tr)}, + } { + coni := build(t, opts...) + require.NotNil(t, coni.cfg.HTTPClient) + assert.Equal(t, tr, coni.cfg.HTTPClient.Transport) + assert.Equal(t, tr, coni.cfg.Transport) + } + }) + + t.Run("WithCloudFetchMinTimeToExpiry sets MinTimeToExpiry", func(t *testing.T) { + coni := build(t, WithCloudFetchMinTimeToExpiry(90*time.Second)) + assert.Equal(t, 90*time.Second, coni.cfg.MinTimeToExpiry) + }) + + t.Run("WithCloudFetchMinTimeToExpiry clamps negative values to 0", func(t *testing.T) { + coni := build(t, WithCloudFetchMinTimeToExpiry(-90*time.Second)) + assert.Equal(t, time.Duration(0), coni.cfg.MinTimeToExpiry) + }) +} + func TestWithQueryTags(t *testing.T) { t.Run("WithQueryTags serializes map into SessionParams QUERY_TAGS", func(t *testing.T) { con, err := NewConnector( diff --git a/driver_e2e_test.go b/driver_e2e_test.go index bf8d4c8c..c0cf8f62 100644 --- a/driver_e2e_test.go +++ b/driver_e2e_test.go @@ -7,10 +7,12 @@ import ( "encoding/json" "fmt" "io" + "net/http" "net/http/httptest" "net/url" "os" "strconv" + "sync" "testing" "time" @@ -710,6 +712,136 @@ func TestE2ECloudFetchExactRowCount(t *testing.T) { require.Equal(t, int64(wantRows), rowCount, "CloudFetch must surface exactly the requested rows, with no Arrow padding") } +// cloudFetchE2EQuery returns a query large enough to force a URL-based +// (CloudFetch) result, as in TestE2ECloudFetchExactRowCount. +func cloudFetchE2EQuery(rows int) string { + return fmt.Sprintf("SELECT id, repeat('x', 64) AS pad FROM range(%d)", rows) +} + +// drainSQLRows reads every row of query through database/sql and returns the +// row count and the first iteration error (rows.Err or a Scan error). +func drainSQLRows(t *testing.T, db *sql.DB, query string) (int64, error) { + t.Helper() + rows, err := db.QueryContext(context.Background(), query) + if err != nil { + return 0, err + } + defer rows.Close() //nolint:errcheck + + var n int64 + for rows.Next() { + var id int64 + var pad string + if scanErr := rows.Scan(&id, &pad); scanErr != nil { + return n, scanErr + } + n++ + } + return n, rows.Err() +} + +// TestE2ECloudFetchMinTimeToExpiryIsHonored covers issue #475 (part 3): +// CloudFetchConfig.MinTimeToExpiry had no ConnOption, so a caller could not +// widen the presigned-link safety margin. With a margin longer than any real +// presigned URL lifetime, every server-issued link must be treated as expired +// and iteration must fail with ErrLinkExpired; without the option the same +// query must drain exactly. +func TestE2ECloudFetchMinTimeToExpiryIsHonored(t *testing.T) { + if testing.Short() { + t.Skip("skipping large CloudFetch drain in -short mode") + } + host, httpPath, token := pecoTestingCreds(t) + + const wantRows = 2000000 + query := cloudFetchE2EQuery(wantRows) + + open := func(extra ...ConnOption) *sql.DB { + opts := append([]ConnOption{ + WithServerHostname(host), + WithPort(443), + WithHTTPPath(httpPath), + WithAccessToken(token), + WithMaxRows(500000), + WithCloudFetch(true), + }, extra...) + connector, err := NewConnector(opts...) + require.NoError(t, err) + return sql.OpenDB(connector) + } + + t.Run("default margin drains exactly", func(t *testing.T) { + db := open() + defer db.Close() //nolint:errcheck + n, err := drainSQLRows(t, db, query) + require.NoError(t, err) + require.Equal(t, int64(wantRows), n) + }) + + t.Run("24h margin treats every link as expired", func(t *testing.T) { + db := open(WithCloudFetchMinTimeToExpiry(24 * time.Hour)) + defer db.Close() //nolint:errcheck + _, err := drainSQLRows(t, db, query) + require.Error(t, err, "a 24h MinTimeToExpiry must reject every server-issued presigned link") + assert.ErrorContains(t, err, dbsqlerr.ErrLinkExpired) + }) +} + +type countingRoundTripper struct { + mu sync.Mutex + hosts []string + next http.RoundTripper +} + +func (c *countingRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) { + c.mu.Lock() + c.hosts = append(c.hosts, req.URL.Hostname()) + c.mu.Unlock() + return c.next.RoundTrip(req) +} + +func (c *countingRoundTripper) seen() []string { + c.mu.Lock() + defer c.mu.Unlock() + return append([]string(nil), c.hosts...) +} + +// TestE2ECloudFetchHTTPClientOption covers issue #475 (part 2): CloudFetch +// downloads must be able to use a caller-supplied *http.Client (so the caller +// can set http.Client.Timeout) WITHOUT replacing the Thrift transport, which +// WithTransport does as a side effect. +func TestE2ECloudFetchHTTPClientOption(t *testing.T) { + if testing.Short() { + t.Skip("skipping large CloudFetch drain in -short mode") + } + host, httpPath, token := pecoTestingCreds(t) + + const wantRows = 2000000 + rt := &countingRoundTripper{next: http.DefaultTransport} + downloadClient := &http.Client{Timeout: 2 * time.Minute, Transport: rt} + + connector, err := NewConnector( + WithServerHostname(host), + WithPort(443), + WithHTTPPath(httpPath), + WithAccessToken(token), + WithMaxRows(500000), + WithCloudFetch(true), + WithCloudFetchHTTPClient(downloadClient), + ) + require.NoError(t, err) + db := sql.OpenDB(connector) + defer db.Close() //nolint:errcheck + + n, err := drainSQLRows(t, db, cloudFetchE2EQuery(wantRows)) + require.NoError(t, err) + require.Equal(t, int64(wantRows), n) + + // The "Thrift transport is not replaced" half of the contract is asserted + // by TestCloudFetchConnOptions (cfg.Transport stays nil); here we only + // check that downloads actually went through the supplied client. + require.NotEmpty(t, rt.seen(), "CloudFetch presigned-URL downloads must go through the supplied http.Client") +} + // TestE2ETimestampNTZTypeName asserts that a zoned TIMESTAMP and a TIMESTAMP_NTZ // in the same result set report distinct database type names. The Databricks // TIMESTAMP_NTZ type is a wall-clock value with no zone; it must be reported as