From 60171fd118f3d4ae7149f192d9089b927ee10a60 Mon Sep 17 00:00:00 2001 From: Rashi Jaiswal Date: Mon, 2 Nov 2026 14:54:32 +0530 Subject: [PATCH] feat: add WithConnectTimeout Signed-off-by: Rashi Jaiswal --- CHANGELOG.md | 1 + connector.go | 16 +++++++ connector_test.go | 14 ++++++ doc.go | 1 + internal/backend/thrift/backend.go | 56 ++++++++++++++++++++++- internal/backend/thrift/backend_test.go | 60 +++++++++++++++++++++++++ internal/config/config.go | 2 + 7 files changed, 148 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 066fac4e..49c19b27 100644 --- a/CHANGELOG.md +++ b/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. diff --git a/connector.go b/connector.go index 4ef687e2..58e3b11e 100644 --- a/connector.go +++ b/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; diff --git a/connector_test.go b/connector_test.go index 3b1382dd..55f1e4d0 100644 --- a/connector_test.go +++ b/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)) diff --git a/doc.go b/doc.go index e78f5a61..3fb49621 100644 --- a/doc.go +++ b/doc.go @@ -87,6 +87,7 @@ Supported functional options include: - WithMaxRows( int): Sets up the max rows fetched per request. Default is 100000. Optional - WithSessionParams( map[string]string): Sets up session parameters including "timezone" and "ansi_mode". Optional - WithTimeout( Duration). Adds timeout (in time.Duration) for the server query execution. Default is no timeout. Optional + - WithConnectTimeout( 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( 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( string). Used to identify partners. Optional - WithCloudFetch (bool). Used to enable cloud fetch for the query execution. Default is true. Optional diff --git a/internal/backend/thrift/backend.go b/internal/backend/thrift/backend.go index 6601fb30..248cb623 100644 --- a/internal/backend/thrift/backend.go +++ b/internal/backend/thrift/backend.go @@ -82,7 +82,7 @@ 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{ @@ -90,7 +90,16 @@ func (b *Backend) OpenSession(ctx context.Context) error { 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 { diff --git a/internal/backend/thrift/backend_test.go b/internal/backend/thrift/backend_test.go index ed76c3bc..f2cddfab 100644 --- a/internal/backend/thrift/backend_test.go +++ b/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) + }) +} diff --git a/internal/config/config.go b/internal/config/config.go index a0ace461..98598b32 100644 --- a/internal/config/config.go +++ b/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,