feat(kernel): add client query timeout

CHANGELOG.md

@@ -1,6 +1,7 @@
 # Release History
 
 ## Unreleased
+- Add `WithClientQueryTimeout` for client-side kernel execution deadlines while preserving the legacy execute contract when the option is omitted.
 - Transparently auto-recover Thrift connections to Reyden / Real-Time warehouses: when a warehouse rejects the default Thrift protocol (SQLSTATE `KP001`), the session is re-opened on the SEA/kernel backend and the warehouse is remembered so later connections skip Thrift. Applies only when no backend was chosen explicitly (`WithUseKernel`). Note: recovery requires a `databricks_kernel` build, since the kernel backend is otherwise not linked in.
 - Parse explicit HTTP(S) server schemes case-insensitively, so mixed-case HTTP is rejected correctly by the SEA/kernel backend.
 

CONNECTION_PARAMETERS.md

@@ -98,6 +98,7 @@ Notes for the SEA/kernel backend:
 |---|---|:---:|:---:|---|---|
 | `maxRows` | `WithMaxRows` | ✅ | ⚠️ | `100000` | Max rows per fetch. On the kernel path the kernel manages paging, so this is accepted but has no effect. |
 | `timeout` | `WithTimeout` | ✅ | ❌ | no timeout | Server-side query timeout, in seconds. On the kernel path use the `STATEMENT_TIMEOUT` session parameter instead. |
+| — | `WithClientQueryTimeout` | ❌ | ✅ | legacy 600s ceiling | Client-side execution deadline through terminal-state mapping. Zero or the maximum `time.Duration` is unlimited. Connector-only; no DSN spelling. |
 | `userAgentEntry` | `WithUserAgentEntry` | ✅ | ✅ | | Identifies your application (partners/ISVs), format `<isv-name+product-name>`. |
 | *(session param)* | `WithSessionParams` | ✅ | ✅ | | Arbitrary server session confs (e.g. `ansi_mode`, `STATEMENT_TIMEOUT`, `QUERY_TAGS`) are forwarded unchanged. |
 | *(via session param)* | `WithQueryTags` | ✅ | ✅ | | Session-level query tags (serialized into `QUERY_TAGS`). |

KERNEL_REV

@@ -1 +1 @@
-21504eae074c7f5b11888a4ebcc7468780b2f844
+bfa6b6003bda6ddd44dbb667fb5dbbc9f4fee79d

README.md

@@ -250,6 +250,7 @@ session parameter on both backends.
 |---|---|---|---|---|
 | `maxRows` | `WithMaxRows` | Thrift only (inert on kernel) | `100000` | Max rows per fetch. On the kernel path the kernel manages paging, so this is accepted but has no effect. |
 | `timeout` | `WithTimeout` | Thrift only | no timeout | Server-side query timeout, in seconds. On the kernel path use the `STATEMENT_TIMEOUT` session parameter instead. |
+| — | `WithClientQueryTimeout` | SEA only | legacy 600s ceiling | Client-side execution deadline. Zero or the maximum `time.Duration` is unlimited. Connector-only; no DSN spelling. |
 | `userAgentEntry` | `WithUserAgentEntry` | Both | | Identifies your application (partners/ISVs), format `<isv-name+product-name>`. |
 | *(session param)* | `WithSessionParams` | Both | | Arbitrary session confs (e.g. `ansi_mode`, `STATEMENT_TIMEOUT`, `QUERY_TAGS`). |
 | *(via session param)* | `WithQueryTags` | Both | | Session-level query tags (serialized into `QUERY_TAGS`). |

connector.go

@@ -165,6 +165,13 @@ func NewConnector(options ...ConnOption) (driver.Connector, error) {
 	for _, opt := range options {
 		opt(cfg)
 	}
+	if err := validateClientQueryTimeout(cfg.ClientQueryTimeout); err != nil {
+		return nil, err
+	}
+	if cfg.ClientQueryTimeout != nil && !cfg.UseKernel {
+		return nil, fmt.Errorf("databricks: WithClientQueryTimeout %w; "+
+			"add WithUseKernel(true) or remove it", dbsqlerr.ErrRequiresKernelBackend)
+	}
 
 	client := client.RetryableClient(cfg)
 
@@ -419,6 +426,21 @@ func WithTimeout(n time.Duration) ConnOption {
 	}
 }
 
+// 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;
+// zero and the maximum time.Duration select unlimited execution. Omitting this
+// option preserves the kernel's legacy 600-second polling ceiling.
+//
+// This connector-only option requires WithUseKernel(true). Negative or
+// out-of-range durations are rejected by NewConnector.
+func WithClientQueryTimeout(timeout time.Duration) ConnOption {
+	return func(c *config.Config) {
+		configured := timeout
+		c.ClientQueryTimeout = &configured
+	}
+}
+
 // Sets the initial catalog name and schema name in the session.
 // Use <select * from foo> instead of <select * from catalog.schema.foo>
 func WithInitialNamespace(catalog, schema string) ConnOption {
@@ -825,6 +847,13 @@ func WithTokenCache(enabled bool) ConnOption {
 // UseKernel nor other backend-selecting options). On success or failure, it
 // returns the backend, session latency, and error.
 func (c *connector) openSessionWithReydenFallback(ctx context.Context) (backend.Backend, int64, error) {
+	// NewConnector enforces this for public construction. Keep the connect-time
+	// guard as defense in depth for package-internal connectors assembled directly.
+	if !c.cfg.UseKernel && c.cfg.ClientQueryTimeout != nil {
+		return nil, 0, fmt.Errorf("databricks: WithClientQueryTimeout %w; "+
+			"add WithUseKernel(true) or remove it", dbsqlerr.ErrRequiresKernelBackend)
+	}
+
 	// Guardrail, checked up front — before the cache pre-check — so the outcome does not depend
 	// on process-global cache state. The experimental WithKernel* options have no Thrift-path
 	// equivalent, so a caller who sets one (a trusted-CA bundle, a hostname-verify skip, a proxy,

connector_test.go

@@ -8,6 +8,8 @@ import (
 
 	"github.com/databricks/databricks-sql-go/auth/pat"
 	"github.com/databricks/databricks-sql-go/auth/tokenprovider"
+	dbsqlerr "github.com/databricks/databricks-sql-go/errors"
+	"github.com/databricks/databricks-sql-go/internal/backend/kernel"
 	"github.com/databricks/databricks-sql-go/internal/client"
 	"github.com/databricks/databricks-sql-go/internal/config"
 	"github.com/golang-jwt/jwt/v5"
@@ -55,6 +57,82 @@ func TestFederatedTokenAuthenticatorPreservesThriftTokenExchange(t *testing.T) {
 	assert.Equal(t, subjectToken, exchange.subjectToken)
 }
 
+func TestWithClientQueryTimeout(t *testing.T) {
+	t.Run("omitted remains distinguishable from explicit zero", func(t *testing.T) {
+		con, err := NewConnector(WithUseKernel(true))
+		require.NoError(t, err)
+		got := con.(*connector).cfg.ClientQueryTimeout
+		assert.Nil(t, got)
+
+		con, err = NewConnector(WithUseKernel(true), WithClientQueryTimeout(0))
+		require.NoError(t, err)
+		got = con.(*connector).cfg.ClientQueryTimeout
+		require.NotNil(t, got)
+		assert.Equal(t, time.Duration(0), *got)
+	})
+
+	t.Run("finite value is stored independently of server timeout", func(t *testing.T) {
+		con, err := NewConnector(
+			WithUseKernel(true),
+			WithTimeout(11*time.Second),
+			WithClientQueryTimeout(2*time.Second),
+		)
+		require.NoError(t, err)
+		cfg := con.(*connector).cfg
+		require.NotNil(t, cfg.ClientQueryTimeout)
+		assert.Equal(t, 2*time.Second, *cfg.ClientQueryTimeout)
+		assert.Equal(t, 11*time.Second, cfg.QueryTimeout)
+	})
+
+	t.Run("maximum duration is the unlimited sentinel", func(t *testing.T) {
+		maximumDuration := time.Duration(1<<63 - 1)
+		con, err := NewConnector(WithUseKernel(true), WithClientQueryTimeout(maximumDuration))
+		require.NoError(t, err)
+		got := con.(*connector).cfg.ClientQueryTimeout
+		require.NotNil(t, got)
+		assert.Equal(t, maximumDuration, *got)
+	})
+
+	t.Run("negative value is rejected during connector construction", func(t *testing.T) {
+		_, err := NewConnector(WithUseKernel(true), WithClientQueryTimeout(-time.Nanosecond))
+		require.Error(t, err)
+		assert.ErrorIs(t, err, dbsqlerr.ErrInvalidKernelConfig)
+		assert.Contains(t, err.Error(), "negative")
+	})
+
+	t.Run("value that rounds beyond the C ABI range is rejected", func(t *testing.T) {
+		maxFinite := time.Duration(kernel.MaxClientQueryTimeoutMilliseconds) * time.Millisecond
+		_, err := NewConnector(WithUseKernel(true), WithClientQueryTimeout(maxFinite+time.Nanosecond))
+		require.Error(t, err)
+		assert.ErrorIs(t, err, dbsqlerr.ErrInvalidKernelConfig)
+		assert.Contains(t, err.Error(), "maximum")
+	})
+
+	t.Run("requires explicit kernel backend", func(t *testing.T) {
+		for _, timeout := range []time.Duration{0, time.Second, time.Duration(1<<63 - 1)} {
+			_, err := NewConnector(WithClientQueryTimeout(timeout))
+			require.Error(t, err)
+			assert.ErrorIs(t, err, dbsqlerr.ErrRequiresKernelBackend)
+		}
+	})
+
+	t.Run("reusing an option does not alias connector configuration", func(t *testing.T) {
+		option := WithClientQueryTimeout(2 * time.Second)
+		first, err := NewConnector(WithUseKernel(true), option)
+		require.NoError(t, err)
+		second, err := NewConnector(WithUseKernel(true), option)
+		require.NoError(t, err)
+
+		firstTimeout := first.(*connector).cfg.ClientQueryTimeout
+		secondTimeout := second.(*connector).cfg.ClientQueryTimeout
+		require.NotNil(t, firstTimeout)
+		require.NotNil(t, secondTimeout)
+		assert.NotSame(t, firstTimeout, secondTimeout)
+		*firstTimeout = 7 * time.Second
+		assert.Equal(t, 2*time.Second, *secondTimeout)
+	})
+}
+
 func TestNewConnector(t *testing.T) {
 	t.Run("Connector initialized with functional options should have all options set", func(t *testing.T) {
 		host := "databricks-host"

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
+  - 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
   - WithMaxDownloadThreads (<num_threads> int). Sets up the max number of concurrent workers for cloud fetch. Default is 10. Optional
@@ -103,6 +104,11 @@ Cancelling a query via context cancellation or timeout is supported.
 	ctx, cancel := context.WithTimeout(context.Background(), 30 * time.Second)
 	defer cancel()
 
+WithClientQueryTimeout is a separate, connector-only kernel option. Its deadline
+covers execution through terminal-state mapping, including the initial execute RPC
+and status polling, but not result materialization or fetching. WithTimeout remains
+the independent server-side Thrift option.
+
 	// Execute query. Query will be cancelled after 30 seconds if still running
 	res, err := db.ExecContext(ctx, "CREATE TABLE example(id int, message string)")
 

errors/errors.go

@@ -74,7 +74,7 @@ var ErrNotSupportedByKernel error = errors.New("not supported by the kernel back
 var ErrKernelNotCompiled error = errors.New("the SEA-via-kernel backend is not compiled into this binary")
 
 // value to be used with errors.Is() to determine that a kernel-only option (e.g.
-// WithKernelTrustedCerts / WithKernelSkipHostnameVerify) was set without
+// WithClientQueryTimeout or WithKernelTrustedCerts) was set without
 // WithUseKernel, so the default (Thrift) backend rejected it rather than
 // connecting with a weaker-than-intended TLS trust store. This is the mirror of
 // ErrNotSupportedByKernel — that sentinel means "the kernel can't honor this
@@ -83,7 +83,7 @@ var ErrKernelNotCompiled error = errors.New("the SEA-via-kernel backend is not c
 var ErrRequiresKernelBackend error = errors.New("requires the SEA-via-kernel backend")
 
 // value to be used with errors.Is() to determine that a kernel-backend option was
-// itself malformed (e.g. a WithKernelProxy URL that does not parse) — as opposed to
+// itself malformed (e.g. a bad proxy URL or client query timeout) — as opposed to
 // unsupported (ErrNotSupportedByKernel) or a transient connect failure. The kernel
 // path validates such options in the Go layer before handing them to the kernel's C
 // ABI, where the failure would otherwise surface as an opaque wrapped string a

internal/backend/kernel/backend.go

@@ -78,8 +78,8 @@ var kernelSessionSeq atomic.Uint64
 // KernelBackend implements backend.Backend over the kernel C ABI. One backend
 // backs one conn, which database/sql serializes to a single goroutine at a time,
 // so the kernel session inherits single-owner-ship and needs no locks; the only
-// concurrency is the per-statement cancel watcher (see operation.go), which
-// touches only the kernel's internal inflight-id slot.
+// concurrency is the detached per-statement cancel watcher (see operation.go),
+// which may briefly outlive a client-timeout return while its cancel RPC ends.
 type KernelBackend struct {
 	cfg       Config
 	session   *C.kernel_session_t

internal/backend/kernel/cgo.go

@@ -152,6 +152,10 @@ func fireCancel(canceller *C.kernel_statement_canceller_t) bool {
 	return bool(dispatched)
 }
 
+func maxClientQueryTimeoutMillisecondsFromC() uint64 {
+	return uint64(C.DATABRICKS_KERNEL_MAX_CLIENT_QUERY_TIMEOUT_MS)
+}
+
 // lastError reads the kernel's thread-local last error and copies its string
 // fields out immediately — the C `char*` fields are valid only until the next
 // FFI call on this thread. Must run on the same OS thread as the failing call;

internal/backend/kernel/config.go

@@ -6,7 +6,10 @@ package kernel
 // dbsql) name kernel.Config and be unit-tested under CGO_ENABLED=0; OpenSession
 // (tagged) is what maps the assembled Config onto the kernel's C setters.
 
-import "time"
+import (
+	"fmt"
+	"time"
+)
 
 // Config is the flat connection config for the kernel backend. The connector
 // fills it from the driver's config so the user-facing options are unchanged.
@@ -26,6 +29,13 @@ type Config struct {
 	// neither unlimited nor an immediate timeout.
 	RequestTimeout time.Duration
 
+	// ClientQueryTimeout is nil when the connector option was omitted. That
+	// distinction is behavioral: nil uses kernel_statement_execute and retains
+	// its legacy 600s polling ceiling, while any non-nil value (including zero)
+	// uses kernel_statement_execute_with_timeout_ms. The pointer is copied while
+	// assembling this Config and each execution snapshots its value.
+	ClientQueryTimeout *time.Duration
+
 	// MaxConnections is the maximum number of idle HTTP connections retained
 	// per host. Zero keeps the kernel default (100).
 	MaxConnections int
@@ -108,6 +118,50 @@ type Config struct {
 	DriverSystemConfiguration *DriverSystemConfiguration
 }
 
+// MaxClientQueryTimeoutMilliseconds is the C ABI's largest finite timeout.
+// Keep it in lockstep with DATABRICKS_KERNEL_MAX_CLIENT_QUERY_TIMEOUT_MS.
+const MaxClientQueryTimeoutMilliseconds uint64 = 9_223_372_036_854
+
+const unlimitedClientQueryTimeout = time.Duration(1<<63 - 1)
+
+// ClientQueryTimeoutMilliseconds converts the public time.Duration to the C ABI
+// value. Zero and time.Duration's maximum are explicit unlimited sentinels.
+// Positive fractional milliseconds round up so they cannot accidentally become
+// unlimited. Values whose rounded form exceeds the C ABI maximum are rejected.
+func ClientQueryTimeoutMilliseconds(timeout time.Duration) (uint64, error) {
+	if timeout < 0 {
+		return 0, fmt.Errorf("client query timeout must not be negative")
+	}
+	if timeout == 0 || timeout == unlimitedClientQueryTimeout {
+		return 0, nil
+	}
+
+	milliseconds := uint64(timeout / time.Millisecond)
+	if timeout%time.Millisecond != 0 {
+		milliseconds++
+	}
+	if milliseconds > MaxClientQueryTimeoutMilliseconds {
+		return 0, fmt.Errorf("client query timeout rounds above the maximum of %d ms", MaxClientQueryTimeoutMilliseconds)
+	}
+	return milliseconds, nil
+}
+
+// configuredClientQueryTimeoutMilliseconds preserves option presence while
+// snapshotting its value for one execution. nil means use the legacy execute
+// entry point; a non-nil pointer (including one containing zero) means use the
+// timeout-aware entry point.
+func configuredClientQueryTimeoutMilliseconds(timeout *time.Duration) (*uint64, error) {
+	if timeout == nil {
+		return nil, nil
+	}
+	configured := *timeout
+	milliseconds, err := ClientQueryTimeoutMilliseconds(configured)
+	if err != nil {
+		return nil, err
+	}
+	return &milliseconds, nil
+}
+
 // RetryConfig is the driver's HTTP retry policy forwarded to the kernel: the
 // backoff-wait bounds, the maximum number of retries after the initial attempt
 // (MaxRetries == 0 disables retries), and the cumulative retry budget across all

internal/backend/kernel/config_test.go

@@ -24,3 +24,49 @@ func TestRequestTimeoutMilliseconds(t *testing.T) {
 		})
 	}
 }
+
+func TestClientQueryTimeoutMilliseconds(t *testing.T) {
+	maxFinite := time.Duration(MaxClientQueryTimeoutMilliseconds) * time.Millisecond
+	maximumDuration := time.Duration(1<<63 - 1)
+	for _, tc := range []struct {
+		name    string
+		timeout time.Duration
+		want    uint64
+		wantErr bool
+	}{
+		{"negative", -time.Nanosecond, 0, true},
+		{"zero is unlimited", 0, 0, false},
+		{"sub-millisecond rounds up", time.Nanosecond, 1, false},
+		{"fractional milliseconds round up", 1500 * time.Microsecond, 2, false},
+		{"whole milliseconds", 12 * time.Second, 12_000, false},
+		{"largest finite value", maxFinite, MaxClientQueryTimeoutMilliseconds, false},
+		{"rounded value above largest finite", maxFinite + time.Nanosecond, 0, true},
+		{"maximum duration is unlimited", maximumDuration, 0, false},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			got, err := ClientQueryTimeoutMilliseconds(tc.timeout)
+			if (err != nil) != tc.wantErr {
+				t.Fatalf("ClientQueryTimeoutMilliseconds(%v) error = %v, wantErr %v", tc.timeout, err, tc.wantErr)
+			}
+			if got != tc.want {
+				t.Errorf("ClientQueryTimeoutMilliseconds(%v) = %d, want %d", tc.timeout, got, tc.want)
+			}
+		})
+	}
+}
+
+func TestConfiguredClientQueryTimeoutMillisecondsPreservesPresenceAndSnapshots(t *testing.T) {
+	if got, err := configuredClientQueryTimeoutMilliseconds(nil); err != nil || got != nil {
+		t.Fatalf("omitted timeout = (%v, %v), want (nil, nil)", got, err)
+	}
+
+	configured := time.Duration(0)
+	got, err := configuredClientQueryTimeoutMilliseconds(&configured)
+	if err != nil || got == nil || *got != 0 {
+		t.Fatalf("explicit zero = (%v, %v), want pointer to 0", got, err)
+	}
+	configured = 5 * time.Second
+	if *got != 0 {
+		t.Errorf("resolved timeout changed after source mutation: got %d, want 0", *got)
+	}
+}