From 647b50ea3d69ac7d448caa6f72a0d37b9816152c Mon Sep 17 00:00:00 2001 From: cat23123 <58714163+cat23123@users.noreply.github.com> Date: Tue, 15 Sep 2026 18:41:01 +0000 Subject: [PATCH 1/5] feat(kernel): add client query timeout Signed-off-by: cat23123 <58714163+cat23123@users.noreply.github.com> --- CHANGELOG.md | 1 + CONNECTION_PARAMETERS.md | 1 + KERNEL_REV | 2 +- README.md | 1 + connector.go | 29 +++++++ connector_test.go | 78 +++++++++++++++++ doc.go | 6 ++ errors/errors.go | 4 +- internal/backend/kernel/backend.go | 4 +- internal/backend/kernel/cgo.go | 4 + internal/backend/kernel/config.go | 56 +++++++++++- internal/backend/kernel/config_test.go | 46 ++++++++++ .../kernel/include/databricks_kernel.h | 26 +++++- internal/backend/kernel/kernel_test.go | 85 +++++++++++++++++++ internal/backend/kernel/operation.go | 79 ++++++++++++++--- internal/config/config.go | 13 ++- internal/config/config_test.go | 9 ++ kernel_backend.go | 4 +- kernel_config.go | 18 ++++ kernel_config_test.go | 31 +++++++ kernel_e2e_test.go | 51 +++++++++++ 21 files changed, 523 insertions(+), 25 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index fee7a63a..87eb7351 100644 --- a/CHANGELOG.md +++ b/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. ## v1.15.1 (2026-09-01) diff --git a/CONNECTION_PARAMETERS.md b/CONNECTION_PARAMETERS.md index 70863dbb..ce1fdb5b 100644 --- a/CONNECTION_PARAMETERS.md +++ b/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 ``. | | *(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`). | diff --git a/KERNEL_REV b/KERNEL_REV index b914d0a7..4ff9cf1b 100644 --- a/KERNEL_REV +++ b/KERNEL_REV @@ -1 +1 @@ -21504eae074c7f5b11888a4ebcc7468780b2f844 +ec639bd48e84c64d708a0b70a532e44ecf83ad7c diff --git a/README.md b/README.md index f4986001..3953d4bb 100644 --- a/README.md +++ b/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 ``. | | *(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`). | diff --git a/connector.go b/connector.go index 742471b1..1909077d 100644 --- a/connector.go +++ b/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 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, diff --git a/connector_test.go b/connector_test.go index e7cef9c9..d5b37800 100644 --- a/connector_test.go +++ b/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" diff --git a/doc.go b/doc.go index 179bc12b..53224d90 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 + - 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 - WithMaxDownloadThreads ( 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)") diff --git a/errors/errors.go b/errors/errors.go index b3d07fa7..81174ad7 100644 --- a/errors/errors.go +++ b/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 diff --git a/internal/backend/kernel/backend.go b/internal/backend/kernel/backend.go index b8580266..a916cb16 100644 --- a/internal/backend/kernel/backend.go +++ b/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 diff --git a/internal/backend/kernel/cgo.go b/internal/backend/kernel/cgo.go index afa23f36..a8a79f6d 100644 --- a/internal/backend/kernel/cgo.go +++ b/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; diff --git a/internal/backend/kernel/config.go b/internal/backend/kernel/config.go index f043082b..062030a5 100644 --- a/internal/backend/kernel/config.go +++ b/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 diff --git a/internal/backend/kernel/config_test.go b/internal/backend/kernel/config_test.go index e1433182..9751322f 100644 --- a/internal/backend/kernel/config_test.go +++ b/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) + } +} diff --git a/internal/backend/kernel/include/databricks_kernel.h b/internal/backend/kernel/include/databricks_kernel.h index 525f2f8b..8100a249 100644 --- a/internal/backend/kernel/include/databricks_kernel.h +++ b/internal/backend/kernel/include/databricks_kernel.h @@ -85,6 +85,9 @@ #include #include +/* Largest finite value accepted by kernel_statement_execute_with_timeout_ms. */ +#define DATABRICKS_KERNEL_MAX_CLIENT_QUERY_TIMEOUT_MS UINT64_C(9223372036854) + #ifdef __cplusplus extern "C" { #endif @@ -705,12 +708,31 @@ KernelStatusCode kernel_statement_set_query_tags(kernel_statement_t* stmt, const char* query_tags); /* - * Wait-for-result execution. On success, `*out` holds an executed handle - * released with `kernel_executed_statement_close`. + * Wait-for-result execution with the existing 600-second polling ceiling and + * no caller-supplied deadline. On success, release `*out` with + * `kernel_executed_statement_close`. */ KernelStatusCode kernel_statement_execute(kernel_statement_t* stmt, kernel_executed_statement_t** out); +/* + * Wait-for-result execution with a client-side deadline in milliseconds. + * Zero selects unlimited execution; values above + * DATABRICKS_KERNEL_MAX_CLIENT_QUERY_TIMEOUT_MS return InvalidArgument. A finite + * deadline starts on entry and covers execution + * through terminal-state mapping, and excludes result materialisation and + * fetching. It is never sent to the server. + * + * Timeout returns KernelStatusCode_Timeout with SQLSTATE HYT00 and leaves + * `*out` NULL. Cleanup is best effort and asynchronous; `stmt` remains reusable. + * + * This function has the same handle invalidation and panic contract as + * kernel_statement_execute. + */ +KernelStatusCode kernel_statement_execute_with_timeout_ms( + kernel_statement_t* stmt, uint64_t timeout_ms, + kernel_executed_statement_t** out); + /* * Submit-and-return (async). DEFERRED in v0: this always returns * `KernelStatusCode_InvalidArgument` (with an explanatory last error) and diff --git a/internal/backend/kernel/kernel_test.go b/internal/backend/kernel/kernel_test.go index d950a3ff..9039db2e 100644 --- a/internal/backend/kernel/kernel_test.go +++ b/internal/backend/kernel/kernel_test.go @@ -9,6 +9,8 @@ import ( "encoding/json" "errors" "os" + "strings" + "sync" "testing" "time" @@ -19,6 +21,89 @@ import ( "github.com/databricks/databricks-sql-go/logger" ) +func TestStatementExecuteTimeoutRouting(t *testing.T) { + err := tryStatementExecuteRoute(nil) + if err == nil || !strings.Contains(err.Error(), "kernel_statement_execute: stmt must not be null") { + t.Fatalf("omitted timeout routed error = %v, want legacy kernel_statement_execute", err) + } + + zero := uint64(0) + err = tryStatementExecuteRoute(&zero) + if err == nil || !strings.Contains(err.Error(), "kernel_statement_execute_with_timeout_ms: stmt must not be null") { + t.Fatalf("explicit zero routed error = %v, want timeout-aware execute", err) + } +} + +func TestClientQueryTimeoutMaximumMatchesCABI(t *testing.T) { + if got := maxClientQueryTimeoutMillisecondsFromC(); got != MaxClientQueryTimeoutMilliseconds { + t.Fatalf("Go client timeout maximum = %d, C ABI maximum = %d", MaxClientQueryTimeoutMilliseconds, got) + } +} + +func TestFinishCancelWatcherPreservesLegacyJoinAndDetachesForClientTimeout(t *testing.T) { + t.Run("client timeout does not wait", func(t *testing.T) { + done := make(chan struct{}) + blockWatcher := make(chan struct{}) + var wg sync.WaitGroup + wg.Add(1) + go func() { + defer wg.Done() + <-done + <-blockWatcher + }() + + freeCalled := make(chan struct{}) + returned := make(chan struct{}) + go func() { + finishCancelWatcher(done, &wg, true, func() { close(freeCalled) }) + close(returned) + }() + + select { + case <-returned: + case <-time.After(time.Second): + t.Fatal("detached client-timeout cleanup waited for the watcher") + } + select { + case <-freeCalled: + t.Fatal("detached cleanup freed a canceller still owned by the watcher") + default: + } + close(blockWatcher) + wg.Wait() + }) + + t.Run("legacy execution still joins", func(t *testing.T) { + done := make(chan struct{}) + var wg sync.WaitGroup + wg.Add(1) + freeCalled := make(chan struct{}) + returned := make(chan struct{}) + go func() { + finishCancelWatcher(done, &wg, false, func() { close(freeCalled) }) + close(returned) + }() + + <-done + select { + case <-returned: + t.Fatal("legacy cleanup returned before the watcher drained") + default: + } + wg.Done() + select { + case <-returned: + case <-time.After(time.Second): + t.Fatal("legacy cleanup did not return after the watcher drained") + } + select { + case <-freeCalled: + default: + t.Fatal("legacy cleanup did not free the canceller") + } + }) +} + // setAuth maps each Auth mode to exactly one kernel_session_config_set_auth_* // value-setter. These are pure config setters (no network), so we can assert the // call succeeds against a freshly allocated config for every mode — exercising the diff --git a/internal/backend/kernel/operation.go b/internal/backend/kernel/operation.go index e59bbe90..e2f355f0 100644 --- a/internal/backend/kernel/operation.go +++ b/internal/backend/kernel/operation.go @@ -29,10 +29,11 @@ import ( // 1. new_statement + set_sql // 2. canceller_new BEFORE execute, so it can observe the server statement id // 3. a watcher goroutine that fires the canceller on ctx.Done() -// 4. the single blocking kernel_statement_execute (inline/CloudFetch and -// long-query polling all happen inside the kernel, invisibly) -// 5. drain the watcher before returning, so a late cancel cannot land on a -// statement that reuses this handle +// 4. the selected synchronous execute entry point (inline/CloudFetch and +// long-query polling all happen inside the kernel) +// 5. on the legacy path, drain the watcher before returning; on the explicit +// client-timeout path, let an in-flight cancel finish in the background so +// it cannot extend the client deadline // // Binds any query parameters (bindParams) before executing; staging statements // are rejected up front by Execute, so none reach here. @@ -42,6 +43,14 @@ func (k *KernelBackend) execute(ctx context.Context, req backend.ExecRequest) (b // debuglog convention (conn.ExecContext logs sql.len=%d). klogCtx(ctx, "Execute sql.len=%d", len(req.Query)) + // Snapshot the immutable connector option once for this execution. Presence + // selects the C entry point, so an explicit zero must remain distinguishable + // from an omitted option. + clientQueryTimeoutMs, err := configuredClientQueryTimeoutMilliseconds(k.cfg.ClientQueryTimeout) + if err != nil { + return &kernelOp{}, fmt.Errorf("kernel: invalid client query timeout: %w", err) + } + // Reject statement text with an interior NUL before touching the kernel: set_sql // takes a NUL-terminated C string with no length, so a NUL would silently // truncate the query (the same reason bound params are guarded by @@ -126,10 +135,17 @@ func (k *KernelBackend) execute(ctx context.Context, req backend.ExecRequest) (b // returns (done) if the RPC never dispatches. done := make(chan struct{}) var watcherWg sync.WaitGroup + watcherOwnsCanceller := clientQueryTimeoutMs != nil && canceller != nil && ctx.Done() != nil if canceller != nil && ctx.Done() != nil { watcherWg.Add(1) go func() { defer watcherWg.Done() + if watcherOwnsCanceller { + // A cancel RPC can retry for minutes. The timeout-aware execute path + // must be free to return while it finishes, so transfer ownership to + // this goroutine and free only after the RPC is quiescent. + defer C.kernel_statement_canceller_free(canceller) + } select { case <-ctx.Done(): case <-done: @@ -170,17 +186,16 @@ func (k *KernelBackend) execute(ctx context.Context, req backend.ExecRequest) (b // The one blocking call. inline vs CloudFetch and long-query polling are all // resolved inside the kernel; Go just waits here. var exec *C.kernel_executed_statement_t - execErr := call(func() C.KernelStatusCode { - return C.kernel_statement_execute(stmt, &exec) - }) + execErr := callStatementExecute(stmt, clientQueryTimeoutMs, &exec) - // Drain the watcher before returning so a late canceller fire cannot land on - // a subsequent statement reusing this handle. - close(done) - watcherWg.Wait() - if canceller != nil { - C.kernel_statement_canceller_free(canceller) - } + // Preserve the legacy join contract when no client timeout was configured. + // On the timeout-aware path the watcher owns the canceller and may outlive + // this call, ensuring a blocked cancel RPC cannot delay the deadline result. + finishCancelWatcher(done, &watcherWg, watcherOwnsCanceller, func() { + if canceller != nil { + C.kernel_statement_canceller_free(canceller) + } + }) op := &kernelOp{backend: k, stmt: stmt, location: k.cfg.Location, decimalAsFloat: k.cfg.DecimalAsFloat} if execErr != nil { @@ -238,6 +253,42 @@ func (k *KernelBackend) execute(ctx context.Context, req backend.ExecRequest) (b return op, nil } +func finishCancelWatcher(done chan struct{}, wg *sync.WaitGroup, detached bool, free func()) { + close(done) + if detached { + return + } + wg.Wait() + free() +} + +// callStatementExecute keeps the compatibility split explicit: nil invokes the +// legacy API (and its existing 600s polling ceiling), while every configured +// value invokes the timeout-aware API. In particular, timeoutMs==0 means +// unlimited execution; it must not fall back to the legacy function. +func callStatementExecute( + stmt *C.kernel_statement_t, + timeoutMs *uint64, + out **C.kernel_executed_statement_t, +) error { + if timeoutMs == nil { + return call(func() C.KernelStatusCode { + return C.kernel_statement_execute(stmt, out) + }) + } + return call(func() C.KernelStatusCode { + return C.kernel_statement_execute_with_timeout_ms(stmt, C.uint64_t(*timeoutMs), out) + }) +} + +// tryStatementExecuteRoute lets tagged tests exercise the real C symbol routing +// with a null statement, avoiding a live session. The selected entry point is +// observable in the kernel's InvalidArgument message. +func tryStatementExecuteRoute(timeoutMs *uint64) error { + var out *C.kernel_executed_statement_t + return callStatementExecute(nil, timeoutMs, &out) +} + // bindParams binds the driver's backend.Param list onto the statement via the // kernel's raw-param bind. Each Param is already stringified with its Databricks // SQL type name; an empty Name is a positional param (ordinal assigned kernel-side diff --git a/internal/config/config.go b/internal/config/config.go index ee7b5b4d..75635927 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -49,6 +49,12 @@ type Config struct { ThriftProtocolVersion cli_service.TProtocolVersion ThriftDebugClientProtocol bool + // ClientQueryTimeout is the connector-only, kernel execution deadline set by + // WithClientQueryTimeout. nil means the option was omitted, which preserves + // the legacy kernel_statement_execute contract. A non-nil zero explicitly + // selects unlimited execution through the timeout-aware C API. + ClientQueryTimeout *time.Duration + // KernelExperimental carries experimental, kernel-backend-only options that // have no equivalent on the default (Thrift) path — currently the richer TLS // surface (a trusted-CA bundle and an independent hostname-skip) the kernel @@ -223,7 +229,7 @@ func (c *Config) DeepCopy() *Config { return nil } - return &Config{ + cp := &Config{ UserConfig: c.UserConfig.DeepCopy(), TLSConfig: c.TLSConfig.Clone(), ArrowConfig: c.ArrowConfig.DeepCopy(), @@ -239,6 +245,11 @@ func (c *Config) DeepCopy() *Config { ThriftDebugClientProtocol: c.ThriftDebugClientProtocol, KernelExperimental: c.KernelExperimental.DeepCopy(), } + if c.ClientQueryTimeout != nil { + timeout := *c.ClientQueryTimeout + cp.ClientQueryTimeout = &timeout + } + return cp } // UserConfig is the set of configurations exposed to users diff --git a/internal/config/config_test.go b/internal/config/config_test.go index 63ad788a..b1e235b0 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -805,6 +805,7 @@ func TestConfig_DeepCopy(t *testing.T) { }) t.Run("copy config with all values", func(t *testing.T) { maxConnections := 37 + clientQueryTimeout := 2 * time.Second cfg := &Config{ UserConfig: UserConfig{}.WithDefaults(), TLSConfig: &tls.Config{MinVersion: tls.VersionTLS12}, @@ -818,6 +819,7 @@ func TestConfig_DeepCopy(t *testing.T) { ThriftTransport: "http", ThriftProtocolVersion: cli_service.TProtocolVersion_SPARK_CLI_SERVICE_PROTOCOL_V8, ThriftDebugClientProtocol: false, + ClientQueryTimeout: &clientQueryTimeout, KernelExperimental: &KernelExperimentalConfig{ TLSTrustedCertsPEM: []byte("ca-bundle"), TLSClientCertPEM: []byte("client-cert"), @@ -836,6 +838,13 @@ func TestConfig_DeepCopy(t *testing.T) { if cfg_copy.KernelExperimental == cfg.KernelExperimental { t.Error("DeepCopy aliased KernelExperimental pointer") } + if cfg_copy.ClientQueryTimeout == cfg.ClientQueryTimeout { + t.Error("DeepCopy aliased ClientQueryTimeout pointer") + } + *cfg_copy.ClientQueryTimeout = 9 * time.Second + if *cfg.ClientQueryTimeout != 2*time.Second { + t.Error("mutating the copy changed the original ClientQueryTimeout") + } if cfg_copy.KernelExperimental.MaxConnections == cfg.KernelExperimental.MaxConnections { t.Error("DeepCopy aliased KernelExperimental MaxConnections pointer") } diff --git a/kernel_backend.go b/kernel_backend.go index e7fad0df..e06b4f23 100644 --- a/kernel_backend.go +++ b/kernel_backend.go @@ -13,8 +13,8 @@ import ( // newKernelBackend builds the SEA-via-kernel backend from the driver config; the // connector opens the session right after, matching the Thrift path. It reads the // same config fields Thrift does and translates them to the kernel's flat -// connection config, so the user-facing options are unchanged — only the routing -// differs. The public API adds nothing beyond WithUseKernel. +// connection config. WithClientQueryTimeout is the one execution-specific public +// option consumed only by this backend; all other routing behavior is unchanged. func newKernelBackend(ctx context.Context, cfg *config.Config) (backend.Backend, error) { // Reject options the kernel path can't honor yet + resolve the auth form. The // validation is pure Go and lives in kernel_config.go (untagged) so its tests — diff --git a/kernel_config.go b/kernel_config.go index 6fd51ec7..2c461634 100644 --- a/kernel_config.go +++ b/kernel_config.go @@ -32,6 +32,17 @@ import ( // resolveKernelAuth's U2M case. const u2mKernelClientID = "databricks-sql-connector" +func validateClientQueryTimeout(timeout *time.Duration) error { + if timeout == nil { + return nil + } + if _, err := kernel.ClientQueryTimeoutMilliseconds(*timeout); err != nil { + return fmt.Errorf("databricks: invalid WithClientQueryTimeout value: %v: %w", + err, dbsqlerr.ErrInvalidKernelConfig) + } + return nil +} + // validateKernelConfig enforces the kernel backend's "nothing silently ignored" // contract: it rejects every option the kernel path can't yet honor with a clear // error (rather than dropping it, which would behave differently than Thrift) and @@ -50,6 +61,9 @@ func validateKernelConfig(cfg *config.Config) (kernel.Auth, error) { } func validateKernelConfigContext(ctx context.Context, cfg *config.Config) (kernel.Auth, error) { + if err := validateClientQueryTimeout(cfg.ClientQueryTimeout); err != nil { + return kernel.Auth{}, err + } // Initial namespace (WithInitialNamespace) is forwarded, not rejected: the // kernel C ABI has no catalog/schema setter, so KernelBackend.OpenSession // selects it post-connect with USE CATALOG / USE SCHEMA. No per-backend handling @@ -163,6 +177,10 @@ func buildKernelConfig(cfg *config.Config, kauth kernel.Auth) kernel.Config { Telemetry: kernelTelemetryConfig(cfg), DriverSystemConfiguration: kernelDriverSystemConfiguration(cfg), } + if cfg.ClientQueryTimeout != nil { + timeout := *cfg.ClientQueryTimeout + kc.ClientQueryTimeout = &timeout + } // TLS: the driver honors TLSConfig only for InsecureSkipVerify (see // internal/client), so map exactly that knob to the kernel. if cfg.TLSConfig != nil && cfg.TLSConfig.InsecureSkipVerify { diff --git a/kernel_config_test.go b/kernel_config_test.go index 318442b0..16390766 100644 --- a/kernel_config_test.go +++ b/kernel_config_test.go @@ -88,6 +88,15 @@ func TestValidateKernelConfig(t *testing.T) { } }) + t.Run("client query timeout accepted and kept separate from server timeout", func(t *testing.T) { + c := baseKernelConfig() + timeout := 2 * time.Second + c.ClientQueryTimeout = &timeout + if _, err := validateKernelConfig(c); err != nil { + t.Errorf("client query timeout should validate on the kernel path, got %v", err) + } + }) + t.Run("PAT resolves to a PAT auth descriptor", func(t *testing.T) { c := baseKernelConfig() // AccessToken = "dapi-x" a, err := validateKernelConfig(c) @@ -494,6 +503,28 @@ func TestKernelConfigFieldsClassified(t *testing.T) { // TestKernelExperimentalFieldsClassified only asserts the disposition map, not the // runtime copy). These run in the default CGO_ENABLED=0 build. func TestBuildKernelConfig(t *testing.T) { + t.Run("client query timeout presence and value forwarded without aliasing", func(t *testing.T) { + c := baseKernelConfig() + kc := buildKernelConfig(c, kernel.Auth{Mode: kernel.AuthPAT, Token: "dapi-x"}) + if kc.ClientQueryTimeout != nil { + t.Fatalf("omitted ClientQueryTimeout forwarded as %v, want nil", *kc.ClientQueryTimeout) + } + + timeout := time.Duration(0) + c.ClientQueryTimeout = &timeout + kc = buildKernelConfig(c, kernel.Auth{Mode: kernel.AuthPAT, Token: "dapi-x"}) + if kc.ClientQueryTimeout == nil || *kc.ClientQueryTimeout != 0 { + t.Fatalf("explicit zero ClientQueryTimeout = %v, want pointer to zero", kc.ClientQueryTimeout) + } + if kc.ClientQueryTimeout == c.ClientQueryTimeout { + t.Fatal("buildKernelConfig aliased the driver ClientQueryTimeout pointer") + } + timeout = 5 * time.Second + if *kc.ClientQueryTimeout != 0 { + t.Errorf("kernel ClientQueryTimeout changed after source mutation: got %v, want 0", *kc.ClientQueryTimeout) + } + }) + t.Run("max connections forwarded", func(t *testing.T) { c := baseKernelConfig() WithKernelMaxConnections(37)(c) diff --git a/kernel_e2e_test.go b/kernel_e2e_test.go index 0b7479be..3a6cf952 100644 --- a/kernel_e2e_test.go +++ b/kernel_e2e_test.go @@ -488,6 +488,57 @@ func TestKernelE2ECancellation(t *testing.T) { t.Logf("cancelled after %v with err=%v", elapsed, err) } +// TestKernelE2EClientQueryTimeout verifies the kernel-owned execution deadline +// returns HYT00 promptly and leaves the same connection reusable. +func TestKernelE2EClientQueryTimeout(t *testing.T) { + const queryTimeout = 5 * time.Second + db := kernelTestDBWith(t, WithClientQueryTimeout(queryTimeout)) + defer db.Close() + db.SetMaxOpenConns(1) + + conn, err := db.Conn(context.Background()) + if err != nil { + t.Fatalf("conn: %v", err) + } + defer conn.Close() + if err := conn.PingContext(context.Background()); err != nil { + t.Fatalf("prime connection: %v", err) + } + + start := time.Now() + rows, err := conn.QueryContext(context.Background(), + "SELECT count(*) FROM range(0, 100000000000) WHERE id % 7 = 0") + elapsed := time.Since(start) + if rows != nil { + rows.Close() + } + if err == nil { + t.Fatal("expected a client query timeout, got nil") + } + var executionErr dbsqlerr.DBExecutionError + if !errors.As(err, &executionErr) { + t.Fatalf("timeout is not a DBExecutionError: %v", err) + } + if executionErr.SqlState() != "HYT00" { + t.Fatalf("timeout SQLSTATE = %q, want HYT00 (err: %v)", executionErr.SqlState(), err) + } + if executionErr.IsRetryable() { + t.Error("client query timeout must not be retryable") + } + if elapsed < queryTimeout-time.Second || elapsed > queryTimeout+10*time.Second { + t.Errorf("timeout returned after %v, want approximately %v", elapsed, queryTimeout) + } + + var got int64 + if err := conn.QueryRowContext(context.Background(), "SELECT 1").Scan(&got); err != nil { + t.Fatalf("reuse connection after timeout: %v", err) + } + if got != 1 { + t.Errorf("SELECT 1 after timeout = %d, want 1", got) + } + t.Logf("client timeout returned after %v; connection reuse succeeded", elapsed) +} + // TestKernelE2EInitialNamespace proves WithInitialNamespace selects the initial // catalog/schema on the kernel session — applied post-connect via USE CATALOG / // USE SCHEMA, since the kernel C ABI has no namespace setter. current_catalog() / From a770a05bf566cd1cb1fe8432ca0aa316d45661dd Mon Sep 17 00:00:00 2001 From: Cathleen Yan <58714163+cathleeny@users.noreply.github.com> Date: Wed, 16 Sep 2026 20:45:57 +0000 Subject: [PATCH 2/5] chore(kernel): update client timeout kernel pin Signed-off-by: Cathleen Yan <58714163+cathleeny@users.noreply.github.com> --- KERNEL_REV | 2 +- internal/backend/kernel/include/databricks_kernel.h | 8 +++++--- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/KERNEL_REV b/KERNEL_REV index 4ff9cf1b..d474f19a 100644 --- a/KERNEL_REV +++ b/KERNEL_REV @@ -1 +1 @@ -ec639bd48e84c64d708a0b70a532e44ecf83ad7c +8c8a5475204f1d1239a290119786abc1adad23d9 diff --git a/internal/backend/kernel/include/databricks_kernel.h b/internal/backend/kernel/include/databricks_kernel.h index 8100a249..4cb46c60 100644 --- a/internal/backend/kernel/include/databricks_kernel.h +++ b/internal/backend/kernel/include/databricks_kernel.h @@ -719,9 +719,11 @@ KernelStatusCode kernel_statement_execute(kernel_statement_t* stmt, * Wait-for-result execution with a client-side deadline in milliseconds. * Zero selects unlimited execution; values above * DATABRICKS_KERNEL_MAX_CLIENT_QUERY_TIMEOUT_MS return InvalidArgument. A finite - * deadline starts on entry and covers execution - * through terminal-state mapping, and excludes result materialisation and - * fetching. It is never sent to the server. + * deadline starts on entry and covers execution through terminal-state mapping, + * and excludes result materialisation and fetching. A status request already in + * flight at the deadline may finish within a fixed five-second grace period; no + * new status request starts after the deadline. The deadline is never sent to + * the server. * * Timeout returns KernelStatusCode_Timeout with SQLSTATE HYT00 and leaves * `*out` NULL. Cleanup is best effort and asynchronous; `stmt` remains reusable. From 7434aabd7fee362e6645bbbcff1dd21889e9838c Mon Sep 17 00:00:00 2001 From: Cathleen Yan <58714163+cathleeny@users.noreply.github.com> Date: Wed, 16 Sep 2026 21:01:22 +0000 Subject: [PATCH 3/5] chore(kernel): refresh client timeout pin Signed-off-by: Cathleen Yan <58714163+cathleeny@users.noreply.github.com> --- KERNEL_REV | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/KERNEL_REV b/KERNEL_REV index d474f19a..61d2dd7b 100644 --- a/KERNEL_REV +++ b/KERNEL_REV @@ -1 +1 @@ -8c8a5475204f1d1239a290119786abc1adad23d9 +4ec0210cfc7c3f6bf5422d6e0134335e48b11f40 From 38b5174a134d3abfd577c387fbcb8fb9fbe60c42 Mon Sep 17 00:00:00 2001 From: Cathleen Yan <58714163+cathleeny@users.noreply.github.com> Date: Wed, 16 Sep 2026 21:21:45 +0000 Subject: [PATCH 4/5] chore(kernel): pin timeout cleanup fix Signed-off-by: Cathleen Yan <58714163+cathleeny@users.noreply.github.com> --- KERNEL_REV | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/KERNEL_REV b/KERNEL_REV index 61d2dd7b..5b074f63 100644 --- a/KERNEL_REV +++ b/KERNEL_REV @@ -1 +1 @@ -4ec0210cfc7c3f6bf5422d6e0134335e48b11f40 +1668d9860fe2d7b47b79c77c64afe2b5c2f5100f From 76885d3d4a0e285f00eb9248d68b6867a11ab3ff Mon Sep 17 00:00:00 2001 From: Cathleen Yan <58714163+cathleeny@users.noreply.github.com> Date: Wed, 16 Sep 2026 21:43:32 +0000 Subject: [PATCH 5/5] build(kernel): update client timeout pin Signed-off-by: Cathleen Yan <58714163+cathleeny@users.noreply.github.com> --- KERNEL_REV | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/KERNEL_REV b/KERNEL_REV index 5b074f63..4383f26e 100644 --- a/KERNEL_REV +++ b/KERNEL_REV @@ -1 +1 @@ -1668d9860fe2d7b47b79c77c64afe2b5c2f5100f +bfa6b6003bda6ddd44dbb667fb5dbbc9f4fee79d