feat: add WithConnectTimeout

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## Unreleased
 
+- Add `WithConnectTimeout`, which bounds session establishment and retries transient OpenSession failures (such as a cold-starting warehouse returning 503) until the deadline. Default (0) keeps a single attempt. Thrift backend only.
 - Upgrade go-oidc to v3.21.0 (go-jose v4.1.6) and x/oauth2 to v0.36.0 to address go-jose/v3 security findings (databricks/databricks-sql-go#526).
 - Reduce kernel-backend JSON rendering allocations for nested Boolean and signed-integer values.
 - **Behavior change:** Native Arrow TIMESTAMP_NTZ results now preserve their wall-clock fields and report `TIMESTAMP_NTZ` on both backends. Returned `time.Time` values use UTC instead of the session location; their represented instant and `.UTC()` value are unchanged (databricks/databricks-sql-go#506). On Thrift, this requires a server-supplied Arrow schema with native timestamps; without it, NTZ still reports `TIMESTAMP` and uses session-location handling. String timestamps (`UseArrowNativeTimestamp=false`, internal) and non-Arrow values also retain legacy parsing.

connector.go

@@ -179,6 +179,11 @@ func NewConnector(options ...ConnOption) (driver.Connector, error) {
 			"add WithUseKernel(true) or remove it", dbsqlerr.ErrRequiresKernelBackend)
 	}
 
+	if cfg.ConnectTimeout != 0 && cfg.UseKernel {
+		return nil, fmt.Errorf("databricks: WithConnectTimeout is not supported on the kernel backend: %w; "+
+			"remove WithUseKernel(true) or WithConnectTimeout", dbsqlerr.ErrInvalidKernelConfig)
+	}
+
 	client := client.RetryableClient(cfg)
 
 	return &connector{cfg: cfg, client: client}, nil
@@ -436,6 +441,17 @@ func WithTimeout(n time.Duration) ConnOption {
 	}
 }
 
+// WithConnectTimeout bounds how long establishing a session (OpenSession) may
+// take and keeps retrying an unavailable warehouse until that deadline, so a
+// cold-starting cluster has time to come up. Default (0) is a single attempt
+// with no added bound, preserving the prior behavior. Applies to the Thrift
+// backend only.
+func WithConnectTimeout(n time.Duration) ConnOption {
+	return func(c *config.Config) {
+		c.ConnectTimeout = n
+	}
+}
+
 // WithClientQueryTimeout sets a client-side deadline for kernel statement
 // execution. It is independent of WithTimeout, which remains a server-side
 // Thrift option. A positive duration is rounded up to the next millisecond;

connector_test.go

@@ -57,6 +57,20 @@ func TestFederatedTokenAuthenticatorPreservesThriftTokenExchange(t *testing.T) {
 	assert.Equal(t, subjectToken, exchange.subjectToken)
 }
 
+func TestWithConnectTimeout(t *testing.T) {
+	con, err := NewConnector()
+	require.NoError(t, err)
+	assert.Equal(t, time.Duration(0), con.(*connector).cfg.ConnectTimeout)
+
+	con, err = NewConnector(WithConnectTimeout(600 * time.Second))
+	require.NoError(t, err)
+	assert.Equal(t, 600*time.Second, con.(*connector).cfg.ConnectTimeout)
+
+	_, err = NewConnector(WithUseKernel(true), WithConnectTimeout(time.Second))
+	require.Error(t, err)
+	assert.ErrorIs(t, err, dbsqlerr.ErrInvalidKernelConfig)
+}
+
 func TestWithClientQueryTimeout(t *testing.T) {
 	t.Run("omitted remains distinguishable from explicit zero", func(t *testing.T) {
 		con, err := NewConnector(WithUseKernel(true))

doc.go

@@ -87,6 +87,7 @@ Supported functional options include:
   - WithMaxRows(<max_rows> int): Sets up the max rows fetched per request. Default is 100000. Optional
   - WithSessionParams(<params_map> map[string]string): Sets up session parameters including "timezone" and "ansi_mode". Optional
   - WithTimeout(<timeout> Duration). Adds timeout (in time.Duration) for the server query execution. Default is no timeout. Optional
+  - WithConnectTimeout(<timeout> Duration). Bounds session establishment and retries transient (429/503) OpenSession failures, such as a cold-starting warehouse, until the deadline. Default is a single attempt. Thrift backend only; rejected with WithUseKernel(true). Optional
   - WithClientQueryTimeout(<timeout> Duration). Bounds client-side statement execution on the kernel backend. Omitted preserves the legacy 600s ceiling; zero or the maximum Duration is unlimited. Requires WithUseKernel(true). Optional
   - WithUserAgentEntry(<isv-name+product-name> string). Used to identify partners. Optional
   - WithCloudFetch (bool). Used to enable cloud fetch for the query execution. Default is true. Optional

internal/backend/thrift/backend.go

@@ -82,15 +82,24 @@ func (b *Backend) OpenSession(ctx context.Context) error {
 	protocolVersion := int64(b.cfg.ThriftProtocolVersion)
 
 	debuglog.Logf(ctx, "thrift.Backend.OpenSession", "sending OpenSession protocolVersion=0x%X catalog=%q schema=%q", protocolVersion, b.cfg.Catalog, b.cfg.Schema)
-	session, err := b.client.OpenSession(ctx, &cli_service.TOpenSessionReq{
+	openReq := &cli_service.TOpenSessionReq{
 		ClientProtocolI64: &protocolVersion,
 		Configuration:     sessionParams,
 		InitialNamespace: &cli_service.TNamespace{
 			CatalogName: catalogName,
 			SchemaName:  schemaName,
 		},
 		CanUseMultipleCatalogs: &b.cfg.CanUseMultipleCatalogs,
-	})
+	}
+
+	// ConnectTimeout unset (<= 0): single attempt, identical to prior behavior.
+	var session *cli_service.TOpenSessionResp
+	var err error
+	if b.cfg.ConnectTimeout <= 0 {
+		session, err = b.client.OpenSession(ctx, openReq)
+	} else {
+		session, err = b.openSessionWithRetry(ctx, openReq)
+	}
 	if err != nil {
 		debuglog.Logf(ctx, "thrift.Backend.OpenSession", "OpenSession failed: %v", err)
 		return dbsqlerrint.NewRequestError(ctx, fmt.Sprintf("error connecting: host=%s port=%d, httpPath=%s", b.cfg.Host, b.cfg.Port, b.cfg.HTTPPath), err)
@@ -103,6 +112,49 @@ func (b *Backend) OpenSession(ctx context.Context) error {
 	return nil
 }
 
+// openSessionWithRetry bounds OpenSession by cfg.ConnectTimeout and retries with
+// exponential backoff while the failure is transient (e.g. a cold-starting
+// warehouse returning HTTP 503), so a slow cold start succeeds instead of
+// failing once the per-request retry budget (RetryMax) is exhausted. These
+// connect retries are independent of the query-level RetryMax. Non-retryable
+// failures (bad token, missing warehouse, ...) fail fast.
+func (b *Backend) openSessionWithRetry(ctx context.Context, req *cli_service.TOpenSessionReq) (*cli_service.TOpenSessionResp, error) {
+	ctx, cancel := context.WithTimeout(ctx, b.cfg.ConnectTimeout)
+	defer cancel()
+
+	wait := b.cfg.RetryWaitMin
+	if wait <= 0 {
+		wait = 1 * time.Second
+	}
+	waitMax := b.cfg.RetryWaitMax
+	if waitMax <= 0 {
+		waitMax = 30 * time.Second
+	}
+
+	log := logger.WithContext("", driverctx.CorrelationIdFromContext(ctx), "")
+	for attempt := 1; ; attempt++ {
+		session, err := b.client.OpenSession(ctx, req)
+		if err == nil {
+			return session, nil
+		}
+		if !errors.Is(err, dbsqlerrint.RetryableError) {
+			return nil, err
+		}
+
+		log.Debug().Err(err).Msgf("connect: attempt %d failed, retrying in %s", attempt, wait)
+		select {
+		case <-ctx.Done():
+			return nil, fmt.Errorf("%w; %w", err, ctx.Err())
+		case <-time.After(wait):
+		}
+
+		wait *= 2
+		if wait > waitMax {
+			wait = waitMax
+		}
+	}
+}
+
 // SessionID returns the formatted server session id (conn.id's value), computed
 // once in OpenSession. Empty until OpenSession succeeds.
 func (b *Backend) SessionID() string {

internal/backend/thrift/backend_test.go

@@ -16,6 +16,7 @@ import (
 	"github.com/databricks/databricks-sql-go/internal/cli_service"
 	"github.com/databricks/databricks-sql-go/internal/client"
 	"github.com/databricks/databricks-sql-go/internal/config"
+	dbsqlerrint "github.com/databricks/databricks-sql-go/internal/errors"
 	"github.com/databricks/databricks-sql-go/internal/thrift_protocol"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -1305,3 +1306,62 @@ func TestBackend_Execute(t *testing.T) {
 		require.Error(t, op.ExecutionError(context.Background(), err))
 	})
 }
+
+func TestBackend_OpenSession_ConnectTimeout(t *testing.T) {
+	t.Parallel()
+	retryable := func() error { return dbsqlerrint.NewRetryableError(fmt.Errorf("503"), "") }
+
+	newBackend := func(timeout time.Duration, open func(context.Context, *cli_service.TOpenSessionReq) (*cli_service.TOpenSessionResp, error)) *Backend {
+		cfg := config.WithDefaults()
+		cfg.ConnectTimeout = timeout
+		cfg.RetryWaitMin = time.Millisecond
+		cfg.RetryWaitMax = 5 * time.Millisecond
+		return NewForTest(&client.TestClient{FnOpenSession: open}, nil, cfg)
+	}
+
+	t.Run("unset timeout makes a single attempt", func(t *testing.T) {
+		var calls int
+		be := newBackend(0, func(context.Context, *cli_service.TOpenSessionReq) (*cli_service.TOpenSessionResp, error) {
+			calls++
+			return nil, retryable()
+		})
+		assert.Error(t, be.OpenSession(context.Background()))
+		assert.Equal(t, 1, calls)
+	})
+
+	t.Run("retries a transient failure until the session opens", func(t *testing.T) {
+		var calls int
+		be := newBackend(5*time.Second, func(context.Context, *cli_service.TOpenSessionReq) (*cli_service.TOpenSessionResp, error) {
+			calls++
+			if calls < 3 {
+				return nil, retryable()
+			}
+			return getTestSession(), nil
+		})
+		assert.NoError(t, be.OpenSession(context.Background()))
+		assert.Equal(t, 3, calls)
+	})
+
+	t.Run("non-retryable failure fails fast", func(t *testing.T) {
+		var calls int
+		be := newBackend(5*time.Second, func(context.Context, *cli_service.TOpenSessionReq) (*cli_service.TOpenSessionResp, error) {
+			calls++
+			return nil, fmt.Errorf("401 unauthorized")
+		})
+		assert.Error(t, be.OpenSession(context.Background()))
+		assert.Equal(t, 1, calls)
+	})
+
+	t.Run("gives up once the deadline passes", func(t *testing.T) {
+		be := newBackend(50*time.Millisecond, func(context.Context, *cli_service.TOpenSessionReq) (*cli_service.TOpenSessionResp, error) {
+			return nil, retryable()
+		})
+		start := time.Now()
+		err := be.OpenSession(context.Background())
+		require.Error(t, err)
+		assert.Less(t, time.Since(start), 5*time.Second)
+		// The error carries both the last retryable error and the deadline.
+		assert.ErrorIs(t, err, dbsqlerrint.RetryableError)
+		assert.ErrorIs(t, err, context.DeadlineExceeded)
+	})
+}

internal/config/config.go

@@ -41,6 +41,7 @@ type Config struct {
 	PollInterval              time.Duration
 	ClientTimeout             time.Duration // max time the http request can last
 	PingTimeout               time.Duration // max time allowed for ping
+	ConnectTimeout            time.Duration // max time allowed to establish a session (0 = single attempt, no bound)
 	CanUseMultipleCatalogs    bool
 	DriverName                string
 	DriverVersion             string
@@ -246,6 +247,7 @@ func (c *Config) DeepCopy() *Config {
 		PollInterval:              c.PollInterval,
 		ClientTimeout:             c.ClientTimeout,
 		PingTimeout:               c.PingTimeout,
+		ConnectTimeout:            c.ConnectTimeout,
 		CanUseMultipleCatalogs:    c.CanUseMultipleCatalogs,
 		DriverName:                c.DriverName,
 		DriverVersion:             c.DriverVersion,