CloudFetch and M2M token exchange are unbounded and uninterruptible (v1.13.0) (#475)

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)

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)

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

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(