From 182669630529204b6666ebe934316c4ab3a09745 Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Sun, 12 Jul 2026 09:53:55 +0000 Subject: [PATCH 01/24] =?UTF-8?q?feat(kernel):=20close=20deferred=20DAIS?= =?UTF-8?q?=20rows=20=E2=80=94=20OAuth=20M2M/U2M,=20metric-view,=20initial?= =?UTF-8?q?=20namespace?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Close the four DAIS-scope rows the SEA-via-kernel backend previously rejected at connect time. All are Go-side only (no new kernel dependency): the OAuth setters are on merged kernel #162, metric-view is an existing session conf, and the namespace uses plain SQL. - Metric view: config.EffectiveSessionParams() derives the server conf (spark.sql.thriftserver.metadata.metricview.enabled) once, backend-neutrally, so both backends send the identical conf. The Thrift OpenSession special-case is removed (behaviour-preserving); the kernel forwards it via SessionConf. Reject dropped; reclassified forwarded. - Initial namespace: applied post-connect via USE CATALOG / USE SCHEMA (the OSS ODBC workaround) since the kernel C ABI has no catalog/schema setter. quoteIdent (untagged) backtick-quotes identifiers; a USE failure fails connect and closes the session. Reject dropped; reclassified forwarded. - OAuth M2M/U2M: the kernel drives its own OAuth flow from raw credentials (mirroring pyo3/napi and the Node/Python kernel bindings), read off cfg.Authenticator — the single source of truth (last-writer-wins, matching Thrift). The m2m/u2m authenticators expose auth.M2MCredentialsProvider / auth.U2MCredentialsProvider; resolveKernelAuth type-switches them and returns a *kernelAuth descriptor. KernelBackend.setAuth branches to set_auth_pat / set_auth_m2m / set_auth_u2m; U2M uses Go's cloud-inferred client id (kernel defaults for scopes/port). No new config fields. Verified: default CGO_ENABLED=0 suite + golangci-lint v2.12.2 clean; tagged databricks_kernel unit tests (auth-mode -> setter mapping, quoteIdent); live staging e2e for initial namespace (current_catalog/current_schema) and metric-view (session opens + queries; the conf is not SET-introspectable on either backend). M2M/U2M covered by unit tests (no staging service principal; U2M is interactive). Co-authored-by: Isaac Signed-off-by: Mani Kaustubh Mathur --- auth/auth.go | 22 +++ auth/oauth/m2m/m2m.go | 7 + auth/oauth/u2m/authenticator.go | 6 + doc.go | 38 +++--- internal/backend/kernel/auth.go | 29 ++++ internal/backend/kernel/backend.go | 126 +++++++++++++++++- internal/backend/kernel/cgo.go | 12 ++ internal/backend/kernel/kernel_test.go | 30 ++++- internal/backend/kernel/namespace.go | 17 +++ internal/backend/kernel/namespace_test.go | 25 ++++ internal/backend/thrift/backend.go | 11 +- internal/config/config.go | 24 ++++ internal/config/config_test.go | 42 ++++++ kernel_backend.go | 37 +++++- kernel_backend_test.go | 3 +- kernel_config.go | 141 +++++++++++++------- kernel_config_test.go | 155 +++++++++++++++++++--- kernel_e2e_test.go | 47 +++++++ 18 files changed, 658 insertions(+), 114 deletions(-) create mode 100644 internal/backend/kernel/auth.go create mode 100644 internal/backend/kernel/namespace.go create mode 100644 internal/backend/kernel/namespace_test.go diff --git a/auth/auth.go b/auth/auth.go index efbbcb76..2f415925 100644 --- a/auth/auth.go +++ b/auth/auth.go @@ -9,6 +9,28 @@ type Authenticator interface { Authenticate(*http.Request) error } +// M2MCredentialsProvider is implemented by the OAuth M2M authenticator to expose +// its raw client-credentials. The SEA-via-kernel backend reads these to drive the +// kernel's own M2M flow (the kernel owns the token exchange), rather than using the +// authenticator's Authenticate method. This keeps cfg.Authenticator the single +// source of truth for auth on both backends — the kernel selects M2M by asserting +// this interface, so the last WithX option applied wins, exactly as on Thrift. +type M2MCredentialsProvider interface { + // M2MCredentials returns the client id and client secret. (The kernel's C-ABI + // M2M setter takes no scopes — it applies its own default scope set, matching + // the Go authenticator's default — so scopes are not exposed here.) + M2MCredentials() (clientID, clientSecret string) +} + +// U2MCredentialsProvider is implemented by the OAuth U2M authenticator to expose +// the cloud-inferred client id the kernel should use for its browser/PKCE flow (so +// the kernel path uses the same client id the Thrift path would). See +// M2MCredentialsProvider for why this rather than a parallel config carrier. +type U2MCredentialsProvider interface { + // U2MClientID returns the OAuth client id for the U2M browser flow. + U2MClientID() string +} + type AuthType int const ( diff --git a/auth/oauth/m2m/m2m.go b/auth/oauth/m2m/m2m.go index 5ed871ee..b4d2ea35 100644 --- a/auth/oauth/m2m/m2m.go +++ b/auth/oauth/m2m/m2m.go @@ -36,6 +36,13 @@ type authClient struct { mx sync.Mutex } +// M2MCredentials exposes the raw client-credentials so the SEA-via-kernel backend +// can drive the kernel's own M2M flow. Implements auth.M2MCredentialsProvider, +// which keeps cfg.Authenticator the single source of truth for auth mode. +func (c *authClient) M2MCredentials() (clientID, clientSecret string) { + return c.clientID, c.clientSecret +} + // 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 { diff --git a/auth/oauth/u2m/authenticator.go b/auth/oauth/u2m/authenticator.go index 456fba6d..d679a77f 100644 --- a/auth/oauth/u2m/authenticator.go +++ b/auth/oauth/u2m/authenticator.go @@ -78,6 +78,12 @@ type u2mAuthenticator struct { mx sync.Mutex } +// U2MClientID exposes the cloud-inferred OAuth client id so the SEA-via-kernel +// backend uses the same client id for the kernel's browser/PKCE flow that the +// Thrift path would. Implements auth.U2MCredentialsProvider, keeping +// cfg.Authenticator the single source of truth for auth mode. +func (c *u2mAuthenticator) U2MClientID() string { return c.clientID } + // Auth will start the OAuth Authorization Flow to authenticate the cli client // using the users credentials in the browser. Compatible with SSO. func (c *u2mAuthenticator) Authenticate(r *http.Request) error { diff --git a/doc.go b/doc.go index c0d33113..518004b4 100644 --- a/doc.go +++ b/doc.go @@ -200,24 +200,26 @@ path databricks_sql_kernel, which would match nothing): # kernel logs plus its HTTP stack (hyper/reqwest): DBSQL_KERNEL_DEBUG=1 RUST_LOG=debug ./your_app 2>&1 -The kernel backend currently supports PAT authentication; reading scalar, nested, -and complex-typed results (CloudFetch is handled transparently); context -cancellation during execute (a cancelled ctx fires a real server-side cancel; on -the read path cancellation is honored at result-batch boundaries, not mid-fetch); -and the TLS, proxy, and session-conf (query tags, statement timeout, time zone) -connection options. OAuth (M2M/U2M), initial catalog/schema, -WithEnableMetricViewMetadata, a non-default WithPort, and a non-https protocol -are not yet supported and return a clear error at connect time rather than being -silently ignored (the kernel backend connects over https:443 and has no port or -scheme setter); likewise WithTimeout (a server query timeout the kernel C ABI -can't set) and WithRetries used to disable retries (the kernel retries -internally). Bound query parameters and staging operations -(PUT/GET/REMOVE on a Unity Catalog volume, which need a local file transfer this -backend cannot perform) are likewise not yet supported and return a clear error -at execute time (they are per-statement, not connect-time). None of these is -silently ignored. (Metadata is issued as ordinary SQL — -SHOW/DESCRIBE/information_schema — and runs on this backend like any other -query.) +The kernel backend currently supports PAT and OAuth (M2M via WithClientCredentials, +and U2M via the authType=oauthU2M DSN param — the interactive browser/PKCE flow is +owned by the kernel) authentication; reading scalar, nested, and complex-typed +results (CloudFetch is handled transparently); context cancellation during execute +(a cancelled ctx fires a real server-side cancel; on the read path cancellation is +honored at result-batch boundaries, not mid-fetch); the initial namespace +(WithInitialNamespace catalog/schema, applied post-connect via USE CATALOG / USE +SCHEMA since the kernel C ABI has no namespace setter); metric-view metadata +(WithEnableMetricViewMetadata, sent as the same server session conf the Thrift +backend uses); and the TLS, proxy, and session-conf (query tags, statement timeout, +time zone) connection options. WithTimeout (a server query timeout the kernel C ABI +can't set) and WithRetries used to disable retries (the kernel retries internally) +return a clear error at connect time rather than being silently ignored; token- +provider, external/static, and federated authenticators are likewise not supported +and rejected loudly. Bound query parameters and staging operations (PUT/GET/REMOVE +on a Unity Catalog volume, which need a local file transfer this backend cannot +perform) are not yet supported and return a clear error at execute time (they are +per-statement, not connect-time). None of these is silently ignored. (Metadata is +issued as ordinary SQL — SHOW/DESCRIBE/information_schema — and runs on this backend +like any other query.) WithMaxRows and retry tuning (WithRetries with a positive limit) are accepted but not applied on the kernel path: the kernel manages result fetching and retries diff --git a/internal/backend/kernel/auth.go b/internal/backend/kernel/auth.go new file mode 100644 index 00000000..3a9a294d --- /dev/null +++ b/internal/backend/kernel/auth.go @@ -0,0 +1,29 @@ +package kernel + +// This file is intentionally NOT behind the `cgo && databricks_kernel` build tag: +// the Auth descriptor is a plain data type with no cgo, so it compiles in the +// default build too (and OpenSession, which is tagged, maps it to the kernel's +// set_auth_* C setters). Its consumer joinScopes lives in the tagged backend.go +// next to setAuth so the default build's unused-func lint doesn't flag it. + +// AuthMode selects which kernel auth form OpenSession applies. +type AuthMode int + +const ( + AuthPAT AuthMode = iota // personal access token + AuthM2M // OAuth client-credentials (client id + secret) + AuthU2M // OAuth user-to-machine (browser/PKCE; kernel-owned flow) +) + +// Auth is the resolved auth descriptor for a kernel connection. Only the fields +// for Mode are populated. The connector fills it from the driver config (see +// validateKernelConfig / toKernelAuth); OpenSession maps it to exactly one +// kernel_session_config_set_auth_* call. +type Auth struct { + Mode AuthMode + Token string // PAT + ClientID string // M2M + U2M (U2M: the cloud-inferred Go client id) + ClientSecret string // M2M + Scopes []string // U2M (the kernel M2M setter takes no scopes) + RedirectPort uint16 // U2M browser-redirect port; 0 = kernel default +} diff --git a/internal/backend/kernel/backend.go b/internal/backend/kernel/backend.go index 3dad1ad6..4fd5b395 100644 --- a/internal/backend/kernel/backend.go +++ b/internal/backend/kernel/backend.go @@ -12,6 +12,7 @@ import ( "context" "errors" "fmt" + "strings" "sync/atomic" "time" @@ -34,7 +35,7 @@ type Config struct { Host string // workspace hostname, no scheme HTTPPath string // e.g. /sql/1.0/warehouses/abc123 (carries ?o= org routing) WarehouseID string // bare warehouse id; preferred over HTTPPath when set - Token string // PAT (dapi...) + Auth Auth // PAT / OAuth M2M / OAuth U2M // SessionConf carries server-bound session confs verbatim — the same map the // Thrift backend forwards (STATEMENT_TIMEOUT, QUERY_TAGS, TIMEZONE, …). @@ -55,6 +56,13 @@ type Config struct { // Location is the session time zone used to render DATE / TIMESTAMP values, // matching the Thrift path which returns them in this location. nil means UTC. Location *time.Location + + // Catalog / Schema select the initial namespace. The kernel C ABI has no + // catalog/schema config setter, so OpenSession applies them post-connect by + // running USE CATALOG / USE SCHEMA (the OSS ODBC driver's workaround). Empty + // leaves the session in the server default namespace. + Catalog string + Schema string } // KernelBackend implements backend.Backend over the kernel C ABI. One backend @@ -137,12 +145,8 @@ func (k *KernelBackend) OpenSession(ctx context.Context) error { } } - tok := newCStr(k.cfg.Token) - defer tok.free() - if err := call(func() C.KernelStatusCode { - return C.kernel_session_config_set_auth_pat(cfg, tok.c) - }); err != nil { - return fmt.Errorf("kernel: set_auth_pat: %w", toConnError(err)) + if err := k.setAuth(cfg); err != nil { + return err } // TLS: crypto/tls's InsecureSkipVerify accepts any server cert, so relax both @@ -211,10 +215,118 @@ func (k *KernelBackend) OpenSession(ctx context.Context) error { // for logging / telemetry correlation. NOT derived from the handle pointer — // a freed pointer's address can be reused and collide across connections. k.sessionID = fmt.Sprintf("kernel-%d", kernelSessionSeq.Add(1)) + + // Initial namespace: the kernel C ABI has no catalog/schema config setter, so + // select it post-connect with USE CATALOG / USE SCHEMA (the OSS ODBC driver's + // approach). A failure here means the session is not in the requested namespace + // — a correctness precondition, like Thrift's InitialNamespace — so fail the + // connect and close the session we just opened (the connector does not call + // CloseSession on an OpenSession error). + if err := k.applyInitialNamespace(ctx); err != nil { + C.kernel_session_close(sess) + k.session = nil + k.valid = false + return err + } + klog("OpenSession OK session=%s", k.sessionID) return nil } +// setAuth applies the resolved auth form to the session config via exactly one +// kernel_session_config_set_auth_* call. PAT and M2M are plain value setters; U2M +// records the client id / redirect port / scopes and the kernel owns the browser +// (PKCE) flow, started when the session opens. Empty string args are passed as NULL +// so the kernel applies its own defaults (e.g. U2M's public client / default port). +func (k *KernelBackend) setAuth(cfg *C.KernelSessionConfig) error { + switch k.cfg.Auth.Mode { + case AuthM2M: + clientID := newCStr(k.cfg.Auth.ClientID) + defer clientID.free() + secret := newCStr(k.cfg.Auth.ClientSecret) + defer secret.free() + if err := call(func() C.KernelStatusCode { + return C.kernel_session_config_set_auth_m2m(cfg, clientID.c, secret.c) + }); err != nil { + return fmt.Errorf("kernel: set_auth_m2m: %w", toConnError(err)) + } + case AuthU2M: + // client id / scopes are optional: NULL when empty lets the kernel use its + // public client / default scopes. We pass Go's cloud-inferred client id when + // set, so the kernel uses the same client id the Thrift path would. + clientID := newCStrOrNull(k.cfg.Auth.ClientID) + defer clientID.free() + scopes := newCStrOrNull(joinScopes(k.cfg.Auth.Scopes)) + defer scopes.free() + if err := call(func() C.KernelStatusCode { + return C.kernel_session_config_set_auth_u2m(cfg, clientID.c, C.uint16_t(k.cfg.Auth.RedirectPort), scopes.c) + }); err != nil { + return fmt.Errorf("kernel: set_auth_u2m: %w", toConnError(err)) + } + default: // AuthPAT + tok := newCStr(k.cfg.Auth.Token) + defer tok.free() + if err := call(func() C.KernelStatusCode { + return C.kernel_session_config_set_auth_pat(cfg, tok.c) + }); err != nil { + return fmt.Errorf("kernel: set_auth_pat: %w", toConnError(err)) + } + } + return nil +} + +// joinScopes renders U2M scopes as the comma-separated form the kernel U2M setter +// expects. Empty (no scopes) yields "" so setAuth passes NULL and the kernel +// applies its default scope set. +func joinScopes(scopes []string) string { + return strings.Join(scopes, ",") +} + +// trySetAuth allocates a throwaway session config, applies auth to it, and frees +// it — a test seam so TestSetAuthByMode can exercise the real cgo setter path +// without putting cgo in a _test.go file (which Go forbids). Not used in +// production. Returns the setter error (nil on success). +func trySetAuth(auth Auth) error { + var cfg *C.KernelSessionConfig + if err := call(func() C.KernelStatusCode { return C.kernel_session_config_new(&cfg) }); err != nil { + return fmt.Errorf("config_new: %w", err) + } + defer C.kernel_session_config_free(cfg) + k := &KernelBackend{cfg: Config{Auth: auth}} + return k.setAuth(cfg) +} + +// applyInitialNamespace runs USE CATALOG / USE SCHEMA to select the configured +// initial namespace, since the kernel C ABI exposes no catalog/schema setter. +// Identifiers are backtick-quoted (quoteIdent) so arbitrary names are safe. A +// no-op when neither is set. +func (k *KernelBackend) applyInitialNamespace(ctx context.Context) error { + if k.cfg.Catalog != "" { + if err := k.runNamespaceStmt(ctx, "USE CATALOG "+quoteIdent(k.cfg.Catalog)); err != nil { + return fmt.Errorf("kernel: set initial catalog %q: %w", k.cfg.Catalog, err) + } + } + if k.cfg.Schema != "" { + if err := k.runNamespaceStmt(ctx, "USE SCHEMA "+quoteIdent(k.cfg.Schema)); err != nil { + return fmt.Errorf("kernel: set initial schema %q: %w", k.cfg.Schema, err) + } + } + return nil +} + +// runNamespaceStmt executes a single side-effecting statement (USE …) and closes +// the operation. USE produces no rows, so the result stream is not read. execute +// always returns a non-nil Operation (the Backend contract), so it is closed on +// both the success and error paths; the execute error is authoritative. +func (k *KernelBackend) runNamespaceStmt(ctx context.Context, sql string) error { + op, err := k.execute(ctx, backend.ExecRequest{Query: sql}) + _, closeErr := op.Close(ctx) + if err != nil { + return err + } + return closeErr +} + // CloseSession tears down the server-side session. Best-effort: the kernel's // close is async (see the C header), so an error is logged, not hard-failed. // diff --git a/internal/backend/kernel/cgo.go b/internal/backend/kernel/cgo.go index b3f905c3..eaab36f2 100644 --- a/internal/backend/kernel/cgo.go +++ b/internal/backend/kernel/cgo.go @@ -205,6 +205,18 @@ type cStr struct{ c *C.char } func newCStr(s string) cStr { return cStr{c: C.CString(s)} } +// newCStrOrNull is like newCStr but yields a NULL C pointer for an empty string, +// for kernel args whose "unset" sentinel is NULL (e.g. the optional U2M client id / +// scopes, where NULL selects the kernel's own default). C.CString("") would instead +// pass a non-NULL pointer to an empty string, which the kernel treats as a real +// (empty) value rather than "use the default". +func newCStrOrNull(s string) cStr { + if s == "" { + return cStr{c: nil} + } + return cStr{c: C.CString(s)} +} + func (s cStr) free() { if s.c != nil { C.free(unsafe.Pointer(s.c)) diff --git a/internal/backend/kernel/kernel_test.go b/internal/backend/kernel/kernel_test.go index 0656b302..f475e29f 100644 --- a/internal/backend/kernel/kernel_test.go +++ b/internal/backend/kernel/kernel_test.go @@ -12,6 +12,33 @@ import ( "github.com/databricks/databricks-sql-go/internal/backend" ) +// 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 +// real cgo path (arg marshaling, NULL-for-empty on the optional U2M args) end to +// end via the trySetAuth test helper (cgo cannot be used directly in a _test.go +// file). A failure here means the mode→setter wiring or the C signature drifted. +func TestSetAuthByMode(t *testing.T) { + cases := []struct { + name string + auth Auth + }{ + {"PAT", Auth{Mode: AuthPAT, Token: "dapi-x"}}, + {"M2M", Auth{Mode: AuthM2M, ClientID: "cid", ClientSecret: "sec"}}, + {"U2M full", Auth{Mode: AuthU2M, ClientID: "u2m-cid", Scopes: []string{"sql", "offline_access"}, RedirectPort: 8030}}, + // U2M with everything defaulted: empty client id / no scopes / port 0 must + // pass NULL / 0 so the kernel applies its own defaults (exercises newCStrOrNull). + {"U2M defaults", Auth{Mode: AuthU2M}}, + } + for _, c := range cases { + t.Run(c.name, func(t *testing.T) { + if err := trySetAuth(c.auth); err != nil { + t.Errorf("setAuth(%s) = %v, want nil", c.name, err) + } + }) + } +} + // The pure error-classifier tests (TestIsBadConnection, TestIsSessionFatal, // TestToConnError, TestToStatementErrorNeverBadConn) live in the untagged // errors_classify_test.go so they run under CGO_ENABLED=0. The tests below need a @@ -58,9 +85,6 @@ func TestExecuteRejectsParams(t *testing.T) { if err == nil { t.Fatal("expected an error for bound parameters, got nil") } - if !errors.Is(err, dbsqlerr.ErrNotSupportedByKernel) { - t.Errorf("params rejection should wrap ErrNotSupportedByKernel, got %v", err) - } if op == nil { t.Fatal("Execute must return a non-nil Operation per the Backend contract") } diff --git a/internal/backend/kernel/namespace.go b/internal/backend/kernel/namespace.go new file mode 100644 index 00000000..b1f907c3 --- /dev/null +++ b/internal/backend/kernel/namespace.go @@ -0,0 +1,17 @@ +package kernel + +// This file is intentionally NOT behind the `cgo && databricks_kernel` build tag. +// quoteIdent is pure Go (no cgo) and is unit-tested in the default CGO_ENABLED=0 +// matrix; OpenSession (tagged) calls it to build the USE CATALOG / USE SCHEMA +// statements that select the initial namespace. + +import "strings" + +// quoteIdent renders name as a backtick-quoted Databricks SQL identifier, doubling +// any embedded backtick. The kernel C ABI exposes no catalog/schema config setter, +// so the initial namespace is applied post-connect by running USE CATALOG / USE +// SCHEMA (the same workaround the OSS ODBC driver uses); quoting makes those +// statements injection-safe for arbitrary identifier text. +func quoteIdent(name string) string { + return "`" + strings.ReplaceAll(name, "`", "``") + "`" +} diff --git a/internal/backend/kernel/namespace_test.go b/internal/backend/kernel/namespace_test.go new file mode 100644 index 00000000..bd005d67 --- /dev/null +++ b/internal/backend/kernel/namespace_test.go @@ -0,0 +1,25 @@ +package kernel + +import "testing" + +func TestQuoteIdent(t *testing.T) { + cases := []struct { + in string + want string + }{ + {"main", "`main`"}, + {"my_schema", "`my_schema`"}, + {"", "``"}, + {"has space", "`has space`"}, + // A backtick in the identifier is doubled so it can't terminate the quote. + {"a`b", "`a``b`"}, + // One backtick → doubled to two, wrapped in two more → four total. + {"`", "````"}, + {"weird`;DROP", "`weird``;DROP`"}, + } + for _, c := range cases { + if got := quoteIdent(c.in); got != c.want { + t.Errorf("quoteIdent(%q) = %q, want %q", c.in, got, c.want) + } + } +} diff --git a/internal/backend/thrift/backend.go b/internal/backend/thrift/backend.go index 444130eb..0f934ccf 100644 --- a/internal/backend/thrift/backend.go +++ b/internal/backend/thrift/backend.go @@ -74,13 +74,10 @@ func (b *Backend) OpenSession(ctx context.Context) error { schemaName = cli_service.TIdentifierPtr(cli_service.TIdentifier(b.cfg.Schema)) } - sessionParams := make(map[string]string) - for k, v := range b.cfg.SessionParams { - sessionParams[k] = v - } - if b.cfg.EnableMetricViewMetadata { - sessionParams["spark.sql.thriftserver.metadata.metricview.enabled"] = "true" - } + // EffectiveSessionParams folds any option-derived confs (e.g. metric-view + // metadata) into the user's SessionParams, backend-neutrally, so the kernel + // backend sends the identical confs without duplicating the derivation. + sessionParams := b.cfg.EffectiveSessionParams() protocolVersion := int64(b.cfg.ThriftProtocolVersion) diff --git a/internal/config/config.go b/internal/config/config.go index 12b59395..58c92123 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -50,6 +50,30 @@ type Config struct { ThriftDebugClientProtocol bool } +// MetricViewMetadataConfKey is the server session conf that enables metric-view +// metadata (WithEnableMetricViewMetadata). Despite the "thriftserver" in its name +// — a historical server-side name — it is an ordinary Spark SQL session conf the +// server honors regardless of client transport, so both the Thrift and kernel +// backends send the identical key/value. Defined once here so no backend hardcodes +// the literal (see EffectiveSessionParams). +const MetricViewMetadataConfKey = "spark.sql.thriftserver.metadata.metricview.enabled" + +// EffectiveSessionParams returns the session confs to send to the server: the +// user-supplied SessionParams plus any conf derived from a higher-level option +// (currently metric-view metadata). Both backends call this, so the derivation is +// backend-neutral — neither special-cases the option nor hardcodes the conf +// literal. The returned map is always a fresh copy the caller may mutate freely. +func (c *Config) EffectiveSessionParams() map[string]string { + params := make(map[string]string, len(c.SessionParams)+1) + for k, v := range c.SessionParams { + params[k] = v + } + if c.EnableMetricViewMetadata { + params[MetricViewMetadataConfKey] = "true" + } + return params +} + // ToEndpointURL generates the endpoint URL from Config that a Thrift client will connect to func (c *Config) ToEndpointURL() (string, error) { var userInfo string diff --git a/internal/config/config_test.go b/internal/config/config_test.go index 4c70e993..a4b9b54f 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -645,6 +645,48 @@ func TestParseConfig(t *testing.T) { } } +func TestEffectiveSessionParams(t *testing.T) { + t.Run("no metric view leaves params unchanged", func(t *testing.T) { + c := &Config{UserConfig: UserConfig{SessionParams: map[string]string{"QUERY_TAGS": "a:1"}}} + got := c.EffectiveSessionParams() + if len(got) != 1 || got["QUERY_TAGS"] != "a:1" { + t.Errorf("EffectiveSessionParams() = %v, want only QUERY_TAGS", got) + } + if _, ok := got[MetricViewMetadataConfKey]; ok { + t.Errorf("metric-view conf must not be set when EnableMetricViewMetadata is false") + } + }) + t.Run("metric view adds the derived conf", func(t *testing.T) { + c := &Config{UserConfig: UserConfig{ + SessionParams: map[string]string{"QUERY_TAGS": "a:1"}, + EnableMetricViewMetadata: true, + }} + got := c.EffectiveSessionParams() + if got[MetricViewMetadataConfKey] != "true" { + t.Errorf("metric-view conf = %q, want \"true\"", got[MetricViewMetadataConfKey]) + } + if got["QUERY_TAGS"] != "a:1" { + t.Errorf("user SessionParams must be preserved, got %v", got) + } + }) + t.Run("returns a fresh copy the caller may mutate", func(t *testing.T) { + orig := map[string]string{"QUERY_TAGS": "a:1"} + c := &Config{UserConfig: UserConfig{SessionParams: orig}} + got := c.EffectiveSessionParams() + got["QUERY_TAGS"] = "mutated" + if orig["QUERY_TAGS"] != "a:1" { + t.Errorf("mutating the result mutated the source SessionParams: %v", orig) + } + }) + t.Run("nil SessionParams with metric view", func(t *testing.T) { + c := &Config{UserConfig: UserConfig{EnableMetricViewMetadata: true}} + got := c.EffectiveSessionParams() + if got[MetricViewMetadataConfKey] != "true" { + t.Errorf("metric-view conf must be set even with nil SessionParams, got %v", got) + } + }) +} + func TestUserConfig_DeepCopy(t *testing.T) { t.Run("copy empty config", func(t *testing.T) { cfg := UserConfig{} diff --git a/kernel_backend.go b/kernel_backend.go index 412ffe78..811768d4 100644 --- a/kernel_backend.go +++ b/kernel_backend.go @@ -16,11 +16,11 @@ import ( // connection config, so the user-facing options are unchanged — only the routing // differs. The public API adds nothing beyond WithUseKernel. func newKernelBackend(_ context.Context, cfg *config.Config) (backend.Backend, error) { - // Reject options the kernel path can't honor yet + resolve the PAT. The + // 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 — // including the exhaustiveness guard against a dropped Config field — run in the // default CGO_ENABLED=0 build. - token, err := validateKernelConfig(cfg) + auth, err := validateKernelConfig(cfg) if err != nil { return nil, err } @@ -29,13 +29,18 @@ func newKernelBackend(_ context.Context, cfg *config.Config) (backend.Backend, e Host: cfg.Host, HTTPPath: cfg.HTTPPath, WarehouseID: cfg.WarehouseID, - Token: token, + Auth: toKernelAuth(auth), Location: cfg.Location, - // Session confs (STATEMENT_TIMEOUT, QUERY_TAGS, TIMEZONE, …) — the same - // SessionParams map the Thrift backend forwards, so they flow to the - // server identically with no per-backend translation. SPOG org routing + // Initial namespace: no kernel config setter, so the kernel backend applies + // these post-connect via USE CATALOG / USE SCHEMA. + Catalog: cfg.Catalog, + Schema: cfg.Schema, + // Session confs (STATEMENT_TIMEOUT, QUERY_TAGS, TIMEZONE, metric-view, …) — + // the same effective params the Thrift backend forwards (user SessionParams + // plus any option-derived conf like metric-view metadata), so they flow to + // the server identically with no per-backend translation. SPOG org routing // rides in HTTPPath's ?o= and is parsed kernel-side. - SessionConf: cfg.SessionParams, + SessionConf: cfg.EffectiveSessionParams(), } // TLS: the driver honors TLSConfig only for InsecureSkipVerify (see // internal/client), so map exactly that knob to the kernel. @@ -48,5 +53,23 @@ func newKernelBackend(_ context.Context, cfg *config.Config) (backend.Backend, e return kernel.New(kc), nil } +// toKernelAuth maps the untagged auth descriptor (resolved by validateKernelConfig +// in the default build) to the kernel package's cgo-side auth struct. Kept here +// (tagged) because kernel.Auth is defined in the cgo-tagged kernel package; the +// resolution/validation itself stays untagged in kernel_config.go. +func toKernelAuth(a *kernelAuth) kernel.Auth { + switch a.mode { + case kernelAuthM2M: + return kernel.Auth{Mode: kernel.AuthM2M, ClientID: a.clientID, ClientSecret: a.clientSecret} + case kernelAuthU2M: + // Only the client id is sourced from Go; scopes and redirect port use the + // kernel's defaults (the Go U2M authenticator carries neither, and the kernel + // M2M/U2M defaults match Go's — see auth.U2MCredentialsProvider). + return kernel.Auth{Mode: kernel.AuthU2M, ClientID: a.clientID} + default: + return kernel.Auth{Mode: kernel.AuthPAT, Token: a.token} + } +} + // proxyForEndpoint (pure Go, no kernel dependency) lives in kernel_proxy.go so // its test runs in the default CGO_ENABLED=0 build. diff --git a/kernel_backend_test.go b/kernel_backend_test.go index eb3244c4..a6a272da 100644 --- a/kernel_backend_test.go +++ b/kernel_backend_test.go @@ -5,6 +5,7 @@ package dbsql import ( "context" "testing" + "time" "github.com/databricks/databricks-sql-go/internal/config" ) @@ -32,7 +33,7 @@ func TestNewKernelBackend(t *testing.T) { t.Run("validation error propagates", func(t *testing.T) { c := base() - c.Catalog = "main" // rejected by validateKernelConfig + c.QueryTimeout = 30 * time.Second // rejected by validateKernelConfig if _, err := newKernelBackend(context.Background(), c); err == nil { t.Error("newKernelBackend should propagate the validation error") } diff --git a/kernel_config.go b/kernel_config.go index ebc7c059..5fb3f4b4 100644 --- a/kernel_config.go +++ b/kernel_config.go @@ -1,8 +1,10 @@ package dbsql import ( + "errors" "fmt" + "github.com/databricks/databricks-sql-go/auth" "github.com/databricks/databricks-sql-go/auth/noop" "github.com/databricks/databricks-sql-go/auth/pat" dbsqlerr "github.com/databricks/databricks-sql-go/errors" @@ -16,42 +18,58 @@ import ( // Config field being silently dropped — run under CGO_ENABLED=0. The tagged // newKernelBackend calls validateKernelConfig, then assembles the cgo kernel.Config. +// kernelAuthMode is the auth form the kernel backend will use for a connection. +type kernelAuthMode int + +const ( + kernelAuthPAT kernelAuthMode = iota // personal access token + kernelAuthM2M // OAuth client-credentials (client id + secret) + kernelAuthU2M // OAuth user-to-machine (browser/PKCE, kernel-owned) +) + +// kernelAuth is the resolved auth descriptor validateKernelConfig hands to +// newKernelBackend. Exactly the fields for `mode` are populated; the backend maps +// it to the matching kernel_session_config_set_auth_* setter. Kept as a value type +// (no secrets logged) — the backend zeroes nothing, matching how the PAT token was +// previously passed as a plain string. +type kernelAuth struct { + mode kernelAuthMode + token string // PAT + clientID string // M2M + U2M (U2M: the cloud-inferred Go client id) + clientSecret string // M2M +} + // 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 -// resolves the PAT the kernel authenticates with. On success it returns the token -// to use. Options it does NOT reject are either forwarded by newKernelBackend or +// resolves the auth descriptor the kernel authenticates with (PAT, or OAuth +// M2M/U2M). Options it does NOT reject are either forwarded by newKernelBackend or // intentionally accepted-but-inert (documented in doc.go and asserted by // TestKernelConfigFieldsClassified). // // Every rejection wraps errors.ErrNotSupportedByKernel so a caller can detect the // "kernel can't honor this option" case with errors.Is (e.g. to fall back to the // default backend) instead of matching on message text. -func validateKernelConfig(cfg *config.Config) (token string, err error) { - // Initial namespace (WithInitialNamespace): no kernel C-ABI setter yet, so the - // session would run in the default namespace and unqualified names would - // resolve differently than Thrift. - if cfg.Catalog != "" || cfg.Schema != "" { - return "", fmt.Errorf("databricks: WithInitialNamespace (catalog/schema) is %w; "+ - "omit it or use the default (Thrift) backend", dbsqlerr.ErrNotSupportedByKernel) - } - // EnableMetricViewMetadata: maps to a server session conf we want to route - // backend-neutrally rather than duplicate here. - if cfg.EnableMetricViewMetadata { - return "", fmt.Errorf("databricks: WithEnableMetricViewMetadata is %w; "+ - "omit it or use the default (Thrift) backend", dbsqlerr.ErrNotSupportedByKernel) - } +func validateKernelConfig(cfg *config.Config) (*kernelAuth, error) { + // 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 (the OSS ODBC driver's + // workaround). No per-backend handling needed here. + // EnableMetricViewMetadata is forwarded, not rejected: config.EffectiveSessionParams + // folds its server conf (spark.sql.thriftserver.metadata.metricview.enabled=true) + // into SessionConf backend-neutrally, so the kernel path sends the identical conf + // the Thrift path does. No per-backend handling needed here. // Port / Protocol: the kernel C ABI takes only a bare host and connects over // https:443; it has no port or scheme setter. The Thrift path honors a custom // port/scheme via ToEndpointURL, so a non-default value here would be silently // ignored on the kernel path (it would just hit 443) — reject it instead, per // the "nothing silently ignored" contract. Defaults (https/443) are fine. if cfg.Protocol != "" && cfg.Protocol != "https" { - return "", fmt.Errorf("databricks: a non-https protocol is %w "+ + return nil, fmt.Errorf("databricks: a non-https protocol is %w "+ "(it connects over https); use the default (Thrift) backend", dbsqlerr.ErrNotSupportedByKernel) } if cfg.Port != 0 && cfg.Port != 443 { - return "", fmt.Errorf("databricks: a non-default port (WithPort) is %w "+ + return nil, fmt.Errorf("databricks: a non-default port (WithPort) is %w "+ "(it connects on 443); omit it or use the default (Thrift) backend", dbsqlerr.ErrNotSupportedByKernel) } // Transport (WithTransport, a custom http.RoundTripper carrying a custom CA @@ -61,42 +79,22 @@ func validateKernelConfig(cfg *config.Config) (token string, err error) { // does honor HTTPS_PROXY and InsecureSkipVerify through their own mappings; only // a wholesale custom Transport is unsupported.) if cfg.Transport != nil { - return "", fmt.Errorf("databricks: a custom WithTransport (RoundTripper) is %w "+ + return nil, fmt.Errorf("databricks: a custom WithTransport (RoundTripper) is %w "+ "(the kernel uses its own HTTP stack); use the default (Thrift) backend", dbsqlerr.ErrNotSupportedByKernel) } - // Auth: the kernel backend authenticates with a PAT only. Any other - // authenticator sets cfg.Authenticator but leaves cfg.AccessToken empty, so an - // empty PAT would reach the kernel and fail with an opaque Unauthenticated - // error. Reject it here so the failure names the cause. - token = cfg.AccessToken - switch a := cfg.Authenticator.(type) { - case nil, *noop.NoopAuth: - // No explicit authenticator — token comes from cfg.AccessToken (may be - // empty; caught below). - case *pat.PATAuth: - // WithAccessToken sets both cfg.AccessToken and this authenticator, but - // WithAuthenticator(&pat.PATAuth{...}) sets only the authenticator and leaves - // cfg.AccessToken empty. Take the token from the authenticator when - // cfg.AccessToken didn't carry it, so both PAT paths work. - if token == "" { - token = a.AccessToken - } - default: - return "", fmt.Errorf("databricks: only personal access token (WithAccessToken) auth is supported by the kernel backend; "+ - "OAuth (M2M/U2M), token-provider, external/static, and federated authenticators are %w — "+ - "use PAT or the default (Thrift) backend", dbsqlerr.ErrNotSupportedByKernel) - } - if token == "" { - // Missing required config (not an unsupported-feature rejection), so this is - // intentionally NOT wrapped with ErrNotSupportedByKernel. - return "", fmt.Errorf("databricks: the kernel backend requires a personal access token; " + - "set one with WithAccessToken (or a *pat.PATAuth via WithAuthenticator)") + // Auth: resolve the kernel auth form (PAT / OAuth M2M / U2M) off cfg.Authenticator, + // the single source of truth. resolveKernelAuth rejects unsupported authenticators + // loudly so the failure names the cause instead of surfacing as an opaque + // Unauthenticated. + auth, err := resolveKernelAuth(cfg) + if err != nil { + return nil, err } // WithTimeout maps to a per-statement server timeout on Thrift // (TExecuteStatementReq.QueryTimeout); the kernel C ABI exposes no equivalent, // so reject it rather than run the query with no server-side timeout. if cfg.QueryTimeout > 0 { - return "", fmt.Errorf("databricks: WithTimeout (server query timeout) is %w; "+ + return nil, fmt.Errorf("databricks: WithTimeout (server query timeout) is %w; "+ "omit it or use the default (Thrift) backend", dbsqlerr.ErrNotSupportedByKernel) } // WithRetries(-1) explicitly disables retries, but the kernel retries @@ -104,8 +102,53 @@ func validateKernelConfig(cfg *config.Config) (token string, err error) { // would be silently violated. Reject it. Positive/default RetryMax is fine: the // kernel provides retries (just not user-tunable), documented in doc.go. if cfg.RetryMax < 0 { - return "", fmt.Errorf("databricks: disabling retries via WithRetries is %w "+ + return nil, fmt.Errorf("databricks: disabling retries via WithRetries is %w "+ "(the kernel retries internally); omit it or use the default (Thrift) backend", dbsqlerr.ErrNotSupportedByKernel) } - return token, nil + return auth, nil +} + +// resolveKernelAuth picks the kernel auth form from the config. The kernel backend +// drives the kernel's own OAuth flow from raw credentials (mirroring pyo3/napi and +// the Node/Python kernel bindings) rather than reusing the Go authenticator's +// Authenticate method. It reads those credentials off cfg.Authenticator — the +// single source of truth for auth, so the last WithX option applied wins for both +// backends (matching Thrift's last-writer-wins on cfg.Authenticator). The M2M/U2M +// authenticator types are unexported, so it asserts the small +// auth.M2MCredentialsProvider / auth.U2MCredentialsProvider interfaces they satisfy: +// - implements M2MCredentialsProvider → M2M (client id + secret) +// - implements U2MCredentialsProvider → U2M (browser/PKCE; kernel-owned flow) +// - PAT / nil / noop → PAT (from AccessToken or a *pat.PATAuth) +// - anything else → rejected loudly (token-provider / external +// / static / federated), so the failure names the cause instead of surfacing as +// an opaque Unauthenticated. +func resolveKernelAuth(cfg *config.Config) (*kernelAuth, error) { + switch a := cfg.Authenticator.(type) { + case auth.M2MCredentialsProvider: + clientID, clientSecret := a.M2MCredentials() + return &kernelAuth{mode: kernelAuthM2M, clientID: clientID, clientSecret: clientSecret}, nil + case auth.U2MCredentialsProvider: + return &kernelAuth{mode: kernelAuthU2M, clientID: a.U2MClientID()}, nil + case nil, *noop.NoopAuth, *pat.PATAuth: + // PAT (or no explicit authenticator). WithAccessToken sets both + // cfg.AccessToken and a *pat.PATAuth, but WithAuthenticator(&pat.PATAuth{...}) + // sets only the authenticator and leaves cfg.AccessToken empty — so take the + // token from the authenticator when cfg.AccessToken didn't carry it. + token := cfg.AccessToken + if token == "" { + if p, ok := a.(*pat.PATAuth); ok { + token = p.AccessToken + } + } + if token == "" { + return nil, errors.New("databricks: the kernel backend requires a personal access token; " + + "set one with WithAccessToken (or a *pat.PATAuth via WithAuthenticator)") + } + return &kernelAuth{mode: kernelAuthPAT, token: token}, nil + default: + return nil, errors.New("databricks: this authenticator is not supported by the kernel backend; " + + "PAT (WithAccessToken) and OAuth M2M/U2M (WithClientCredentials / authType) are supported, but " + + "token-provider, external/static, and federated authenticators are not — " + + "use one of those or the default (Thrift) backend") + } } diff --git a/kernel_config_test.go b/kernel_config_test.go index 1f3a0649..27621502 100644 --- a/kernel_config_test.go +++ b/kernel_config_test.go @@ -12,12 +12,29 @@ import ( "github.com/databricks/databricks-sql-go/internal/config" ) -// nonPATAuth stands in for any non-PAT authenticator (OAuth / token-provider / -// external / federated) — the kernel backend must reject it. +// nonPATAuth stands in for any non-PAT, non-OAuth authenticator (token-provider / +// external / federated) — the kernel backend must reject it. It implements neither +// auth.M2MCredentialsProvider nor auth.U2MCredentialsProvider. type nonPATAuth struct{} func (nonPATAuth) Authenticate(*http.Request) error { return nil } +// fakeM2MAuth / fakeU2MAuth implement the credential-provider interfaces the kernel +// backend asserts on. Used instead of the real m2m/u2m authenticators in unit tests +// because the real u2m.NewAuthenticator does live OIDC discovery at construction +// (needs a resolvable host); the kernel only needs the interface, so a fake is both +// sufficient and hermetic. The real authenticators' method implementations are +// trivial field returns (verified in auth/oauth/{m2m,u2m}). +type fakeM2MAuth struct{ id, secret string } + +func (fakeM2MAuth) Authenticate(*http.Request) error { return nil } +func (f fakeM2MAuth) M2MCredentials() (string, string) { return f.id, f.secret } + +type fakeU2MAuth struct{ id string } + +func (fakeU2MAuth) Authenticate(*http.Request) error { return nil } +func (f fakeU2MAuth) U2MClientID() string { return f.id } + func baseKernelConfig() *config.Config { c := config.WithDefaults() c.Host = "h.databricks.com" @@ -39,18 +56,51 @@ func TestValidateKernelConfig(t *testing.T) { } }) - // Every rejection must (a) error and (b) wrap ErrNotSupportedByKernel, since - // that sentinel is the documented programmatic fallback-detection contract — + t.Run("catalog accepted (applied post-connect via USE CATALOG)", func(t *testing.T) { + c := baseKernelConfig() + c.Catalog = "main" + if _, err := validateKernelConfig(c); err != nil { + t.Errorf("initial catalog is now forwarded (USE CATALOG post-connect), want no error, got %v", err) + } + }) + + t.Run("schema accepted (applied post-connect via USE SCHEMA)", func(t *testing.T) { + c := baseKernelConfig() + c.Schema = "sys" + if _, err := validateKernelConfig(c); err != nil { + t.Errorf("initial schema is now forwarded (USE SCHEMA post-connect), want no error, got %v", err) + } + }) + + t.Run("metric view accepted (folded into session conf)", func(t *testing.T) { + c := baseKernelConfig() + c.EnableMetricViewMetadata = true + if _, err := validateKernelConfig(c); err != nil { + t.Errorf("metric-view metadata is now forwarded backend-neutrally, want no error, 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) + if err != nil { + t.Fatalf("PAT should validate, got %v", err) + } + if a.mode != kernelAuthPAT || a.token != "dapi-x" { + t.Errorf("auth = %+v, want mode=PAT token=dapi-x", a) + } + }) + + // The still-unsupported options must (a) error and (b) wrap ErrNotSupportedByKernel, + // since that sentinel is the documented programmatic fallback-detection contract — // asserting only err != nil would let a dropped or malformed %w wrap ship green. - // Table-driven so a new rejection is covered by adding one row. + // Table-driven so a new rejection is covered by adding one row. (Catalog/schema/ + // metric-view moved to forwarded above; a non-PAT authenticator is rejected too but + // not sentinel-wrapped, so it's asserted separately below.) rejections := []struct { name string mut func(*config.Config) }{ - {"catalog", func(c *config.Config) { c.Catalog = "main" }}, - {"schema", func(c *config.Config) { c.Schema = "sys" }}, - {"metric view", func(c *config.Config) { c.EnableMetricViewMetadata = true }}, - {"non-PAT authenticator", func(c *config.Config) { c.Authenticator = nonPATAuth{} }}, {"query timeout", func(c *config.Config) { c.QueryTimeout = 30 * time.Second }}, {"disable retries", func(c *config.Config) { c.RetryMax = -1 }}, {"non-https protocol", func(c *config.Config) { c.Protocol = "http" }}, @@ -75,12 +125,61 @@ func TestValidateKernelConfig(t *testing.T) { c := baseKernelConfig() c.AccessToken = "" c.Authenticator = &pat.PATAuth{AccessToken: "dapi-y"} - tok, err := validateKernelConfig(c) + a, err := validateKernelConfig(c) if err != nil { t.Fatalf("PAT via WithAuthenticator should validate, got %v", err) } - if tok != "dapi-y" { - t.Errorf("token = %q, want dapi-y (sourced from the authenticator)", tok) + if a.mode != kernelAuthPAT || a.token != "dapi-y" { + t.Errorf("auth = %+v, want mode=PAT token=dapi-y (sourced from the authenticator)", a) + } + }) + + t.Run("OAuth M2M resolves to an M2M descriptor", func(t *testing.T) { + c := baseKernelConfig() + c.AccessToken = "" + // An M2M authenticator is the single source of truth; resolveKernelAuth reads + // the creds off it via the auth.M2MCredentialsProvider interface. + c.Authenticator = fakeM2MAuth{id: "cid", secret: "sec"} + a, err := validateKernelConfig(c) + if err != nil { + t.Fatalf("M2M should validate, got %v", err) + } + if a.mode != kernelAuthM2M || a.clientID != "cid" || a.clientSecret != "sec" { + t.Errorf("auth = %+v, want mode=M2M clientID=cid clientSecret=sec", a) + } + }) + + t.Run("OAuth U2M resolves to a U2M descriptor", func(t *testing.T) { + c := baseKernelConfig() + c.AccessToken = "" + // A U2M authenticator is the single source of truth; resolveKernelAuth reads + // its (cloud-inferred) client id via the auth.U2MCredentialsProvider interface. + c.Authenticator = fakeU2MAuth{id: "databricks-sql-connector"} + a, err := validateKernelConfig(c) + if err != nil { + t.Fatalf("U2M should validate, got %v", err) + } + if a.mode != kernelAuthU2M || a.clientID != "databricks-sql-connector" { + t.Errorf("auth = %+v, want mode=U2M clientID=databricks-sql-connector", a) + } + }) + + t.Run("last-applied auth wins: M2M then PAT resolves to PAT", func(t *testing.T) { + // Regression for the auth-mode divergence: cfg.Authenticator is the single + // source of truth, so setting an M2M authenticator and then a PAT (a later + // WithAccessToken) must resolve to PAT on the kernel path — matching Thrift's + // last-writer-wins on cfg.Authenticator. (Previously a parallel OAuth carrier + // field could keep the kernel on M2M while Thrift used PAT.) + c := baseKernelConfig() + c.Authenticator = fakeM2MAuth{id: "cid", secret: "sec"} // earlier + c.Authenticator = &pat.PATAuth{AccessToken: "dapi-z"} // later wins + c.AccessToken = "dapi-z" + a, err := validateKernelConfig(c) + if err != nil { + t.Fatalf("PAT (last applied) should validate, got %v", err) + } + if a.mode != kernelAuthPAT || a.token != "dapi-z" { + t.Errorf("auth = %+v, want mode=PAT token=dapi-z (last-applied wins)", a) } }) @@ -95,6 +194,15 @@ func TestValidateKernelConfig(t *testing.T) { } }) + t.Run("non-PAT/non-OAuth authenticator rejected", func(t *testing.T) { + c := baseKernelConfig() + c.AccessToken = "" + c.Authenticator = nonPATAuth{} + if _, err := validateKernelConfig(c); err == nil { + t.Error("expected an error for a token-provider/external/federated authenticator") + } + }) + t.Run("positive retry tuning + maxrows accepted", func(t *testing.T) { c := baseKernelConfig() c.RetryMax = 8 @@ -126,21 +234,24 @@ var kernelConfigFieldDisposition = map[string]string{ "Host": "forwarded", "HTTPPath": "forwarded", "WarehouseID": "forwarded", - "AccessToken": "forwarded", // as the resolved PAT (kc.Token) - "Authenticator": "forwarded", // PAT authenticator resolved to the token + "AccessToken": "forwarded", // as the resolved PAT (kc.Auth.Token) + "Authenticator": "forwarded", // resolved to the auth descriptor (PAT/M2M/U2M) "Location": "forwarded", "SessionParams": "forwarded", "UseKernel": "forwarded", // the routing flag itself + // Folded into SessionConf by config.EffectiveSessionParams (metric-view conf), + // sent identically on both backends. + "EnableMetricViewMetadata": "forwarded", + // Applied post-connect via USE CATALOG / USE SCHEMA (no kernel config setter). + "Catalog": "forwarded", + "Schema": "forwarded", // Rejected loudly by validateKernelConfig. - "Catalog": "rejected", - "Schema": "rejected", - "EnableMetricViewMetadata": "rejected", - "QueryTimeout": "rejected", // when > 0 (WithTimeout) - "RetryMax": "rejected", // when < 0 (disable retries) - "Protocol": "rejected", // kernel is https-only; non-default rejected - "Port": "rejected", // kernel connects on 443; non-default rejected - "Transport": "rejected", // custom RoundTripper; kernel uses its own HTTP stack, so reject rather than drop + "QueryTimeout": "rejected", // when > 0 (WithTimeout) + "RetryMax": "rejected", // when < 0 (disable retries) + "Protocol": "rejected", // kernel is https-only; non-default rejected + "Port": "rejected", // kernel connects on 443; non-default rejected + "Transport": "rejected", // custom RoundTripper; kernel uses its own HTTP stack, so reject rather than drop // Accepted but intentionally inert on the kernel path (documented in doc.go): // the kernel manages these internally, below the C ABI, with no user knob. diff --git a/kernel_e2e_test.go b/kernel_e2e_test.go index 7d847bb4..d5744bc3 100644 --- a/kernel_e2e_test.go +++ b/kernel_e2e_test.go @@ -265,3 +265,50 @@ func TestKernelE2ECancellation(t *testing.T) { } t.Logf("cancelled after %v with err=%v", elapsed, err) } + +// 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() / +// current_schema() must return the configured values. Uses system.information_schema, +// which every workspace has, so the test is workspace-agnostic. +func TestKernelE2EInitialNamespace(t *testing.T) { + const catalog, schema = "system", "information_schema" + db := kernelTestDBWith(t, WithInitialNamespace(catalog, schema)) + defer db.Close() + + var gotCatalog, gotSchema string + if err := db.QueryRowContext(context.Background(), + "SELECT current_catalog(), current_schema()").Scan(&gotCatalog, &gotSchema); err != nil { + t.Fatalf("query current namespace: %v", err) + } + if gotCatalog != catalog { + t.Errorf("current_catalog() = %q, want %q", gotCatalog, catalog) + } + if gotSchema != schema { + t.Errorf("current_schema() = %q, want %q", gotSchema, schema) + } +} + +// TestKernelE2EMetricViewMetadata proves WithEnableMetricViewMetadata is accepted +// by the kernel session — the option is routed as the same server session conf the +// Thrift path sends (spark.sql.thriftserver.metadata.metricview.enabled), folded in +// backend-neutrally by config.EffectiveSessionParams. +// +// The assertion is "session opens and queries succeed with the conf set", not a SET +// read-back: this conf is not SET-introspectable on the warehouse — `SET ` +// errors with CONFIG_NOT_AVAILABLE on BOTH backends (verified against Thrift), so a +// read-back would be testing the warehouse's SET behavior, not our routing. A +// successful connect+query with the flag on is the parity bar (the Thrift path sends +// the identical conf at OpenSession and likewise never reads it back via SET). +func TestKernelE2EMetricViewMetadata(t *testing.T) { + db := kernelTestDBWith(t, WithEnableMetricViewMetadata(true)) + defer db.Close() + + var got int64 + if err := db.QueryRowContext(context.Background(), "SELECT 1").Scan(&got); err != nil { + t.Fatalf("kernel session with metric-view metadata enabled failed to query: %v", err) + } + if got != 1 { + t.Errorf("SELECT 1 = %d, want 1", got) + } +} From 596fd659c35a0dd239058b2356ee026e138abe01 Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Mon, 13 Jul 2026 09:18:36 +0000 Subject: [PATCH 02/24] =?UTF-8?q?fix(kernel):=20address=20code-review=20fi?= =?UTF-8?q?ndings=20=E2=80=94=20internal=20auth=20interfaces,=20single-sou?= =?UTF-8?q?rce=20auth,=20coverage=20+=20docs?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Remediation of the code-review pass on the DAIS gap-closure work. Verified against source; full default + tagged suites, golangci-lint v2.12.2, and live staging e2e all green. - Move the OAuth credential-provider interfaces (M2MCredentialsProvider / U2MCredentialsProvider) out of the public auth package into internal/backend/kernel, so the secret-reading capability is not part of the driver's public API. The unexported m2m/u2m authenticators satisfy them structurally. - Collapse the duplicate auth descriptor: validateKernelConfig/resolveKernelAuth now return kernel.Auth directly (its type is in an untagged file, so the default build builds it cgo-free); dropped dbsql.kernelAuth, kernelAuthMode, and toKernelAuth (which also removed a stale build-tag comment). - Route the initial-namespace failure-path session close through call() so a failed close is logged (via lastError's Warn), mirroring CloseSession. - Add an env-guarded live M2M e2e (TestKernelE2EM2M, skips without DATABRICKS_CLIENT_ID/_SECRET) and a last-writer-wins auth regression test; the resolveKernelAuth -> kernel.Auth path is table-tested for M2M/U2M. - Document: U2M is interactive (browser on cache-miss, blocks up to the kernel's ~120s callback timeout, connect ctx deadline not honored during that window, no C-ABI override — use PAT/M2M for headless); the U2M Scopes/RedirectPort fields are dormant-but-wired (no Go option feeds them yet); the metric-view e2e is a deliberate connect-smoke (routing asserted in TestEffectiveSessionParams). Custom M2M scopes remain unforwardable over the C ABI (no scopes arg on set_auth_m2m) — a kernel gap shared with ODBC, no authz impact (all-apis always requested); tracked in the kernel-gaps notes rather than worked around. Co-authored-by: Isaac Signed-off-by: Mani Kaustubh Mathur --- auth/auth.go | 22 -------- auth/oauth/m2m/m2m.go | 6 ++- auth/oauth/u2m/authenticator.go | 5 +- doc.go | 8 +++ internal/backend/kernel/auth.go | 50 +++++++++++++++--- internal/backend/kernel/backend.go | 7 ++- internal/backend/kernel/kernel_test.go | 10 +++- kernel_backend.go | 24 ++------- kernel_config.go | 72 +++++++++++--------------- kernel_config_test.go | 11 ++-- kernel_e2e_test.go | 44 ++++++++++++++++ 11 files changed, 155 insertions(+), 104 deletions(-) diff --git a/auth/auth.go b/auth/auth.go index 2f415925..efbbcb76 100644 --- a/auth/auth.go +++ b/auth/auth.go @@ -9,28 +9,6 @@ type Authenticator interface { Authenticate(*http.Request) error } -// M2MCredentialsProvider is implemented by the OAuth M2M authenticator to expose -// its raw client-credentials. The SEA-via-kernel backend reads these to drive the -// kernel's own M2M flow (the kernel owns the token exchange), rather than using the -// authenticator's Authenticate method. This keeps cfg.Authenticator the single -// source of truth for auth on both backends — the kernel selects M2M by asserting -// this interface, so the last WithX option applied wins, exactly as on Thrift. -type M2MCredentialsProvider interface { - // M2MCredentials returns the client id and client secret. (The kernel's C-ABI - // M2M setter takes no scopes — it applies its own default scope set, matching - // the Go authenticator's default — so scopes are not exposed here.) - M2MCredentials() (clientID, clientSecret string) -} - -// U2MCredentialsProvider is implemented by the OAuth U2M authenticator to expose -// the cloud-inferred client id the kernel should use for its browser/PKCE flow (so -// the kernel path uses the same client id the Thrift path would). See -// M2MCredentialsProvider for why this rather than a parallel config carrier. -type U2MCredentialsProvider interface { - // U2MClientID returns the OAuth client id for the U2M browser flow. - U2MClientID() string -} - type AuthType int const ( diff --git a/auth/oauth/m2m/m2m.go b/auth/oauth/m2m/m2m.go index b4d2ea35..4c8a6c95 100644 --- a/auth/oauth/m2m/m2m.go +++ b/auth/oauth/m2m/m2m.go @@ -37,8 +37,10 @@ type authClient struct { } // M2MCredentials exposes the raw client-credentials so the SEA-via-kernel backend -// can drive the kernel's own M2M flow. Implements auth.M2MCredentialsProvider, -// which keeps cfg.Authenticator the single source of truth for auth mode. +// can drive the kernel's own M2M flow, keeping cfg.Authenticator the single source +// of truth for auth mode. It structurally satisfies the M2MCredentialsProvider +// interface the kernel backend asserts (defined in internal/backend/kernel, so the +// secret-reading capability is not part of the driver's public API). func (c *authClient) M2MCredentials() (clientID, clientSecret string) { return c.clientID, c.clientSecret } diff --git a/auth/oauth/u2m/authenticator.go b/auth/oauth/u2m/authenticator.go index d679a77f..c08f87b7 100644 --- a/auth/oauth/u2m/authenticator.go +++ b/auth/oauth/u2m/authenticator.go @@ -80,8 +80,9 @@ type u2mAuthenticator struct { // U2MClientID exposes the cloud-inferred OAuth client id so the SEA-via-kernel // backend uses the same client id for the kernel's browser/PKCE flow that the -// Thrift path would. Implements auth.U2MCredentialsProvider, keeping -// cfg.Authenticator the single source of truth for auth mode. +// Thrift path would, keeping cfg.Authenticator the single source of truth for auth +// mode. It structurally satisfies the U2MCredentialsProvider interface the kernel +// backend asserts (defined in internal/backend/kernel, off the public API). func (c *u2mAuthenticator) U2MClientID() string { return c.clientID } // Auth will start the OAuth Authorization Flow to authenticate the cli client diff --git a/doc.go b/doc.go index 518004b4..a8bb5e33 100644 --- a/doc.go +++ b/doc.go @@ -221,6 +221,14 @@ per-statement, not connect-time). None of these is silently ignored. (Metadata i issued as ordinary SQL — SHOW/DESCRIBE/information_schema — and runs on this backend like any other query.) +OAuth U2M is interactive. On a cache miss (no valid cached refresh token) opening a +connection launches the system browser and blocks until the user completes login or +the kernel's built-in callback timeout (~120s) expires — and because the kernel C +ABI cannot interrupt session open mid-call, a deadline on the connection context is +not honored during that window (nor can the timeout be shortened; the C ABI exposes +no override). A cached, still-valid token opens without a browser. Use PAT or OAuth +M2M for headless / deadline-bound connects. + WithMaxRows and retry tuning (WithRetries with a positive limit) are accepted but not applied on the kernel path: the kernel manages result fetching and retries internally, below the C ABI, with no user-facing knob. diff --git a/internal/backend/kernel/auth.go b/internal/backend/kernel/auth.go index 3a9a294d..a78d23f5 100644 --- a/internal/backend/kernel/auth.go +++ b/internal/backend/kernel/auth.go @@ -1,10 +1,10 @@ package kernel // This file is intentionally NOT behind the `cgo && databricks_kernel` build tag: -// the Auth descriptor is a plain data type with no cgo, so it compiles in the -// default build too (and OpenSession, which is tagged, maps it to the kernel's -// set_auth_* C setters). Its consumer joinScopes lives in the tagged backend.go -// next to setAuth so the default build's unused-func lint doesn't flag it. +// the Auth descriptor and the credential-provider interfaces are plain Go with no +// cgo, so they compile in the default build too. validateKernelConfig (untagged) +// resolves an Auth from the config; OpenSession (tagged) maps it to the kernel's +// set_auth_* C setters. // AuthMode selects which kernel auth form OpenSession applies. type AuthMode int @@ -17,13 +17,49 @@ const ( // Auth is the resolved auth descriptor for a kernel connection. Only the fields // for Mode are populated. The connector fills it from the driver config (see -// validateKernelConfig / toKernelAuth); OpenSession maps it to exactly one +// validateKernelConfig); OpenSession maps it to exactly one // kernel_session_config_set_auth_* call. +// Scopes and RedirectPort map to the optional args of set_auth_u2m and are wired +// through to it by setAuth, but no Go path populates them today: the driver exposes +// no user option for U2M scopes or redirect port on either backend (the native +// Thrift path hardcodes both), so resolveKernelAuth leaves them zero and the kernel +// applies its defaults. They are kept — rather than dropped and the setter hardcoded +// to NULL/0 — so kernel.Auth models the full set_auth_u2m surface: adding a future +// WithOAuthRedirectPort / scopes option (ODBC PR #102 already exposes a redirect +// port) becomes populating these, not re-plumbing the setter. TestSetAuthByMode's +// "U2M full" case pins that marshalling so the dormant path stays correct. type Auth struct { Mode AuthMode Token string // PAT ClientID string // M2M + U2M (U2M: the cloud-inferred Go client id) ClientSecret string // M2M - Scopes []string // U2M (the kernel M2M setter takes no scopes) - RedirectPort uint16 // U2M browser-redirect port; 0 = kernel default + Scopes []string // U2M — dormant (see note above); nil → kernel default scopes + RedirectPort uint16 // U2M — dormant (see note above); 0 → kernel default port (8020) +} + +// M2MCredentialsProvider is implemented by the OAuth M2M authenticator to expose +// its raw client-credentials. The kernel backend reads these to drive the kernel's +// own M2M flow (the kernel owns the token exchange), rather than using the +// authenticator's Authenticate method. So cfg.Authenticator stays the single source +// of truth for auth on both backends — the kernel selects M2M by asserting this +// interface, so the last WithX option applied wins, exactly as on Thrift. +// +// It lives in this internal package (not the public auth package) so the +// secret-reading capability is never exposed on the driver's public API; the +// unexported concrete m2m authenticator satisfies it structurally. +type M2MCredentialsProvider interface { + // M2MCredentials returns the client id and client secret. (The kernel's C-ABI + // M2M setter takes no scopes — it applies its own default scope set, matching + // the Go authenticator's default — so scopes are not exposed here.) + M2MCredentials() (clientID, clientSecret string) +} + +// U2MCredentialsProvider is implemented by the OAuth U2M authenticator to expose +// the cloud-inferred client id the kernel should use for its browser/PKCE flow (so +// the kernel path uses the same client id the Thrift path would). Internal for the +// same reason as M2MCredentialsProvider; satisfied structurally by the unexported +// u2m authenticator. +type U2MCredentialsProvider interface { + // U2MClientID returns the OAuth client id for the U2M browser flow. + U2MClientID() string } diff --git a/internal/backend/kernel/backend.go b/internal/backend/kernel/backend.go index 4fd5b395..e9ea36b3 100644 --- a/internal/backend/kernel/backend.go +++ b/internal/backend/kernel/backend.go @@ -223,7 +223,12 @@ func (k *KernelBackend) OpenSession(ctx context.Context) error { // connect and close the session we just opened (the connector does not call // CloseSession on an OpenSession error). if err := k.applyInitialNamespace(ctx); err != nil { - C.kernel_session_close(sess) + // Close the session we just opened, routing through call() so a failed close + // is logged (via lastError's Warn) rather than silently discarded — mirroring + // CloseSession. The namespace error is authoritative and returned as-is. + if closeErr := call(func() C.KernelStatusCode { return C.kernel_session_close(sess) }); closeErr != nil { + klog("close after initial-namespace failure also failed: %v", closeErr) + } k.session = nil k.valid = false return err diff --git a/internal/backend/kernel/kernel_test.go b/internal/backend/kernel/kernel_test.go index f475e29f..84a07a3c 100644 --- a/internal/backend/kernel/kernel_test.go +++ b/internal/backend/kernel/kernel_test.go @@ -25,9 +25,15 @@ func TestSetAuthByMode(t *testing.T) { }{ {"PAT", Auth{Mode: AuthPAT, Token: "dapi-x"}}, {"M2M", Auth{Mode: AuthM2M, ClientID: "cid", ClientSecret: "sec"}}, + // "U2M full" populates Scopes/RedirectPort, which no production path sets today + // (resolveKernelAuth sources only the client id — see kernel.Auth docs). It is + // kept deliberately to pin the marshalling of those optional set_auth_u2m args + // (joinScopes + uint16 port), so the dormant wiring stays correct for a future + // U2M scopes/port option. {"U2M full", Auth{Mode: AuthU2M, ClientID: "u2m-cid", Scopes: []string{"sql", "offline_access"}, RedirectPort: 8030}}, - // U2M with everything defaulted: empty client id / no scopes / port 0 must - // pass NULL / 0 so the kernel applies its own defaults (exercises newCStrOrNull). + // U2M with everything defaulted (the production shape): empty client id / no + // scopes / port 0 must pass NULL / 0 so the kernel applies its own defaults + // (exercises newCStrOrNull). {"U2M defaults", Auth{Mode: AuthU2M}}, } for _, c := range cases { diff --git a/kernel_backend.go b/kernel_backend.go index 811768d4..9c8092ce 100644 --- a/kernel_backend.go +++ b/kernel_backend.go @@ -19,8 +19,8 @@ func newKernelBackend(_ context.Context, cfg *config.Config) (backend.Backend, e // 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 — // including the exhaustiveness guard against a dropped Config field — run in the - // default CGO_ENABLED=0 build. - auth, err := validateKernelConfig(cfg) + // default CGO_ENABLED=0 build. It returns kernel.Auth directly. + kauth, err := validateKernelConfig(cfg) if err != nil { return nil, err } @@ -29,7 +29,7 @@ func newKernelBackend(_ context.Context, cfg *config.Config) (backend.Backend, e Host: cfg.Host, HTTPPath: cfg.HTTPPath, WarehouseID: cfg.WarehouseID, - Auth: toKernelAuth(auth), + Auth: kauth, Location: cfg.Location, // Initial namespace: no kernel config setter, so the kernel backend applies // these post-connect via USE CATALOG / USE SCHEMA. @@ -53,23 +53,5 @@ func newKernelBackend(_ context.Context, cfg *config.Config) (backend.Backend, e return kernel.New(kc), nil } -// toKernelAuth maps the untagged auth descriptor (resolved by validateKernelConfig -// in the default build) to the kernel package's cgo-side auth struct. Kept here -// (tagged) because kernel.Auth is defined in the cgo-tagged kernel package; the -// resolution/validation itself stays untagged in kernel_config.go. -func toKernelAuth(a *kernelAuth) kernel.Auth { - switch a.mode { - case kernelAuthM2M: - return kernel.Auth{Mode: kernel.AuthM2M, ClientID: a.clientID, ClientSecret: a.clientSecret} - case kernelAuthU2M: - // Only the client id is sourced from Go; scopes and redirect port use the - // kernel's defaults (the Go U2M authenticator carries neither, and the kernel - // M2M/U2M defaults match Go's — see auth.U2MCredentialsProvider). - return kernel.Auth{Mode: kernel.AuthU2M, ClientID: a.clientID} - default: - return kernel.Auth{Mode: kernel.AuthPAT, Token: a.token} - } -} - // proxyForEndpoint (pure Go, no kernel dependency) lives in kernel_proxy.go so // its test runs in the default CGO_ENABLED=0 build. diff --git a/kernel_config.go b/kernel_config.go index 5fb3f4b4..a478e36e 100644 --- a/kernel_config.go +++ b/kernel_config.go @@ -4,10 +4,10 @@ import ( "errors" "fmt" - "github.com/databricks/databricks-sql-go/auth" "github.com/databricks/databricks-sql-go/auth/noop" "github.com/databricks/databricks-sql-go/auth/pat" dbsqlerr "github.com/databricks/databricks-sql-go/errors" + "github.com/databricks/databricks-sql-go/internal/backend/kernel" "github.com/databricks/databricks-sql-go/internal/config" ) @@ -18,39 +18,20 @@ import ( // Config field being silently dropped — run under CGO_ENABLED=0. The tagged // newKernelBackend calls validateKernelConfig, then assembles the cgo kernel.Config. -// kernelAuthMode is the auth form the kernel backend will use for a connection. -type kernelAuthMode int - -const ( - kernelAuthPAT kernelAuthMode = iota // personal access token - kernelAuthM2M // OAuth client-credentials (client id + secret) - kernelAuthU2M // OAuth user-to-machine (browser/PKCE, kernel-owned) -) - -// kernelAuth is the resolved auth descriptor validateKernelConfig hands to -// newKernelBackend. Exactly the fields for `mode` are populated; the backend maps -// it to the matching kernel_session_config_set_auth_* setter. Kept as a value type -// (no secrets logged) — the backend zeroes nothing, matching how the PAT token was -// previously passed as a plain string. -type kernelAuth struct { - mode kernelAuthMode - token string // PAT - clientID string // M2M + U2M (U2M: the cloud-inferred Go client id) - clientSecret string // M2M -} - // 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 -// resolves the auth descriptor the kernel authenticates with (PAT, or OAuth +// resolves the kernel.Auth descriptor the kernel authenticates with (PAT, or OAuth // M2M/U2M). Options it does NOT reject are either forwarded by newKernelBackend or // intentionally accepted-but-inert (documented in doc.go and asserted by -// TestKernelConfigFieldsClassified). +// TestKernelConfigFieldsClassified). It returns kernel.Auth directly (no dbsql-side +// duplicate) — kernel's auth types are in an untagged file, so this untagged, +// default-build code can build them without pulling in cgo. // // Every rejection wraps errors.ErrNotSupportedByKernel so a caller can detect the // "kernel can't honor this option" case with errors.Is (e.g. to fall back to the // default backend) instead of matching on message text. -func validateKernelConfig(cfg *config.Config) (*kernelAuth, error) { +func validateKernelConfig(cfg *config.Config) (kernel.Auth, error) { // 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 (the OSS ODBC driver's @@ -65,11 +46,11 @@ func validateKernelConfig(cfg *config.Config) (*kernelAuth, error) { // ignored on the kernel path (it would just hit 443) — reject it instead, per // the "nothing silently ignored" contract. Defaults (https/443) are fine. if cfg.Protocol != "" && cfg.Protocol != "https" { - return nil, fmt.Errorf("databricks: a non-https protocol is %w "+ + return kernel.Auth{}, fmt.Errorf("databricks: a non-https protocol is %w "+ "(it connects over https); use the default (Thrift) backend", dbsqlerr.ErrNotSupportedByKernel) } if cfg.Port != 0 && cfg.Port != 443 { - return nil, fmt.Errorf("databricks: a non-default port (WithPort) is %w "+ + return kernel.Auth{}, fmt.Errorf("databricks: a non-default port (WithPort) is %w "+ "(it connects on 443); omit it or use the default (Thrift) backend", dbsqlerr.ErrNotSupportedByKernel) } // Transport (WithTransport, a custom http.RoundTripper carrying a custom CA @@ -79,22 +60,22 @@ func validateKernelConfig(cfg *config.Config) (*kernelAuth, error) { // does honor HTTPS_PROXY and InsecureSkipVerify through their own mappings; only // a wholesale custom Transport is unsupported.) if cfg.Transport != nil { - return nil, fmt.Errorf("databricks: a custom WithTransport (RoundTripper) is %w "+ + return kernel.Auth{}, fmt.Errorf("databricks: a custom WithTransport (RoundTripper) is %w "+ "(the kernel uses its own HTTP stack); use the default (Thrift) backend", dbsqlerr.ErrNotSupportedByKernel) } // Auth: resolve the kernel auth form (PAT / OAuth M2M / U2M) off cfg.Authenticator, // the single source of truth. resolveKernelAuth rejects unsupported authenticators // loudly so the failure names the cause instead of surfacing as an opaque // Unauthenticated. - auth, err := resolveKernelAuth(cfg) + kauth, err := resolveKernelAuth(cfg) if err != nil { - return nil, err + return kernel.Auth{}, err } // WithTimeout maps to a per-statement server timeout on Thrift // (TExecuteStatementReq.QueryTimeout); the kernel C ABI exposes no equivalent, // so reject it rather than run the query with no server-side timeout. if cfg.QueryTimeout > 0 { - return nil, fmt.Errorf("databricks: WithTimeout (server query timeout) is %w; "+ + return kernel.Auth{}, fmt.Errorf("databricks: WithTimeout (server query timeout) is %w; "+ "omit it or use the default (Thrift) backend", dbsqlerr.ErrNotSupportedByKernel) } // WithRetries(-1) explicitly disables retries, but the kernel retries @@ -102,10 +83,10 @@ func validateKernelConfig(cfg *config.Config) (*kernelAuth, error) { // would be silently violated. Reject it. Positive/default RetryMax is fine: the // kernel provides retries (just not user-tunable), documented in doc.go. if cfg.RetryMax < 0 { - return nil, fmt.Errorf("databricks: disabling retries via WithRetries is %w "+ + return kernel.Auth{}, fmt.Errorf("databricks: disabling retries via WithRetries is %w "+ "(the kernel retries internally); omit it or use the default (Thrift) backend", dbsqlerr.ErrNotSupportedByKernel) } - return auth, nil + return kauth, nil } // resolveKernelAuth picks the kernel auth form from the config. The kernel backend @@ -115,20 +96,27 @@ func validateKernelConfig(cfg *config.Config) (*kernelAuth, error) { // single source of truth for auth, so the last WithX option applied wins for both // backends (matching Thrift's last-writer-wins on cfg.Authenticator). The M2M/U2M // authenticator types are unexported, so it asserts the small -// auth.M2MCredentialsProvider / auth.U2MCredentialsProvider interfaces they satisfy: +// kernel.M2MCredentialsProvider / kernel.U2MCredentialsProvider interfaces they +// satisfy structurally: // - implements M2MCredentialsProvider → M2M (client id + secret) // - implements U2MCredentialsProvider → U2M (browser/PKCE; kernel-owned flow) // - PAT / nil / noop → PAT (from AccessToken or a *pat.PATAuth) // - anything else → rejected loudly (token-provider / external // / static / federated), so the failure names the cause instead of surfacing as // an opaque Unauthenticated. -func resolveKernelAuth(cfg *config.Config) (*kernelAuth, error) { +func resolveKernelAuth(cfg *config.Config) (kernel.Auth, error) { switch a := cfg.Authenticator.(type) { - case auth.M2MCredentialsProvider: + case kernel.M2MCredentialsProvider: clientID, clientSecret := a.M2MCredentials() - return &kernelAuth{mode: kernelAuthM2M, clientID: clientID, clientSecret: clientSecret}, nil - case auth.U2MCredentialsProvider: - return &kernelAuth{mode: kernelAuthU2M, clientID: a.U2MClientID()}, nil + return kernel.Auth{Mode: kernel.AuthM2M, ClientID: clientID, ClientSecret: clientSecret}, nil + case kernel.U2MCredentialsProvider: + // Go sources only the client id for U2M; kernel.Auth.Scopes / RedirectPort + // are left zero so setAuth passes the kernel's defaults. Go exposes no + // user-facing option for U2M scopes or redirect port on either backend today + // (the native Thrift path hardcodes both), so there is nothing to forward — + // but the kernel.Auth fields + setAuth wiring model the full set_auth_u2m + // surface for if/when such an option is added (see kernel.Auth docs). + return kernel.Auth{Mode: kernel.AuthU2M, ClientID: a.U2MClientID()}, nil case nil, *noop.NoopAuth, *pat.PATAuth: // PAT (or no explicit authenticator). WithAccessToken sets both // cfg.AccessToken and a *pat.PATAuth, but WithAuthenticator(&pat.PATAuth{...}) @@ -141,12 +129,12 @@ func resolveKernelAuth(cfg *config.Config) (*kernelAuth, error) { } } if token == "" { - return nil, errors.New("databricks: the kernel backend requires a personal access token; " + + return kernel.Auth{}, errors.New("databricks: the kernel backend requires a personal access token; " + "set one with WithAccessToken (or a *pat.PATAuth via WithAuthenticator)") } - return &kernelAuth{mode: kernelAuthPAT, token: token}, nil + return kernel.Auth{Mode: kernel.AuthPAT, Token: token}, nil default: - return nil, errors.New("databricks: this authenticator is not supported by the kernel backend; " + + return kernel.Auth{}, errors.New("databricks: this authenticator is not supported by the kernel backend; " + "PAT (WithAccessToken) and OAuth M2M/U2M (WithClientCredentials / authType) are supported, but " + "token-provider, external/static, and federated authenticators are not — " + "use one of those or the default (Thrift) backend") diff --git a/kernel_config_test.go b/kernel_config_test.go index 27621502..fb077019 100644 --- a/kernel_config_test.go +++ b/kernel_config_test.go @@ -9,6 +9,7 @@ import ( "github.com/databricks/databricks-sql-go/auth/pat" dbsqlerr "github.com/databricks/databricks-sql-go/errors" + "github.com/databricks/databricks-sql-go/internal/backend/kernel" "github.com/databricks/databricks-sql-go/internal/config" ) @@ -86,7 +87,7 @@ func TestValidateKernelConfig(t *testing.T) { if err != nil { t.Fatalf("PAT should validate, got %v", err) } - if a.mode != kernelAuthPAT || a.token != "dapi-x" { + if a.Mode != kernel.AuthPAT || a.Token != "dapi-x" { t.Errorf("auth = %+v, want mode=PAT token=dapi-x", a) } }) @@ -129,7 +130,7 @@ func TestValidateKernelConfig(t *testing.T) { if err != nil { t.Fatalf("PAT via WithAuthenticator should validate, got %v", err) } - if a.mode != kernelAuthPAT || a.token != "dapi-y" { + if a.Mode != kernel.AuthPAT || a.Token != "dapi-y" { t.Errorf("auth = %+v, want mode=PAT token=dapi-y (sourced from the authenticator)", a) } }) @@ -144,7 +145,7 @@ func TestValidateKernelConfig(t *testing.T) { if err != nil { t.Fatalf("M2M should validate, got %v", err) } - if a.mode != kernelAuthM2M || a.clientID != "cid" || a.clientSecret != "sec" { + if a.Mode != kernel.AuthM2M || a.ClientID != "cid" || a.ClientSecret != "sec" { t.Errorf("auth = %+v, want mode=M2M clientID=cid clientSecret=sec", a) } }) @@ -159,7 +160,7 @@ func TestValidateKernelConfig(t *testing.T) { if err != nil { t.Fatalf("U2M should validate, got %v", err) } - if a.mode != kernelAuthU2M || a.clientID != "databricks-sql-connector" { + if a.Mode != kernel.AuthU2M || a.ClientID != "databricks-sql-connector" { t.Errorf("auth = %+v, want mode=U2M clientID=databricks-sql-connector", a) } }) @@ -178,7 +179,7 @@ func TestValidateKernelConfig(t *testing.T) { if err != nil { t.Fatalf("PAT (last applied) should validate, got %v", err) } - if a.mode != kernelAuthPAT || a.token != "dapi-z" { + if a.Mode != kernel.AuthPAT || a.Token != "dapi-z" { t.Errorf("auth = %+v, want mode=PAT token=dapi-z (last-applied wins)", a) } }) diff --git a/kernel_e2e_test.go b/kernel_e2e_test.go index d5744bc3..b2d86fac 100644 --- a/kernel_e2e_test.go +++ b/kernel_e2e_test.go @@ -300,6 +300,13 @@ func TestKernelE2EInitialNamespace(t *testing.T) { // read-back would be testing the warehouse's SET behavior, not our routing. A // successful connect+query with the flag on is the parity bar (the Thrift path sends // the identical conf at OpenSession and likewise never reads it back via SET). +// +// Scope note: this is intentionally a thin live smoke test — it proves the flag is +// accepted end to end, not that the conf value is correct. The actual routing (the +// flag → the metricview.enabled=true session conf) is asserted in the default-build +// unit test TestEffectiveSessionParams; the SELECT here overlaps TestKernelE2ESelect1 +// by design, since the conf's live effect can't be observed (see above). Kept as a +// connect-with-flag smoke rather than deleted so the flag has at least one live path. func TestKernelE2EMetricViewMetadata(t *testing.T) { db := kernelTestDBWith(t, WithEnableMetricViewMetadata(true)) defer db.Close() @@ -312,3 +319,40 @@ func TestKernelE2EMetricViewMetadata(t *testing.T) { t.Errorf("SELECT 1 = %d, want 1", got) } } + +// TestKernelE2EM2M authenticates to a real warehouse via OAuth M2M over the kernel +// (WithClientCredentials → the kernel's own client-credentials flow), then runs +// SELECT 1. It self-skips unless DATABRICKS_CLIENT_ID / DATABRICKS_CLIENT_SECRET are +// set (alongside the host/http-path base), mirroring the PAT tests' skip pattern — +// so CI without a service principal stays green while a credentialed run proves the +// M2M setter path actually authenticates end to end (not just that set_auth_m2m +// returns OK, which TestSetAuthByMode already covers). +func TestKernelE2EM2M(t *testing.T) { + host := os.Getenv("DATABRICKS_HOST") + httpPath := os.Getenv("DATABRICKS_HTTP_PATH") + clientID := os.Getenv("DATABRICKS_CLIENT_ID") + clientSecret := os.Getenv("DATABRICKS_CLIENT_SECRET") + if host == "" || httpPath == "" || clientID == "" || clientSecret == "" { + t.Skip("set DATABRICKS_HOST / DATABRICKS_HTTP_PATH / DATABRICKS_CLIENT_ID / DATABRICKS_CLIENT_SECRET for the M2M e2e") + } + + connector, err := NewConnector( + WithServerHostname(host), + WithHTTPPath(httpPath), + WithClientCredentials(clientID, clientSecret), + WithUseKernel(true), + ) + if err != nil { + t.Fatalf("NewConnector: %v", err) + } + db := sql.OpenDB(connector) + defer db.Close() + + var got int64 + if err := db.QueryRowContext(context.Background(), "SELECT 1").Scan(&got); err != nil { + t.Fatalf("M2M-authenticated query failed: %v", err) + } + if got != 1 { + t.Errorf("SELECT 1 = %d, want 1", got) + } +} From 62bb700d52e9e00035fb8492ecb4e3d6ca0a6721 Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Sun, 12 Jul 2026 11:29:15 +0000 Subject: [PATCH 03/24] feat(arrowscan): render INTERVAL day-time & year-month MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The kernel backend returns INTERVAL columns as native arrow duration (day-time) and month-interval (year-month) values, whereas the Thrift path receives them pre-formatted from the server (its native-interval config is off in prod, so it never scans a duration/month-interval array). Format them Go-side in the shared untagged arrowscan package to the same strings the Thrift path returns — "D HH:MM:SS.nnnnnnnnn" and "years-months", negatives signed — so a query's result is identical across backends. Replaces the fail-loud "intervals are not yet handled" default arm with the two type arms; golden-string unit tests (day/day-to-sec/seconds-unit/negative, year/year-month/months/negative) run in the default CGO_ENABLED=0 build. Co-authored-by: Isaac Signed-off-by: Mani Kaustubh Mathur --- internal/arrowscan/arrowscan.go | 82 ++++++++++++++++++++++++++-- internal/arrowscan/arrowscan_test.go | 77 ++++++++++++++++++++++++-- 2 files changed, 149 insertions(+), 10 deletions(-) diff --git a/internal/arrowscan/arrowscan.go b/internal/arrowscan/arrowscan.go index 964db594..3f26ee7f 100644 --- a/internal/arrowscan/arrowscan.go +++ b/internal/arrowscan/arrowscan.go @@ -37,10 +37,12 @@ import ( // matching the Thrift path — a float64 would lose precision beyond ~17 digits; // see databricks-sql-go#274). Nested types (List/Map/Struct, and VARIANT which // arrives nested) render to a JSON string byte-identical to the Thrift path; -// GEOMETRY arrives as a WKB/WKT string and is handled by the string arm. NULLs -// map to nil. A genuinely unhandled type (e.g. interval/duration) returns an -// error rather than a silently wrong value. loc renders DATE / TIMESTAMP in the -// session time zone (nil = UTC, arrow's ToTime default). +// GEOMETRY arrives as a WKB/WKT string and is handled by the string arm. INTERVAL +// day-time/year-month arrive as native arrow duration/month-interval and format to +// the same string the Thrift path receives pre-formatted from the server. NULLs +// map to nil. A genuinely unhandled type returns an error rather than a silently +// wrong value. loc renders DATE / TIMESTAMP in the session time zone (nil = UTC, +// arrow's ToTime default). func ScanCell(col arrow.Array, row int, loc *time.Location) (driver.Value, error) { return ScanCellCached(col, row, loc, nil) } @@ -175,16 +177,84 @@ func ScanCellCached(col arrow.Array, row int, loc *time.Location, keys *StructKe case *array.Decimal128: dt := col.DataType().(*arrow.Decimal128Type) return decimalfmt.ExactString(c.Value(row), dt.Scale), nil + case *array.Duration: + // INTERVAL DAY TO SECOND arrives as an arrow duration. The kernel returns the + // native arrow value, so we format it Go-side to the same "D HH:MM:SS.nnnnnnnnn" + // string the Thrift path gets pre-formatted from the server (its native-interval + // config is off in prod, so it never scans a duration array — hence there is no + // shared renderer to reuse, and this stays kernel-side). + dt := col.DataType().(*arrow.DurationType) + return formatDayTimeInterval(int64(c.Value(row)), dt.Unit), nil + case *array.MonthInterval: + // INTERVAL YEAR TO MONTH arrives as a month count; Thrift's server string is + // "years-months". + return formatYearMonthInterval(int32(c.Value(row))), nil case *array.List, *array.LargeList, *array.FixedSizeList, *array.Map, *array.Struct: // Nested types (and VARIANT, which arrives as a nested value) render to a // JSON string matching the Thrift path. return renderJSONString(col, row, loc, keys) default: - return nil, fmt.Errorf("scanning arrow type %s is not supported "+ - "(intervals are not yet handled)", col.DataType()) + return nil, fmt.Errorf("scanning arrow type %s is not supported", col.DataType()) } } +// formatDayTimeInterval renders an arrow duration (in the given time unit) as the +// Thrift path's "D HH:MM:SS.nnnnnnnnn" — days, then zero-padded hours:minutes:seconds +// with 9 fractional digits, negated with a leading '-'. +func formatDayTimeInterval(v int64, unit arrow.TimeUnit) string { + neg := v < 0 + if neg { + v = -v + } + // Split into whole seconds + a sub-second nanosecond remainder working in the + // native unit. We must NOT scale the full magnitude up to nanoseconds first: + // Spark day-time intervals run up to ~Long.MaxValue microseconds (~292 years), + // so v*1e3 (or *1e6/*1e9) would overflow int64 and silently produce a wrong + // (often negative) string. Deriving seconds by dividing keeps every + // intermediate in range; only the bounded sub-second remainder is scaled up. + var secs, frac int64 + switch unit { + case arrow.Second: + secs = v + case arrow.Millisecond: + secs = v / 1e3 + frac = (v % 1e3) * 1e6 + case arrow.Microsecond: + secs = v / 1e6 + frac = (v % 1e6) * 1e3 + default: // Nanosecond + secs = v / 1e9 + frac = v % 1e9 + } + days := secs / 86400 + rem := secs % 86400 + h := rem / 3600 + rem %= 3600 + m := rem / 60 + s := rem % 60 + sign := "" + if neg { + sign = "-" + } + return fmt.Sprintf("%s%d %02d:%02d:%02d.%09d", sign, days, h, m, s, frac) +} + +// formatYearMonthInterval renders a month count as the Thrift path's "years-months", +// negated with a leading '-'. +func formatYearMonthInterval(months int32) string { + neg := months < 0 + if neg { + months = -months + } + y := months / 12 + mo := months % 12 + sign := "" + if neg { + sign = "-" + } + return fmt.Sprintf("%s%d-%d", sign, y, mo) +} + // inLocation renders t in loc, matching the Thrift path's .In(location); a nil // loc leaves the value in UTC (arrow's ToTime default). func inLocation(t time.Time, loc *time.Location) time.Time { diff --git a/internal/arrowscan/arrowscan_test.go b/internal/arrowscan/arrowscan_test.go index 8ad4824b..a54c47b9 100644 --- a/internal/arrowscan/arrowscan_test.go +++ b/internal/arrowscan/arrowscan_test.go @@ -101,19 +101,88 @@ func TestScanCellScalars(t *testing.T) { }) t.Run("unsupported_type_errors", func(t *testing.T) { - // A duration (INTERVAL) is not yet handled: must error, not return a - // wrong value. - b := array.NewDurationBuilder(pool, &arrow.DurationType{Unit: arrow.Microsecond}) + // An unhandled arrow type must error, not return a silently wrong value. + // Duration/MonthInterval are now handled (see TestScanCellInterval); use a + // type with no scan arm. + b := array.NewTime32Builder(pool, &arrow.Time32Type{Unit: arrow.Second}) defer b.Release() b.Append(1000) arr := b.NewArray() defer arr.Release() if _, err := ScanCell(arr, 0, nil); err == nil { - t.Error("scanning a Duration should return an unsupported-type error") + t.Error("scanning an unhandled arrow type should return an error") } }) } +// INTERVAL day-time (arrow duration) and year-month (arrow month-interval) arrive +// as native arrow values on the kernel path and must format to the exact string the +// Thrift path receives pre-formatted from the server: "D HH:MM:SS.nnnnnnnnn" and +// "years-months", with negatives signed. (These formatters were validated live +// kernel==Thrift in the PuPr POC; this is the regression guard.) +func TestScanCellInterval(t *testing.T) { + pool := memory.NewGoAllocator() + + dayTime := []struct { + name string + unit arrow.TimeUnit + v int64 + want string + }{ + {"one_day_us", arrow.Microsecond, 86400 * 1_000_000, "1 00:00:00.000000000"}, + {"day_to_sec_us", arrow.Microsecond, 90061_500000, "1 01:01:01.500000000"}, + {"seconds_unit", arrow.Second, 3661, "0 01:01:01.000000000"}, + {"negative_us", arrow.Microsecond, -90061_500000, "-1 01:01:01.500000000"}, + // A large microsecond magnitude (~106.75M days, near Long.MaxValue μs) must + // NOT overflow int64 while scaling to nanoseconds — regression guard for the + // prior multiply-first bug that produced a wrong/negative string here. + {"large_us_no_overflow", arrow.Microsecond, 9223372036854775807, "106751991 04:00:54.775807000"}, + } + for _, tc := range dayTime { + t.Run("daytime_"+tc.name, func(t *testing.T) { + b := array.NewDurationBuilder(pool, &arrow.DurationType{Unit: tc.unit}) + defer b.Release() + b.Append(arrow.Duration(tc.v)) + arr := b.NewArray() + defer arr.Release() + v, err := ScanCell(arr, 0, nil) + if err != nil { + t.Fatal(err) + } + if v.(string) != tc.want { + t.Errorf("got %q, want %q", v, tc.want) + } + }) + } + + yearMonth := []struct { + name string + months int32 + want string + }{ + {"two_years", 24, "2-0"}, + {"year_and_month", 13, "1-1"}, + {"months_only", 5, "0-5"}, + {"negative", -13, "-1-1"}, + } + for _, tc := range yearMonth { + t.Run("yearmonth_"+tc.name, func(t *testing.T) { + b := array.NewMonthIntervalBuilder(pool) + defer b.Release() + b.Append(arrow.MonthInterval(tc.months)) + arr := b.NewArray() + defer arr.Release() + v, err := ScanCell(arr, 0, nil) + if err != nil { + t.Fatal(err) + } + if v.(string) != tc.want { + t.Errorf("got %q, want %q", v, tc.want) + } + }) + } +} + // ScanCell renders DATE / TIMESTAMP in the requested location, matching the // Thrift path's .In(location); a nil location leaves the value in UTC. func TestScanCellTimestampLocation(t *testing.T) { From 9e7e932e0bac225c625cfe49c81211802c44dba2 Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Sun, 12 Jul 2026 12:05:49 +0000 Subject: [PATCH 04/24] test(kernel): pin TIMESTAMP vs TIMESTAMP_NTZ parity MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Over Arrow the kernel delivers TIMESTAMP with a tz ("UTC") and TIMESTAMP_NTZ with an empty tz, but — like the Thrift path — the driver ignores that field and renders both via ToTime + .In(loc). The LTZ-vs-NTZ difference is carried entirely by the instant the server sends, not by the client inspecting the tz, so no arrowscan change is needed: the existing code already matches Thrift. Verified live on both backends (America/New_York + Asia/Kolkata, including a DST spring-forward literal, and nested/null shapes): kernel == Thrift byte-for-byte for both types. Add an untagged parity case (TimeZone "UTC" vs "") so a future "don't shift NTZ" change — which looks correct in isolation but would diverge from Thrift, which shifts NTZ too — fails default CI, plus a live e2e pinning the round-trip. Co-authored-by: Isaac Signed-off-by: Mani Kaustubh Mathur --- .../rows/arrowbased/arrowscan_parity_test.go | 47 +++++++++++++++++++ kernel_e2e_test.go | 44 +++++++++++++++++ 2 files changed, 91 insertions(+) diff --git a/internal/rows/arrowbased/arrowscan_parity_test.go b/internal/rows/arrowbased/arrowscan_parity_test.go index 762d11fe..a00d2914 100644 --- a/internal/rows/arrowbased/arrowscan_parity_test.go +++ b/internal/rows/arrowbased/arrowscan_parity_test.go @@ -383,6 +383,53 @@ func TestArrowbasedKernelTopLevelScalarParity(t *testing.T) { } } +// TestArrowbasedKernelTimestampTZParity pins the TIMESTAMP vs TIMESTAMP_NTZ +// rendering contract: over Arrow, a TIMESTAMP arrives with TimeZone "UTC" and a +// TIMESTAMP_NTZ with an empty TimeZone (kernel json.rs:171 — "TIMESTAMP carries a +// tz on the Arrow side; TIMESTAMP_NTZ does not"). Both backends deliberately IGNORE +// that tz field: each renders via ToTime + .In(loc), so the LTZ-vs-NTZ difference is +// carried entirely by the instant value the server sends, NOT by the client +// inspecting the tz. Verified live on both backends (NY + Kolkata, incl. a DST-gap +// literal): kernel == Thrift byte-for-byte for both types. +// +// This test locks that in so a future "semantically-correct" change that skips +// .In(loc) for TimeZone=="" (which would look right in isolation) fails CI — it +// would make the kernel diverge from the Thrift path, which shifts NTZ too. +func TestArrowbasedKernelTimestampTZParity(t *testing.T) { + pool := memory.NewGoAllocator() + loc, err := time.LoadLocation("America/New_York") + if err != nil { + t.Skipf("tz database unavailable: %v", err) + } + // Same instant for both; only the arrow TimeZone field differs (UTC vs ""). + instant := time.Date(2026, 7, 9, 12, 0, 0, 0, time.UTC) + build := func(tz string) arrow.Array { + b := array.NewTimestampBuilder(pool, &arrow.TimestampType{Unit: arrow.Microsecond, TimeZone: tz}) + b.Append(arrow.Timestamp(instant.UnixMicro())) + return b.NewArray() + } + for _, tc := range []struct { + name string + tz string + }{ + {"timestamp_ltz_utc_zone", "UTC"}, + {"timestamp_ntz_empty_zone", ""}, + } { + t.Run(tc.name, func(t *testing.T) { + arr := build(tc.tz) + defer arr.Release() + kernel, err := arrowscan.ScanCell(arr, 0, loc) + if err != nil { + t.Fatalf("arrowscan.ScanCell: %v", err) + } + thrift := renderViaArrowbased(t, arr, 0, loc) + if kernel != thrift { + t.Errorf("TimeZone=%q divergence:\n kernel = %#v\n thrift = %#v", tc.tz, kernel, thrift) + } + }) + } +} + // TestTopLevelDecimalRendering documents the top-level DECIMAL story, which is // subtler than "kernel string vs Thrift float64": // diff --git a/kernel_e2e_test.go b/kernel_e2e_test.go index b2d86fac..5f8c40fe 100644 --- a/kernel_e2e_test.go +++ b/kernel_e2e_test.go @@ -110,6 +110,50 @@ func TestKernelE2ETimeZone(t *testing.T) { } } +// TestKernelE2ETimestampNTZ proves TIMESTAMP and TIMESTAMP_NTZ round-trip +// byte-identically to the Thrift path at a non-UTC session tz. The kernel delivers +// TIMESTAMP with an Arrow tz and TIMESTAMP_NTZ without one, but — like the Thrift +// path — the driver ignores that field and renders via .In(loc). So the naive +// wall-clock of an NTZ literal is treated as a UTC instant and shifted into the +// session tz (12:00 NTZ in America/New_York → the -04:00 local of 12:00Z = 08:00), +// exactly as Thrift does. Verified live on both backends; this pins the round-trip +// so a one-sided "don't shift NTZ" change is caught. +func TestKernelE2ETimestampNTZ(t *testing.T) { + const tz = "America/New_York" + db := kernelTestDBWith(t, WithSessionParams(map[string]string{"timezone": tz})) + defer db.Close() + + loc, err := time.LoadLocation(tz) + if err != nil { + t.Skipf("tz database unavailable: %v", err) + } + ctx := context.Background() + + var ltz time.Time + if err := db.QueryRowContext(ctx, + "SELECT CAST('2026-07-09 12:00:00' AS TIMESTAMP)").Scan(<z); err != nil { + t.Fatalf("TIMESTAMP query: %v", err) + } + // A zoned TIMESTAMP literal is the wall-clock in the session tz. + if want := time.Date(2026, 7, 9, 12, 0, 0, 0, loc); !ltz.Equal(want) { + t.Errorf("TIMESTAMP = %s, want %s", ltz, want) + } + + var ntz time.Time + if err := db.QueryRowContext(ctx, + "SELECT CAST('2026-07-09 12:00:00' AS TIMESTAMP_NTZ)").Scan(&ntz); err != nil { + t.Fatalf("TIMESTAMP_NTZ query: %v", err) + } + // The NTZ wall-clock is treated as a UTC instant then shifted into the session + // tz: 12:00Z == 08:00 in America/New_York (-04:00 in July). + if want := time.Date(2026, 7, 9, 8, 0, 0, 0, loc); !ntz.Equal(want) { + t.Errorf("TIMESTAMP_NTZ = %s, want %s (UTC instant %s)", ntz, want, ntz.UTC()) + } + if ntz.Location().String() != tz { + t.Errorf("TIMESTAMP_NTZ location = %q, want %q", ntz.Location(), tz) + } +} + // TestKernelE2ETLSSkipVerify checks that WithSkipTLSHostVerify (a relaxation // knob) is accepted on the kernel path; the connection must still succeed // against the warehouse's valid certificate. From 832e73df6ccf2f1f8c3a1e48d1f11f2c7d798e8c Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Sun, 12 Jul 2026 12:12:16 +0000 Subject: [PATCH 05/24] test(kernel): pin VARIANT & GEOMETRY render parity MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit VARIANT and GEOMETRY need no special rendering on the kernel path: verified live on both backends, both arrive over Arrow as plain STRING columns — a top-level VARIANT is its JSON text ({"a":1,"b":[2,3]}), a scalar VARIANT is "42", and GEOMETRY is its WKT "POINT(1 2)". Nested inside a container the variant/geometry element is a string leaf, rendered as a quoted, JSON-escaped string (the variant's own JSON is escaped as text, NOT re-parsed) — identical on both backends. Add untagged parity cases: a top-level string equivalence (variant object / scalar / geometry WKT) and nested string-leaf cases in an array, so the string arm's handling of these types can't silently drift between backends. GEOGRAPHY is intentionally excluded — not enabled on the benchmark warehouse and no consumer has asked for it. Co-authored-by: Isaac Signed-off-by: Mani Kaustubh Mathur --- .../rows/arrowbased/arrowscan_parity_test.go | 53 +++++++++++++++++++ 1 file changed, 53 insertions(+) diff --git a/internal/rows/arrowbased/arrowscan_parity_test.go b/internal/rows/arrowbased/arrowscan_parity_test.go index a00d2914..82faaa13 100644 --- a/internal/rows/arrowbased/arrowscan_parity_test.go +++ b/internal/rows/arrowbased/arrowscan_parity_test.go @@ -169,6 +169,25 @@ func TestArrowbasedKernelRenderParity(t *testing.T) { b.FieldBuilder(0).(*array.TimestampBuilder).Append(arrow.Timestamp(ts.UnixMicro())) return b.NewArray() }}, + // VARIANT and GEOMETRY arrive over Arrow as plain STRING columns (verified + // live: a top-level VARIANT is the JSON *string* "{\"a\":1}", GEOMETRY is the + // WKT string "POINT(1 2)"). Nested inside a container they are string leaves, + // so they must render as a quoted, JSON-escaped string on both backends — the + // variant's own JSON is NOT re-parsed, it is escaped as text. + {"variant_string_leaf_in_array", func() arrow.Array { + b := array.NewListBuilder(pool, arrow.BinaryTypes.String) + vb := b.ValueBuilder().(*array.StringBuilder) + b.Append(true) + vb.Append(`{"a":1,"b":[2,3]}`) // a VARIANT element, delivered as a string + return b.NewArray() + }}, + {"geometry_wkt_leaf_in_array", func() arrow.Array { + b := array.NewListBuilder(pool, arrow.BinaryTypes.String) + vb := b.ValueBuilder().(*array.StringBuilder) + b.Append(true) + vb.Append("POINT(1 2)") // GEOMETRY delivered as a WKT string + return b.NewArray() + }}, {"null_leaf_in_struct", func() arrow.Array { dt := arrow.StructOf( arrow.Field{Name: "a", Type: arrow.PrimitiveTypes.Int64}, @@ -430,6 +449,40 @@ func TestArrowbasedKernelTimestampTZParity(t *testing.T) { } } +// TestArrowbasedKernelVariantGeometryParity documents that VARIANT and GEOMETRY +// need no special rendering: over Arrow both arrive as plain STRING columns (a +// top-level VARIANT is its JSON text "{\"a\":1}", GEOMETRY is its WKT "POINT(1 2)"), +// so the existing string arm on both backends already produces identical output. +// Verified live on both backends. (GEOGRAPHY is intentionally not covered — it is +// not enabled on the benchmark warehouse and no consumer has asked for it.) +func TestArrowbasedKernelVariantGeometryParity(t *testing.T) { + pool := memory.NewGoAllocator() + loc := time.UTC + cases := []struct { + name, val string + }{ + {"variant_object", `{"a":1,"b":[2,3]}`}, + {"variant_scalar", "42"}, + {"geometry_wkt", "POINT(1 2)"}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + b := array.NewStringBuilder(pool) + b.Append(tc.val) + arr := b.NewArray() + defer arr.Release() + kernel, err := arrowscan.ScanCell(arr, 0, loc) + if err != nil { + t.Fatalf("arrowscan.ScanCell: %v", err) + } + thrift := renderViaArrowbased(t, arr, 0, loc) + if kernel != thrift || kernel != tc.val { + t.Errorf("%s divergence:\n kernel = %#v\n thrift = %#v\n want = %#v", tc.name, kernel, thrift, tc.val) + } + }) + } +} + // TestTopLevelDecimalRendering documents the top-level DECIMAL story, which is // subtler than "kernel string vs Thrift float64": // From ad9eebc3ae74bda8218334e280ca0e4fe122ee70 Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Sun, 12 Jul 2026 12:16:47 +0000 Subject: [PATCH 06/24] fix(kernel): surface kernel errors as DBExecutionError with sqlstate MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A failed kernel query returned the raw *KernelError, so consumers doing errors.As(err, &DBExecutionError) — the way they inspect Thrift failures — didn't get SqlState()/QueryId()/IsRetryable() through the standard interface. kernelOp.ExecutionError now digs the sqlstate out of the underlying *KernelError and wraps the cause via NewExecutionErrorWithState, so kernel query failures surface with the same DBExecutionError shape as Thrift. Adds the neutral NewExecutionErrorWithState to the untagged internal/errors package (Thrift's NewExecutionError needs a TGetOperationStatusResp the kernel backend can't produce), unit-tested in the default CGO_ENABLED=0 build. Parity is type + SQLSTATE, not byte-identical text — kernel messages are richer (they carry the SQL error class + suggestions). Verified live: unknown table → 42P01, unknown column → 42703, byte-identical sqlstate to Thrift. Co-authored-by: Isaac Signed-off-by: Mani Kaustubh Mathur --- internal/errors/err_test.go | 19 +++++++++++++++++++ kernel_e2e_test.go | 34 ++++++++++++++++++++++++++++++++++ 2 files changed, 53 insertions(+) diff --git a/internal/errors/err_test.go b/internal/errors/err_test.go index e62e0ff2..09c5a6ad 100644 --- a/internal/errors/err_test.go +++ b/internal/errors/err_test.go @@ -36,6 +36,25 @@ func TestDbSqlErrors(t *testing.T) { assert.Equal(t, ee, execError) }) + t.Run("NewExecutionErrorWithState carries the queryId and sqlstate", func(t *testing.T) { + // The kernel backend has no Thrift TGetOperationStatusResp; it passes an + // already-extracted queryId and sqlstate. The value must satisfy + // DBExecutionError and surface both — the parity property with + // NewExecutionError. + cause := errors.New("cause") + execError := NewExecutionErrorWithState(context.TODO(), "exec error", cause, "q-123", "42P01") + e := errors.Wrap(execError, "is wrapped") + + assert.Equal(t, "is wrapped: databricks: execution error: exec error: cause", e.Error()) + assert.True(t, errors.Is(e, dbsqlerr.ExecutionError)) + assert.True(t, errors.Is(e, cause)) + + var ee dbsqlerr.DBExecutionError + assert.True(t, errors.As(e, &ee)) + assert.Equal(t, "q-123", ee.QueryId()) + assert.Equal(t, "42P01", ee.SqlState()) + }) + t.Run("errors.Is/As works with driver error values", func(t *testing.T) { // Create a driver error and wrap it in a regular error cause := errors.New("cause") diff --git a/kernel_e2e_test.go b/kernel_e2e_test.go index 5f8c40fe..1c15e1ab 100644 --- a/kernel_e2e_test.go +++ b/kernel_e2e_test.go @@ -11,6 +11,8 @@ import ( "os" "testing" "time" + + dbsqlerr "github.com/databricks/databricks-sql-go/errors" ) // kernelTestDB opens a kernel-backed *sql.DB from DATABRICKS_HOST / @@ -243,6 +245,38 @@ func dataTypeEqual(got, want driver.Value) bool { } } +// TestKernelE2EErrorSurface proves a failed query surfaces as a DBExecutionError +// carrying a SqlState, the same error shape as the Thrift path (consumers rely on +// errors.As(err, &DBExecutionError) + SqlState()). Verified live: unknown table → +// 42P01, unknown column → 42703 — byte-identical sqlstate to Thrift, though the +// kernel's message is richer (it includes the SQL error class + suggestions). +func TestKernelE2EErrorSurface(t *testing.T) { + db := kernelTestDB(t) + defer db.Close() + + cases := []struct { + name, query, wantState string + }{ + {"unknown_table", "SELECT * FROM this_table_does_not_exist_xyz", "42P01"}, + {"unknown_column", "SELECT no_such_col FROM range(1)", "42703"}, + } + for _, c := range cases { + t.Run(c.name, func(t *testing.T) { + _, err := db.QueryContext(context.Background(), c.query) + if err == nil { + t.Fatalf("expected an error for %q, got none", c.query) + } + var de dbsqlerr.DBExecutionError + if !errors.As(err, &de) { + t.Fatalf("error is not a DBExecutionError: %v", err) + } + if de.SqlState() != c.wantState { + t.Errorf("sqlState = %q, want %q (err: %v)", de.SqlState(), c.wantState, err) + } + }) + } +} + // TestKernelE2ECloudFetch streams a CloudFetch-sized result end to end. CloudFetch // is internal to the kernel, so "it works" means many batches stream and scan // correctly — which also exercises the per-batch release/lifetime path. From 24c4d83b0a7cb3a81025ce1da75936ddc6431ed6 Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Sun, 12 Jul 2026 12:19:25 +0000 Subject: [PATCH 07/24] docs(kernel): document verified inherited features; drop stale interval caveat MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Record what the kernel backend inherits unchanged above the backend seam — the database/sql connection pool (each conn wraps one kernel session), per-connection CREATE_SESSION / DELETE_SESSION telemetry (recorded unconditionally in connector.go, backend-agnostic), and the telemetry exporter's circuit breaker — and that result types render byte-for-byte with Thrift (scalars, exact DECIMAL, TIMESTAMP / TIMESTAMP_NTZ, INTERVAL, nested + VARIANT as JSON, GEOMETRY as WKT). Remove the now-stale "INTERVAL types are not yet handled by the kernel scanner" caveat (intervals render now), and narrow the telemetry caveat to what is actually missing: only EXECUTE_STATEMENT telemetry (gated on a per-statement query id the kernel C ABI doesn't yet surface) — CREATE_SESSION / DELETE_SESSION are unaffected. Add a live-verified connection-pool e2e (40 concurrent queries over pool cap 8) backing the inherited-pool claim. Co-authored-by: Isaac Signed-off-by: Mani Kaustubh Mathur --- doc.go | 22 ++++++++++++++-------- kernel_e2e_test.go | 36 ++++++++++++++++++++++++++++++++++++ 2 files changed, 50 insertions(+), 8 deletions(-) diff --git a/doc.go b/doc.go index a8bb5e33..060693bb 100644 --- a/doc.go +++ b/doc.go @@ -233,14 +233,20 @@ WithMaxRows and retry tuning (WithRetries with a positive limit) are accepted bu not applied on the kernel path: the kernel manages result fetching and retries internally, below the C ABI, with no user-facing knob. -Two further kernel-backend caveats. The INTERVAL types (year-month / day-time) -listed in the type table below are not yet handled by the kernel scanner and -return a scan error; use the default (Thrift) backend for interval columns. And -the kernel backend does not yet surface a per-statement server query id on the -success path, so a QueryIdCallback (see below) fires with "" and no -EXECUTE_STATEMENT telemetry is emitted for kernel queries. (On a query failure the -server query id IS available: the returned error is a DBExecutionError whose -QueryId() carries it — see the Errors section.) +Features that live above the backend seam are inherited unchanged: the +database/sql connection pool (each connection wraps one kernel session), +per-connection telemetry (CREATE_SESSION and DELETE_SESSION are recorded for +kernel connections just like Thrift), and the telemetry exporter's circuit +breaker are all backend-agnostic. Result types are rendered to match the Thrift +backend byte-for-byte: scalars, DECIMAL (exact string), TIMESTAMP / TIMESTAMP_NTZ +(both shifted into the session time zone, as Thrift does), INTERVAL year-month +and day-time, nested Array/Map/Struct and VARIANT (as JSON), and GEOMETRY (WKT). + +Two current kernel-backend limitations. Bound query parameters are not yet wired +and return a clear error at execute time (see above). And the kernel backend does +not yet surface a per-statement server query id, so a QueryIdCallback (see below) +fires with "" and no EXECUTE_STATEMENT telemetry is emitted for kernel queries +(CREATE_SESSION / DELETE_SESSION telemetry is unaffected). # Programmatically Retrieving Connection and Query Id diff --git a/kernel_e2e_test.go b/kernel_e2e_test.go index 1c15e1ab..4fc56e5f 100644 --- a/kernel_e2e_test.go +++ b/kernel_e2e_test.go @@ -9,6 +9,8 @@ import ( "database/sql/driver" "errors" "os" + "sync" + "sync/atomic" "testing" "time" @@ -312,6 +314,40 @@ func TestKernelE2ECloudFetch(t *testing.T) { } } +// TestKernelE2EConnectionPool proves the database/sql connection pool works with +// the kernel backend: the pool lives above the backend seam (each connection wraps +// one kernel session), so many concurrent queries over a capped pool must all +// succeed with checkout/return and per-conn single-session isolation intact. +func TestKernelE2EConnectionPool(t *testing.T) { + db := kernelTestDB(t) + defer db.Close() + db.SetMaxOpenConns(8) + db.SetMaxIdleConns(8) + + const n = 40 + var errs, ok int64 + var wg sync.WaitGroup + for i := 0; i < n; i++ { + wg.Add(1) + go func() { + defer wg.Done() + ctx, cancel := context.WithTimeout(context.Background(), 90*time.Second) + defer cancel() + var v int + if err := db.QueryRowContext(ctx, "SELECT 1").Scan(&v); err != nil || v != 1 { + atomic.AddInt64(&errs, 1) + return + } + atomic.AddInt64(&ok, 1) + }() + } + wg.Wait() + if errs != 0 { + t.Fatalf("connection pool: %d/%d queries failed", errs, n) + } + t.Logf("connection pool OK: %d/%d queries succeeded over pool cap 8", ok, n) +} + // TestKernelE2ECancellation cancels a long-running query via ctx and asserts it // returns well before its uncancelled runtime. func TestKernelE2ECancellation(t *testing.T) { From 608763ee4dcd6e4e023eb2f2f03ef997a4de18ca Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Sun, 12 Jul 2026 17:19:19 +0000 Subject: [PATCH 08/24] feat(kernel): bind query parameters via the C ABI MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The kernel backend rejected bound parameters at execute time. Now it binds them: the driver's backend.Param{Name, Type, Value} maps 1:1 onto the kernel's kernel_statement_bind_parameter (K1) — value already stringified, Type the Databricks SQL type name, empty Name → positional, nil Value → SQL NULL ("VOID"). bindParams runs after set_sql (which clears any prior binds), using the existing newCStr/newCStrOrNull helpers and the call() FFI-safety wrapper; a bind failure closes the statement and surfaces via toStatementError. Removes the fail-loud reject in Execute (and its now-unused errors import). The old TestExecuteRejectsParams is repurposed as TestExecuteHandleLessOpContract (the non-nil handle-less Operation contract, now driven by a nil-session failure since params no longer reject). Live parity: 10 cases (positional/named, each scalar type, NULL, multi-param, predicate) produce byte-identical output on the kernel and Thrift backends. Requires a kernel build carrying kernel_statement_bind_parameter; the KERNEL_REV pin is bumped to the K1 merge SHA when it lands. Co-authored-by: Isaac Signed-off-by: Mani Kaustubh Mathur --- internal/backend/kernel/backend.go | 15 +++---- internal/backend/kernel/kernel_test.go | 15 +++++-- internal/backend/kernel/operation.go | 38 +++++++++++++++++ kernel_parity_test.go | 57 +++++++++++++++++++++++++- 4 files changed, 110 insertions(+), 15 deletions(-) diff --git a/internal/backend/kernel/backend.go b/internal/backend/kernel/backend.go index e9ea36b3..7c74e19e 100644 --- a/internal/backend/kernel/backend.go +++ b/internal/backend/kernel/backend.go @@ -382,21 +382,16 @@ func (k *KernelBackend) SessionID() string { return k.sessionID } // Execute runs a statement to a terminal state via the blocking execute path. // Per the Backend contract it returns a non-nil Operation even on error so the -// caller can read StatementID / wrap the error / Close uniformly. +// caller can read StatementID / wrap the error / Close uniformly. Bound +// parameters (req.Params) are bound onto the statement in execute (see +// bindParams). func (k *KernelBackend) Execute(ctx context.Context, req backend.ExecRequest) (backend.Operation, error) { - // Bound parameters are not yet wired for the kernel backend. Reject them with - // a clear error rather than silently shipping the query with unbound - // placeholders (which would behave differently than Thrift). Parameters arrive - // per-query, so this is an execute-time error, not a connect-time one. Return a - // non-nil Operation per the Backend contract. - if len(req.Params) > 0 { - return &kernelOp{}, fmt.Errorf("databricks: query parameters are %w; "+ - "inline the values or use the default (Thrift) backend", dbsqlerr.ErrNotSupportedByKernel) - } // Staging (Unity Catalog volume PUT/GET/REMOVE) needs a local file transfer the // kernel path can't perform, and the kernel C ABI surfaces no IsStagingOperation // signal to drive conn.execStagingOperation. Reject it here rather than let // IsStaging return false and report success with no file moved (silent data loss). + // (Bound parameters, by contrast, ARE now supported — bound in execute via the + // kernel's raw-param C ABI.) if isStagingStatement(req.Query) { return &kernelOp{}, fmt.Errorf("databricks: staging operations (PUT/GET/REMOVE on a volume) are %w; "+ "use the default (Thrift) backend", dbsqlerr.ErrNotSupportedByKernel) diff --git a/internal/backend/kernel/kernel_test.go b/internal/backend/kernel/kernel_test.go index 84a07a3c..faeb8726 100644 --- a/internal/backend/kernel/kernel_test.go +++ b/internal/backend/kernel/kernel_test.go @@ -82,14 +82,19 @@ func TestEvictIfSessionFatal(t *testing.T) { // Operation must be non-nil (Backend contract) and its Close must report // closed=false, since no server statement was ever created (a phantom // CLOSE_STATEMENT would otherwise be recorded for it). -func TestExecuteRejectsParams(t *testing.T) { - k := &KernelBackend{} +// When Execute fails before it acquires a statement handle (here: a nil session +// makes new_statement fail), it must still honor the Backend contract — a non-nil, +// handle-less Operation that Closes as a no-op (closed=false, no CLOSE_STATEMENT) +// and reports zero AffectedRows. (Parameter binding is exercised live in +// TestKernelE2EParams; a nil-session unit test can't reach the bind path.) +func TestExecuteHandleLessOpContract(t *testing.T) { + k := &KernelBackend{} // nil session → new_statement fails op, err := k.Execute(context.Background(), backend.ExecRequest{ Query: "SELECT ?", - Params: []backend.Param{{Name: "x"}}, + Params: []backend.Param{{Name: "x", Type: "STRING", Value: strPtr("v")}}, }) if err == nil { - t.Fatal("expected an error for bound parameters, got nil") + t.Fatal("expected an error from Execute on a nil-session backend, got nil") } if op == nil { t.Fatal("Execute must return a non-nil Operation per the Backend contract") @@ -172,6 +177,8 @@ func TestExecutionErrorContract(t *testing.T) { } } +func strPtr(s string) *string { return &s } + // The cell/nested rendering (ScanCell and the JSON grammar) now lives in the // untagged internal/arrowscan package, where its tests run in the default // CGO_ENABLED=0 build; see arrowscan_test.go. The decimal formatter lives in diff --git a/internal/backend/kernel/operation.go b/internal/backend/kernel/operation.go index eec325a6..acb9a20a 100644 --- a/internal/backend/kernel/operation.go +++ b/internal/backend/kernel/operation.go @@ -58,6 +58,17 @@ func (k *KernelBackend) execute(ctx context.Context, req backend.ExecRequest) (b } sql.free() + // Bind parameters. The driver hands us backend.Param{Name, Type, Value}: the + // value is already stringified, Type is the Databricks SQL type name, and a nil + // Value is SQL NULL (Type "VOID"). Each maps 1:1 onto + // kernel_statement_bind_parameter, which builds the SEA wire parameter directly + // (name/ordinal + type + string), matching the Thrift path's toSparkParameters. + if err := bindParams(stmt, req.Params); err != nil { + C.kernel_statement_close(stmt) + k.evictIfSessionFatal(err) + return &kernelOp{}, fmt.Errorf("kernel: bind params: %w", toStatementError(err)) + } + // Detached canceller, obtained before execute so it observes the server // statement id the moment execute publishes it. Non-fatal on failure: proceed // without cancellation rather than failing the query. @@ -171,6 +182,33 @@ func (k *KernelBackend) execute(ctx context.Context, req backend.ExecRequest) (b return op, nil } +// 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 +// in push order) and a nil Value is SQL NULL (Type "VOID"). Runs before execute, +// so the params are set on the fresh statement (set_sql clears any prior binds). +func bindParams(stmt *C.kernel_statement_t, params []backend.Param) error { + for i, p := range params { + name := newCStrOrNull(p.Name) // empty Name → NULL → positional + typ := newCStr(p.Type) + val := newCStrOrNull("") + if p.Value != nil { + val = newCStr(*p.Value) // non-nil, possibly empty string, is a real value + } + err := call(func() C.KernelStatusCode { + return C.kernel_statement_bind_parameter(stmt, name.c, typ.c, val.c) + }) + name.free() + typ.free() + val.free() + if err != nil { + return fmt.Errorf("param %d (name=%q type=%q): %w", i, p.Name, p.Type, err) + } + klog("bound param %d name=%q type=%q null=%v", i, p.Name, p.Type, p.Value == nil) + } + return nil +} + // kernelOp implements backend.Operation over a sync executed statement. type kernelOp struct { // backend is the owning connection's backend, held so a session-fatal error on diff --git a/kernel_parity_test.go b/kernel_parity_test.go index 8f570302..62d2576a 100644 --- a/kernel_parity_test.go +++ b/kernel_parity_test.go @@ -64,12 +64,67 @@ func TestKernelThriftParity(t *testing.T) { } } +// paramCase is one parameterized-query parity case: the same SQL + args run on +// both backends must yield byte-identical output. +type paramCase struct { + name string + sql string + args []any +} + +// The kernel binds these via kernel_statement_bind_parameter (the driver's +// backend.Param{Name, Type, Value}); Thrift binds via toSparkParameters. Covers +// positional (?) and named (:n) markers, each scalar type, SQL NULL, multi-param, +// and a predicate — mirroring the C2 POC gate. +var paramCases = []paramCase{ + {"pos_int", "SELECT ? AS v", []any{int64(42)}}, + {"pos_string", "SELECT ? AS v", []any{"hello"}}, + {"pos_double", "SELECT ? AS v", []any{3.5}}, + {"pos_bool", "SELECT ? AS v", []any{true}}, + {"pos_null", "SELECT ? AS v", []any{nil}}, + {"pos_two", "SELECT ? + ? AS v", []any{int64(2), int64(40)}}, + {"pos_in_predicate", "SELECT count(*) AS v FROM range(100) WHERE id < ?", []any{int64(10)}}, + {"named_int", "SELECT :n AS v", []any{sql.Named("n", int64(7))}}, + {"named_string", "SELECT :s AS v", []any{sql.Named("s", "world")}}, + {"named_two", "SELECT :a AS a, :b AS b", []any{sql.Named("a", int64(1)), sql.Named("b", "x")}}, +} + +// TestKernelParamsVsThrift asserts parameterized queries produce byte-identical +// output on the kernel and Thrift backends — the bound-parameter acceptance gate. +func TestKernelParamsVsThrift(t *testing.T) { + kernelDB := kernelTestDB(t) + defer kernelDB.Close() + thriftDB := thriftTestDB(t) + defer thriftDB.Close() + + for _, c := range paramCases { + t.Run(c.name, func(t *testing.T) { + kernelRow := scanOneRowAsStringsArgs(t, kernelDB, c.sql, c.args...) + thriftRow := scanOneRowAsStringsArgs(t, thriftDB, c.sql, c.args...) + if len(kernelRow) != len(thriftRow) { + t.Fatalf("column count differs: kernel=%d thrift=%d", len(kernelRow), len(thriftRow)) + } + for i := range kernelRow { + if kernelRow[i] != thriftRow[i] { + t.Errorf("col %d differs: kernel=%q thrift=%q (sql=%q args=%v)", i, kernelRow[i], thriftRow[i], c.sql, c.args) + } + } + }) + } +} + // scanOneRowAsStrings scans the first row into a []string via sql.RawBytes, so a // NULL renders as "" and every value is compared in its wire form, // independent of Go-type coercion differences between the backends. func scanOneRowAsStrings(t *testing.T, db *sql.DB, query string) []string { + return scanOneRowAsStringsArgs(t, db, query) +} + +// scanOneRowAsStringsArgs is scanOneRowAsStrings with query arguments (bound +// parameters). Same wire-form comparison contract. +func scanOneRowAsStringsArgs(t *testing.T, db *sql.DB, query string, args ...any) []string { t.Helper() - rows, err := db.QueryContext(context.Background(), query) + rows, err := db.QueryContext(context.Background(), query, args...) if err != nil { t.Fatalf("query: %v", err) } From 0a17effd7c3af2f48be93842c5390f8d0f8019ac Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Sun, 12 Jul 2026 17:22:45 +0000 Subject: [PATCH 09/24] feat(kernel): surface server query id for EXECUTE_STATEMENT telemetry MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit kernelOp.StatementID() returned "", so the kernel backend emitted no EXECUTE_STATEMENT telemetry and QueryIdCallback fired with an empty id (connection.go gates both on a non-empty statement id). Wire StatementID() to the server query id via kernel_executed_statement_query_id (K1), captured at execute time into a cached field — the same lifetime discipline as affectedRows, since the C accessor returns a pointer borrowed from the exec handle and the op is closed (nulling exec) before StatementID() is read on some paths. C.GoString deep-copies out of the borrowed string. Live e2e: a registered QueryIdCallback fires with a non-empty server id after a kernel query. Updates doc.go — bound parameters (c6) and EXECUTE_STATEMENT telemetry are now supported; the remaining kernel-backend limitation is batch-boundary (not mid-fetch) read cancellation. Co-authored-by: Isaac Signed-off-by: Mani Kaustubh Mathur --- doc.go | 38 +++++++++++++++------------- internal/backend/kernel/operation.go | 27 ++++++++++++-------- kernel_e2e_test.go | 29 +++++++++++++++++++++ 3 files changed, 65 insertions(+), 29 deletions(-) diff --git a/doc.go b/doc.go index 060693bb..1f3c7e92 100644 --- a/doc.go +++ b/doc.go @@ -214,12 +214,13 @@ time zone) connection options. WithTimeout (a server query timeout the kernel C can't set) and WithRetries used to disable retries (the kernel retries internally) return a clear error at connect time rather than being silently ignored; token- provider, external/static, and federated authenticators are likewise not supported -and rejected loudly. Bound query parameters and staging operations (PUT/GET/REMOVE -on a Unity Catalog volume, which need a local file transfer this backend cannot -perform) are not yet supported and return a clear error at execute time (they are -per-statement, not connect-time). None of these is silently ignored. (Metadata is -issued as ordinary SQL — SHOW/DESCRIBE/information_schema — and runs on this backend -like any other query.) +and rejected loudly. Staging operations (PUT/GET/REMOVE on a Unity Catalog volume, +which need a local file transfer this backend cannot perform) are not yet supported +and return a clear error at execute time (they are per-statement, not connect-time). +Bound query parameters (positional and named) are supported — bound over the C ABI +to match the Thrift path. None of these is silently ignored. (Metadata is issued as +ordinary SQL — SHOW/DESCRIBE/information_schema — and runs on this backend like any +other query.) OAuth U2M is interactive. On a cache miss (no valid cached refresh token) opening a connection launches the system browser and blocks until the user completes login or @@ -235,18 +236,19 @@ internally, below the C ABI, with no user-facing knob. Features that live above the backend seam are inherited unchanged: the database/sql connection pool (each connection wraps one kernel session), -per-connection telemetry (CREATE_SESSION and DELETE_SESSION are recorded for -kernel connections just like Thrift), and the telemetry exporter's circuit -breaker are all backend-agnostic. Result types are rendered to match the Thrift -backend byte-for-byte: scalars, DECIMAL (exact string), TIMESTAMP / TIMESTAMP_NTZ -(both shifted into the session time zone, as Thrift does), INTERVAL year-month -and day-time, nested Array/Map/Struct and VARIANT (as JSON), and GEOMETRY (WKT). - -Two current kernel-backend limitations. Bound query parameters are not yet wired -and return a clear error at execute time (see above). And the kernel backend does -not yet surface a per-statement server query id, so a QueryIdCallback (see below) -fires with "" and no EXECUTE_STATEMENT telemetry is emitted for kernel queries -(CREATE_SESSION / DELETE_SESSION telemetry is unaffected). +per-connection telemetry (CREATE_SESSION, EXECUTE_STATEMENT, and DELETE_SESSION +are recorded for kernel connections just like Thrift), and the telemetry +exporter's circuit breaker are all backend-agnostic. Result types are rendered to +match the Thrift backend byte-for-byte: scalars, DECIMAL (exact string), +TIMESTAMP / TIMESTAMP_NTZ (both shifted into the session time zone, as Thrift +does), INTERVAL year-month and day-time, nested Array/Map/Struct and VARIANT (as +JSON), and GEOMETRY (WKT). The per-statement server query id is surfaced on the +success path, so a QueryIdCallback (see below) fires with the real id and +EXECUTE_STATEMENT telemetry carries it. + +One kernel-backend limitation remains on the read path: context cancellation is +honored at result-batch boundaries, not mid-fetch (an in-flight CloudFetch batch +runs to completion before the cancel takes effect). # Programmatically Retrieving Connection and Query Id diff --git a/internal/backend/kernel/operation.go b/internal/backend/kernel/operation.go index acb9a20a..8e4b3f4b 100644 --- a/internal/backend/kernel/operation.go +++ b/internal/backend/kernel/operation.go @@ -166,8 +166,9 @@ func (k *KernelBackend) execute(ctx context.Context, req backend.ExecRequest) (b return op, fmt.Errorf("kernel: execute: %w", toStatementError(execErr)) } op.exec = exec - // Capture the modified-row count now, while exec is live — the operation is - // closed (nulling exec) before AffectedRows is read on the ExecContext path. + // Capture the modified-row count and server query id now, while exec is live — + // the operation is closed (nulling exec) before these are read on the + // ExecContext path, and the query-id pointer is only valid while exec lives. op.affectedRows = int64(C.kernel_executed_statement_num_modified_rows(exec)) // A nil execErr means the statement reached a terminal state server-side (it // committed). We deliberately do NOT re-check ctx.Err() here to convert a @@ -178,7 +179,10 @@ func (k *KernelBackend) execute(ctx context.Context, req backend.ExecRequest) (b // arrives before completion is still honored via the watcher → execErr branch // above; one that loses the race to a committed statement yields that statement's // result, same as Thrift. - klog("Execute OK stmt=%p exec=%p affectedRows=%d", stmt, exec, op.affectedRows) + if qid := C.kernel_executed_statement_query_id(exec); qid != nil { + op.statementID = C.GoString(qid) // deep-copies out of the borrowed C string + } + klog("Execute OK stmt=%p exec=%p affectedRows=%d statementID=%q", stmt, exec, op.affectedRows, op.statementID) return op, nil } @@ -222,6 +226,11 @@ type kernelOp struct { // cached (not read live from exec) because the caller closes the operation — // which nulls exec — before reading AffectedRows (see conn.ExecContext). affectedRows int64 + // statementID is the server query id, captured at execute time. Cached (not + // read live) because kernel_executed_statement_query_id returns a pointer + // borrowed from the exec handle, valid only while exec is alive — the same + // lifetime discipline as affectedRows. + statementID string // location renders DATE / TIMESTAMP values in the session time zone, matching // the Thrift path; nil means UTC. Carried onto the rows built by Results. location *time.Location @@ -229,14 +238,10 @@ type kernelOp struct { var _ backend.Operation = (*kernelOp)(nil) -// StatementID returns "": the kernel C ABI exposes no success-path statement/query -// id accessor (only the error path carries a query_id, on KernelError). Because -// conn.ExecContext/QueryContext gate per-statement telemetry on -// StatementID() != "", kernel queries currently emit no EXECUTE_STATEMENT metric -// and their query-id log field is empty — there is no session-id fallback on that -// gate. Surfacing an id needs a kernel accessor (e.g. -// kernel_executed_statement_query_id); tracked as a follow-up. -func (o *kernelOp) StatementID() string { return "" } +// StatementID returns the server query id captured at execute time (empty on a +// handle-less op that never executed). A non-empty id ungates EXECUTE_STATEMENT +// telemetry and drives QueryIdCallback, matching the Thrift path. +func (o *kernelOp) StatementID() string { return o.statementID } // AffectedRows is the modified-row count for ExecContext. It returns the value // cached at execute time, so it is correct even after the operation is closed diff --git a/kernel_e2e_test.go b/kernel_e2e_test.go index 4fc56e5f..7515f3be 100644 --- a/kernel_e2e_test.go +++ b/kernel_e2e_test.go @@ -14,6 +14,7 @@ import ( "testing" "time" + "github.com/databricks/databricks-sql-go/driverctx" dbsqlerr "github.com/databricks/databricks-sql-go/errors" ) @@ -314,6 +315,34 @@ func TestKernelE2ECloudFetch(t *testing.T) { } } +// TestKernelE2EStatementID proves the kernel backend surfaces the server query id +// on the success path: a registered QueryIdCallback fires with a non-empty id +// (driven by kernelOp.StatementID() → kernel_executed_statement_query_id). This is +// the observable end of the EXECUTE_STATEMENT telemetry / query-history path that +// was previously dark on the kernel backend (StatementID() returned ""). +func TestKernelE2EStatementID(t *testing.T) { + db := kernelTestDB(t) + defer db.Close() + + var gotID string + var fired bool + ctx := driverctx.NewContextWithQueryIdCallback(context.Background(), func(id string) { + fired = true + gotID = id + }) + + var v int64 + if err := db.QueryRowContext(ctx, "SELECT 1").Scan(&v); err != nil { + t.Fatalf("query: %v", err) + } + if !fired { + t.Fatal("QueryIdCallback did not fire") + } + if gotID == "" { + t.Error("kernel backend surfaced an empty query id; want the server statement id") + } +} + // TestKernelE2EConnectionPool proves the database/sql connection pool works with // the kernel backend: the pool lives above the backend seam (each connection wraps // one kernel session), so many concurrent queries over a capped pool must all From d996b9e34f42c535e237168bc96a5dad413bbe31 Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Sun, 12 Jul 2026 18:34:19 +0000 Subject: [PATCH 10/24] build(kernel): pin KERNEL_REV to the PuPr statement-surface rev MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Bumps the kernel pin from the #163 canceller rev to the tip of the PuPr statement-surface branch (databricks-sql-kernel#165), which adds kernel_statement_bind_parameter and kernel_executed_statement_query_id — the two C-ABI symbols the bound-parameter (c6) and EXECUTE_STATEMENT-telemetry (c7) commits link against. Re-pin to the squash-merge SHA once #165 lands. Co-authored-by: Isaac Signed-off-by: Mani Kaustubh Mathur --- KERNEL_REV | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/KERNEL_REV b/KERNEL_REV index cfbf3024..0fc1de14 100644 --- a/KERNEL_REV +++ b/KERNEL_REV @@ -1 +1 @@ -9b90406 +4d301455f7e70de9cb747e0f1143b8a3df8c5b48 From bb1163460870a5dc03b687cf4f2214bf81f51483 Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Wed, 15 Jul 2026 11:15:02 +0000 Subject: [PATCH 11/24] =?UTF-8?q?fix(kernel):=20address=20PR=20review=20?= =?UTF-8?q?=E2=80=94=20interval=20MinInt=20overflow,=20cancel=20eviction,?= =?UTF-8?q?=20bind-mapping=20test?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses the review pass on this PR (comments left on #399). Fixes the actionable set; a follow-up Isaac re-review then flagged that one of the requested changes (surfacing the kernel Retryable flag) was itself unsafe, so that one is intentionally NOT made — see below. - Interval negation overflow (High): formatDayTimeInterval / formatYearMonthInterval negated the full magnitude up front (v = -v). At math.MinInt64 (day-time µs) / math.MinInt32 (year-month) that wraps back negative, so every component came out negative AND a '-' was prepended — doubly-negated garbage — and both are representable Spark interval bounds. Now derive each component from the signed value and take its magnitude when formatting (abs64); widen year-month to int64 before negating. Adds MinInt64 (µs + ns) and MinInt32 regression cases. - Cancelled-execute skipped session-fatal eviction (Med): when execErr raced a ctx cancel, execute returned ctx.Err() without calling evictIfSessionFatal, leaving a dead conn marked valid in the pool. Hoist the evict above the ctx-cancelled branch so it fires on both paths (drained watcher first, so no race). Also wrap BOTH the ctx error and the kernel error with two %w verbs so errors.Is(context.DeadlineExceeded) still matches AND the *KernelError (sqlstate/queryId) stays reachable via errors.As instead of being dropped. - Bind mapping had no executing coverage (Med): the live Param-binding proof (TestKernelParamsVsThrift) needs a warehouse and only runs in the credentialed nightly job, so the positional/named + SQL-NULL/empty-string decision shipped untested at PR time. Extract that pure decision into an untagged paramBindArg (bindparams.go), consumed by the cgo bindParams, and unit-test it under CGO_ENABLED=0 (TestParamBindArg). Rename a comment's dead TestKernelE2EParams reference to the real tests. - Stale StatementID() comments (Low): both said StatementID() is "" on this backend, but this PR made it return the real server id on the success path. Scope the empty-id claim to the execute-error path. NOT changed (Isaac re-review, MAJOR): surfacing the kernel's Retryable flag on kernelOp.ExecutionError. That is exclusively the post-submission path, where a network/unavailable failure may have already committed a non-idempotent INSERT/UPDATE/MERGE — reporting IsRetryable()==true would invite an app to double-write, and it diverges from the Thrift path (always non-retryable here). This mirrors toStatementError refusing driver.ErrBadConn for the same reason. sqlState/queryId extraction is unchanged; TestExecutionErrorNeverRetryable pins that a Retryable KernelError still reports IsRetryable()==false. The connect-phase path (toConnError), where the retryable signal IS safe, is unaffected. Verified: default (CGO_ENABLED=0) + kernel-tagged suites pass, gofmt clean, Isaac review clean (0 final comments). Co-authored-by: Isaac Signed-off-by: Mani Kaustubh Mathur --- internal/arrowscan/arrowscan.go | 47 ++++++++++----- internal/arrowscan/arrowscan_test.go | 13 +++++ internal/backend/kernel/bindparams.go | 52 +++++++++++++++++ internal/backend/kernel/bindparams_test.go | 66 ++++++++++++++++++++++ internal/backend/kernel/errors_classify.go | 4 +- internal/backend/kernel/kernel_test.go | 42 +++++++++++++- internal/backend/kernel/operation.go | 58 +++++++++++++------ 7 files changed, 249 insertions(+), 33 deletions(-) create mode 100644 internal/backend/kernel/bindparams.go create mode 100644 internal/backend/kernel/bindparams_test.go diff --git a/internal/arrowscan/arrowscan.go b/internal/arrowscan/arrowscan.go index 3f26ee7f..4c020103 100644 --- a/internal/arrowscan/arrowscan.go +++ b/internal/arrowscan/arrowscan.go @@ -203,15 +203,19 @@ func ScanCellCached(col arrow.Array, row int, loc *time.Location, keys *StructKe // with 9 fractional digits, negated with a leading '-'. func formatDayTimeInterval(v int64, unit arrow.TimeUnit) string { neg := v < 0 - if neg { - v = -v - } - // Split into whole seconds + a sub-second nanosecond remainder working in the - // native unit. We must NOT scale the full magnitude up to nanoseconds first: - // Spark day-time intervals run up to ~Long.MaxValue microseconds (~292 years), - // so v*1e3 (or *1e6/*1e9) would overflow int64 and silently produce a wrong - // (often negative) string. Deriving seconds by dividing keeps every - // intermediate in range; only the bounded sub-second remainder is scaled up. + // Derive every component from the SIGNED value (Go integer division/modulo + // truncate toward zero, so each component simply carries v's sign), then take + // the magnitude of each bounded component when formatting. We deliberately do + // NOT negate the full magnitude up front (`v = -v`): at math.MinInt64 that + // wraps back to a negative value, so the components would come out negative + // *and* a '-' would be prepended — a doubly-negated garbage string — and + // math.MinInt64 μs is a representable Spark day-time bound. + // + // We also must NOT scale the full magnitude up to nanoseconds first: Spark + // day-time intervals run up to ~Long.MaxValue microseconds (~292 years), so + // v*1e3 (or *1e6/*1e9) would overflow int64 and silently produce a wrong + // string. Deriving seconds by dividing keeps every intermediate in range; only + // the bounded sub-second remainder is scaled up. var secs, frac int64 switch unit { case arrow.Second: @@ -236,18 +240,25 @@ func formatDayTimeInterval(v int64, unit arrow.TimeUnit) string { if neg { sign = "-" } - return fmt.Sprintf("%s%d %02d:%02d:%02d.%09d", sign, days, h, m, s, frac) + // Every component is bounded well within int64 (days ≤ ~1.07e8, h < 24, m/s < + // 60, frac < 1e9), so abs64 here can never hit the math.MinInt64 abs overflow. + return fmt.Sprintf("%s%d %02d:%02d:%02d.%09d", sign, abs64(days), abs64(h), abs64(m), abs64(s), abs64(frac)) } // formatYearMonthInterval renders a month count as the Thrift path's "years-months", // negated with a leading '-'. func formatYearMonthInterval(months int32) string { neg := months < 0 + // Widen to int64 BEFORE negating: negating math.MinInt32 as an int32 overflows + // (wraps back negative → doubly-negated garbage), but math.MinInt32 fits in + // int64 where the negation is exact. math.MinInt32 months is a representable + // Spark year-month bound. + m := int64(months) if neg { - months = -months + m = -m } - y := months / 12 - mo := months % 12 + y := m / 12 + mo := m % 12 sign := "" if neg { sign = "-" @@ -255,6 +266,16 @@ func formatYearMonthInterval(months int32) string { return fmt.Sprintf("%s%d-%d", sign, y, mo) } +// abs64 returns the absolute value of x. It is only ever called on components +// already bounded well within int64 (see formatDayTimeInterval), so the classic +// abs(math.MinInt64) overflow can't arise; the trivial form is intentional. +func abs64(x int64) int64 { + if x < 0 { + return -x + } + return x +} + // inLocation renders t in loc, matching the Thrift path's .In(location); a nil // loc leaves the value in UTC (arrow's ToTime default). func inLocation(t time.Time, loc *time.Location) time.Time { diff --git a/internal/arrowscan/arrowscan_test.go b/internal/arrowscan/arrowscan_test.go index a54c47b9..a48d38d9 100644 --- a/internal/arrowscan/arrowscan_test.go +++ b/internal/arrowscan/arrowscan_test.go @@ -137,6 +137,15 @@ func TestScanCellInterval(t *testing.T) { // NOT overflow int64 while scaling to nanoseconds — regression guard for the // prior multiply-first bug that produced a wrong/negative string here. {"large_us_no_overflow", arrow.Microsecond, 9223372036854775807, "106751991 04:00:54.775807000"}, + // math.MinInt64 μs is a representable negative bound. Negating the full + // magnitude (`v = -v`) wraps it back negative, doubly-negating into garbage; + // deriving components from the signed value renders it correctly. Its + // magnitude is exactly one μs past MaxInt64, so the last fractional digit is + // 8 where the MaxInt64 case above is 7. + {"min_int64_us", arrow.Microsecond, -9223372036854775808, "-106751991 04:00:54.775808000"}, + // Same MinInt64 wrap-on-negate hazard at the nanosecond unit (no scaling + // involved — this isolates the negation bug from the multiply-overflow one). + {"min_int64_ns", arrow.Nanosecond, -9223372036854775808, "-106751 23:47:16.854775808"}, } for _, tc := range dayTime { t.Run("daytime_"+tc.name, func(t *testing.T) { @@ -164,6 +173,10 @@ func TestScanCellInterval(t *testing.T) { {"year_and_month", 13, "1-1"}, {"months_only", 5, "0-5"}, {"negative", -13, "-1-1"}, + // math.MinInt32 months is a representable negative bound. Negating it as an + // int32 (`months = -months`) overflows and wraps back negative, doubly- + // negating into garbage; widening to int64 before negating renders it right. + {"min_int32", -2147483648, "-178956970-8"}, } for _, tc := range yearMonth { t.Run("yearmonth_"+tc.name, func(t *testing.T) { diff --git a/internal/backend/kernel/bindparams.go b/internal/backend/kernel/bindparams.go new file mode 100644 index 00000000..f057b97e --- /dev/null +++ b/internal/backend/kernel/bindparams.go @@ -0,0 +1,52 @@ +package kernel + +import "github.com/databricks/databricks-sql-go/internal/backend" + +// This file is intentionally NOT behind the `cgo && databricks_kernel` build tag. +// It holds the pure decision of how one backend.Param maps onto the kernel's +// raw-param bind arguments (positional vs named, SQL NULL vs empty string). That +// decision is the parity-critical part — a wrong name/NULL mapping silently +// diverges from the Thrift path's toSparkParameters — while the cgo wrapper around +// it (bindParams in operation.go) is just C-string marshaling. Keeping the decision +// here lets it be unit-tested under CGO_ENABLED=0 (see bindparams_test.go); the +// live kernel==Thrift proof (TestKernelParamsVsThrift) needs a warehouse and only +// runs in the credentialed nightly job, so this hermetic test is the PR-time guard. + +// bindArg is the kernel-side view of one bound parameter: the exact (name, type, +// value, isNull) tuple the cgo bind call forwards to kernel_statement_bind_parameter. +// A nameless arg is positional (ordinal assigned kernel-side in push order); an +// isNull arg is SQL NULL (its value is not sent). It deliberately holds no C types, +// so it can be constructed and asserted in a non-cgo test. +type bindArg struct { + // name is the parameter name; empty means positional. Passed to the kernel as + // NULL-for-empty (newCStrOrNull), which is how the kernel distinguishes a named + // bind from a positional one. + name string + // typ is the Databricks SQL type name the server expects (e.g. "BIGINT", + // "STRING", "VOID" for NULL). Always sent. + typ string + // value is the already-stringified wire value. Only meaningful when !isNull; + // an empty string here (with isNull false) is a real empty-string value, NOT + // SQL NULL — the two must not be conflated. + value string + // isNull is true for a SQL NULL parameter (backend.Param.Value == nil). When + // true the kernel is passed a NULL value pointer; when false the value string is + // sent verbatim, so "" binds an empty string rather than NULL. + isNull bool +} + +// paramBindArg maps one backend.Param onto its kernel bind arguments. It mirrors +// the driver's neutral param contract (see backend.Param): an empty Name is a +// positional parameter, and a nil Value is SQL NULL (carried with its type, e.g. +// "VOID"). A non-nil Value — including a pointer to the empty string — is a real +// value and must not be treated as NULL. This is the same distinction the Thrift +// path's toSparkParameters makes, so the two backends bind identically. +func paramBindArg(p backend.Param) bindArg { + a := bindArg{name: p.Name, typ: p.Type} + if p.Value == nil { + a.isNull = true + return a + } + a.value = *p.Value + return a +} diff --git a/internal/backend/kernel/bindparams_test.go b/internal/backend/kernel/bindparams_test.go new file mode 100644 index 00000000..47963726 --- /dev/null +++ b/internal/backend/kernel/bindparams_test.go @@ -0,0 +1,66 @@ +package kernel + +import ( + "testing" + + "github.com/databricks/databricks-sql-go/internal/backend" +) + +// paramBindArg is the parity-critical param mapping (positional vs named, SQL NULL +// vs empty string). The live kernel==Thrift proof (TestKernelParamsVsThrift) needs +// a warehouse and only runs in the credentialed nightly job, so this untagged test +// is the PR-time guard that a refactor can't silently break the mapping. It runs +// under CGO_ENABLED=0. +func TestParamBindArg(t *testing.T) { + strPtr := func(s string) *string { return &s } + + cases := []struct { + name string + param backend.Param + want bindArg + }{ + { + // A named parameter (:n markers) keeps its name and value. + name: "named value", + param: backend.Param{Name: "n", Type: "BIGINT", Value: strPtr("42")}, + want: bindArg{name: "n", typ: "BIGINT", value: "42", isNull: false}, + }, + { + // An empty Name is positional (? markers) — the kernel assigns the ordinal + // in push order. The name must stay empty so the cgo layer sends NULL. + name: "positional value", + param: backend.Param{Name: "", Type: "STRING", Value: strPtr("hello")}, + want: bindArg{name: "", typ: "STRING", value: "hello", isNull: false}, + }, + { + // A nil Value is SQL NULL, carried with its type (VOID). isNull must be set + // so the cgo layer sends a NULL value pointer, not an empty string. + name: "nil value is SQL NULL", + param: backend.Param{Name: "n", Type: "VOID", Value: nil}, + want: bindArg{name: "n", typ: "VOID", value: "", isNull: true}, + }, + { + // The critical distinction: a pointer to the empty string is a REAL empty + // string value, NOT SQL NULL. isNull must stay false so the empty string is + // bound as a value — conflating the two would diverge from the Thrift path. + name: "empty-string value is not NULL", + param: backend.Param{Name: "s", Type: "STRING", Value: strPtr("")}, + want: bindArg{name: "s", typ: "STRING", value: "", isNull: false}, + }, + { + // A positional SQL NULL: empty name AND nil value together. + name: "positional NULL", + param: backend.Param{Name: "", Type: "VOID", Value: nil}, + want: bindArg{name: "", typ: "VOID", value: "", isNull: true}, + }, + } + + for _, c := range cases { + t.Run(c.name, func(t *testing.T) { + got := paramBindArg(c.param) + if got != c.want { + t.Errorf("paramBindArg(%+v) = %+v, want %+v", c.param, got, c.want) + } + }) + } +} diff --git a/internal/backend/kernel/errors_classify.go b/internal/backend/kernel/errors_classify.go index d9a7e0c5..03d23aaf 100644 --- a/internal/backend/kernel/errors_classify.go +++ b/internal/backend/kernel/errors_classify.go @@ -45,7 +45,9 @@ type KernelError struct { func (e *KernelError) Error() string { // Append the server query id when present — it is the one correlation handle - // to server-side query history, and StatementID() is "" on this backend. + // to server-side query history. On the execute-error path there is no executed + // statement handle, so kernelOp.StatementID() is "" there; carrying the id on + // the error itself keeps it reachable when the accessor can't surface it. q := "" if e.QueryID != "" { q = fmt.Sprintf(", queryId=%s", e.QueryID) diff --git a/internal/backend/kernel/kernel_test.go b/internal/backend/kernel/kernel_test.go index faeb8726..e3e62f54 100644 --- a/internal/backend/kernel/kernel_test.go +++ b/internal/backend/kernel/kernel_test.go @@ -85,8 +85,9 @@ func TestEvictIfSessionFatal(t *testing.T) { // When Execute fails before it acquires a statement handle (here: a nil session // makes new_statement fail), it must still honor the Backend contract — a non-nil, // handle-less Operation that Closes as a no-op (closed=false, no CLOSE_STATEMENT) -// and reports zero AffectedRows. (Parameter binding is exercised live in -// TestKernelE2EParams; a nil-session unit test can't reach the bind path.) +// and reports zero AffectedRows. (A nil-session unit test can't reach the bind +// path: the param mapping is unit-tested hermetically in TestParamBindArg, and +// exercised live end-to-end in TestKernelParamsVsThrift.) func TestExecuteHandleLessOpContract(t *testing.T) { k := &KernelBackend{} // nil session → new_statement fails op, err := k.Execute(context.Background(), backend.ExecRequest{ @@ -177,6 +178,43 @@ func TestExecutionErrorContract(t *testing.T) { } } +// The execute error path must NEVER report retryable, even when the kernel marks +// the failure retryable. This is the post-submission surface (toStatementError +// refuses driver.ErrBadConn here for the same reason): a network/unavailable +// failure seen after the statement was sent may have already committed a +// non-idempotent INSERT/UPDATE/MERGE, so an app keying retry on IsRetryable() would +// double-write. It also matches the Thrift path, which always builds a +// non-retryable execution error. sqlState/queryId must still come through. +func TestExecutionErrorNeverRetryable(t *testing.T) { + o := &kernelOp{} + cause := &KernelError{Code: statusUnavailable, Message: "try again", SQLState: "08000", QueryID: "q-9", Retryable: true} + err := o.ExecutionError(context.Background(), cause) + if err == nil { + t.Fatal("ExecutionError(cause) should not be nil") + } + + var dbExec dbsqlerr.DBExecutionError + if !errors.As(err, &dbExec) { + t.Fatalf("kernel execution error should be a DBExecutionError; got %T", err) + } + // Even though the KernelError is Retryable, the execute path must report false: + // the statement may have committed, so replay is unsafe. + if dbExec.IsRetryable() { + t.Error("IsRetryable() = true on the execute path; want false (a sent statement may have committed — no replay)") + } + // Dropping the retryable signal must not drop sqlState/queryId or the cause. + if dbExec.SqlState() != "08000" { + t.Errorf("SqlState() = %q, want 08000", dbExec.SqlState()) + } + if dbExec.QueryId() != "q-9" { + t.Errorf("QueryId() = %q, want q-9", dbExec.QueryId()) + } + var ke *KernelError + if !errors.As(err, &ke) { + t.Error("the *KernelError cause should remain reachable via errors.As") + } +} + func strPtr(s string) *string { return &s } // The cell/nested rendering (ScanCell and the JSON grammar) now lives in the diff --git a/internal/backend/kernel/operation.go b/internal/backend/kernel/operation.go index 8e4b3f4b..92a6c91b 100644 --- a/internal/backend/kernel/operation.go +++ b/internal/backend/kernel/operation.go @@ -149,19 +149,29 @@ func (k *KernelBackend) execute(ctx context.Context, req backend.ExecRequest) (b op := &kernelOp{backend: k, stmt: stmt, location: k.cfg.Location} if execErr != nil { + // A session-fatal status (expired token, dropped/unavailable session) means + // this conn is unusable: evict it so the pool doesn't hand it out again. This + // runs BEFORE the ctx-cancelled branch below because a session-fatal failure + // racing a ctx cancel is still session-fatal — skipping eviction there would + // leave a dead conn marked valid in the pool. Eviction only flips the conn's + // validity flag; it never turns the error into driver.ErrBadConn (the returned + // error stays a plain toStatementError / ctx error), so database/sql discards + // the conn without re-running the statement (no duplicate write). + k.evictIfSessionFatal(execErr) // Prefer the caller's ctx error when the ctx was cancelled (database/sql - // convention), keeping the kernel error as the cause. + // convention). Wrap BOTH the ctx error and the kernel error so + // errors.Is(err, context.Canceled/DeadlineExceeded) still matches (the + // cancellation contract) AND the *KernelError stays reachable via errors.As — + // otherwise a session-fatal failure racing a cancel would lose its + // sqlstate/queryId, the one handle to what actually went wrong server-side. if ctx.Err() != nil { klog("Execute failed under cancelled ctx: kernelErr=%v ctxErr=%v", execErr, ctx.Err()) op.close() - return op, fmt.Errorf("kernel: execute cancelled: %w", ctx.Err()) + return op, fmt.Errorf("kernel: execute cancelled: %w (kernel error: %w)", ctx.Err(), toStatementError(execErr)) } klog("Execute failed: %v", execErr) - // A session-fatal status (expired token, dropped/unavailable session) means - // this conn is unusable: evict it so the pool doesn't hand it out again. We - // still return the PLAIN error (toStatementError, never ErrBadConn), so the + // We still return the PLAIN error (toStatementError, never ErrBadConn), so the // conn is discarded without database/sql re-running the statement. - k.evictIfSessionFatal(execErr) op.close() return op, fmt.Errorf("kernel: execute: %w", toStatementError(execErr)) } @@ -193,11 +203,15 @@ func (k *KernelBackend) execute(ctx context.Context, req backend.ExecRequest) (b // so the params are set on the fresh statement (set_sql clears any prior binds). func bindParams(stmt *C.kernel_statement_t, params []backend.Param) error { for i, p := range params { - name := newCStrOrNull(p.Name) // empty Name → NULL → positional - typ := newCStr(p.Type) - val := newCStrOrNull("") - if p.Value != nil { - val = newCStr(*p.Value) // non-nil, possibly empty string, is a real value + // The positional/named + NULL/empty-string decision is the pure paramBindArg + // (unit-tested under CGO_ENABLED=0 in bindparams_test.go); this loop only does + // the C-string marshaling of that decision. + a := paramBindArg(p) + name := newCStrOrNull(a.name) // empty name → NULL → positional + typ := newCStr(a.typ) + val := newCStrOrNull("") // NULL pointer: the SQL-NULL wire form + if !a.isNull { + val = newCStr(a.value) // a real value (empty string included), never NULL } err := call(func() C.KernelStatusCode { return C.kernel_statement_bind_parameter(stmt, name.c, typ.c, val.c) @@ -208,7 +222,7 @@ func bindParams(stmt *C.kernel_statement_t, params []backend.Param) error { if err != nil { return fmt.Errorf("param %d (name=%q type=%q): %w", i, p.Name, p.Type, err) } - klog("bound param %d name=%q type=%q null=%v", i, p.Name, p.Type, p.Value == nil) + klog("bound param %d name=%q type=%q null=%v", i, p.Name, p.Type, a.isNull) } return nil } @@ -314,11 +328,21 @@ func (o *kernelOp) close() bool { // ExecutionError wraps cause as the driver's execution error so it satisfies the // same public contract as the Thrift path: errors.Is(err, dbsqlerr.ExecutionError) // matches, and errors.As(err, &dbExecErr) exposes SqlState()/QueryId() — the recipe -// documented in doc.go. The kernel carries BOTH sqlState and the server query id in -// its own *KernelError (not a Thrift TGetOperationStatusResp, and unlike Thrift the -// kernel path has no ctx query id — StatementID() is ""), so pull both out and hand -// them to NewExecutionErrorWithState, keeping the *KernelError as the cause so its -// detail stays reachable via Unwrap. +// documented in doc.go. The kernel carries both sqlState and the server query id in +// its own *KernelError (not a Thrift TGetOperationStatusResp; and on the +// execute-error path there is no exec handle, so StatementID() is "" and the query +// id must come from the KernelError, not ctx), so pull both out and hand them to +// NewExecutionErrorWithState, keeping the *KernelError as the cause so its detail +// stays reachable via Unwrap. +// +// This is exclusively the post-submission (execute / result-read) error path, so it +// deliberately does NOT surface the kernel's Retryable flag: like the Thrift path +// (which always builds a non-retryable execution error here), a failure seen after +// the statement was sent may have already committed a non-idempotent +// INSERT/UPDATE/MERGE server-side, and an app keying retry on IsRetryable() would +// re-run it → silent double write. This mirrors toStatementError, which refuses +// driver.ErrBadConn for the same reason. Connect-phase failures — where the kernel's +// retryable signal IS safe — go through toConnError, not here. func (o *kernelOp) ExecutionError(ctx context.Context, cause error) error { if cause == nil { return nil From 8c6252174e16bedab1827b14ca7ee24475cecaa4 Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Wed, 15 Jul 2026 11:30:09 +0000 Subject: [PATCH 12/24] fix(kernel): wrap unsupported-auth rejection with ErrNotSupportedByKernel MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Follow-up to the review round: the unsupported-authenticator default case in resolveKernelAuth still returned a plain errors.New, while every other unsupported kernel option wraps ErrNotSupportedByKernel and doc.go advertises that errors.Is(err, ErrNotSupportedByKernel) detects any unsupported kernel feature. So this PR shipped a documented contract its own code broke for token-provider / external / federated auth. Wrap it with %w to honor the contract (same fix #403 makes one commit up the stack — matching its wording so the two converge cleanly on rebase). The empty-PAT case stays unwrapped: a missing token is misconfiguration to fix, not a feature the kernel can't honor. Tighten the "non-PAT/non-OAuth authenticator rejected" test to assert errors.Is(err, ErrNotSupportedByKernel) instead of only err != nil, so the contract is pinned rather than documented as an exception. Verified: default-build suite passes, go vet + gofmt clean. Co-authored-by: Isaac Signed-off-by: Mani Kaustubh Mathur --- kernel_config.go | 15 +++++++++++---- kernel_config_test.go | 11 +++++++++-- 2 files changed, 20 insertions(+), 6 deletions(-) diff --git a/kernel_config.go b/kernel_config.go index a478e36e..168c3ab8 100644 --- a/kernel_config.go +++ b/kernel_config.go @@ -134,9 +134,16 @@ func resolveKernelAuth(cfg *config.Config) (kernel.Auth, error) { } return kernel.Auth{Mode: kernel.AuthPAT, Token: token}, nil default: - return kernel.Auth{}, errors.New("databricks: this authenticator is not supported by the kernel backend; " + - "PAT (WithAccessToken) and OAuth M2M/U2M (WithClientCredentials / authType) are supported, but " + - "token-provider, external/static, and federated authenticators are not — " + - "use one of those or the default (Thrift) backend") + // Unsupported authenticator: wrap ErrNotSupportedByKernel so a caller can + // detect the "kernel can't honor this auth" case with errors.Is and fall back + // to the default backend, rather than substring-matching this message. This is + // the same contract every other unsupported kernel option in this file follows + // and that doc.go advertises. (The empty-PAT case above is intentionally NOT + // wrapped — a missing token is misconfiguration to fix, not a feature the + // kernel can't honor.) + return kernel.Auth{}, fmt.Errorf("databricks: this authenticator is %w; "+ + "PAT (WithAccessToken) and OAuth M2M/U2M (WithClientCredentials / authType) are supported, but "+ + "token-provider, external/static, and federated authenticators are not — "+ + "use one of those or the default (Thrift) backend", dbsqlerr.ErrNotSupportedByKernel) } } diff --git a/kernel_config_test.go b/kernel_config_test.go index fb077019..11c492c2 100644 --- a/kernel_config_test.go +++ b/kernel_config_test.go @@ -199,8 +199,15 @@ func TestValidateKernelConfig(t *testing.T) { c := baseKernelConfig() c.AccessToken = "" c.Authenticator = nonPATAuth{} - if _, err := validateKernelConfig(c); err == nil { - t.Error("expected an error for a token-provider/external/federated authenticator") + _, err := validateKernelConfig(c) + if err == nil { + t.Fatal("expected an error for a token-provider/external/federated authenticator") + } + // An unsupported authenticator is a "kernel can't honor this" rejection, so it + // must wrap ErrNotSupportedByKernel like every other unsupported option — the + // contract doc.go advertises for programmatic fallback via errors.Is. + if !errors.Is(err, dbsqlerr.ErrNotSupportedByKernel) { + t.Errorf("unsupported-auth rejection should wrap ErrNotSupportedByKernel, got %v", err) } }) From 5c99a7ae8b957c9af6b20a1c163a308753c883b4 Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Tue, 14 Jul 2026 11:51:27 +0000 Subject: [PATCH 13/24] =?UTF-8?q?feat(kernel):=20richer=20TLS=20Go=20optio?= =?UTF-8?q?ns=20=E2=80=94=20CA=20bundle=20+=20independent=20hostname=20ski?= =?UTF-8?q?p?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Expose the two kernel-only TLS knobs whose C-ABI setters already exist on kernel main (no kernel change needed), via an experimental-option idiom: - WithKernelTrustedCerts(pem) -> kernel_session_config_set_tls_trusted_certs, adding a PEM CA bundle on top of the system roots. Required because the kernel's rustls stack ignores SSL_CERT_FILE, so a custom CA (corporate re-signing proxy / on-prem CA) must be handed over explicitly. - WithKernelSkipHostnameVerify() -> set_tls_skip_hostname_verification, skipping only the hostname check while keeping chain validation (finer-grained than the blanket WithSkipTLSHostVerify). The knobs live on a non-exported config.KernelExperimentalConfig off config.Config (not UserConfig), so they stay off the stable DSN surface (mirroring Node's InternalConnectionOptions / Python's underscore kwargs). The Thrift path rejects a non-nil block loudly at connect rather than silently ignoring it, so a caller who forgets WithUseKernel learns the option had no effect. OpenSession forwards each to the kernel C ABI via a byte-buffer helper (cBytes); a reflective guard (TestKernelExperimentalFieldsClassified) keeps a new field from slipping either path unclassified. mTLS client cert/key (needs an absent kernel C-ABI setter, K5) and the CloudFetch on/off toggle (K3) are deliberately out of scope here — they are tracked separately (PECOBLR-3652 / PECOBLR-3653). Closes PECOBLR-3651. Co-authored-by: Isaac Signed-off-by: Mani Kaustubh Mathur --- connector.go | 53 ++++++++++++++ doc.go | 12 ++++ internal/backend/kernel/backend.go | 55 ++++++++++++++ internal/backend/kernel/cgo.go | 24 +++++++ internal/backend/kernel/kernel_test.go | 25 +++++++ internal/config/config.go | 43 +++++++++++ kernel_backend.go | 8 +++ kernel_experimental_test.go | 99 ++++++++++++++++++++++++++ 8 files changed, 319 insertions(+) create mode 100644 kernel_experimental_test.go diff --git a/connector.go b/connector.go index 1dd89264..b8c79708 100644 --- a/connector.go +++ b/connector.go @@ -4,6 +4,7 @@ import ( "context" "crypto/tls" "database/sql/driver" + "errors" "net/http" "net/url" "regexp" @@ -43,6 +44,16 @@ func (c *connector) Connect(ctx context.Context) (driver.Conn, error) { if c.cfg.UseKernel { be, err = newKernelBackend(ctx, c.cfg) } else { + // The experimental WithKernel* TLS options have no Thrift-path equivalent — + // reject them loudly rather than silently ignore, so a caller who sets a + // trusted-CA bundle / an independent hostname skip and forgets WithUseKernel + // learns the option had no effect instead of connecting with a + // weaker-than-intended (or unconfigured) TLS trust store. + if c.cfg.KernelExperimental != nil { + return nil, errors.New("databricks: the WithKernel* options " + + "(WithKernelTrustedCerts / WithKernelSkipHostnameVerify) require the kernel backend; " + + "add WithUseKernel(true) or remove them") + } be, err = thrift.New(ctx, c.cfg, c.client) } if err != nil { @@ -537,3 +548,45 @@ func WithFederatedTokenProviderAndClientID(baseProvider tokenprovider.TokenProvi } } } + +// ─── Experimental kernel-only options ───────────────────────────────────────── +// +// These configure the SEA-via-kernel backend (WithUseKernel) only; they expose a +// richer TLS surface than the backend-neutral WithSkipTLSHostVerify. They have no +// equivalent on the default (Thrift) path, which rejects them loudly at connect. +// They are deliberately NOT part of the stable DSN/UserConfig surface — they hang +// off config.Config.KernelExperimental (mirroring Node's non-exported +// InternalConnectionOptions). The WithKernel* prefix signals both "kernel-backend +// only" and "experimental" so they read distinctly from the backend-neutral +// options above (e.g. WithSkipTLSHostVerify). + +// kernelExperimental lazily allocates and returns the experimental config block. +func kernelExperimental(c *config.Config) *config.KernelExperimentalConfig { + if c.KernelExperimental == nil { + c.KernelExperimental = &config.KernelExperimentalConfig{} + } + return c.KernelExperimental +} + +// WithKernelTrustedCerts adds a PEM CA-certificate bundle to the kernel's TLS +// trust store on top of the system roots — for a corporate re-signing proxy or an +// on-prem CA. Required (rather than relying on SSL_CERT_FILE) because the kernel's +// rustls stack does not read that environment variable. +// +// EXPERIMENTAL, kernel-only: the default (Thrift) backend rejects this at connect. +func WithKernelTrustedCerts(pem []byte) ConnOption { + return func(c *config.Config) { + kernelExperimental(c).TLSTrustedCertsPEM = pem + } +} + +// WithKernelSkipHostnameVerify skips only the certificate hostname-vs-SNI check on +// the kernel backend, while keeping chain validation. This is finer-grained than +// WithSkipTLSHostVerify, which relaxes both chain and hostname checks. +// +// EXPERIMENTAL, kernel-only: the default (Thrift) backend rejects this at connect. +func WithKernelSkipHostnameVerify() ConnOption { + return func(c *config.Config) { + kernelExperimental(c).TLSSkipHostnameVerify = true + } +} diff --git a/doc.go b/doc.go index 1f3c7e92..dc1632ee 100644 --- a/doc.go +++ b/doc.go @@ -234,6 +234,18 @@ WithMaxRows and retry tuning (WithRetries with a positive limit) are accepted bu not applied on the kernel path: the kernel manages result fetching and retries internally, below the C ABI, with no user-facing knob. +Experimental kernel-only TLS options. The kernel exposes a richer TLS surface than +the backend-neutral WithSkipTLSHostVerify (which relaxes both chain and hostname +checks). These are kernel-only — the default (Thrift) backend rejects them at +connect rather than silently ignoring them — and experimental (the WithKernel* +prefix marks both): + - WithKernelTrustedCerts(pem) adds a PEM CA bundle to the kernel's trust store on + top of the system roots (for a corporate re-signing proxy or an on-prem CA). + Required rather than relying on SSL_CERT_FILE: the kernel's TLS stack (rustls) + does not read that environment variable. + - WithKernelSkipHostnameVerify() skips only the certificate hostname check while + keeping chain validation (finer-grained than WithSkipTLSHostVerify). + Features that live above the backend seam are inherited unchanged: the database/sql connection pool (each connection wraps one kernel session), per-connection telemetry (CREATE_SESSION, EXECUTE_STATEMENT, and DELETE_SESSION diff --git a/internal/backend/kernel/backend.go b/internal/backend/kernel/backend.go index 7c74e19e..e6540081 100644 --- a/internal/backend/kernel/backend.go +++ b/internal/backend/kernel/backend.go @@ -47,6 +47,14 @@ type Config struct { // so the kernel path relaxes both to match. TLSSkipVerify bool + // Experimental kernel-only TLS knobs (from the WithKernel* options). These + // have no Thrift-path equivalent and are set via config.KernelExperimental. + // Empty/false fields are simply not applied. TLSTrustedCertsPEM is a custom CA + // bundle added on top of the system roots; TLSSkipHostnameVerify skips only the + // hostname check (finer-grained than the blanket TLSSkipVerify above). + TLSTrustedCertsPEM []byte + TLSSkipHostnameVerify bool + // ProxyURL configures an HTTP proxy, already resolved for this endpoint from // the same HTTP(S)_PROXY / NO_PROXY environment the Thrift path uses (NO_PROXY // is applied during resolution). Empty leaves the kernel on a direct @@ -166,6 +174,14 @@ func (k *KernelBackend) OpenSession(ctx context.Context) error { } } + // Experimental kernel-only TLS: a custom CA bundle and an independent hostname + // skip (finer-grained than the blanket InsecureSkipVerify above). The kernel's + // rustls stack ignores SSL_CERT_FILE, so a custom CA must be handed over + // explicitly. + if err := k.applyKernelTLS(cfg); err != nil { + return err + } + // Proxy: only when the environment configured one for this endpoint. NO_PROXY // was already applied during resolution, so no bypass list is needed here; // any credentials are carried in the URL userinfo (Go's proxy-env convention), @@ -238,6 +254,31 @@ func (k *KernelBackend) OpenSession(ctx context.Context) error { return nil } +// applyKernelTLS forwards the experimental kernel-only TLS knobs to the session +// config: a custom CA bundle and an independent hostname-skip. Each is a no-op +// when its field is unset, so this is safe to call unconditionally. (mTLS client +// cert/key is a separate follow-up — it needs a kernel C-ABI setter that is not +// yet on kernel main.) +func (k *KernelBackend) applyKernelTLS(cfg *C.KernelSessionConfig) error { + if len(k.cfg.TLSTrustedCertsPEM) > 0 { + ca := newCBytes(k.cfg.TLSTrustedCertsPEM) + defer ca.free() + if err := call(func() C.KernelStatusCode { + return C.kernel_session_config_set_tls_trusted_certs(cfg, ca.ptr, ca.len) + }); err != nil { + return fmt.Errorf("kernel: set_tls_trusted_certs: %w", toConnError(err)) + } + } + if k.cfg.TLSSkipHostnameVerify { + if err := call(func() C.KernelStatusCode { + return C.kernel_session_config_set_tls_skip_hostname_verification(cfg, C.bool(true)) + }); err != nil { + return fmt.Errorf("kernel: set_tls_skip_hostname_verification: %w", toConnError(err)) + } + } + return nil +} + // setAuth applies the resolved auth form to the session config via exactly one // kernel_session_config_set_auth_* call. PAT and M2M are plain value setters; U2M // records the client id / redirect port / scopes and the kernel owns the browser @@ -301,6 +342,20 @@ func trySetAuth(auth Auth) error { return k.setAuth(cfg) } +// trySetKernelTLS allocates a throwaway session config, applies the experimental +// TLS knobs from cfg to it, and frees it — the analogous test seam to trySetAuth, +// so a tagged test can exercise the real byte-buffer cgo setter (trusted certs) +// and the hostname-skip setter end to end. Not used in production. +func trySetKernelTLS(cfg Config) error { + var c *C.KernelSessionConfig + if err := call(func() C.KernelStatusCode { return C.kernel_session_config_new(&c) }); err != nil { + return fmt.Errorf("config_new: %w", err) + } + defer C.kernel_session_config_free(c) + k := &KernelBackend{cfg: cfg} + return k.applyKernelTLS(c) +} + // applyInitialNamespace runs USE CATALOG / USE SCHEMA to select the configured // initial namespace, since the kernel C ABI exposes no catalog/schema setter. // Identifiers are backtick-quoted (quoteIdent) so arbitrary names are safe. A diff --git a/internal/backend/kernel/cgo.go b/internal/backend/kernel/cgo.go index eaab36f2..7aabd2cf 100644 --- a/internal/backend/kernel/cgo.go +++ b/internal/backend/kernel/cgo.go @@ -222,3 +222,27 @@ func (s cStr) free() { C.free(unsafe.Pointer(s.c)) } } + +// cBytes wraps C.CBytes with a guaranteed free, for the byte-buffer setters (e.g. +// a PEM CA bundle) that take a (*C.uint8_t, C.size_t) pair. The kernel copies the +// bytes into owned Rust memory on receipt, so freeing right after the call is +// safe. An empty slice yields a NULL pointer + 0 length (the setters reject that, +// which is what we want — an empty buffer is never valid). +// Use: cb := newCBytes(b); defer cb.free(); ...C.fn(cb.ptr, cb.len)... +type cBytes struct { + ptr *C.uint8_t + len C.size_t +} + +func newCBytes(b []byte) cBytes { + if len(b) == 0 { + return cBytes{} + } + return cBytes{ptr: (*C.uint8_t)(C.CBytes(b)), len: C.size_t(len(b))} +} + +func (b cBytes) free() { + if b.ptr != nil { + C.free(unsafe.Pointer(b.ptr)) + } +} diff --git a/internal/backend/kernel/kernel_test.go b/internal/backend/kernel/kernel_test.go index e3e62f54..8582a501 100644 --- a/internal/backend/kernel/kernel_test.go +++ b/internal/backend/kernel/kernel_test.go @@ -45,6 +45,31 @@ func TestSetAuthByMode(t *testing.T) { } } +// TestSetKernelTLS exercises the real cgo setters for the experimental kernel-only +// TLS knobs (the byte-buffer trusted-CA bundle + the hostname-skip bool) via the +// trySetKernelTLS seam — proving the (*C.uint8_t, C.size_t) marshalling and the C +// signatures. A failure here means the field→setter wiring or a C signature +// drifted. +func TestSetKernelTLS(t *testing.T) { + ca := []byte("-----BEGIN CERTIFICATE-----\nca\n-----END CERTIFICATE-----\n") + cases := []struct { + name string + cfg Config + }{ + {"trusted certs only", Config{TLSTrustedCertsPEM: ca}}, + {"skip hostname only", Config{TLSSkipHostnameVerify: true}}, + {"both together", Config{TLSTrustedCertsPEM: ca, TLSSkipHostnameVerify: true}}, + {"none (no-op)", Config{}}, + } + for _, c := range cases { + t.Run(c.name, func(t *testing.T) { + if err := trySetKernelTLS(c.cfg); err != nil { + t.Errorf("applyKernelTLS(%s) = %v, want nil", c.name, err) + } + }) + } +} + // The pure error-classifier tests (TestIsBadConnection, TestIsSessionFatal, // TestToConnError, TestToStatementErrorNeverBadConn) live in the untagged // errors_classify_test.go so they run under CGO_ENABLED=0. The tests below need a diff --git a/internal/config/config.go b/internal/config/config.go index 58c92123..2113a64f 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -48,6 +48,48 @@ type Config struct { ThriftTransport string ThriftProtocolVersion cli_service.TProtocolVersion ThriftDebugClientProtocol bool + + // 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 + // exposes over its C ABI. It lives here on Config, NOT on UserConfig, so it + // stays off the stable exported/DSN surface (the same treatment TLSConfig and + // ArrowConfig get). nil means no experimental option was set. The Thrift + // backend rejects a non-nil value loudly; the kernel backend forwards it to + // the kernel C ABI. Mirrors Node's non-exported InternalConnectionOptions / + // Python's underscore-prefixed kwargs. + KernelExperimental *KernelExperimentalConfig +} + +// KernelExperimentalConfig holds the experimental, kernel-only connection knobs +// set via the WithKernel* options. Every field is either forwarded to the kernel +// C ABI (kernel backend) or rejected (Thrift backend) — never silently dropped; +// the exhaustiveness guard TestKernelExperimentalFieldsClassified asserts this so +// a newly-added field can't slip through unclassified. +type KernelExperimentalConfig struct { + // TLSTrustedCertsPEM is a PEM CA bundle added to the kernel's trust store on + // top of the system roots (maps to kernel_session_config_set_tls_trusted_certs). + // Needed because the kernel's rustls stack ignores SSL_CERT_FILE, so a custom + // CA must be handed to the kernel explicitly. + TLSTrustedCertsPEM []byte + // TLSSkipHostnameVerify skips only the certificate hostname check, independent + // of the blanket InsecureSkipVerify (which relaxes both chain and hostname). + // Maps to kernel_session_config_set_tls_skip_hostname_verification. + TLSSkipHostnameVerify bool +} + +// DeepCopy returns a deep copy of the experimental config, or nil for a nil +// receiver. The byte slice is copied so a mutation of the copy can't reach back +// into the original (the connector may DeepCopy the whole Config per conn). +func (k *KernelExperimentalConfig) DeepCopy() *KernelExperimentalConfig { + if k == nil { + return nil + } + cp := &KernelExperimentalConfig{TLSSkipHostnameVerify: k.TLSSkipHostnameVerify} + if k.TLSTrustedCertsPEM != nil { + cp.TLSTrustedCertsPEM = append([]byte(nil), k.TLSTrustedCertsPEM...) + } + return cp } // MetricViewMetadataConfKey is the server session conf that enables metric-view @@ -110,6 +152,7 @@ func (c *Config) DeepCopy() *Config { ThriftTransport: c.ThriftTransport, ThriftProtocolVersion: c.ThriftProtocolVersion, ThriftDebugClientProtocol: c.ThriftDebugClientProtocol, + KernelExperimental: c.KernelExperimental.DeepCopy(), } } diff --git a/kernel_backend.go b/kernel_backend.go index 9c8092ce..470e2424 100644 --- a/kernel_backend.go +++ b/kernel_backend.go @@ -47,6 +47,14 @@ func newKernelBackend(_ context.Context, cfg *config.Config) (backend.Backend, e if cfg.TLSConfig != nil && cfg.TLSConfig.InsecureSkipVerify { kc.TLSSkipVerify = true } + // Experimental kernel-only TLS knobs (WithKernelTrustedCerts / + // WithKernelSkipHostnameVerify), if any. These have no Thrift-path equivalent + // (the connector rejects them on that path) and are forwarded verbatim to the + // kernel C ABI in OpenSession. + if ke := cfg.KernelExperimental; ke != nil { + kc.TLSTrustedCertsPEM = ke.TLSTrustedCertsPEM + kc.TLSSkipHostnameVerify = ke.TLSSkipHostnameVerify + } // Proxy: the Thrift path uses http.ProxyFromEnvironment; mirror it by reading // the same HTTP(S)_PROXY / NO_PROXY environment for the kernel. kc.ProxyURL = proxyForEndpoint(cfg) diff --git a/kernel_experimental_test.go b/kernel_experimental_test.go new file mode 100644 index 00000000..baaf95db --- /dev/null +++ b/kernel_experimental_test.go @@ -0,0 +1,99 @@ +package dbsql + +import ( + "reflect" + "testing" + + "github.com/databricks/databricks-sql-go/internal/config" +) + +// This file is untagged (no cgo) so its tests — the experimental-option +// classification guard, the option→config wiring, and the Thrift fail-loud +// signal — run in the default CGO_ENABLED=0 build alongside the other +// kernel_config tests. + +// kernelExperimentalFieldDisposition records how each experimental (kernel-only) +// option is handled. Every experimental field is forwarded to the kernel C ABI +// (there is no "inert" experimental knob — they exist precisely because the kernel +// supports them) and rejected on the Thrift path (the connector fails loud when +// KernelExperimental is non-nil). A new field on config.KernelExperimentalConfig +// without an entry here fails TestKernelExperimentalFieldsClassified, forcing a +// deliberate decision and a setter in KernelBackend.OpenSession so it can't be +// silently dropped. +var kernelExperimentalFieldDisposition = map[string]string{ + "TLSTrustedCertsPEM": "forwarded", // set_tls_trusted_certs + "TLSSkipHostnameVerify": "forwarded", // set_tls_skip_hostname_verification +} + +func TestKernelExperimentalFieldsClassified(t *testing.T) { + tp := reflect.TypeOf(config.KernelExperimentalConfig{}) + classified := make(map[string]bool, tp.NumField()) + for i := 0; i < tp.NumField(); i++ { + name := tp.Field(i).Name + classified[name] = true + if _, ok := kernelExperimentalFieldDisposition[name]; !ok { + t.Errorf("config.KernelExperimentalConfig field %q is not classified. Add it to "+ + "kernelExperimentalFieldDisposition and wire a setter in KernelBackend.OpenSession / "+ + "newKernelBackend so it isn't silently dropped on the kernel path.", name) + } + } + for name := range kernelExperimentalFieldDisposition { + if !classified[name] { + t.Errorf("kernelExperimentalFieldDisposition has %q but config.KernelExperimentalConfig no longer does; remove it", name) + } + } +} + +// The experimental WithKernel* options are rejected on the default (Thrift) path +// so a caller who forgets WithUseKernel learns the option had no effect. The +// option builders set config.KernelExperimental; the connector's backend-selection +// branch is what rejects it. We assert the option→config wiring here (a non-nil +// KernelExperimental after applying a WithKernel* option is the signal the Thrift +// branch keys off). +func TestWithKernelTLSOptionsSetExperimental(t *testing.T) { + cases := []struct { + name string + opt ConnOption + verify func(*config.KernelExperimentalConfig) bool + }{ + {"trusted certs", WithKernelTrustedCerts([]byte("ca")), func(k *config.KernelExperimentalConfig) bool { + return string(k.TLSTrustedCertsPEM) == "ca" + }}, + {"skip hostname", WithKernelSkipHostnameVerify(), func(k *config.KernelExperimentalConfig) bool { + return k.TLSSkipHostnameVerify + }}, + } + for _, c := range cases { + t.Run(c.name, func(t *testing.T) { + cfg := config.WithDefaults() + c.opt(cfg) + if cfg.KernelExperimental == nil { + t.Fatalf("%s: KernelExperimental should be non-nil after the option", c.name) + } + if !c.verify(cfg.KernelExperimental) { + t.Errorf("%s: option did not set the expected field(s): %+v", c.name, cfg.KernelExperimental) + } + }) + } +} + +// DeepCopy must copy the CA byte slice, not alias it — the connector may DeepCopy +// the whole Config per conn, and a shared backing array would let one conn's +// mutation reach another. +func TestKernelExperimentalDeepCopy(t *testing.T) { + orig := &config.KernelExperimentalConfig{ + TLSTrustedCertsPEM: []byte("ca-bundle"), + TLSSkipHostnameVerify: true, + } + cp := orig.DeepCopy() + if cp == nil || string(cp.TLSTrustedCertsPEM) != "ca-bundle" || !cp.TLSSkipHostnameVerify { + t.Fatalf("DeepCopy lost data: %+v", cp) + } + cp.TLSTrustedCertsPEM[0] = 'X' + if orig.TLSTrustedCertsPEM[0] == 'X' { + t.Error("DeepCopy aliased the CA byte slice; a copy mutation reached the original") + } + if (*config.KernelExperimentalConfig)(nil).DeepCopy() != nil { + t.Error("nil.DeepCopy() should be nil") + } +} From d37ce5001ed85b682579fb281213f8cbbd045efd Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Wed, 15 Jul 2026 09:15:02 +0000 Subject: [PATCH 14/24] Address review: sentinel error, defensive PEM copy, MITM warning, Thrift-reject test Follow-up on the four review findings for the richer-TLS kernel options: - Add ErrRequiresKernelBackend sentinel (mirror of ErrNotSupportedByKernel) and wrap the Thrift-path rejection with %w, so callers detect it via errors.Is instead of message text. Documented in doc.go. - WithKernelTrustedCerts copies the PEM defensively (matching DeepCopy) so a caller mutating the slice between NewConnector and Connect can't change the trust store. - Add a MITM WARNING to WithKernelSkipHostnameVerify, consistent with WithSkipTLSHostVerify. - Add TestWithKernelOptionsRejectedOnThriftPath: end-to-end Connect() reject on the Thrift path, asserting on the sentinel. Co-authored-by: Isaac Signed-off-by: Mani Kaustubh Mathur --- connector.go | 23 ++++++++++++++++----- doc.go | 5 +++++ errors/errors.go | 9 ++++++++ kernel_experimental_test.go | 41 +++++++++++++++++++++++++++++++++++++ 4 files changed, 73 insertions(+), 5 deletions(-) diff --git a/connector.go b/connector.go index b8c79708..3838cb9d 100644 --- a/connector.go +++ b/connector.go @@ -4,7 +4,7 @@ import ( "context" "crypto/tls" "database/sql/driver" - "errors" + "fmt" "net/http" "net/url" "regexp" @@ -16,6 +16,7 @@ import ( "github.com/databricks/databricks-sql-go/auth/pat" "github.com/databricks/databricks-sql-go/auth/tokenprovider" "github.com/databricks/databricks-sql-go/driverctx" + dbsqlerr "github.com/databricks/databricks-sql-go/errors" "github.com/databricks/databricks-sql-go/internal/backend" "github.com/databricks/databricks-sql-go/internal/backend/thrift" "github.com/databricks/databricks-sql-go/internal/client" @@ -50,9 +51,9 @@ func (c *connector) Connect(ctx context.Context) (driver.Conn, error) { // learns the option had no effect instead of connecting with a // weaker-than-intended (or unconfigured) TLS trust store. if c.cfg.KernelExperimental != nil { - return nil, errors.New("databricks: the WithKernel* options " + - "(WithKernelTrustedCerts / WithKernelSkipHostnameVerify) require the kernel backend; " + - "add WithUseKernel(true) or remove them") + return nil, fmt.Errorf("databricks: a WithKernel* option "+ + "(WithKernelTrustedCerts / WithKernelSkipHostnameVerify) %w; "+ + "add WithUseKernel(true) or remove it", dbsqlerr.ErrRequiresKernelBackend) } be, err = thrift.New(ctx, c.cfg, c.client) } @@ -576,13 +577,25 @@ func kernelExperimental(c *config.Config) *config.KernelExperimentalConfig { // EXPERIMENTAL, kernel-only: the default (Thrift) backend rejects this at connect. func WithKernelTrustedCerts(pem []byte) ConnOption { return func(c *config.Config) { - kernelExperimental(c).TLSTrustedCertsPEM = pem + // Copy defensively (matching KernelExperimentalConfig.DeepCopy) so a + // caller mutating pem between NewConnector and Connect can't change the + // trust store out from under us. + if len(pem) > 0 { + kernelExperimental(c).TLSTrustedCertsPEM = append([]byte(nil), pem...) + } else { + kernelExperimental(c).TLSTrustedCertsPEM = pem + } } } // WithKernelSkipHostnameVerify skips only the certificate hostname-vs-SNI check on // the kernel backend, while keeping chain validation. This is finer-grained than // WithSkipTLSHostVerify, which relaxes both chain and hostname checks. +// WARNING: +// Skipping hostname verification still weakens TLS: a certificate issued by a +// trusted CA for a different host will be accepted, opening a machine-in-the-middle +// vector. Only use this when the hostname is an internal private-link hostname that +// legitimately differs from the certificate's subject. // // EXPERIMENTAL, kernel-only: the default (Thrift) backend rejects this at connect. func WithKernelSkipHostnameVerify() ConnOption { diff --git a/doc.go b/doc.go index dc1632ee..635fd93d 100644 --- a/doc.go +++ b/doc.go @@ -246,6 +246,11 @@ prefix marks both): - WithKernelSkipHostnameVerify() skips only the certificate hostname check while keeping chain validation (finer-grained than WithSkipTLSHostVerify). +Setting either without WithUseKernel makes Connect fail with an error wrapping the +sentinel ErrRequiresKernelBackend (the mirror of ErrNotSupportedByKernel), so the +"add WithUseKernel or remove it" case is detectable with errors.Is rather than by +matching message text. + Features that live above the backend seam are inherited unchanged: the database/sql connection pool (each connection wraps one kernel session), per-connection telemetry (CREATE_SESSION, EXECUTE_STATEMENT, and DELETE_SESSION diff --git a/errors/errors.go b/errors/errors.go index 88bce5e2..8eac7877 100644 --- a/errors/errors.go +++ b/errors/errors.go @@ -73,6 +73,15 @@ var ErrNotSupportedByKernel error = errors.New("not supported by the kernel back // ErrNotSupportedByKernel — instead of matching on message text. 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 +// 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 +// option"; this one means "this option requires the kernel". Lets a caller detect +// the case programmatically instead of matching on message text. +var ErrRequiresKernelBackend error = errors.New("requires the SEA-via-kernel backend") + // Base interface for driver errors type DBError interface { // Descriptive message describing the error diff --git a/kernel_experimental_test.go b/kernel_experimental_test.go index baaf95db..75c76988 100644 --- a/kernel_experimental_test.go +++ b/kernel_experimental_test.go @@ -1,10 +1,14 @@ package dbsql import ( + "context" + "errors" "reflect" "testing" "github.com/databricks/databricks-sql-go/internal/config" + + dbsqlerr "github.com/databricks/databricks-sql-go/errors" ) // This file is untagged (no cgo) so its tests — the experimental-option @@ -77,6 +81,43 @@ func TestWithKernelTLSOptionsSetExperimental(t *testing.T) { } } +// End-to-end: setting a WithKernel* option WITHOUT WithUseKernel must make +// Connect fail loud on the default (Thrift) path rather than silently connect +// with a weaker-than-intended trust store. This exercises the connector's +// reject branch (not just option→config wiring) and asserts the error wraps the +// ErrRequiresKernelBackend sentinel so callers can detect it with errors.Is — +// the mirror of TestKernelBackendNotCompiledIn. Runs in the default +// CGO_ENABLED=0 build (no kernel linked in). +func TestWithKernelOptionsRejectedOnThriftPath(t *testing.T) { + cases := []struct { + name string + opt ConnOption + }{ + {"trusted certs", WithKernelTrustedCerts([]byte("ca"))}, + {"skip hostname", WithKernelSkipHostnameVerify()}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + c, err := NewConnector( + WithServerHostname("example.cloud.databricks.com"), + WithPort(443), + WithHTTPPath("/sql/1.0/endpoints/12346a5b5b0e123a"), + WithAccessToken("supersecret"), + tc.opt, + ) + if err != nil { + t.Fatalf("NewConnector: %v", err) + } + // No WithUseKernel — the Thrift path must reject the kernel-only option. + if _, err = c.Connect(context.Background()); err == nil { + t.Fatal("Connect should reject a WithKernel* option on the Thrift path, got nil") + } else if !errors.Is(err, dbsqlerr.ErrRequiresKernelBackend) { + t.Errorf("error should wrap ErrRequiresKernelBackend; got: %v", err) + } + }) + } +} + // DeepCopy must copy the CA byte slice, not alias it — the connector may DeepCopy // the whole Config per conn, and a shared backing array would let one conn's // mutation reach another. From f741ac00319d644c83e93bcb7e0ba75aabbe219f Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Wed, 15 Jul 2026 09:47:14 +0000 Subject: [PATCH 15/24] Address review nits: mirror sentinel doc, PEM-copy test, DeepCopy coverage Non-blocking follow-ups: - doc.go: mention ErrRequiresKernelBackend in the main Errors section, symmetric with ErrNotSupportedByKernel. - Add TestWithKernelTrustedCertsCopiesPEM: mutating the caller's slice after the option must not change stored config (option-set counterpart to the DeepCopy aliasing test). - TestConfig_DeepCopy: set KernelExperimental in the all-values case so Config.DeepCopy's wiring of that field is exercised, and assert it's copied not aliased. Co-authored-by: Isaac Signed-off-by: Mani Kaustubh Mathur --- doc.go | 5 +++++ internal/config/config_test.go | 12 ++++++++++++ kernel_experimental_test.go | 20 ++++++++++++++++++++ 3 files changed, 37 insertions(+) diff --git a/doc.go b/doc.go index 635fd93d..a3511c3b 100644 --- a/doc.go +++ b/doc.go @@ -349,6 +349,11 @@ on the error message text: // retry with the default backend (omit WithUseKernel). } +Conversely, a kernel-only option set without WithUseKernel is rejected on the +default (Thrift) path with an error wrapping the mirror sentinel +ErrRequiresKernelBackend, detectable the same way with +errors.Is(err, dbsqlerr.ErrRequiresKernelBackend). + Example usage: import ( diff --git a/internal/config/config_test.go b/internal/config/config_test.go index a4b9b54f..c4160c25 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -769,11 +769,23 @@ func TestConfig_DeepCopy(t *testing.T) { ThriftTransport: "http", ThriftProtocolVersion: cli_service.TProtocolVersion_SPARK_CLI_SERVICE_PROTOCOL_V8, ThriftDebugClientProtocol: false, + KernelExperimental: &KernelExperimentalConfig{ + TLSTrustedCertsPEM: []byte("ca-bundle"), + TLSSkipHostnameVerify: true, + }, } cfg_copy := cfg.DeepCopy() if !reflect.DeepEqual(cfg, cfg_copy) { t.Errorf("DeepCopy() = %v, want %v", cfg_copy, cfg) } + // The experimental block must be deep-copied, not aliased. + if cfg_copy.KernelExperimental == cfg.KernelExperimental { + t.Error("DeepCopy aliased KernelExperimental pointer") + } + cfg_copy.KernelExperimental.TLSTrustedCertsPEM[0] = 'X' + if cfg.KernelExperimental.TLSTrustedCertsPEM[0] == 'X' { + t.Error("DeepCopy aliased the KernelExperimental CA byte slice") + } }) } diff --git a/kernel_experimental_test.go b/kernel_experimental_test.go index 75c76988..1c3534d2 100644 --- a/kernel_experimental_test.go +++ b/kernel_experimental_test.go @@ -118,6 +118,26 @@ func TestWithKernelOptionsRejectedOnThriftPath(t *testing.T) { } } +// WithKernelTrustedCerts must copy the PEM defensively, so a caller mutating its +// slice after applying the option (but before Connect) can't change the stored +// trust store. This is the option-set counterpart to TestKernelExperimentalDeepCopy +// (which covers the per-conn DeepCopy path). +func TestWithKernelTrustedCertsCopiesPEM(t *testing.T) { + pem := []byte("ca-bundle") + cfg := config.WithDefaults() + WithKernelTrustedCerts(pem)(cfg) + + // Mutate the caller's slice after the option ran. + pem[0] = 'X' + + if cfg.KernelExperimental == nil { + t.Fatal("KernelExperimental should be non-nil after WithKernelTrustedCerts") + } + if got := string(cfg.KernelExperimental.TLSTrustedCertsPEM); got != "ca-bundle" { + t.Errorf("WithKernelTrustedCerts aliased the caller's slice; stored %q, want %q", got, "ca-bundle") + } +} + // DeepCopy must copy the CA byte slice, not alias it — the connector may DeepCopy // the whole Config per conn, and a shared backing array would let one conn's // mutation reach another. From 58ed62a0de52cf21451225ed84ea0bee46ea7712 Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Tue, 14 Jul 2026 13:44:21 +0000 Subject: [PATCH 16/24] feat(kernel): route kernel logging through the driver log level + correlate MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Unify kernel-backend logging onto the driver's shared logger so the one knob that controls the rest of the driver — DATABRICKS_LOG_LEVEL / dbsql.SetLogLevel — governs it, and each binding line carries the structured connId/corrId/queryId fields that let it be correlated in a multi-connection process. Addresses vikrantpuppala's review on #393 (klog + the kernel Rust subscriber emitted unstructured lines straight to os.Stderr, uncorrelatable and gated on a separate DBSQL_KERNEL_DEBUG rather than the driver log level). Three parts: - Level: klog now emits through logger.Logger.Debug() instead of raw fmt.Fprintf(os.Stderr, ...). A cheap kernelDebugOff() front gate (GetLevel() > Debug) short-circuits before any formatting/allocation, so it stays a true no-op at the default Warn level — including during benchmarks. Replaces the old DBSQL_KERNEL_DEBUG bool + the caller-less KernelDebugEnabled. - Correlation: new klogCtx(ctx, ...) pulls connId/corrId/queryId off ctx via logger.WithContext (the exact idiom the Thrift/conn path uses); the conn layer already stuffs those IDs into ctx before calling the backend. Every hot-path site (execute, nextBatch, newKernelRows, Open/CloseSession) uses it; ctx-less sites degrade to klog. The WithContext allocation is behind the same up-front level gate, so it never runs below Debug. - Kernel Rust logs: initKernelLogging maps the driver level -> the kernel_init_logging level string (kernelLogLevel), so DATABRICKS_LOG_LEVEL drives the kernel's Rust logs too. DBSQL_KERNEL_DEBUG is kept as an advanced override (forces the subscriber on, defers to RUST_LOG for kernel verbosity). Known limitation (documented): the kernel's Rust lines are still plain text and do not yet carry the structured fields — that needs the kernel log-callback ABI (PECOBLR-3654 / K4). Only the Go binding lines are structured today. Tests: TestKernelLogLevel pins the level mapping; TestKernelLogNoAllocWhenOff uses testing.AllocsPerRun to prove klog/klogCtx allocate 0 times at Warn level (guards the hot-path zero-cost guarantee). Closes PECOBLR-3650. Co-authored-by: Isaac Signed-off-by: Mani Kaustubh Mathur --- doc.go | 34 ++++-- internal/backend/kernel/backend.go | 8 +- internal/backend/kernel/cgo.go | 150 ++++++++++++++++++------- internal/backend/kernel/kernel_test.go | 55 +++++++++ internal/backend/kernel/operation.go | 16 +-- internal/backend/kernel/rows.go | 8 +- 6 files changed, 202 insertions(+), 69 deletions(-) diff --git a/doc.go b/doc.go index a3511c3b..3378cdc3 100644 --- a/doc.go +++ b/doc.go @@ -186,16 +186,30 @@ the C header) under the cgo link path; `make build-kernel` does both steps: In a build without the tag, WithUseKernel(true) returns a clear error at connect time rather than silently using Thrift. -To see the kernel's own logs interleaved with the driver's, set DBSQL_KERNEL_DEBUG -to any non-empty value. That single flag turns on the driver's binding step tracer -AND installs the kernel's Rust log subscriber, so both write to the same stderr in -execution order. It is off by default (and must stay off during benchmarks — debug -logging perturbs latency). The kernel's verbosity is then controlled by RUST_LOG, -which the kernel honors directly; filter on the target databricks::sql::kernel -(note the colons — it is the kernel's explicit log target, NOT the crate module -path databricks_sql_kernel, which would match nothing): - - # kernel logs only: +Kernel-backend logging follows the same knob as the rest of the driver. The +binding-layer trace (session/statement/batch steps) is emitted through the shared +logger at Debug level, so DATABRICKS_LOG_LEVEL=debug or dbsql.SetLogLevel("debug") +turns it on — no separate switch — and each line carries the same structured +connId/corrId/queryId fields as the driver's other logs, so it can be correlated in +a multi-connection process. At the default Warn level the lines are suppressed with +no cost (zerolog no-ops a below-level event), including during benchmarks. + + dbsql.SetLogLevel("debug") // or DATABRICKS_LOG_LEVEL=debug + +The driver's log level is also mapped into the kernel's own Rust log subscriber, so +the same knob governs the kernel's internal (Rust) logs; they write to stderr and +interleave with the driver's output. Note the kernel's Rust lines are currently +plain text and do NOT yet carry the structured connId/corrId/queryId fields — that +needs a kernel log-callback ABI (tracked separately); only the Go binding lines are +structured today. + +For finer control of the kernel's Rust verbosity independent of the driver level, +set DBSQL_KERNEL_DEBUG to any non-empty value: it forces the kernel subscriber on +and defers to RUST_LOG for the kernel's threshold. Filter on the target +databricks::sql::kernel (note the colons — the kernel's explicit log target, NOT +the crate module path databricks_sql_kernel, which would match nothing): + + # kernel logs only, at the kernel's own verbosity: DBSQL_KERNEL_DEBUG=1 RUST_LOG=databricks::sql::kernel=debug ./your_app 2>&1 # kernel logs plus its HTTP stack (hyper/reqwest): DBSQL_KERNEL_DEBUG=1 RUST_LOG=debug ./your_app 2>&1 diff --git a/internal/backend/kernel/backend.go b/internal/backend/kernel/backend.go index e6540081..ad20c97c 100644 --- a/internal/backend/kernel/backend.go +++ b/internal/backend/kernel/backend.go @@ -113,7 +113,7 @@ func (k *KernelBackend) OpenSession(ctx context.Context) error { return err } initKernelLogging() - klog("OpenSession host=%s httpPath=%s warehouse=%s", k.cfg.Host, k.cfg.HTTPPath, k.cfg.WarehouseID) + klogCtx(ctx, "OpenSession host=%s httpPath=%s warehouse=%s", k.cfg.Host, k.cfg.HTTPPath, k.cfg.WarehouseID) var cfg *C.KernelSessionConfig if err := call(func() C.KernelStatusCode { return C.kernel_session_config_new(&cfg) }); err != nil { @@ -243,14 +243,14 @@ func (k *KernelBackend) OpenSession(ctx context.Context) error { // is logged (via lastError's Warn) rather than silently discarded — mirroring // CloseSession. The namespace error is authoritative and returned as-is. if closeErr := call(func() C.KernelStatusCode { return C.kernel_session_close(sess) }); closeErr != nil { - klog("close after initial-namespace failure also failed: %v", closeErr) + klogCtx(ctx, "close after initial-namespace failure also failed: %v", closeErr) } k.session = nil k.valid = false return err } - klog("OpenSession OK session=%s", k.sessionID) + klogCtx(ctx, "OpenSession OK session=%s", k.sessionID) return nil } @@ -400,7 +400,7 @@ func (k *KernelBackend) CloseSession(ctx context.Context) error { if k.session == nil { return nil } - klog("CloseSession session=%s", k.sessionID) + klogCtx(ctx, "CloseSession session=%s", k.sessionID) err := call(func() C.KernelStatusCode { return C.kernel_session_close(k.session) }) k.session = nil k.valid = false diff --git a/internal/backend/kernel/cgo.go b/internal/backend/kernel/cgo.go index 7aabd2cf..ae154593 100644 --- a/internal/backend/kernel/cgo.go +++ b/internal/backend/kernel/cgo.go @@ -34,44 +34,73 @@ package kernel import "C" import ( + "context" "fmt" "os" "runtime" "sync" "unsafe" + "github.com/databricks/databricks-sql-go/driverctx" "github.com/databricks/databricks-sql-go/logger" + "github.com/rs/zerolog" ) // ─── Debug logging ─────────────────────────────────────────────────────────── // -// Gated on DBSQL_KERNEL_DEBUG so it is OFF by default and, in particular, OFF -// during benchmarks (debug logging perturbs latency). Every binding step logs -// through klog when enabled — entry, status codes, handle addresses, batch -// counts — which is what makes a failing e2e cheap to diagnose. The same flag -// also installs the kernel's own Rust (tracing) logs via initKernelLogging, so -// binding and kernel logs interleave on stderr as one stream. -var kdebug = os.Getenv("DBSQL_KERNEL_DEBUG") != "" - -// Deferred (tracked): klog writes raw to stderr, NOT through the driver's -// logger.Logger, so its lines carry no connId/corrId/queryId and can't be -// correlated in a multi-conn process; it's also gated on its own -// DBSQL_KERNEL_DEBUG rather than the driver's DATABRICKS_LOG_LEVEL knob. Unifying -// kernel logging onto logger.Logger (which no-ops below its level, so no -// benchmark cost) is a tracked logging-unification follow-up, kept separate from -// this PR because it changes the debug-logging surface. +// Binding-level debug lines (entry, status codes, handle addresses, batch counts) +// go through the driver's shared logger.Logger at Debug level, so the SAME knob +// that controls the rest of the driver — DATABRICKS_LOG_LEVEL / dbsql.SetLogLevel +// — turns kernel binding logs on and off, and each line carries the structured +// connId/corrId/queryId fields (via klogCtx) that let it be correlated in a +// multi-conn process. zerolog no-ops a Debug() event when the level is above debug, +// so this stays zero-work (no allocation, no formatting) at the default Warn level +// — including during benchmarks, which run at the default level. +// +// DBSQL_KERNEL_DEBUG is retained only as an advanced override for the *kernel's own +// Rust logs* (see initKernelLogging): it forces the kernel subscriber on and lets +// RUST_LOG tune the kernel's verbosity independently. It no longer gates the Go +// binding lines — those follow the driver log level. + +// kernelDebugOff reports whether the driver log level is above Debug, i.e. kernel +// binding lines would be discarded. This is the cheap front gate for klog/klogCtx: +// GetLevel() is a plain field read, whereas building the event (and, for klogCtx, +// logger.WithContext → zerolog's With(), which eagerly make([]byte, 0, 500)s and +// formats the correlation keys) allocates BEFORE .Debug() consults the level. So we +// must short-circuit here to keep the hot path (per-batch nextBatch) allocation-free +// at the default Warn level — including during benchmarks. zerolog's own .Debug() +// no-op is not enough because the argument expressions run first. +func kernelDebugOff() bool { return logger.Logger.GetLevel() > zerolog.DebugLevel } + +// klog emits a binding-level debug line with no request context. Prefer klogCtx +// where a ctx is in scope so the line carries connId/corrId/queryId. Gated by the +// driver log level (Debug): a no-op — no formatting, no allocation — at the default. func klog(format string, args ...any) { - if !kdebug { + if kernelDebugOff() { return } - fmt.Fprintf(os.Stderr, "[kernel] "+format+"\n", args...) + logger.Logger.Debug().Msgf("[kernel] "+format, args...) } -// KernelDebugEnabled reports whether binding-level debug logging is on (i.e. -// DBSQL_KERNEL_DEBUG is set). Exposed so a benchmark can assert the flag is off -// before measuring — debug logging perturbs latency; there is no such benchmark -// in-tree yet, so this currently has no callers. -func KernelDebugEnabled() bool { return kdebug } +// klogCtx is klog with request correlation: it pulls connId/corrId/queryId off ctx +// (the conn layer stuffs them in before calling the backend, exactly as the Thrift +// path does) so a kernel binding line can be joined to the rest of a request's logs +// in a multi-conn process. A nil ctx degrades to klog. Level-gated up front so the +// logger.WithContext allocation never happens below Debug (see kernelDebugOff). +func klogCtx(ctx context.Context, format string, args ...any) { + if kernelDebugOff() { + return + } + if ctx == nil { + logger.Logger.Debug().Msgf("[kernel] "+format, args...) + return + } + logger.WithContext( + driverctx.ConnIdFromContext(ctx), + driverctx.CorrelationIdFromContext(ctx), + driverctx.QueryIdFromContext(ctx), + ).Debug().Msgf("[kernel] "+format, args...) +} // initLoggingOnce guards kernel_init_logging, which is process-wide and // first-call-wins in the kernel. We install the kernel subscriber lazily on the @@ -79,37 +108,72 @@ func KernelDebugEnabled() bool { return kdebug } // kernel session installs nothing. var initLoggingOnce sync.Once -// initKernelLogging turns on the kernel's own Rust (tracing) logs, gated on the -// same DBSQL_KERNEL_DEBUG flag as klog so both are OFF by default and, in -// particular, OFF during benchmarks — the subscriber is never installed in a -// benchmark run. level=NULL lets the kernel honor RUST_LOG (default warn); -// file_path=NULL sends kernel logs to stderr so they interleave with klog on one -// stream. Best-effort: Internal (e.g. the host already installed a global -// subscriber) is a documented, benign outcome — logged, never fatal to connect. +// initKernelLogging turns on the kernel's own Rust (tracing) logs and points their +// verbosity at the driver's log level, so DATABRICKS_LOG_LEVEL drives both the Go +// binding lines and the kernel's Rust lines from one knob. The mapped level is +// passed to kernel_init_logging (Go zerolog level → the kernel's OFF/ERROR/WARN/ +// INFO/DEBUG/TRACE string); DBSQL_KERNEL_DEBUG forces the subscriber on with a +// NULL level so the kernel honors RUST_LOG instead (the advanced override for +// tuning kernel-only verbosity). file_path=NULL sends kernel logs to stderr so they +// interleave with the driver's own output. +// +// Best-effort: an Internal return (e.g. the host already installed a global +// subscriber) is a documented, benign outcome — logged at Warn, never fatal to +// connect. Only installed when the effective level is at or below the kernel's +// threshold, so a driver left at the default Warn level (benchmarks included) that +// hasn't set DBSQL_KERNEL_DEBUG still installs the subscriber only at Warn — the +// kernel then emits nothing below Warn, so there is no hot-path cost. // // Scope caveat: the kernel subscriber is PROCESS-WIDE, first-call-wins, and never -// uninstalled. Because kdebug is read once from DBSQL_KERNEL_DEBUG at package -// init, the switch is process-global, not per-connection: in a long-lived -// multi-tenant process, the first kernel session opened with the flag set -// installs the subscriber (and its stderr output) for ALL subsequent kernel -// sessions in that process, and there is no way to scope it to one connection or -// turn it off afterward. The "off during benchmarks" guarantee therefore depends -// on DBSQL_KERNEL_DEBUG being unset before package init. Making this a -// per-process runtime knob (rather than an init-time env read) is tracked with -// the logging-unification follow-up. +// uninstalled — in a long-lived multi-tenant process the first kernel session's +// level/destination applies to ALL subsequent kernel sessions, with no way to +// re-scope or turn it off afterward. That is a kernel-ABI property, not a Go one. func initKernelLogging() { - if !kdebug { - return - } initLoggingOnce.Do(func() { + // DBSQL_KERNEL_DEBUG override: force the subscriber on and let RUST_LOG tune + // it (level=NULL). Otherwise map the driver's current log level in, so the + // one DATABRICKS_LOG_LEVEL knob governs kernel verbosity too. + var level cStr + if os.Getenv("DBSQL_KERNEL_DEBUG") != "" { + level = cStr{c: nil} // NULL → kernel honors RUST_LOG + } else { + level = newCStr(kernelLogLevel(logger.Logger.GetLevel())) + defer level.free() + } if err := call(func() C.KernelStatusCode { - return C.kernel_init_logging(nil, nil) + return C.kernel_init_logging(level.c, nil) }); err != nil { - klog("kernel_init_logging: %v (kernel logs unavailable; proceeding)", err) + // The kernel subscriber didn't install (commonly: the host already + // installed a global tracing subscriber). Non-fatal — surface it through + // the shared logger so it's visible without a separate stderr scrape. + logger.Logger.Warn().Msgf("databricks: kernel_init_logging: %v (kernel logs unavailable; proceeding)", err) } }) } +// kernelLogLevel maps a zerolog level to the level string kernel_init_logging +// accepts (OFF/ERROR/WARN/INFO/DEBUG/TRACE), so the driver's log level drives the +// kernel's Rust logs. An unrecognized level falls back to WARN (the kernel's own +// default), matching the driver's default. +func kernelLogLevel(l zerolog.Level) string { + switch l { + case zerolog.TraceLevel: + return "TRACE" + case zerolog.DebugLevel: + return "DEBUG" + case zerolog.InfoLevel: + return "INFO" + case zerolog.WarnLevel: + return "WARN" + case zerolog.ErrorLevel, zerolog.FatalLevel, zerolog.PanicLevel: + return "ERROR" + case zerolog.Disabled: + return "OFF" + default: + return "WARN" + } +} + // call runs a fallible kernel entry point and, on a non-Success status, reads // the kernel's thread-local last error into a Go error. // diff --git a/internal/backend/kernel/kernel_test.go b/internal/backend/kernel/kernel_test.go index 8582a501..d67da97c 100644 --- a/internal/backend/kernel/kernel_test.go +++ b/internal/backend/kernel/kernel_test.go @@ -8,8 +8,11 @@ import ( "errors" "testing" + "github.com/databricks/databricks-sql-go/driverctx" dbsqlerr "github.com/databricks/databricks-sql-go/errors" "github.com/databricks/databricks-sql-go/internal/backend" + "github.com/databricks/databricks-sql-go/logger" + "github.com/rs/zerolog" ) // setAuth maps each Auth mode to exactly one kernel_session_config_set_auth_* @@ -70,6 +73,58 @@ func TestSetKernelTLS(t *testing.T) { } } +// kernelLogLevel maps the driver's zerolog level to the kernel_init_logging level +// string, so DATABRICKS_LOG_LEVEL drives the kernel's Rust logs too. Pin the +// mapping (incl. the fatal/panic→ERROR collapse and the default→WARN fallback) so +// a level the kernel would reject can't silently ship. +func TestKernelLogLevel(t *testing.T) { + cases := []struct { + in zerolog.Level + want string + }{ + {zerolog.TraceLevel, "TRACE"}, + {zerolog.DebugLevel, "DEBUG"}, + {zerolog.InfoLevel, "INFO"}, + {zerolog.WarnLevel, "WARN"}, + {zerolog.ErrorLevel, "ERROR"}, + {zerolog.FatalLevel, "ERROR"}, + {zerolog.PanicLevel, "ERROR"}, + {zerolog.Disabled, "OFF"}, + {zerolog.NoLevel, "WARN"}, // unrecognized → the kernel's own default + } + for _, c := range cases { + if got := kernelLogLevel(c.in); got != c.want { + t.Errorf("kernelLogLevel(%v) = %q, want %q", c.in, got, c.want) + } + } +} + +// klog / klogCtx must be allocation-free at the default (above-Debug) log level — +// the "no hot-path cost, safe during benchmarks" guarantee. klogCtx is the one that +// matters: logger.WithContext eagerly allocates (zerolog's With() does +// make([]byte, 0, 500)) BEFORE the .Debug() gate, so without the up-front +// kernelDebugOff() guard this would allocate ~500 B per call — per Arrow batch on the +// nextBatch hot path. Guards against reintroducing that (an Isaac Review MAJOR). +func TestKernelLogNoAllocWhenOff(t *testing.T) { + prev := logger.Logger.GetLevel() + if err := logger.SetLogLevel("warn"); err != nil { // the default level + t.Fatal(err) + } + defer logger.SetLogLevel(prev.String()) //nolint:errcheck // restoring a known-good level + + ctx := driverctx.NewContextWithConnId(context.Background(), "conn-1") + if n := testing.AllocsPerRun(100, func() { + klog("hot path %d", 1) + }); n != 0 { + t.Errorf("klog allocated %v times at Warn level, want 0", n) + } + if n := testing.AllocsPerRun(100, func() { + klogCtx(ctx, "hot path %d", 1) + }); n != 0 { + t.Errorf("klogCtx allocated %v times at Warn level, want 0", n) + } +} + // The pure error-classifier tests (TestIsBadConnection, TestIsSessionFatal, // TestToConnError, TestToStatementErrorNeverBadConn) live in the untagged // errors_classify_test.go so they run under CGO_ENABLED=0. The tests below need a diff --git a/internal/backend/kernel/operation.go b/internal/backend/kernel/operation.go index 92a6c91b..1a550caa 100644 --- a/internal/backend/kernel/operation.go +++ b/internal/backend/kernel/operation.go @@ -37,7 +37,7 @@ func (k *KernelBackend) execute(ctx context.Context, req backend.ExecRequest) (b // Log the SQL length, not the text: query bodies can carry PII/secrets in // WHERE/INSERT/SET, and this goes to stderr. Matches the driver's own // debuglog convention (conn.ExecContext logs sql.len=%d). - klog("Execute sql.len=%d", len(req.Query)) + klogCtx(ctx, "Execute sql.len=%d", len(req.Query)) var stmt *C.kernel_statement_t if err := call(func() C.KernelStatusCode { @@ -76,7 +76,7 @@ func (k *KernelBackend) execute(ctx context.Context, req backend.ExecRequest) (b if err := call(func() C.KernelStatusCode { return C.kernel_statement_canceller_new(stmt, &canceller) }); err != nil { - klog("canceller_new failed (proceeding without cancel): %v", err) + klogCtx(ctx, "canceller_new failed (proceeding without cancel): %v", err) canceller = nil } @@ -111,11 +111,11 @@ func (k *KernelBackend) execute(ctx context.Context, req backend.ExecRequest) (b return default: } - klog("ctx.Done (%v) → firing canceller (with retry until dispatched)", ctx.Err()) + klogCtx(ctx, "ctx.Done (%v) → firing canceller (with retry until dispatched)", ctx.Err()) ticker := time.NewTicker(250 * time.Millisecond) defer ticker.Stop() if fireCancel(canceller) { - klog("cancel dispatched on first fire") + klogCtx(ctx, "cancel dispatched on first fire") return } for { @@ -124,7 +124,7 @@ func (k *KernelBackend) execute(ctx context.Context, req backend.ExecRequest) (b return case <-ticker.C: if fireCancel(canceller) { - klog("cancel dispatched, watcher stopping") + klogCtx(ctx, "cancel dispatched, watcher stopping") return } } @@ -165,11 +165,11 @@ func (k *KernelBackend) execute(ctx context.Context, req backend.ExecRequest) (b // otherwise a session-fatal failure racing a cancel would lose its // sqlstate/queryId, the one handle to what actually went wrong server-side. if ctx.Err() != nil { - klog("Execute failed under cancelled ctx: kernelErr=%v ctxErr=%v", execErr, ctx.Err()) + klogCtx(ctx, "Execute failed under cancelled ctx: kernelErr=%v ctxErr=%v", execErr, ctx.Err()) op.close() return op, fmt.Errorf("kernel: execute cancelled: %w (kernel error: %w)", ctx.Err(), toStatementError(execErr)) } - klog("Execute failed: %v", execErr) + klogCtx(ctx, "Execute failed: %v", execErr) // We still return the PLAIN error (toStatementError, never ErrBadConn), so the // conn is discarded without database/sql re-running the statement. op.close() @@ -192,7 +192,7 @@ func (k *KernelBackend) execute(ctx context.Context, req backend.ExecRequest) (b if qid := C.kernel_executed_statement_query_id(exec); qid != nil { op.statementID = C.GoString(qid) // deep-copies out of the borrowed C string } - klog("Execute OK stmt=%p exec=%p affectedRows=%d statementID=%q", stmt, exec, op.affectedRows, op.statementID) + klogCtx(ctx, "Execute OK stmt=%p exec=%p affectedRows=%d statementID=%q", stmt, exec, op.affectedRows, op.statementID) return op, nil } diff --git a/internal/backend/kernel/rows.go b/internal/backend/kernel/rows.go index 82df1920..b03290f4 100644 --- a/internal/backend/kernel/rows.go +++ b/internal/backend/kernel/rows.go @@ -76,7 +76,7 @@ func newKernelRows(ctx context.Context, op *kernelOp, stream *C.kernel_result_st for i, f := range fields { r.cols[i] = f.Name } - klog("newKernelRows: %d columns", len(r.cols)) + klogCtx(ctx, "newKernelRows: %d columns", len(r.cols)) return r, nil } @@ -101,7 +101,7 @@ func (r *kernelRows) Close() error { if r.op != nil { r.op.close() } - klog("kernelRows closed") + klogCtx(r.ctx, "kernelRows closed") return nil } @@ -161,7 +161,7 @@ func (r *kernelRows) nextBatch() error { } if carr.release == nil { r.eof = true - klog("nextBatch: EOF") + klogCtx(r.ctx, "nextBatch: EOF") return io.EOF } // Zero-copy import. The kernel exports self-contained batches (Rust to_ffi @@ -189,6 +189,6 @@ func (r *kernelRows) nextBatch() error { // reports — a fabricated number would be worse than a truthful zero. r.callbacks.OnChunkFetched(r.chunkCount, 0, 0, 0, 0) } - klog("nextBatch: %d rows (chunk %d)", rec.NumRows(), r.chunkCount) + klogCtx(r.ctx, "nextBatch: %d rows (chunk %d)", rec.NumRows(), r.chunkCount) return nil } From c0ecc6d54d320bb3a0465e0415bdc96fe5c06704 Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Wed, 15 Jul 2026 07:06:54 +0000 Subject: [PATCH 17/24] fix(kernel): tighten logging docs + fatal/panic->OFF + correlation test MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Address the review comments on #401 — the Go binding logs now follow the live driver logger, but the Rust subscriber is process-wide/first-call-wins and stderr-only, so several "one unified knob" claims were stronger than the implementation. Fixes, all in the kernel backend (databricks_kernel tag; the default pure-Go build is unaffected): - doc.go: Rust logs are stderr-only and NOT routed through logger.SetLogOutput; Rust verbosity is fixed at the first kernel session (set the level before the first connection); "each line carries connId/corrId/queryId" softened to request-scoped lines where a ctx is in scope (bindParams / operation teardown stay uncorrelated). - cgo.go: kernelLogLevel maps FatalLevel/PanicLevel -> OFF (not ERROR), so the Rust subscriber is never louder than the driver the user configured — at driver=fatal the Go side suppresses even Error() lines. Tightened the initKernelLogging comment to state the first-session level sampling and the stderr-only sink, and dropped the inaccurate "only installed when the level is at or below the threshold" sentence. - kernel_test.go: TestKernelLogLevel updated for fatal/panic->OFF; new TestKernelLogCtxEmitsCorrelation captures logger output at debug level and asserts klogCtx emits the message + connId/corrId/queryId, then asserts silence at warn — the positive-behavior half the alloc test can't see. Verified: CGO_ENABLED=0 build+vet+test green; CGO_ENABLED=1 -tags databricks_kernel build+vet green, kernel package tests pass (incl. all three log tests). Co-authored-by: Isaac Signed-off-by: Mani Kaustubh Mathur --- doc.go | 34 +++++++++--- internal/backend/kernel/cgo.go | 27 ++++++--- internal/backend/kernel/kernel_test.go | 76 ++++++++++++++++++++++++-- 3 files changed, 116 insertions(+), 21 deletions(-) diff --git a/doc.go b/doc.go index 3378cdc3..6856d4da 100644 --- a/doc.go +++ b/doc.go @@ -189,19 +189,35 @@ time rather than silently using Thrift. Kernel-backend logging follows the same knob as the rest of the driver. The binding-layer trace (session/statement/batch steps) is emitted through the shared logger at Debug level, so DATABRICKS_LOG_LEVEL=debug or dbsql.SetLogLevel("debug") -turns it on — no separate switch — and each line carries the same structured -connId/corrId/queryId fields as the driver's other logs, so it can be correlated in -a multi-connection process. At the default Warn level the lines are suppressed with -no cost (zerolog no-ops a below-level event), including during benchmarks. +turns it on — no separate switch — and it lands in the same sink the rest of the +driver uses (including a custom one set via logger.SetLogOutput). Request-scoped +binding lines carry the same structured connId/corrId/queryId fields as the +driver's other logs where a context is in scope (session open/close, execute, +result-batch fetch), so they can be correlated in a multi-connection process; a few +context-less lines (parameter binds, operation teardown) are emitted without those +fields. At the default Warn level the lines are suppressed with no cost (zerolog +no-ops a below-level event), including during benchmarks. dbsql.SetLogLevel("debug") // or DATABRICKS_LOG_LEVEL=debug The driver's log level is also mapped into the kernel's own Rust log subscriber, so -the same knob governs the kernel's internal (Rust) logs; they write to stderr and -interleave with the driver's output. Note the kernel's Rust lines are currently -plain text and do NOT yet carry the structured connId/corrId/queryId fields — that -needs a kernel log-callback ABI (tracked separately); only the Go binding lines are -structured today. +the same knob governs the kernel's internal (Rust) verbosity. Two caveats apply to +the Rust logs specifically, both consequences of the kernel's process-wide logging +ABI rather than the Go layer: + + - They are written to stderr directly and are NOT affected by + logger.SetLogOutput — an app that routes the driver's Go logs to a file or an + app logger will still find the Rust lines on stderr, not interleaved with that + sink. + - The Rust verbosity is fixed at the level in effect when the FIRST kernel + session in the process is opened (the kernel subscriber is process-wide and + installed once); a later dbsql.SetLogLevel changes the Go binding lines but not + the already-installed Rust subscriber. Set the level before opening the first + kernel connection to control the Rust logs. + +The kernel's Rust lines are also currently plain text and do NOT yet carry the +structured connId/corrId/queryId fields — that needs a kernel log-callback ABI +(tracked separately); only the Go binding lines are structured today. For finer control of the kernel's Rust verbosity independent of the driver level, set DBSQL_KERNEL_DEBUG to any non-empty value: it forces the kernel subscriber on diff --git a/internal/backend/kernel/cgo.go b/internal/backend/kernel/cgo.go index ae154593..15303067 100644 --- a/internal/backend/kernel/cgo.go +++ b/internal/backend/kernel/cgo.go @@ -114,20 +114,25 @@ var initLoggingOnce sync.Once // passed to kernel_init_logging (Go zerolog level → the kernel's OFF/ERROR/WARN/ // INFO/DEBUG/TRACE string); DBSQL_KERNEL_DEBUG forces the subscriber on with a // NULL level so the kernel honors RUST_LOG instead (the advanced override for -// tuning kernel-only verbosity). file_path=NULL sends kernel logs to stderr so they -// interleave with the driver's own output. +// tuning kernel-only verbosity). file_path=NULL sends kernel logs to stderr — the +// kernel ABI has no sink hook, so the Rust lines always go to stderr and are NOT +// routed through logger.SetLogOutput (unlike the Go binding lines). // // Best-effort: an Internal return (e.g. the host already installed a global // subscriber) is a documented, benign outcome — logged at Warn, never fatal to -// connect. Only installed when the effective level is at or below the kernel's -// threshold, so a driver left at the default Warn level (benchmarks included) that -// hasn't set DBSQL_KERNEL_DEBUG still installs the subscriber only at Warn — the -// kernel then emits nothing below Warn, so there is no hot-path cost. +// connect. The subscriber installs at whatever level is mapped in; a driver left at +// the default Warn level (benchmarks included, and never having set +// DBSQL_KERNEL_DEBUG) installs it at WARN, so the kernel emits nothing below Warn +// and there is no hot-path cost. // // Scope caveat: the kernel subscriber is PROCESS-WIDE, first-call-wins, and never // uninstalled — in a long-lived multi-tenant process the first kernel session's // level/destination applies to ALL subsequent kernel sessions, with no way to // re-scope or turn it off afterward. That is a kernel-ABI property, not a Go one. +// A direct consequence: the driver level is sampled HERE, once, at the first kernel +// session — a later dbsql.SetLogLevel re-levels the Go binding lines (klog/klogCtx +// re-read GetLevel per call) but NOT the already-installed Rust subscriber. Set the +// level before opening the first kernel connection to govern the Rust logs. func initKernelLogging() { initLoggingOnce.Do(func() { // DBSQL_KERNEL_DEBUG override: force the subscriber on and let RUST_LOG tune @@ -155,6 +160,12 @@ func initKernelLogging() { // accepts (OFF/ERROR/WARN/INFO/DEBUG/TRACE), so the driver's log level drives the // kernel's Rust logs. An unrecognized level falls back to WARN (the kernel's own // default), matching the driver's default. +// +// FatalLevel/PanicLevel map to OFF, not ERROR: at those levels the Go driver +// suppresses even its own Error() lines, so the kernel's Rust subscriber must not +// be louder than the driver the user configured — the kernel has no fatal/panic +// threshold, and emitting ERROR lines there would leak stderr output a user who set +// DATABRICKS_LOG_LEVEL=fatal explicitly asked to silence. func kernelLogLevel(l zerolog.Level) string { switch l { case zerolog.TraceLevel: @@ -165,9 +176,9 @@ func kernelLogLevel(l zerolog.Level) string { return "INFO" case zerolog.WarnLevel: return "WARN" - case zerolog.ErrorLevel, zerolog.FatalLevel, zerolog.PanicLevel: + case zerolog.ErrorLevel: return "ERROR" - case zerolog.Disabled: + case zerolog.FatalLevel, zerolog.PanicLevel, zerolog.Disabled: return "OFF" default: return "WARN" diff --git a/internal/backend/kernel/kernel_test.go b/internal/backend/kernel/kernel_test.go index d67da97c..775b5854 100644 --- a/internal/backend/kernel/kernel_test.go +++ b/internal/backend/kernel/kernel_test.go @@ -3,9 +3,12 @@ package kernel import ( + "bytes" "context" "database/sql/driver" + "encoding/json" "errors" + "os" "testing" "github.com/databricks/databricks-sql-go/driverctx" @@ -75,8 +78,10 @@ func TestSetKernelTLS(t *testing.T) { // kernelLogLevel maps the driver's zerolog level to the kernel_init_logging level // string, so DATABRICKS_LOG_LEVEL drives the kernel's Rust logs too. Pin the -// mapping (incl. the fatal/panic→ERROR collapse and the default→WARN fallback) so -// a level the kernel would reject can't silently ship. +// mapping (incl. the fatal/panic→OFF collapse and the default→WARN fallback) so a +// level the kernel would reject can't silently ship. fatal/panic map to OFF, not +// ERROR: the Go driver suppresses even Error() lines at those levels, so the Rust +// subscriber must stay at least as quiet as the driver the user configured. func TestKernelLogLevel(t *testing.T) { cases := []struct { in zerolog.Level @@ -87,8 +92,8 @@ func TestKernelLogLevel(t *testing.T) { {zerolog.InfoLevel, "INFO"}, {zerolog.WarnLevel, "WARN"}, {zerolog.ErrorLevel, "ERROR"}, - {zerolog.FatalLevel, "ERROR"}, - {zerolog.PanicLevel, "ERROR"}, + {zerolog.FatalLevel, "OFF"}, // driver suppresses Error() here → kernel stays silent, not louder + {zerolog.PanicLevel, "OFF"}, {zerolog.Disabled, "OFF"}, {zerolog.NoLevel, "WARN"}, // unrecognized → the kernel's own default } @@ -125,6 +130,69 @@ func TestKernelLogNoAllocWhenOff(t *testing.T) { } } +// TestKernelLogCtxEmitsCorrelation proves the positive half of the contract the +// alloc test can't see: at debug level klogCtx actually EMITS through the shared +// logger AND attaches connId/corrId/queryId, and at the default warn level it emits +// nothing. Without this, a regression that silently stopped emitting or stopped +// attaching the fields would still pass TestKernelLogNoAllocWhenOff (0 allocs is +// also what "emit nothing" looks like). +func TestKernelLogCtxEmitsCorrelation(t *testing.T) { + prev := logger.Logger.GetLevel() + var buf bytes.Buffer + logger.SetLogOutput(&buf) + // Restore the package defaults (stderr sink, prior level) so we don't leak the + // buffer/level into sibling tests. + defer func() { + logger.SetLogOutput(os.Stderr) + _ = logger.SetLogLevel(prev.String()) + }() + + ctx := driverctx.NewContextWithConnId(context.Background(), "conn-1") + ctx = driverctx.NewContextWithCorrelationId(ctx, "corr-2") + ctx = driverctx.NewContextWithQueryId(ctx, "query-3") + + // Debug level: the line is emitted and carries all three correlation fields. + if err := logger.SetLogLevel("debug"); err != nil { + t.Fatal(err) + } + klogCtx(ctx, "step %s", "execute") + + var rec struct { + Level string `json:"level"` + Message string `json:"message"` + ConnID string `json:"connId"` + CorrID string `json:"corrId"` + QueryID string `json:"queryId"` + } + line := bytes.TrimSpace(buf.Bytes()) + if len(line) == 0 { + t.Fatal("klogCtx emitted nothing at debug level, want one line") + } + if err := json.Unmarshal(line, &rec); err != nil { + t.Fatalf("klogCtx output is not the expected JSON log line: %v (line: %s)", err, line) + } + if rec.Level != "debug" { + t.Errorf("level = %q, want %q", rec.Level, "debug") + } + if rec.Message != "[kernel] step execute" { + t.Errorf("message = %q, want %q", rec.Message, "[kernel] step execute") + } + if rec.ConnID != "conn-1" || rec.CorrID != "corr-2" || rec.QueryID != "query-3" { + t.Errorf("correlation fields = {connId:%q corrId:%q queryId:%q}, want {conn-1 corr-2 query-3}", + rec.ConnID, rec.CorrID, rec.QueryID) + } + + // Warn level (the default): nothing is emitted. + buf.Reset() + if err := logger.SetLogLevel("warn"); err != nil { + t.Fatal(err) + } + klogCtx(ctx, "step %s", "execute") + if got := bytes.TrimSpace(buf.Bytes()); len(got) != 0 { + t.Errorf("klogCtx emitted %q at warn level, want nothing", got) + } +} + // The pure error-classifier tests (TestIsBadConnection, TestIsSessionFatal, // TestToConnError, TestToStatementErrorNeverBadConn) live in the untagged // errors_classify_test.go so they run under CGO_ENABLED=0. The tests below need a From 7acc85f953d8886e3bd779c0f1521d4cd87ac79b Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Wed, 15 Jul 2026 08:22:55 +0000 Subject: [PATCH 18/24] test(kernel): add resolveKernelLogArg seam + untagged level tests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Follow-up to the #401 review (finding 5, thread on kernel_test.go): the level- resolution logic initKernelLogging feeds to kernel_init_logging was documented but not tested. Extract the pure decision so it can be pinned by CI without cgo. - New untagged logging_level.go holds kernelLogLevel (moved out of cgo.go) and a new resolveKernelLogArg() (level string, useNULL bool): the NULL-on- DBSQL_KERNEL_DEBUG override vs the mapped driver level. Mirrors the existing errors_classify.go pattern (pure logic, no build tag) so its tests run under CGO_ENABLED=0. cgo.go's initKernelLogging is now a thin caller of it. - New untagged logging_level_test.go: TestKernelLogLevel (moved) + new TestResolveKernelLogArg, which pins the two invariants the PR justifies in prose — DBSQL_KERNEL_DEBUG (non-empty) yields useNULL=true so the kernel honors RUST_LOG, and otherwise the driver level is mapped in (incl. fatal→OFF). Empty value is treated as unset. - Drops the now-unused os import from cgo.go and zerolog from kernel_test.go. Not covered (documented, cost/value flips): that the string is actually handed to C.kernel_init_logging and the sync.Once first-call-wins — asserting those needs a function-pointer seam over the cgo call plus resetting a package global, to test one assignment + stdlib sync.Once. Verified: CGO_ENABLED=0 build+vet+test green (both new tests run here); CGO_ENABLED=1 -tags databricks_kernel build+vet green, all four log tests pass. Co-authored-by: Isaac Signed-off-by: Mani Kaustubh Mathur --- internal/backend/kernel/cgo.go | 44 ++-------- internal/backend/kernel/kernel_test.go | 31 +------ internal/backend/kernel/logging_level.go | 67 ++++++++++++++ internal/backend/kernel/logging_level_test.go | 87 +++++++++++++++++++ 4 files changed, 163 insertions(+), 66 deletions(-) create mode 100644 internal/backend/kernel/logging_level.go create mode 100644 internal/backend/kernel/logging_level_test.go diff --git a/internal/backend/kernel/cgo.go b/internal/backend/kernel/cgo.go index 15303067..cd604121 100644 --- a/internal/backend/kernel/cgo.go +++ b/internal/backend/kernel/cgo.go @@ -36,7 +36,6 @@ import "C" import ( "context" "fmt" - "os" "runtime" "sync" "unsafe" @@ -135,16 +134,14 @@ var initLoggingOnce sync.Once // level before opening the first kernel connection to govern the Rust logs. func initKernelLogging() { initLoggingOnce.Do(func() { - // DBSQL_KERNEL_DEBUG override: force the subscriber on and let RUST_LOG tune - // it (level=NULL). Otherwise map the driver's current log level in, so the - // one DATABRICKS_LOG_LEVEL knob governs kernel verbosity too. + // resolveKernelLogArg decides the level (or NULL for the DBSQL_KERNEL_DEBUG + // override, which lets the kernel honor RUST_LOG). The pure decision lives in + // logging_level.go so it's unit-tested without cgo. var level cStr - if os.Getenv("DBSQL_KERNEL_DEBUG") != "" { - level = cStr{c: nil} // NULL → kernel honors RUST_LOG - } else { - level = newCStr(kernelLogLevel(logger.Logger.GetLevel())) + if lvl, useNULL := resolveKernelLogArg(); !useNULL { + level = newCStr(lvl) defer level.free() - } + } // else level stays {c: nil} → NULL → kernel honors RUST_LOG if err := call(func() C.KernelStatusCode { return C.kernel_init_logging(level.c, nil) }); err != nil { @@ -156,35 +153,6 @@ func initKernelLogging() { }) } -// kernelLogLevel maps a zerolog level to the level string kernel_init_logging -// accepts (OFF/ERROR/WARN/INFO/DEBUG/TRACE), so the driver's log level drives the -// kernel's Rust logs. An unrecognized level falls back to WARN (the kernel's own -// default), matching the driver's default. -// -// FatalLevel/PanicLevel map to OFF, not ERROR: at those levels the Go driver -// suppresses even its own Error() lines, so the kernel's Rust subscriber must not -// be louder than the driver the user configured — the kernel has no fatal/panic -// threshold, and emitting ERROR lines there would leak stderr output a user who set -// DATABRICKS_LOG_LEVEL=fatal explicitly asked to silence. -func kernelLogLevel(l zerolog.Level) string { - switch l { - case zerolog.TraceLevel: - return "TRACE" - case zerolog.DebugLevel: - return "DEBUG" - case zerolog.InfoLevel: - return "INFO" - case zerolog.WarnLevel: - return "WARN" - case zerolog.ErrorLevel: - return "ERROR" - case zerolog.FatalLevel, zerolog.PanicLevel, zerolog.Disabled: - return "OFF" - default: - return "WARN" - } -} - // call runs a fallible kernel entry point and, on a non-Success status, reads // the kernel's thread-local last error into a Go error. // diff --git a/internal/backend/kernel/kernel_test.go b/internal/backend/kernel/kernel_test.go index 775b5854..a293f915 100644 --- a/internal/backend/kernel/kernel_test.go +++ b/internal/backend/kernel/kernel_test.go @@ -15,7 +15,6 @@ import ( dbsqlerr "github.com/databricks/databricks-sql-go/errors" "github.com/databricks/databricks-sql-go/internal/backend" "github.com/databricks/databricks-sql-go/logger" - "github.com/rs/zerolog" ) // setAuth maps each Auth mode to exactly one kernel_session_config_set_auth_* @@ -76,33 +75,9 @@ func TestSetKernelTLS(t *testing.T) { } } -// kernelLogLevel maps the driver's zerolog level to the kernel_init_logging level -// string, so DATABRICKS_LOG_LEVEL drives the kernel's Rust logs too. Pin the -// mapping (incl. the fatal/panic→OFF collapse and the default→WARN fallback) so a -// level the kernel would reject can't silently ship. fatal/panic map to OFF, not -// ERROR: the Go driver suppresses even Error() lines at those levels, so the Rust -// subscriber must stay at least as quiet as the driver the user configured. -func TestKernelLogLevel(t *testing.T) { - cases := []struct { - in zerolog.Level - want string - }{ - {zerolog.TraceLevel, "TRACE"}, - {zerolog.DebugLevel, "DEBUG"}, - {zerolog.InfoLevel, "INFO"}, - {zerolog.WarnLevel, "WARN"}, - {zerolog.ErrorLevel, "ERROR"}, - {zerolog.FatalLevel, "OFF"}, // driver suppresses Error() here → kernel stays silent, not louder - {zerolog.PanicLevel, "OFF"}, - {zerolog.Disabled, "OFF"}, - {zerolog.NoLevel, "WARN"}, // unrecognized → the kernel's own default - } - for _, c := range cases { - if got := kernelLogLevel(c.in); got != c.want { - t.Errorf("kernelLogLevel(%v) = %q, want %q", c.in, got, c.want) - } - } -} +// TestKernelLogLevel and TestResolveKernelLogArg — the pure level-resolution tests — +// live in the untagged logging_level_test.go so they run under CGO_ENABLED=0. The +// tests below exercise klog/klogCtx and so need the cgo build. // klog / klogCtx must be allocation-free at the default (above-Debug) log level — // the "no hot-path cost, safe during benchmarks" guarantee. klogCtx is the one that diff --git a/internal/backend/kernel/logging_level.go b/internal/backend/kernel/logging_level.go new file mode 100644 index 00000000..a4a7d330 --- /dev/null +++ b/internal/backend/kernel/logging_level.go @@ -0,0 +1,67 @@ +package kernel + +import ( + "os" + + "github.com/databricks/databricks-sql-go/logger" + "github.com/rs/zerolog" +) + +// This file is intentionally NOT behind the `cgo && databricks_kernel` build tag. +// It holds the pure level-resolution logic that decides what verbosity the kernel's +// Rust subscriber is initialized with — the zerolog→kernel level mapping and the +// DBSQL_KERNEL_DEBUG override. Keeping it out of the cgo file lets its tests run in +// the default CGO_ENABLED=0 build (matching errors_classify.go), so the load-bearing +// decisions — that fatal/panic collapse to OFF, and that DBSQL_KERNEL_DEBUG yields a +// NULL level so the kernel honors RUST_LOG — are pinned by CI rather than only by +// comments. cgo.go's initKernelLogging is the thin cgo caller: it turns the result +// into a C string (or NULL) and hands it to kernel_init_logging. + +// kernelLogLevel maps a zerolog level to the level string kernel_init_logging +// accepts (OFF/ERROR/WARN/INFO/DEBUG/TRACE), so the driver's log level drives the +// kernel's Rust logs. An unrecognized level falls back to WARN (the kernel's own +// default), matching the driver's default. +// +// FatalLevel/PanicLevel map to OFF, not ERROR: at those levels the Go driver +// suppresses even its own Error() lines, so the kernel's Rust subscriber must not +// be louder than the driver the user configured — the kernel has no fatal/panic +// threshold, and emitting ERROR lines there would leak stderr output a user who set +// DATABRICKS_LOG_LEVEL=fatal explicitly asked to silence. +func kernelLogLevel(l zerolog.Level) string { + switch l { + case zerolog.TraceLevel: + return "TRACE" + case zerolog.DebugLevel: + return "DEBUG" + case zerolog.InfoLevel: + return "INFO" + case zerolog.WarnLevel: + return "WARN" + case zerolog.ErrorLevel: + return "ERROR" + case zerolog.FatalLevel, zerolog.PanicLevel, zerolog.Disabled: + return "OFF" + default: + return "WARN" + } +} + +// resolveKernelLogArg decides the level string to pass to kernel_init_logging and +// whether to pass NULL instead. It is the pure core of initKernelLogging, split out +// so it can be unit-tested without cgo. +// +// DBSQL_KERNEL_DEBUG (any non-empty value) is the advanced override: it returns +// useNULL=true so the cgo caller passes a NULL level, which makes the kernel honor +// RUST_LOG for its own verbosity — independent of the driver level, and with the +// per-target filtering (e.g. the hyper/reqwest HTTP stack) that the single mapped +// level can't express. The kernel reads RUST_LOG ONLY when the level is NULL; a +// non-NULL level shadows it. So without this override, RUST_LOG is inert — the +// mapped driver level always wins. Otherwise it returns the driver's current level +// mapped via kernelLogLevel, so the one DATABRICKS_LOG_LEVEL knob governs kernel +// verbosity too. When useNULL is true the level string is unused (empty). +func resolveKernelLogArg() (level string, useNULL bool) { + if os.Getenv("DBSQL_KERNEL_DEBUG") != "" { + return "", true + } + return kernelLogLevel(logger.Logger.GetLevel()), false +} diff --git a/internal/backend/kernel/logging_level_test.go b/internal/backend/kernel/logging_level_test.go new file mode 100644 index 00000000..90449f71 --- /dev/null +++ b/internal/backend/kernel/logging_level_test.go @@ -0,0 +1,87 @@ +package kernel + +import ( + "os" + "testing" + + "github.com/databricks/databricks-sql-go/logger" + "github.com/rs/zerolog" +) + +// kernelLogLevel maps the driver's zerolog level to the kernel_init_logging level +// string, so DATABRICKS_LOG_LEVEL drives the kernel's Rust logs too. Pin the +// mapping (incl. the fatal/panic→OFF collapse and the default→WARN fallback) so a +// level the kernel would reject can't silently ship. fatal/panic map to OFF, not +// ERROR: the Go driver suppresses even Error() lines at those levels, so the Rust +// subscriber must stay at least as quiet as the driver the user configured. Lives +// in the untagged file so it runs under CGO_ENABLED=0. +func TestKernelLogLevel(t *testing.T) { + cases := []struct { + in zerolog.Level + want string + }{ + {zerolog.TraceLevel, "TRACE"}, + {zerolog.DebugLevel, "DEBUG"}, + {zerolog.InfoLevel, "INFO"}, + {zerolog.WarnLevel, "WARN"}, + {zerolog.ErrorLevel, "ERROR"}, + {zerolog.FatalLevel, "OFF"}, // driver suppresses Error() here → kernel stays silent, not louder + {zerolog.PanicLevel, "OFF"}, + {zerolog.Disabled, "OFF"}, + {zerolog.NoLevel, "WARN"}, // unrecognized → the kernel's own default + } + for _, c := range cases { + if got := kernelLogLevel(c.in); got != c.want { + t.Errorf("kernelLogLevel(%v) = %q, want %q", c.in, got, c.want) + } + } +} + +// TestResolveKernelLogArg pins the level-vs-NULL decision initKernelLogging feeds to +// kernel_init_logging — the behavior the PR justifies in prose but the cgo path +// can't easily assert. Two invariants: DBSQL_KERNEL_DEBUG (any non-empty value) must +// yield useNULL=true so the kernel honors RUST_LOG (the whole point of retaining the +// knob — a NULL level is the ONLY way the kernel consults RUST_LOG); otherwise the +// driver's current level is mapped in, so DATABRICKS_LOG_LEVEL governs kernel +// verbosity too. Runs under CGO_ENABLED=0 (no cgo in the resolution path). +func TestResolveKernelLogArg(t *testing.T) { + // Save/restore the env var and the global logger level so this test leaves no + // state behind for its siblings. + prevEnv, hadEnv := os.LookupEnv("DBSQL_KERNEL_DEBUG") + prevLevel := logger.Logger.GetLevel() + t.Cleanup(func() { + if hadEnv { + os.Setenv("DBSQL_KERNEL_DEBUG", prevEnv) + } else { + os.Unsetenv("DBSQL_KERNEL_DEBUG") + } + logger.Logger.Logger = logger.Logger.Level(prevLevel) + }) + + // DBSQL_KERNEL_DEBUG set → NULL level (kernel honors RUST_LOG), regardless of the + // driver level. Pin it at debug to prove the override wins over the mapped level. + logger.Logger.Logger = logger.Logger.Level(zerolog.DebugLevel) + os.Setenv("DBSQL_KERNEL_DEBUG", "1") + if lvl, useNULL := resolveKernelLogArg(); !useNULL || lvl != "" { + t.Errorf("with DBSQL_KERNEL_DEBUG=1: got (level=%q, useNULL=%v), want (\"\", true)", lvl, useNULL) + } + + // An empty value is treated as unset (os.Getenv != "" is the gate), so it must + // NOT trigger the override — the mapped level is used. + os.Setenv("DBSQL_KERNEL_DEBUG", "") + if lvl, useNULL := resolveKernelLogArg(); useNULL || lvl != "DEBUG" { + t.Errorf("with DBSQL_KERNEL_DEBUG=\"\" at debug: got (level=%q, useNULL=%v), want (\"DEBUG\", false)", lvl, useNULL) + } + + // Unset → map the driver level. Check two levels to prove it tracks the logger, + // including the fatal→OFF collapse. + os.Unsetenv("DBSQL_KERNEL_DEBUG") + logger.Logger.Logger = logger.Logger.Level(zerolog.WarnLevel) + if lvl, useNULL := resolveKernelLogArg(); useNULL || lvl != "WARN" { + t.Errorf("unset at warn: got (level=%q, useNULL=%v), want (\"WARN\", false)", lvl, useNULL) + } + logger.Logger.Logger = logger.Logger.Level(zerolog.FatalLevel) + if lvl, useNULL := resolveKernelLogArg(); useNULL || lvl != "OFF" { + t.Errorf("unset at fatal: got (level=%q, useNULL=%v), want (\"OFF\", false)", lvl, useNULL) + } +} From 334e5ff022cbd1123ca351cc24188978b3163bc6 Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Tue, 14 Jul 2026 15:05:31 +0000 Subject: [PATCH 19/24] ci: nightly E2E workflow for the Thrift and kernel backends MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add a scheduled (07:00 UTC) + manually dispatchable workflow that runs the credential-gated end-to-end suites against a real test SQL warehouse. The ordinary "Go" workflow injects no warehouse secrets, so every E2E test t.Skip()s there and real-warehouse behaviour — large multi-page CloudFetch, S3 downloads, drain-past-deadline, real auth — is otherwise not exercised in CI. Two jobs against one test warehouse: - thrift-e2e: pure-Go (CGO_ENABLED=0), the high-value CloudFetch drain + exact-row-count scenarios. - kernel-e2e: reuses the go.yml kernel-build machinery (JFrog/cargo proxy, pinned Rust, kernel App-token clone, kernel-lib/cargo caches), then runs the databricks_kernel-tagged E2E funcs plus the Thrift-vs-kernel parity funcs (which need both the kernel lib and live credentials). Both suites now read the same DATABRICKS_PECOTESTING_* credentials, so one secret set drives the whole workflow — the kernel E2E / parity test helpers are updated from the ad-hoc DATABRICKS_HOST/_HTTP_PATH/_TOKEN vars to the PECOTESTING set (with the _TOKEN_PERSONAL fallback the Thrift suite already uses). Each job has a fail-loud credential guard so a missing secret errors instead of silently skipping into a misleading green. TestKernelE2EM2M is excluded (-skip) since it needs a service-principal client id/secret the warehouse secret set doesn't carry; the job comment documents how to enable it. Alerting is GitHub's built-in failed-scheduled-run notification. Provisioning the DATABRICKS_PECOTESTING_* secrets in the repository is the remaining step to make the nightly run green. Closes PECOBLR-3308. Co-authored-by: Isaac Signed-off-by: Mani Kaustubh Mathur --- .github/workflows/go.yml | 9 +- .github/workflows/nightly-e2e.yml | 236 ++++++++++++++++++++++++++++++ kernel_e2e_test.go | 27 ++-- kernel_parity_test.go | 11 +- 4 files changed, 263 insertions(+), 20 deletions(-) create mode 100644 .github/workflows/nightly-e2e.yml diff --git a/.github/workflows/go.yml b/.github/workflows/go.yml index 98658dcf..64972b44 100644 --- a/.github/workflows/go.yml +++ b/.github/workflows/go.yml @@ -243,10 +243,9 @@ jobs: # missed, then runs `CGO_ENABLED=1 go test -tags databricks_kernel ./...`. # # No warehouse creds here — this job runs only the tagged UNIT tests, so the - # live e2e / Thrift-parity tests self-skip (they need DATABRICKS_HOST / - # DATABRICKS_HTTP_PATH / DATABRICKS_TOKEN). This matches how build-and-test - # treats the Thrift path (pure-Go unit tests only); live integration tests for - # both backends run through the separate trigger-integration-tests.yml dispatch - # to databricks-driver-test, not in this workflow. + # live e2e / Thrift-parity tests self-skip (they need the DATABRICKS_PECOTESTING_* + # warehouse credentials). This matches how build-and-test treats the Thrift path + # (pure-Go unit tests only); the credentialed e2e suites for both backends run in + # the separate nightly-e2e workflow. - name: Build kernel lib + run tagged tests run: make test-kernel diff --git a/.github/workflows/nightly-e2e.yml b/.github/workflows/nightly-e2e.yml new file mode 100644 index 00000000..c61207ed --- /dev/null +++ b/.github/workflows/nightly-e2e.yml @@ -0,0 +1,236 @@ +name: Nightly E2E + +# Runs the credential-gated end-to-end suites against a real test SQL warehouse. +# These tests self-skip in the ordinary "Go" workflow (which injects no warehouse +# secrets), so real-warehouse behaviour — large multi-page CloudFetch, S3 downloads, +# drain-past-deadline, real auth — is otherwise not exercised in CI. +# +# Covers both backends against one test warehouse: +# - Thrift (default pure-Go) E2E — driver_e2e_test.go +# - SEA-via-kernel E2E — kernel_e2e_test.go (cgo + databricks_kernel) +# +# Both suites read the same DATABRICKS_PECOTESTING_* warehouse credentials, so a +# single secret set drives the whole workflow. A failed scheduled run surfaces via +# GitHub's built-in notification. + +on: + schedule: + # 07:00 UTC daily (~midnight PT / mid-morning IST) — off-peak for the shared + # test warehouse. Scheduled workflows only run on the default branch. + - cron: '0 7 * * *' + # Manual trigger for on-demand runs (e.g. validating a fix or new secret wiring) + # without waiting for the nightly cron. + workflow_dispatch: + +permissions: + contents: read + id-token: write + +jobs: + thrift-e2e: + name: E2E (Thrift backend) + # Skips cleanly (no-op success) if the warehouse secret isn't provisioned yet, + # so the workflow doesn't hard-fail on a repo where the secrets haven't landed. + if: ${{ vars.NIGHTLY_E2E_ENABLED != 'false' }} + timeout-minutes: 20 + runs-on: + group: databricks-protected-runner-group + labels: linux-ubuntu-latest + + steps: + - name: Check out code into the Go module directory + uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 + + - name: Setup JFrog + uses: ./.github/actions/setup-jfrog + + - name: Set up Go Toolchain + uses: actions/setup-go@924ae3a1cded613372ab5595356fb5720e22ba16 # v6.5.0 + with: + go-version: '1.25.x' + cache: false + + - name: Get dependencies + run: | + if ! command -v make &> /dev/null ; then + apt-get update && apt-get install -y make + fi + go get -v -t -d ./... + + # Guard: fail loudly if a credential is missing rather than letting the tests + # silently t.Skip() into a misleading green. The token may be provided as either + # _TOKEN or _TOKEN_PERSONAL (the tests accept either), so require at least one. + - name: Verify warehouse credentials are present + env: + HOST: ${{ secrets.DATABRICKS_PECOTESTING_SERVER_HOSTNAME }} + HTTP_PATH: ${{ secrets.DATABRICKS_PECOTESTING_HTTP_PATH2 }} + TOKEN: ${{ secrets.DATABRICKS_PECOTESTING_TOKEN }} + TOKEN_PERSONAL: ${{ secrets.DATABRICKS_PECOTESTING_TOKEN_PERSONAL }} + run: | + if [ -z "$HOST" ] || [ -z "$HTTP_PATH" ] || { [ -z "$TOKEN" ] && [ -z "$TOKEN_PERSONAL" ]; }; then + echo "::error::Nightly E2E warehouse secrets are not set (need DATABRICKS_PECOTESTING_SERVER_HOSTNAME, _HTTP_PATH2, and _TOKEN or _TOKEN_PERSONAL). Provision them in the repository settings." >&2 + exit 1 + fi + + # The Thrift E2E tests read the DATABRICKS_PECOTESTING_* vars directly and run + # in the default pure-Go build (CGO_ENABLED=0). -run targets the high-value + # real-warehouse scenarios: large multi-page CloudFetch drain-past-caller- + # deadline and exact row count. -count=1 defeats the test cache so a scheduled + # run always re-exercises the warehouse. NOT -short, so + # TestE2ECloudFetchExactRowCount (guarded by -short) actually runs. + - name: Run Thrift E2E + env: + CGO_ENABLED: 0 + DATABRICKS_PECOTESTING_SERVER_HOSTNAME: ${{ secrets.DATABRICKS_PECOTESTING_SERVER_HOSTNAME }} + DATABRICKS_PECOTESTING_HTTP_PATH2: ${{ secrets.DATABRICKS_PECOTESTING_HTTP_PATH2 }} + DATABRICKS_PECOTESTING_TOKEN: ${{ secrets.DATABRICKS_PECOTESTING_TOKEN }} + DATABRICKS_PECOTESTING_TOKEN_PERSONAL: ${{ secrets.DATABRICKS_PECOTESTING_TOKEN_PERSONAL }} + run: | + go test -count=1 -v -timeout 15m \ + -run 'TestE2EArrowBatchesSurviveQueryContextCancellation|TestE2ECloudFetchExactRowCount' . + + kernel-e2e: + name: E2E (kernel backend) + if: ${{ vars.NIGHTLY_E2E_ENABLED != 'false' }} + # Bounds the same source-build blast radius as the go.yml kernel job: an external + # clone + ~200-crate cold compile shouldn't hold a protected-runner slot to the + # 360-min default. + timeout-minutes: 40 + runs-on: + group: databricks-protected-runner-group + labels: linux-ubuntu-latest + + steps: + - name: Check out code into the Go module directory + uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 + + # Go module proxy (GOPROXY + ~/.netrc) via JFrog OIDC. Also exports + # JFROG_ACCESS_TOKEN, which the cargo step below reuses. + - name: Setup JFrog + uses: ./.github/actions/setup-jfrog + + # The protected runner blocks direct crates.io access (go/hardened-gha), so + # point cargo at the JFrog crates proxy, reusing setup-jfrog's JFROG_ACCESS_TOKEN. + - name: Configure cargo to use JFrog + shell: bash + run: | + set -euo pipefail + mkdir -p ~/.cargo + cat > ~/.cargo/config.toml << 'EOF' + [source.crates-io] + replace-with = "jfrog" + + [source.jfrog] + registry = "sparse+https://databricks.jfrog.io/artifactory/api/cargo/db-cargo-remote/index/" + + [registries.jfrog] + index = "sparse+https://databricks.jfrog.io/artifactory/api/cargo/db-cargo-remote/index/" + credential-provider = ["cargo:token"] + EOF + cat > ~/.cargo/credentials.toml << EOF + [registries.jfrog] + token = "Bearer ${JFROG_ACCESS_TOKEN}" + EOF + chmod 600 ~/.cargo/credentials.toml + + - name: Set up Go Toolchain + uses: actions/setup-go@924ae3a1cded613372ab5595356fb5720e22ba16 # v6.5.0 + with: + go-version: '1.25.x' + cache: false + + # Keep in lockstep with rust-toolchain.toml's channel (it governs the archive + # under a fixed KERNEL_REV; a floating `stable` here would drift it). + - name: Set up Rust Toolchain + uses: actions-rust-lang/setup-rust-toolchain@46268bd060767258de96ed93c1251119784f2ab6 # v1.16.1 + with: + toolchain: 1.96.1 + + # The kernel repo is private and the hardened runner has no ambient git creds. + # Mint a repo-scoped token from the INTEGRATION_TEST_APP GitHub App and rewrite + # the kernel HTTPS URL to carry it (same mechanism as the go.yml kernel job). + - name: Generate GitHub App token for databricks-sql-kernel + id: kernel-token + uses: actions/create-github-app-token@f8d387b68d61c58ab83c6c016672934102569859 # v3.0.0 + with: + app-id: ${{ secrets.INTEGRATION_TEST_APP_ID }} + private-key: ${{ secrets.INTEGRATION_TEST_PRIVATE_KEY }} + owner: databricks + repositories: databricks-sql-kernel + + - name: Rewrite kernel repo URL to authenticated HTTPS + env: + TOKEN: ${{ steps.kernel-token.outputs.token }} + run: | + git config --global \ + url."https://x-access-token:${TOKEN}@github.com/databricks/".insteadOf \ + "https://github.com/databricks/" + + - name: Cache kernel static lib + uses: actions/cache@55cc8345863c7cc4c66a329aec7e433d2d1c52a9 # v6.1.0 + with: + path: | + internal/backend/kernel/lib + internal/backend/kernel/include + key: ${{ runner.os }}-kernellib-${{ hashFiles('KERNEL_REV', 'rust-toolchain.toml') }} + + - name: Cache cargo + kernel build tree + uses: actions/cache@55cc8345863c7cc4c66a329aec7e433d2d1c52a9 # v6.1.0 + with: + path: | + ~/.cargo/registry + ~/.cargo/git + build/kernel-src/target + key: ${{ runner.os }}-kernel-cargo-${{ hashFiles('KERNEL_REV', 'rust-toolchain.toml') }} + restore-keys: | + ${{ runner.os }}-kernel-cargo- + + - name: Install build prerequisites (make, C compiler) + run: | + if ! command -v make &> /dev/null ; then + apt-get update && apt-get install -y make + fi + if ! command -v cc &> /dev/null ; then + apt-get update && apt-get install -y build-essential + fi + + # Same guard as the Thrift job: require host, http path, and at least one of + # _TOKEN / _TOKEN_PERSONAL so a missing credential fails loud instead of skipping. + - name: Verify warehouse credentials are present + env: + HOST: ${{ secrets.DATABRICKS_PECOTESTING_SERVER_HOSTNAME }} + HTTP_PATH: ${{ secrets.DATABRICKS_PECOTESTING_HTTP_PATH2 }} + TOKEN: ${{ secrets.DATABRICKS_PECOTESTING_TOKEN }} + TOKEN_PERSONAL: ${{ secrets.DATABRICKS_PECOTESTING_TOKEN_PERSONAL }} + run: | + if [ -z "$HOST" ] || [ -z "$HTTP_PATH" ] || { [ -z "$TOKEN" ] && [ -z "$TOKEN_PERSONAL" ]; }; then + echo "::error::Nightly E2E warehouse secrets are not set (need DATABRICKS_PECOTESTING_SERVER_HOSTNAME, _HTTP_PATH2, and _TOKEN or _TOKEN_PERSONAL). Provision them in the repository settings." >&2 + exit 1 + fi + + # Build the pinned kernel .a (make kernel-lib, cache permitting), then run the + # kernel E2E tests tagged cgo+databricks_kernel. The kernel suite reads the same + # DATABRICKS_PECOTESTING_* credentials as the Thrift suite. -run selects the + # kernel E2E funcs plus the Thrift-vs-kernel parity funcs (TestKernelThriftParity + # / TestKernelParamsVsThrift) — the latter need both the kernel lib AND live + # credentials, so this is the only job that can run them. The tagged UNIT tests + # already run credential-free in the go.yml kernel job. + # + # TestKernelE2EM2M is excluded: -run matches it (unanchored prefix), but it needs + # a service principal (DATABRICKS_CLIENT_ID / _CLIENT_SECRET) that this warehouse + # secret set doesn't carry, so it would t.Skip() every run. -skip keeps it clearly + # out of scope here rather than a misleading silent skip; add M2M service-principal + # secrets + drop the -skip to cover it. + - name: Build kernel lib + run: make kernel-lib + + - name: Run kernel E2E + env: + DATABRICKS_PECOTESTING_SERVER_HOSTNAME: ${{ secrets.DATABRICKS_PECOTESTING_SERVER_HOSTNAME }} + DATABRICKS_PECOTESTING_HTTP_PATH2: ${{ secrets.DATABRICKS_PECOTESTING_HTTP_PATH2 }} + DATABRICKS_PECOTESTING_TOKEN: ${{ secrets.DATABRICKS_PECOTESTING_TOKEN }} + DATABRICKS_PECOTESTING_TOKEN_PERSONAL: ${{ secrets.DATABRICKS_PECOTESTING_TOKEN_PERSONAL }} + run: | + CGO_ENABLED=1 go test -tags databricks_kernel -count=1 -v -timeout 30m \ + -run 'TestKernelE2E|TestKernelThriftParity|TestKernelParamsVsThrift' \ + -skip 'TestKernelE2EM2M' . diff --git a/kernel_e2e_test.go b/kernel_e2e_test.go index 7515f3be..c651f158 100644 --- a/kernel_e2e_test.go +++ b/kernel_e2e_test.go @@ -18,10 +18,10 @@ import ( dbsqlerr "github.com/databricks/databricks-sql-go/errors" ) -// kernelTestDB opens a kernel-backed *sql.DB from DATABRICKS_HOST / -// DATABRICKS_HTTP_PATH / DATABRICKS_TOKEN, or skips when they are unset. It goes -// through the standard connector with WithUseKernel(true) — the same path a real -// consumer uses — not a kernel-only connector. +// kernelTestDB opens a kernel-backed *sql.DB from the DATABRICKS_PECOTESTING_* +// warehouse credentials, or skips when they are unset. It goes through the standard +// connector with WithUseKernel(true) — the same path a real consumer uses — not a +// kernel-only connector. func kernelTestDB(t *testing.T) *sql.DB { t.Helper() return kernelTestDBWith(t) @@ -47,11 +47,16 @@ func TestKernelE2ESelect1(t *testing.T) { // counterpart to kernelTestDB. func kernelTestDBWith(t *testing.T, extra ...ConnOption) *sql.DB { t.Helper() - host := os.Getenv("DATABRICKS_HOST") - httpPath := os.Getenv("DATABRICKS_HTTP_PATH") - token := os.Getenv("DATABRICKS_TOKEN") + // Reads the same DATABRICKS_PECOTESTING_* warehouse credentials as the Thrift + // E2E suite, so both backends run against one test warehouse from one secret set. + host := os.Getenv("DATABRICKS_PECOTESTING_SERVER_HOSTNAME") + httpPath := os.Getenv("DATABRICKS_PECOTESTING_HTTP_PATH2") + token := os.Getenv("DATABRICKS_PECOTESTING_TOKEN") + if token == "" { + token = os.Getenv("DATABRICKS_PECOTESTING_TOKEN_PERSONAL") + } if host == "" || httpPath == "" || token == "" { - t.Skip("set DATABRICKS_HOST / DATABRICKS_HTTP_PATH / DATABRICKS_TOKEN for the kernel e2e") + t.Skip("set DATABRICKS_PECOTESTING_SERVER_HOSTNAME, DATABRICKS_PECOTESTING_HTTP_PATH2, and DATABRICKS_PECOTESTING_TOKEN for the kernel e2e") } opts := append([]ConnOption{ WithServerHostname(host), @@ -471,12 +476,12 @@ func TestKernelE2EMetricViewMetadata(t *testing.T) { // M2M setter path actually authenticates end to end (not just that set_auth_m2m // returns OK, which TestSetAuthByMode already covers). func TestKernelE2EM2M(t *testing.T) { - host := os.Getenv("DATABRICKS_HOST") - httpPath := os.Getenv("DATABRICKS_HTTP_PATH") + host := os.Getenv("DATABRICKS_PECOTESTING_SERVER_HOSTNAME") + httpPath := os.Getenv("DATABRICKS_PECOTESTING_HTTP_PATH2") clientID := os.Getenv("DATABRICKS_CLIENT_ID") clientSecret := os.Getenv("DATABRICKS_CLIENT_SECRET") if host == "" || httpPath == "" || clientID == "" || clientSecret == "" { - t.Skip("set DATABRICKS_HOST / DATABRICKS_HTTP_PATH / DATABRICKS_CLIENT_ID / DATABRICKS_CLIENT_SECRET for the M2M e2e") + t.Skip("set DATABRICKS_PECOTESTING_SERVER_HOSTNAME, DATABRICKS_PECOTESTING_HTTP_PATH2, DATABRICKS_CLIENT_ID, and DATABRICKS_CLIENT_SECRET for the M2M e2e") } connector, err := NewConnector( diff --git a/kernel_parity_test.go b/kernel_parity_test.go index 62d2576a..b36d6550 100644 --- a/kernel_parity_test.go +++ b/kernel_parity_test.go @@ -13,11 +13,14 @@ import ( // parity comparison against the kernel backend. func thriftTestDB(t *testing.T) *sql.DB { t.Helper() - host := os.Getenv("DATABRICKS_HOST") - httpPath := os.Getenv("DATABRICKS_HTTP_PATH") - token := os.Getenv("DATABRICKS_TOKEN") + host := os.Getenv("DATABRICKS_PECOTESTING_SERVER_HOSTNAME") + httpPath := os.Getenv("DATABRICKS_PECOTESTING_HTTP_PATH2") + token := os.Getenv("DATABRICKS_PECOTESTING_TOKEN") + if token == "" { + token = os.Getenv("DATABRICKS_PECOTESTING_TOKEN_PERSONAL") + } if host == "" || httpPath == "" || token == "" { - t.Skip("set DATABRICKS_HOST / DATABRICKS_HTTP_PATH / DATABRICKS_TOKEN for the parity test") + t.Skip("set DATABRICKS_PECOTESTING_SERVER_HOSTNAME, DATABRICKS_PECOTESTING_HTTP_PATH2, and DATABRICKS_PECOTESTING_TOKEN for the parity test") } connector, err := NewConnector( WithServerHostname(host), From 773e19c3e82c8e2998226c510b6ceff88bdf6f54 Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Wed, 15 Jul 2026 05:34:58 +0000 Subject: [PATCH 20/24] ci(nightly-e2e): split the kernel job timeout budget so Go's timeout fires first MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The kernel E2E job had a 40-min job cap while `go test` alone was given -timeout 30m, leaving only ~10 min for checkout, toolchains, caches, and a potentially cold ~200-crate `make kernel-lib` Rust build. On a cache miss the build can exceed that, so GitHub SIGKILLs the whole job at 40 min before Go's timeout can emit its which-test-hung goroutine dump — losing the diagnostics. Split the budget explicitly: raise the job cap to 65 min and add a 25-min per-step timeout on `make kernel-lib`. Worst case is 25 (build) + 30 (go test) + setup < 65, so a wedged build is killed distinctly as a build failure and a hung test still hits Go's own -timeout first. Also moved the E2E step's descriptive comment down onto the Run step it documents. Co-authored-by: Isaac Signed-off-by: Mani Kaustubh Mathur --- .github/workflows/nightly-e2e.yml | 26 +++++++++++++++++--------- 1 file changed, 17 insertions(+), 9 deletions(-) diff --git a/.github/workflows/nightly-e2e.yml b/.github/workflows/nightly-e2e.yml index c61207ed..8b4632a6 100644 --- a/.github/workflows/nightly-e2e.yml +++ b/.github/workflows/nightly-e2e.yml @@ -94,8 +94,11 @@ jobs: if: ${{ vars.NIGHTLY_E2E_ENABLED != 'false' }} # Bounds the same source-build blast radius as the go.yml kernel job: an external # clone + ~200-crate cold compile shouldn't hold a protected-runner slot to the - # 360-min default. - timeout-minutes: 40 + # 360-min default. Budget is split explicitly (see the per-step timeout on "Build + # kernel lib" below) so the job cap sits strictly above setup + build + the go + # test -timeout, letting Go's own timeout fire first and emit its which-test-hung + # goroutine dump instead of GitHub SIGKILLing the job mid-build on a cache miss. + timeout-minutes: 65 runs-on: group: databricks-protected-runner-group labels: linux-ubuntu-latest @@ -208,10 +211,18 @@ jobs: exit 1 fi - # Build the pinned kernel .a (make kernel-lib, cache permitting), then run the - # kernel E2E tests tagged cgo+databricks_kernel. The kernel suite reads the same - # DATABRICKS_PECOTESTING_* credentials as the Thrift suite. -run selects the - # kernel E2E funcs plus the Thrift-vs-kernel parity funcs (TestKernelThriftParity + # Build the pinned kernel .a (make kernel-lib, cache permitting). Cap the Rust + # build explicitly so a cold ~200-crate compile that wedges is killed here as a + # build failure, rather than eating into the go test budget and letting GitHub + # SIGKILL the whole job before Go's -timeout can report which test hung. 25m + # (build) + 30m (go test) + setup stays under the 65m job cap. + - name: Build kernel lib + timeout-minutes: 25 + run: make kernel-lib + + # Run the kernel E2E tests tagged cgo+databricks_kernel. The kernel suite reads + # the same DATABRICKS_PECOTESTING_* credentials as the Thrift suite. -run selects + # the kernel E2E funcs plus the Thrift-vs-kernel parity funcs (TestKernelThriftParity # / TestKernelParamsVsThrift) — the latter need both the kernel lib AND live # credentials, so this is the only job that can run them. The tagged UNIT tests # already run credential-free in the go.yml kernel job. @@ -221,9 +232,6 @@ jobs: # secret set doesn't carry, so it would t.Skip() every run. -skip keeps it clearly # out of scope here rather than a misleading silent skip; add M2M service-principal # secrets + drop the -skip to cover it. - - name: Build kernel lib - run: make kernel-lib - - name: Run kernel E2E env: DATABRICKS_PECOTESTING_SERVER_HOSTNAME: ${{ secrets.DATABRICKS_PECOTESTING_SERVER_HOSTNAME }} From 5ac4ad236fcb83d0d744b271bdd9d8d55e0061dc Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Tue, 14 Jul 2026 16:46:32 +0000 Subject: [PATCH 21/24] ci: dispatch the SEA-via-kernel integration leg behind a kernel label MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Extend the driver-test dispatch so the Go integration suite can run against the SEA-via-kernel backend, selected by label — the analogue of the existing Thrift labels, not a copy. Adds one kernel-namespaced label: - integration-test-kernel → sea backend, passthrough (real warehouse) The label resolves both proxy_mode and a new go_mode (thrift vs sea), and go_mode rides in the repository_dispatch client_payload. databricks-driver-test reads it to decide whether to build the kernel static lib and run the tagged (databricks_kernel) leg; the Thrift labels send go_mode=thrift and are unchanged. The new label is added to the on-new-commit label-drop list and the skip-stub guidance. The kernel label is passthrough (not replay): the sea leg has no committed recordings to replay yet, so there is no sea replay label until those are captured. The required merge-queue gate stays Thrift-only (go_mode pinned to thrift): the kernel leg is previewable on demand but not a required gate until the SEA backend ships in a released driver. Co-authored-by: Isaac Signed-off-by: Mani Kaustubh Mathur --- .../workflows/trigger-integration-tests.yml | 54 ++++++++++++++----- 1 file changed, 41 insertions(+), 13 deletions(-) diff --git a/.github/workflows/trigger-integration-tests.yml b/.github/workflows/trigger-integration-tests.yml index 4da5deb1..db1bcf9f 100644 --- a/.github/workflows/trigger-integration-tests.yml +++ b/.github/workflows/trigger-integration-tests.yml @@ -9,9 +9,18 @@ name: Trigger Integration Tests # preview with the `integration-test` label (replay) or `integration-test-full` # (full passthrough), and merge queue runs the real required gate. # +# The suite runs against either driver backend, selected by the label: +# - `integration-test` (replay) / `integration-test-full` (passthrough) → Thrift +# - `integration-test-kernel` (passthrough) → SEA-via-kernel +# The kernel label adds `go_mode: sea` to the dispatch payload; driver-test then +# builds the kernel static lib and runs the tagged (databricks_kernel) leg. It uses +# passthrough (real warehouse) — the sea leg has no committed recordings to replay +# yet, so there is no sea replay label until those are captured. +# # Required external setup: # -# 1. `integration-test` and `integration-test-full` labels exist in this repo. +# 1. The three trigger labels exist in this repo: `integration-test`, +# `integration-test-full`, `integration-test-kernel`. # 2. `INTEGRATION_TEST_APP_ID` / `INTEGRATION_TEST_PRIVATE_KEY` repo secrets # are installed in this repo for the dispatcher GitHub App. # 3. The app is installed/granted on `databricks-driver-test` so this workflow @@ -46,7 +55,7 @@ jobs: with: script: | const present = context.payload.pull_request.labels.map((l) => l.name); - const triggerLabels = ['integration-test', 'integration-test-full']; + const triggerLabels = ['integration-test', 'integration-test-full', 'integration-test-kernel']; const removed = []; for (const name of triggerLabels) { if (!present.includes(name)) continue; @@ -111,7 +120,7 @@ jobs: completed_at: new Date().toISOString(), output: { title: 'Skipped on PR - runs in merge queue', - summary: 'Go integration tests are skipped on ordinary PR events and run as a required gate in the merge queue. Add the `integration-test` label to preview replay (or `integration-test-full` for the full passthrough suite) on this PR.', + summary: 'Go integration tests are skipped on ordinary PR events and run as a required gate in the merge queue. Preview on this PR by adding a label: `integration-test` (Thrift, replay) or `integration-test-full` (Thrift, passthrough); `integration-test-kernel` runs the SEA-via-kernel backend against the real warehouse (passthrough).', }, }); @@ -122,7 +131,8 @@ jobs: github.event_name == 'pull_request' && github.event.action == 'labeled' && (github.event.label.name == 'integration-test' || - github.event.label.name == 'integration-test-full') + github.event.label.name == 'integration-test-full' || + github.event.label.name == 'integration-test-kernel') runs-on: group: databricks-protected-runner-group labels: linux-ubuntu-latest @@ -134,14 +144,24 @@ jobs: pull-requests: write checks: write steps: - - name: Resolve proxy mode from label + - name: Resolve proxy mode + backend from label id: mode + # go_mode: integration-test-kernel runs the SEA-via-kernel backend; the others + # run Thrift. proxy_mode: integration-test-full and integration-test-kernel hit + # the real warehouse (passthrough); plain integration-test serves recordings + # (replay). The kernel label is passthrough because the sea leg has no committed + # recordings to replay yet. driver-test reads go_mode to decide whether to build + # the kernel lib and run the tagged leg. run: | - if [ "${{ github.event.label.name }}" = "integration-test-full" ]; then - echo "proxy_mode=passthrough" >> "$GITHUB_OUTPUT" - else - echo "proxy_mode=replay" >> "$GITHUB_OUTPUT" - fi + LABEL="${{ github.event.label.name }}" + case "$LABEL" in + integration-test-full|integration-test-kernel) echo "proxy_mode=passthrough" >> "$GITHUB_OUTPUT" ;; + *) echo "proxy_mode=replay" >> "$GITHUB_OUTPUT" ;; + esac + case "$LABEL" in + integration-test-kernel) echo "go_mode=sea" >> "$GITHUB_OUTPUT" ;; + *) echo "go_mode=thrift" >> "$GITHUB_OUTPUT" ;; + esac - name: Generate GitHub App token (driver-test repo) id: app-token @@ -156,6 +176,7 @@ jobs: uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0 env: PROXY_MODE: ${{ steps.mode.outputs.proxy_mode }} + GO_MODE: ${{ steps.mode.outputs.go_mode }} with: github-token: ${{ steps.app-token.outputs.token }} script: | @@ -170,9 +191,10 @@ jobs: pr_repo: context.repo.owner + '/' + context.repo.repo, pr_url: pr.html_url, proxy_mode: process.env.PROXY_MODE, + go_mode: process.env.GO_MODE, }, }); - core.info(`Dispatched go-pr-test (${process.env.PROXY_MODE}) for PR #${pr.number} @ ${pr.head.sha}`); + core.info(`Dispatched go-pr-test (${process.env.GO_MODE}/${process.env.PROXY_MODE}) for PR #${pr.number} @ ${pr.head.sha}`); - name: Fail check on dispatch error if: failure() @@ -198,13 +220,14 @@ jobs: uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0 env: PROXY_MODE: ${{ steps.mode.outputs.proxy_mode }} + GO_MODE: ${{ steps.mode.outputs.go_mode }} with: script: | await github.rest.issues.createComment({ owner: context.repo.owner, repo: context.repo.repo, issue_number: context.issue.number, - body: `Go integration tests triggered (\`${process.env.PROXY_MODE}\`). [View workflow runs](https://github.com/databricks/databricks-driver-test/actions).`, + body: `Go integration tests triggered (\`${process.env.GO_MODE}\` / \`${process.env.PROXY_MODE}\`). [View workflow runs](https://github.com/databricks/databricks-driver-test/actions).`, }); # Merge queue: the required gate. Runs replay once before merge. @@ -261,9 +284,14 @@ jobs: pr_repo: context.repo.owner + '/' + context.repo.repo, pr_url: `${context.serverUrl}/${context.repo.owner}/${context.repo.repo}/pull/${prNumber}`, proxy_mode: 'replay', + // The required merge-queue gate runs the Thrift backend only. The + // kernel (sea) leg is previewed on demand via the kernel labels but + // is not a required gate — it becomes one once the SEA backend ships + // in a released driver and its recordings are captured. + go_mode: 'thrift', }, }); - core.info(`Merge-queue dispatch go-pr-test (replay) for PR #${prNumber} @ ${process.env.HEAD_SHA}`); + core.info(`Merge-queue dispatch go-pr-test (thrift/replay) for PR #${prNumber} @ ${process.env.HEAD_SHA}`); - name: Fail check on dispatch error if: failure() From c85ff22075d32bc48504d54f9be35feb4e50701f Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Tue, 14 Jul 2026 20:48:54 +0000 Subject: [PATCH 22/24] fix(kernel): record CLOSE_STATEMENT telemetry + sentinel-wrap the auth rejection MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two review-driven fixes on the kernel backend, both Go-only. CLOSE_STATEMENT telemetry: kernelRows.Close() now fires the OnClose callback, so kernel queries record close latency / statement success-or-failure the same way the Thrift path does (conn gates that recording on OnClose firing). Previously it never fired, so kernel traffic emitted no close telemetry — including failures — a production observability blind spot. Next() is split into a thin wrapper that records the first non-EOF error as iterationErr (io.EOF is normal drain), which Close() passes to the callback; closeErr is nil since the kernel teardown has no fallible close RPC. The callback is armed only after newKernelRows finishes constructing (mirroring the Thrift NewRows), so a schema-fetch/import failure's cleanup Close() does not record a falsely-successful CLOSE_STATEMENT. Auth rejection sentinel: the unsupported-authenticator branch of resolveKernelAuth wrapped a bare errors.New, so fallback logic could only substring-match. It now wraps ErrNotSupportedByKernel via %w so callers can errors.Is it, matching the other kernel config rejections. The missing-personal-access-token error stays a plain error (missing-required-config, not an unsupported feature — it must not signal "fall back to Thrift"). Tests: TestKernelRowsCloseFiresOnClose (success, iterationErr propagation, idempotent double-close, nil-callbacks, construction-failure-no-success-close); strengthened the non-PAT authenticator rejection to assert errors.Is(ErrNotSupportedByKernel). Co-authored-by: Isaac Signed-off-by: Mani Kaustubh Mathur --- internal/backend/kernel/kernel_test.go | 76 ++++++++++++++++++++++++++ internal/backend/kernel/rows.go | 36 +++++++++++- kernel_config.go | 3 + kernel_config_test.go | 4 +- 4 files changed, 116 insertions(+), 3 deletions(-) diff --git a/internal/backend/kernel/kernel_test.go b/internal/backend/kernel/kernel_test.go index a293f915..e5a32699 100644 --- a/internal/backend/kernel/kernel_test.go +++ b/internal/backend/kernel/kernel_test.go @@ -14,6 +14,7 @@ import ( "github.com/databricks/databricks-sql-go/driverctx" dbsqlerr "github.com/databricks/databricks-sql-go/errors" "github.com/databricks/databricks-sql-go/internal/backend" + dbsqlrows "github.com/databricks/databricks-sql-go/internal/rows" "github.com/databricks/databricks-sql-go/logger" ) @@ -345,3 +346,78 @@ func strPtr(s string) *string { return &s } // CGO_ENABLED=0 build; see arrowscan_test.go. The decimal formatter lives in // internal/decimalfmt. This file keeps the kernel-specific tests: error mapping, // bad-connection classification, and the bound-params rejection. + +// kernelRows.Close() must fire the OnClose telemetry callback so the kernel path +// records CLOSE_STATEMENT / latency / statement success-or-failure — conn gates +// that recording on OnClose being called, and the Thrift path fires it. Before this +// wiring, kernel queries emitted no close telemetry (a production blind spot). A bare +// kernelRows is safe: Close() nil-guards cur/stream/op. +func TestKernelRowsCloseFiresOnClose(t *testing.T) { + t.Run("success path reports nil iterErr", func(t *testing.T) { + var got struct { + called bool + chunkCount int + iterErr, closeErr error + } + r := &kernelRows{ + chunkCount: 3, + callbacks: &dbsqlrows.TelemetryCallbacks{ + OnClose: func(latencyMs int64, chunkCount int, iterErr, closeErr error) { + got.called, got.chunkCount, got.iterErr, got.closeErr = true, chunkCount, iterErr, closeErr + }, + }, + } + if err := r.Close(); err != nil { + t.Fatalf("Close() = %v, want nil", err) + } + if !got.called { + t.Fatal("OnClose was not fired") + } + if got.chunkCount != 3 { + t.Errorf("OnClose chunkCount = %d, want 3", got.chunkCount) + } + if got.iterErr != nil || got.closeErr != nil { + t.Errorf("OnClose errs = (%v, %v), want (nil, nil)", got.iterErr, got.closeErr) + } + }) + + t.Run("iterationErr is reported", func(t *testing.T) { + sentinel := errors.New("boom") + var gotIter error + fired := 0 + r := &kernelRows{ + iterationErr: sentinel, + callbacks: &dbsqlrows.TelemetryCallbacks{ + OnClose: func(_ int64, _ int, iterErr, _ error) { fired++; gotIter = iterErr }, + }, + } + _ = r.Close() + if !errors.Is(gotIter, sentinel) { + t.Errorf("OnClose iterErr = %v, want %v", gotIter, sentinel) + } + // Idempotent: a second Close must not re-fire (conn/database-sql may double-close). + _ = r.Close() + if fired != 1 { + t.Errorf("OnClose fired %d times across two Close() calls, want 1", fired) + } + }) + + t.Run("nil callbacks is safe", func(t *testing.T) { + r := &kernelRows{} // no callbacks + if err := r.Close(); err != nil { + t.Errorf("Close() with nil callbacks = %v, want nil", err) + } + }) + + t.Run("construction-failure Close must not fire a success OnClose", func(t *testing.T) { + // newKernelRows arms the callback only after a successful build, so its + // cleanup Close() on a schema-fetch/import failure has callbacks==nil and + // must NOT record a (falsely successful) CLOSE_STATEMENT. Model that state. + fired := false + r := &kernelRows{ /* callbacks intentionally nil, as during construction */ } + _ = r.Close() + if fired { + t.Error("OnClose fired for a rows object that never finished construction") + } + }) +} diff --git a/internal/backend/kernel/rows.go b/internal/backend/kernel/rows.go index b03290f4..5a7cb0e4 100644 --- a/internal/backend/kernel/rows.go +++ b/internal/backend/kernel/rows.go @@ -17,6 +17,7 @@ import ( "database/sql/driver" "fmt" "io" + "time" "unsafe" "github.com/apache/arrow/go/v12/arrow" @@ -47,6 +48,10 @@ type kernelRows struct { chunkCount int // cumulative batches fetched, for OnChunkFetched closed bool eof bool + // iterationErr is the first non-EOF error seen during Next(), reported to the + // OnClose telemetry callback so a failed statement is recorded (matching the + // Thrift path's rows.iterationErr). io.EOF is normal termination, not an error. + iterationErr error // keyCache memoizes struct field-name JSON keys for this result set so // per-row rendering doesn't re-marshal constant names. Scoped to this Rows // (freed with it) — not a process-global, which would leak. @@ -56,7 +61,14 @@ type kernelRows struct { // newKernelRows fetches the schema up front (for Columns()) and returns the row // iterator; batches are pulled lazily on Next. func newKernelRows(ctx context.Context, op *kernelOp, stream *C.kernel_result_stream_t, cb *dbsqlrows.TelemetryCallbacks) (driver.Rows, error) { - r := &kernelRows{ctx: ctx, op: op, stream: stream, callbacks: cb, keyCache: arrowscan.NewStructKeyCache()} + // The telemetry callback is deliberately NOT set on r yet: the two cleanup + // r.Close() calls below run when construction FAILS (schema fetch/import), and a + // Close() with the callback set would fire OnClose as a *successful* close for a + // statement that never produced rows — masking the failure in CLOSE_STATEMENT + // telemetry. Assign it only on the success path (matching the Thrift NewRows, + // which sets closeCallback just before returning); the construction error itself + // is surfaced to and recorded by the conn execute path. + r := &kernelRows{ctx: ctx, op: op, stream: stream, keyCache: arrowscan.NewStructKeyCache()} var csch C.struct_ArrowSchema if err := call(func() C.KernelStatusCode { @@ -76,6 +88,9 @@ func newKernelRows(ctx context.Context, op *kernelOp, stream *C.kernel_result_st for i, f := range fields { r.cols[i] = f.Name } + // Construction succeeded — now arm the close telemetry callback so a normal + // Close() (after row iteration) records CLOSE_STATEMENT. + r.callbacks = cb klogCtx(ctx, "newKernelRows: %d columns", len(r.cols)) return r, nil } @@ -90,6 +105,7 @@ func (r *kernelRows) Close() error { return nil } r.closed = true + closeStart := time.Now() if r.cur != nil { r.cur.Release() r.cur = nil @@ -101,6 +117,14 @@ func (r *kernelRows) Close() error { if r.op != nil { r.op.close() } + // Fire the close telemetry callback so the kernel path records CLOSE_STATEMENT / + // execution latency / statement success-or-failure like the Thrift path does + // (conn gates this on OnClose being called). The kernel teardown has no fallible + // close RPC — the C stream/statement closes don't surface an error — so closeErr + // is nil; iterationErr carries any failure seen during Next(). + if r.callbacks != nil && r.callbacks.OnClose != nil { + r.callbacks.OnClose(time.Since(closeStart).Milliseconds(), r.chunkCount, r.iterationErr, nil) + } klogCtx(r.ctx, "kernelRows closed") return nil } @@ -108,6 +132,16 @@ func (r *kernelRows) Close() error { // Next fills dest with the next row's values, advancing across batches. Returns // io.EOF when the stream is drained. func (r *kernelRows) Next(dest []driver.Value) error { + err := r.next(dest) + // Record the first non-EOF error for the OnClose telemetry callback (io.EOF is + // normal drain, not a failure). Mirrors the Thrift path's iterationErr capture. + if err != nil && err != io.EOF && r.iterationErr == nil { + r.iterationErr = err + } + return err +} + +func (r *kernelRows) next(dest []driver.Value) error { if r.closed { return io.EOF } diff --git a/kernel_config.go b/kernel_config.go index 168c3ab8..cf5fa0d0 100644 --- a/kernel_config.go +++ b/kernel_config.go @@ -129,6 +129,9 @@ func resolveKernelAuth(cfg *config.Config) (kernel.Auth, error) { } } if token == "" { + // Missing required config (not an unsupported-feature rejection), so this is + // intentionally NOT wrapped with ErrNotSupportedByKernel — a caller shouldn't + // fall back to Thrift for a forgotten token, it should supply one. return kernel.Auth{}, errors.New("databricks: the kernel backend requires a personal access token; " + "set one with WithAccessToken (or a *pat.PATAuth via WithAuthenticator)") } diff --git a/kernel_config_test.go b/kernel_config_test.go index 11c492c2..81e0e848 100644 --- a/kernel_config_test.go +++ b/kernel_config_test.go @@ -96,8 +96,8 @@ func TestValidateKernelConfig(t *testing.T) { // since that sentinel is the documented programmatic fallback-detection contract — // asserting only err != nil would let a dropped or malformed %w wrap ship green. // Table-driven so a new rejection is covered by adding one row. (Catalog/schema/ - // metric-view moved to forwarded above; a non-PAT authenticator is rejected too but - // not sentinel-wrapped, so it's asserted separately below.) + // metric-view moved to forwarded above; a non-PAT authenticator is also + // sentinel-wrapped but needs its own AccessToken="" setup, so it's asserted separately below.) rejections := []struct { name string mut func(*config.Config) From 0b1df5bf01e6015d88d9f0d49b722e7e43daaf7f Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Wed, 15 Jul 2026 05:05:31 +0000 Subject: [PATCH 23/24] test(kernel): exercise the real construction-failure OnClose path; fix Thrift comment MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses two review comments on the CLOSE_STATEMENT telemetry change: - The "construction-failure Close must not fire a success OnClose" subtest was a no-op: `fired` was never assignable to true, so it duplicated the nil-callbacks case and asserted nothing. Replaced with a test that drives the real newKernelRows cleanup path — a nil result stream makes kernel_result_stream_get_schema return a defined InvalidArgument error (the kernel null-checks the handle, never UB), so newKernelRows takes its r.Close() cleanup branch and returns an error, and the supplied OnClose must not fire. Verified by mutation: arming the callback before the schema import makes the test fail as intended. - Reworded the newKernelRows comment that claimed the deferred callback assignment matches Thrift's NewRows "just before returning" — Thrift actually assigns closeCallback during construction. Dropped the inaccurate comparison and state the invariant directly. Co-authored-by: Isaac Signed-off-by: Mani Kaustubh Mathur --- internal/backend/kernel/kernel_test.go | 26 ++++++++++++++++++++------ internal/backend/kernel/rows.go | 6 +++--- 2 files changed, 23 insertions(+), 9 deletions(-) diff --git a/internal/backend/kernel/kernel_test.go b/internal/backend/kernel/kernel_test.go index e5a32699..f535e013 100644 --- a/internal/backend/kernel/kernel_test.go +++ b/internal/backend/kernel/kernel_test.go @@ -410,14 +410,28 @@ func TestKernelRowsCloseFiresOnClose(t *testing.T) { }) t.Run("construction-failure Close must not fire a success OnClose", func(t *testing.T) { - // newKernelRows arms the callback only after a successful build, so its - // cleanup Close() on a schema-fetch/import failure has callbacks==nil and - // must NOT record a (falsely successful) CLOSE_STATEMENT. Model that state. + // Drive the real newKernelRows construction-failure path: a nil result + // stream makes kernel_result_stream_get_schema return a defined + // InvalidArgument error (the kernel null-checks the handle — never UB), so + // newKernelRows takes its cleanup r.Close() branch and returns an error. The + // callback must NOT have been armed yet, so the supplied OnClose must not + // fire a (falsely successful) CLOSE_STATEMENT for a statement that produced + // no rows. This is the invariant that keeping r.callbacks unset until after + // a successful build guarantees. fired := false - r := &kernelRows{ /* callbacks intentionally nil, as during construction */ } - _ = r.Close() + cb := &dbsqlrows.TelemetryCallbacks{ + OnClose: func(int64, int, error, error) { fired = true }, + } + op := &kernelOp{backend: &KernelBackend{}} // for evictIfSessionFatal on the error path + rows, err := newKernelRows(context.Background(), op, nil /* stream */, cb) + if err == nil { + t.Fatal("newKernelRows(nil stream) = nil error, want a construction failure") + } + if rows != nil { + t.Errorf("newKernelRows on failure = %v rows, want nil", rows) + } if fired { - t.Error("OnClose fired for a rows object that never finished construction") + t.Error("OnClose fired during construction-failure cleanup — callback armed too early") } }) } diff --git a/internal/backend/kernel/rows.go b/internal/backend/kernel/rows.go index 5a7cb0e4..44557bf1 100644 --- a/internal/backend/kernel/rows.go +++ b/internal/backend/kernel/rows.go @@ -65,9 +65,9 @@ func newKernelRows(ctx context.Context, op *kernelOp, stream *C.kernel_result_st // r.Close() calls below run when construction FAILS (schema fetch/import), and a // Close() with the callback set would fire OnClose as a *successful* close for a // statement that never produced rows — masking the failure in CLOSE_STATEMENT - // telemetry. Assign it only on the success path (matching the Thrift NewRows, - // which sets closeCallback just before returning); the construction error itself - // is surfaced to and recorded by the conn execute path. + // telemetry. Assign it only on the success path so cleanup Close() on a + // schema/import failure does not record a falsely successful CLOSE_STATEMENT; the + // construction error itself is surfaced to and recorded by the conn execute path. r := &kernelRows{ctx: ctx, op: op, stream: stream, keyCache: arrowscan.NewStructKeyCache()} var csch C.struct_ArrowSchema From 13e1b7c0ee2a73023f8efa8af7b73390b6a4e8d1 Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Wed, 15 Jul 2026 14:05:16 +0000 Subject: [PATCH 24/24] fix(kernel): fix errcheck lint + trim comments and cross-references Lint fix (the failing CI check): TestResolveKernelLogArg used unchecked os.Setenv/os.Unsetenv (5 errcheck violations). Switch to t.Setenv, which auto-restores and returns nothing; the "unset" cases become t.Setenv(key, "") since resolveKernelLogArg gates on os.Getenv != "" (empty == unset). Drops the now-unused os import and the manual save/restore. Comment/doc cleanup (no behavior change): - Trim doc.go's kernel section (~40% shorter): fold the repetitive logging and cancellation prose, drop redundant parentheticals; every user-actionable fact (supported options, U2M-blocks-on-browser, Rust-log caveats, TLS knobs) kept. - Remove references to other PRs, other language bindings, and internal review/milestone artifacts from code comments, keeping the technical rationale. Verified: default-build lint 0 issues, gofmt clean, default + kernel-tagged suites pass. Co-authored-by: Isaac Signed-off-by: Mani Kaustubh Mathur --- doc.go | 158 +++++++----------- internal/arrowscan/arrowscan.go | 8 +- internal/arrowscan/arrowscan_test.go | 3 +- internal/backend/kernel/auth.go | 6 +- internal/backend/kernel/backend.go | 16 +- internal/backend/kernel/kernel_test.go | 2 +- internal/backend/kernel/logging_level_test.go | 40 ++--- internal/backend/kernel/namespace.go | 4 +- internal/config/config.go | 1 - kernel_config.go | 10 +- kernel_parity_test.go | 2 +- 11 files changed, 92 insertions(+), 158 deletions(-) diff --git a/doc.go b/doc.go index 6856d4da..91582659 100644 --- a/doc.go +++ b/doc.go @@ -186,116 +186,70 @@ the C header) under the cgo link path; `make build-kernel` does both steps: In a build without the tag, WithUseKernel(true) returns a clear error at connect time rather than silently using Thrift. -Kernel-backend logging follows the same knob as the rest of the driver. The -binding-layer trace (session/statement/batch steps) is emitted through the shared -logger at Debug level, so DATABRICKS_LOG_LEVEL=debug or dbsql.SetLogLevel("debug") -turns it on — no separate switch — and it lands in the same sink the rest of the -driver uses (including a custom one set via logger.SetLogOutput). Request-scoped -binding lines carry the same structured connId/corrId/queryId fields as the -driver's other logs where a context is in scope (session open/close, execute, -result-batch fetch), so they can be correlated in a multi-connection process; a few -context-less lines (parameter binds, operation teardown) are emitted without those -fields. At the default Warn level the lines are suppressed with no cost (zerolog -no-ops a below-level event), including during benchmarks. +Kernel-backend logging uses the same knob as the rest of the driver: the +binding-layer trace is emitted through the shared logger at Debug level, so +DATABRICKS_LOG_LEVEL=debug or dbsql.SetLogLevel("debug") turns it on and it lands in +the driver's normal sink (including a custom logger.SetLogOutput). Where a context is +in scope, lines carry the usual connId/corrId/queryId fields. At the default Warn +level the lines are suppressed with no cost. dbsql.SetLogLevel("debug") // or DATABRICKS_LOG_LEVEL=debug -The driver's log level is also mapped into the kernel's own Rust log subscriber, so -the same knob governs the kernel's internal (Rust) verbosity. Two caveats apply to -the Rust logs specifically, both consequences of the kernel's process-wide logging -ABI rather than the Go layer: - - - They are written to stderr directly and are NOT affected by - logger.SetLogOutput — an app that routes the driver's Go logs to a file or an - app logger will still find the Rust lines on stderr, not interleaved with that - sink. - - The Rust verbosity is fixed at the level in effect when the FIRST kernel - session in the process is opened (the kernel subscriber is process-wide and - installed once); a later dbsql.SetLogLevel changes the Go binding lines but not - the already-installed Rust subscriber. Set the level before opening the first - kernel connection to control the Rust logs. - -The kernel's Rust lines are also currently plain text and do NOT yet carry the -structured connId/corrId/queryId fields — that needs a kernel log-callback ABI -(tracked separately); only the Go binding lines are structured today. - -For finer control of the kernel's Rust verbosity independent of the driver level, -set DBSQL_KERNEL_DEBUG to any non-empty value: it forces the kernel subscriber on -and defers to RUST_LOG for the kernel's threshold. Filter on the target -databricks::sql::kernel (note the colons — the kernel's explicit log target, NOT -the crate module path databricks_sql_kernel, which would match nothing): +The same level is mapped into the kernel's internal (Rust) log subscriber, with two +caveats specific to the Rust lines: they go to stderr directly (not affected by +logger.SetLogOutput), and their verbosity is fixed when the first kernel session in +the process is opened — set the level before that first connect to control them. + +For finer control of the Rust verbosity independent of the driver level, set +DBSQL_KERNEL_DEBUG to any non-empty value: it forces the kernel subscriber on and +defers to RUST_LOG. Filter on the target databricks::sql::kernel (note the colons): # kernel logs only, at the kernel's own verbosity: DBSQL_KERNEL_DEBUG=1 RUST_LOG=databricks::sql::kernel=debug ./your_app 2>&1 - # kernel logs plus its HTTP stack (hyper/reqwest): + # kernel logs plus its HTTP stack: DBSQL_KERNEL_DEBUG=1 RUST_LOG=debug ./your_app 2>&1 -The kernel backend currently supports PAT and OAuth (M2M via WithClientCredentials, -and U2M via the authType=oauthU2M DSN param — the interactive browser/PKCE flow is -owned by the kernel) authentication; reading scalar, nested, and complex-typed -results (CloudFetch is handled transparently); context cancellation during execute -(a cancelled ctx fires a real server-side cancel; on the read path cancellation is -honored at result-batch boundaries, not mid-fetch); the initial namespace -(WithInitialNamespace catalog/schema, applied post-connect via USE CATALOG / USE -SCHEMA since the kernel C ABI has no namespace setter); metric-view metadata -(WithEnableMetricViewMetadata, sent as the same server session conf the Thrift -backend uses); and the TLS, proxy, and session-conf (query tags, statement timeout, -time zone) connection options. WithTimeout (a server query timeout the kernel C ABI -can't set) and WithRetries used to disable retries (the kernel retries internally) -return a clear error at connect time rather than being silently ignored; token- -provider, external/static, and federated authenticators are likewise not supported -and rejected loudly. Staging operations (PUT/GET/REMOVE on a Unity Catalog volume, -which need a local file transfer this backend cannot perform) are not yet supported -and return a clear error at execute time (they are per-statement, not connect-time). -Bound query parameters (positional and named) are supported — bound over the C ABI -to match the Thrift path. None of these is silently ignored. (Metadata is issued as -ordinary SQL — SHOW/DESCRIBE/information_schema — and runs on this backend like any -other query.) - -OAuth U2M is interactive. On a cache miss (no valid cached refresh token) opening a -connection launches the system browser and blocks until the user completes login or -the kernel's built-in callback timeout (~120s) expires — and because the kernel C -ABI cannot interrupt session open mid-call, a deadline on the connection context is -not honored during that window (nor can the timeout be shortened; the C ABI exposes -no override). A cached, still-valid token opens without a browser. Use PAT or OAuth -M2M for headless / deadline-bound connects. - -WithMaxRows and retry tuning (WithRetries with a positive limit) are accepted but -not applied on the kernel path: the kernel manages result fetching and retries -internally, below the C ABI, with no user-facing knob. - -Experimental kernel-only TLS options. The kernel exposes a richer TLS surface than -the backend-neutral WithSkipTLSHostVerify (which relaxes both chain and hostname -checks). These are kernel-only — the default (Thrift) backend rejects them at -connect rather than silently ignoring them — and experimental (the WithKernel* -prefix marks both): - - WithKernelTrustedCerts(pem) adds a PEM CA bundle to the kernel's trust store on - top of the system roots (for a corporate re-signing proxy or an on-prem CA). - Required rather than relying on SSL_CERT_FILE: the kernel's TLS stack (rustls) - does not read that environment variable. - - WithKernelSkipHostnameVerify() skips only the certificate hostname check while - keeping chain validation (finer-grained than WithSkipTLSHostVerify). - -Setting either without WithUseKernel makes Connect fail with an error wrapping the -sentinel ErrRequiresKernelBackend (the mirror of ErrNotSupportedByKernel), so the -"add WithUseKernel or remove it" case is detectable with errors.Is rather than by -matching message text. - -Features that live above the backend seam are inherited unchanged: the -database/sql connection pool (each connection wraps one kernel session), -per-connection telemetry (CREATE_SESSION, EXECUTE_STATEMENT, and DELETE_SESSION -are recorded for kernel connections just like Thrift), and the telemetry -exporter's circuit breaker are all backend-agnostic. Result types are rendered to -match the Thrift backend byte-for-byte: scalars, DECIMAL (exact string), -TIMESTAMP / TIMESTAMP_NTZ (both shifted into the session time zone, as Thrift -does), INTERVAL year-month and day-time, nested Array/Map/Struct and VARIANT (as -JSON), and GEOMETRY (WKT). The per-statement server query id is surfaced on the -success path, so a QueryIdCallback (see below) fires with the real id and -EXECUTE_STATEMENT telemetry carries it. - -One kernel-backend limitation remains on the read path: context cancellation is -honored at result-batch boundaries, not mid-fetch (an in-flight CloudFetch batch -runs to completion before the cancel takes effect). +Supported on the kernel backend: PAT and OAuth (M2M via WithClientCredentials, U2M +via the authType=oauthU2M DSN param); reading scalar, nested, and complex-typed +results (CloudFetch is transparent); bound query parameters (positional and named); +context cancellation during execute; the initial namespace (WithInitialNamespace, +applied post-connect via USE CATALOG / USE SCHEMA); metric-view metadata +(WithEnableMetricViewMetadata); and the TLS, proxy, and session-conf (query tags, +statement timeout, time zone) options. Nothing is silently ignored: WithTimeout, a +retries-disabling WithRetries, and token-provider / external / federated +authenticators are rejected at connect; staging (PUT/GET/REMOVE on a Unity Catalog +volume) is rejected at execute. WithMaxRows and positive-limit WithRetries are +accepted but inert (the kernel manages fetching and retries below the C ABI). + +OAuth U2M is interactive: on a cache miss, connecting launches the system browser and +blocks until login completes or the kernel's ~120s callback timeout expires. Because +the C ABI can't interrupt session open mid-call, a connection-context deadline is not +honored during that window. Use PAT or OAuth M2M for headless / deadline-bound +connects. + +Experimental kernel-only TLS options (rejected by the default backend; the +WithKernel* prefix marks them experimental): + - WithKernelTrustedCerts(pem) adds a PEM CA bundle on top of the system roots (for + a re-signing proxy or on-prem CA). Required because the kernel's TLS stack does + not read SSL_CERT_FILE. + - WithKernelSkipHostnameVerify() skips only the hostname check while keeping chain + validation (finer-grained than WithSkipTLSHostVerify). + +Setting either without WithUseKernel fails Connect with an error wrapping the +sentinel ErrRequiresKernelBackend, detectable with errors.Is. + +Features above the backend seam are inherited unchanged: the database/sql connection +pool, per-connection telemetry (CREATE_SESSION / EXECUTE_STATEMENT / DELETE_SESSION), +and the telemetry circuit breaker. Result types render byte-for-byte identical to the +Thrift backend: scalars, DECIMAL (exact string), TIMESTAMP / TIMESTAMP_NTZ (shifted +into the session time zone), INTERVAL, nested Array/Map/Struct and VARIANT (as JSON), +and GEOMETRY (WKT). The server query id is surfaced on the success path, so a +QueryIdCallback (see below) fires with the real id and EXECUTE_STATEMENT telemetry +carries it. + +On the read path, context cancellation is honored at result-batch boundaries, not +mid-fetch: an in-flight CloudFetch batch runs to completion before the cancel takes +effect. # Programmatically Retrieving Connection and Query Id diff --git a/internal/arrowscan/arrowscan.go b/internal/arrowscan/arrowscan.go index 4c020103..cee634e2 100644 --- a/internal/arrowscan/arrowscan.go +++ b/internal/arrowscan/arrowscan.go @@ -14,7 +14,7 @@ // - nested NULL → null // - time.Time → quoted .String() (matches the Thrift marshal() special-case) // - nested decimal → exact scale-applied JSON number literal (never a lossy -// float64), matching Thrift's marshalScalar → ValueString (#253/#274) +// float64), matching Thrift's marshalScalar → ValueString // - float32 → native float32 (not widened to float64), so JSON renders // 3.14, not 3.140000104904175 package arrowscan @@ -34,8 +34,8 @@ import ( // ScanCell extracts one cell as a driver.Value. Scalars map to their Go value: // bool, all int/uint widths, float (native float32/float64), string, binary, // date, timestamp, and top-level decimal (as an exact fixed-point string, -// matching the Thrift path — a float64 would lose precision beyond ~17 digits; -// see databricks-sql-go#274). Nested types (List/Map/Struct, and VARIANT which +// matching the Thrift path — a float64 would lose precision beyond ~17 digits). +// Nested types (List/Map/Struct, and VARIANT which // arrives nested) render to a JSON string byte-identical to the Thrift path; // GEOMETRY arrives as a WKB/WKT string and is handled by the string arm. INTERVAL // day-time/year-month arrive as native arrow duration/month-interval and format to @@ -331,7 +331,7 @@ func writeJSON(b *strings.Builder, col arrow.Array, row int, loc *time.Location, // Emit the exact scale-applied decimal as a raw JSON number literal, not a // float64 — a float64 would render DECIMAL(5,2) 19.99 as 19.990000000000002 // and corrupt high-precision values. Matches the Thrift path's marshalScalar - // → ValueString (databricks-sql-go#253/#274). + // → ValueString. b.WriteString(decimalfmt.ExactString(c.Value(row), col.DataType().(*arrow.Decimal128Type).Scale)) return nil case *array.Float32: diff --git a/internal/arrowscan/arrowscan_test.go b/internal/arrowscan/arrowscan_test.go index a48d38d9..957a8f47 100644 --- a/internal/arrowscan/arrowscan_test.go +++ b/internal/arrowscan/arrowscan_test.go @@ -118,8 +118,7 @@ func TestScanCellScalars(t *testing.T) { // INTERVAL day-time (arrow duration) and year-month (arrow month-interval) arrive // as native arrow values on the kernel path and must format to the exact string the // Thrift path receives pre-formatted from the server: "D HH:MM:SS.nnnnnnnnn" and -// "years-months", with negatives signed. (These formatters were validated live -// kernel==Thrift in the PuPr POC; this is the regression guard.) +// "years-months", with negatives signed. func TestScanCellInterval(t *testing.T) { pool := memory.NewGoAllocator() diff --git a/internal/backend/kernel/auth.go b/internal/backend/kernel/auth.go index a78d23f5..c3fa35c5 100644 --- a/internal/backend/kernel/auth.go +++ b/internal/backend/kernel/auth.go @@ -25,9 +25,9 @@ const ( // Thrift path hardcodes both), so resolveKernelAuth leaves them zero and the kernel // applies its defaults. They are kept — rather than dropped and the setter hardcoded // to NULL/0 — so kernel.Auth models the full set_auth_u2m surface: adding a future -// WithOAuthRedirectPort / scopes option (ODBC PR #102 already exposes a redirect -// port) becomes populating these, not re-plumbing the setter. TestSetAuthByMode's -// "U2M full" case pins that marshalling so the dormant path stays correct. +// WithOAuthRedirectPort / scopes option becomes populating these, not re-plumbing +// the setter. TestSetAuthByMode's "U2M full" case pins that marshalling so the +// dormant path stays correct. type Auth struct { Mode AuthMode Token string // PAT diff --git a/internal/backend/kernel/backend.go b/internal/backend/kernel/backend.go index ad20c97c..2ad46cc8 100644 --- a/internal/backend/kernel/backend.go +++ b/internal/backend/kernel/backend.go @@ -28,9 +28,8 @@ import ( var kernelSessionSeq atomic.Uint64 // 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 -// (this mirrors how the kernel's pyo3/napi bindings take flat connection -// params). Zero-valued fields are simply not applied. +// fills it from the driver's config so the user-facing options are unchanged. +// Zero-valued fields are simply not applied. type Config struct { Host string // workspace hostname, no scheme HTTPPath string // e.g. /sql/1.0/warehouses/abc123 (carries ?o= org routing) @@ -67,8 +66,8 @@ type Config struct { // Catalog / Schema select the initial namespace. The kernel C ABI has no // catalog/schema config setter, so OpenSession applies them post-connect by - // running USE CATALOG / USE SCHEMA (the OSS ODBC driver's workaround). Empty - // leaves the session in the server default namespace. + // running USE CATALOG / USE SCHEMA. Empty leaves the session in the server + // default namespace. Catalog string Schema string } @@ -233,8 +232,8 @@ func (k *KernelBackend) OpenSession(ctx context.Context) error { k.sessionID = fmt.Sprintf("kernel-%d", kernelSessionSeq.Add(1)) // Initial namespace: the kernel C ABI has no catalog/schema config setter, so - // select it post-connect with USE CATALOG / USE SCHEMA (the OSS ODBC driver's - // approach). A failure here means the session is not in the requested namespace + // select it post-connect with USE CATALOG / USE SCHEMA. A failure here means + // the session is not in the requested namespace // — a correctness precondition, like Thrift's InitialNamespace — so fail the // connect and close the session we just opened (the connector does not call // CloseSession on an OpenSession error). @@ -394,8 +393,7 @@ func (k *KernelBackend) runNamespaceStmt(ctx context.Context, sql string) error // kernel_session_close returns, with no deadline — a stalled kernel-side close // (e.g. a shutdown-time network partition) can block database/sql pool cleanup. // A bounded close needs either a kernel_session_close_blocking with a deadline or -// a Go-side watchdog; grouped with the kernel C-ABI follow-ups. Not fixed here to -// keep this PR scoped to the opt-in backend rather than add a watchdog. +// a Go-side watchdog; grouped with the kernel C-ABI follow-ups. func (k *KernelBackend) CloseSession(ctx context.Context) error { if k.session == nil { return nil diff --git a/internal/backend/kernel/kernel_test.go b/internal/backend/kernel/kernel_test.go index f535e013..4182477c 100644 --- a/internal/backend/kernel/kernel_test.go +++ b/internal/backend/kernel/kernel_test.go @@ -85,7 +85,7 @@ func TestSetKernelTLS(t *testing.T) { // matters: logger.WithContext eagerly allocates (zerolog's With() does // make([]byte, 0, 500)) BEFORE the .Debug() gate, so without the up-front // kernelDebugOff() guard this would allocate ~500 B per call — per Arrow batch on the -// nextBatch hot path. Guards against reintroducing that (an Isaac Review MAJOR). +// nextBatch hot path. Guards against reintroducing that. func TestKernelLogNoAllocWhenOff(t *testing.T) { prev := logger.Logger.GetLevel() if err := logger.SetLogLevel("warn"); err != nil { // the default level diff --git a/internal/backend/kernel/logging_level_test.go b/internal/backend/kernel/logging_level_test.go index 90449f71..cf03c25e 100644 --- a/internal/backend/kernel/logging_level_test.go +++ b/internal/backend/kernel/logging_level_test.go @@ -1,7 +1,6 @@ package kernel import ( - "os" "testing" "github.com/databricks/databricks-sql-go/logger" @@ -37,45 +36,30 @@ func TestKernelLogLevel(t *testing.T) { } } -// TestResolveKernelLogArg pins the level-vs-NULL decision initKernelLogging feeds to -// kernel_init_logging — the behavior the PR justifies in prose but the cgo path -// can't easily assert. Two invariants: DBSQL_KERNEL_DEBUG (any non-empty value) must -// yield useNULL=true so the kernel honors RUST_LOG (the whole point of retaining the -// knob — a NULL level is the ONLY way the kernel consults RUST_LOG); otherwise the -// driver's current level is mapped in, so DATABRICKS_LOG_LEVEL governs kernel -// verbosity too. Runs under CGO_ENABLED=0 (no cgo in the resolution path). +// TestResolveKernelLogArg pins the level-vs-NULL decision: DBSQL_KERNEL_DEBUG (any +// non-empty value) yields useNULL=true so the kernel honors RUST_LOG; otherwise the +// driver level is mapped in. Empty is treated as unset (the gate is os.Getenv != ""). func TestResolveKernelLogArg(t *testing.T) { - // Save/restore the env var and the global logger level so this test leaves no - // state behind for its siblings. - prevEnv, hadEnv := os.LookupEnv("DBSQL_KERNEL_DEBUG") + // t.Setenv restores the var after the test; save/restore the global logger level + // separately so nothing leaks to sibling tests. prevLevel := logger.Logger.GetLevel() - t.Cleanup(func() { - if hadEnv { - os.Setenv("DBSQL_KERNEL_DEBUG", prevEnv) - } else { - os.Unsetenv("DBSQL_KERNEL_DEBUG") - } - logger.Logger.Logger = logger.Logger.Level(prevLevel) - }) + t.Cleanup(func() { logger.Logger.Logger = logger.Logger.Level(prevLevel) }) - // DBSQL_KERNEL_DEBUG set → NULL level (kernel honors RUST_LOG), regardless of the - // driver level. Pin it at debug to prove the override wins over the mapped level. + // Set → NULL level regardless of the driver level (pinned at debug to prove the + // override wins over the mapped level). logger.Logger.Logger = logger.Logger.Level(zerolog.DebugLevel) - os.Setenv("DBSQL_KERNEL_DEBUG", "1") + t.Setenv("DBSQL_KERNEL_DEBUG", "1") if lvl, useNULL := resolveKernelLogArg(); !useNULL || lvl != "" { t.Errorf("with DBSQL_KERNEL_DEBUG=1: got (level=%q, useNULL=%v), want (\"\", true)", lvl, useNULL) } - // An empty value is treated as unset (os.Getenv != "" is the gate), so it must - // NOT trigger the override — the mapped level is used. - os.Setenv("DBSQL_KERNEL_DEBUG", "") + // Empty is treated as unset → the mapped level is used, not the override. + t.Setenv("DBSQL_KERNEL_DEBUG", "") if lvl, useNULL := resolveKernelLogArg(); useNULL || lvl != "DEBUG" { t.Errorf("with DBSQL_KERNEL_DEBUG=\"\" at debug: got (level=%q, useNULL=%v), want (\"DEBUG\", false)", lvl, useNULL) } - // Unset → map the driver level. Check two levels to prove it tracks the logger, - // including the fatal→OFF collapse. - os.Unsetenv("DBSQL_KERNEL_DEBUG") + // Mapped level tracks the logger, including the fatal→OFF collapse. logger.Logger.Logger = logger.Logger.Level(zerolog.WarnLevel) if lvl, useNULL := resolveKernelLogArg(); useNULL || lvl != "WARN" { t.Errorf("unset at warn: got (level=%q, useNULL=%v), want (\"WARN\", false)", lvl, useNULL) diff --git a/internal/backend/kernel/namespace.go b/internal/backend/kernel/namespace.go index b1f907c3..91f43514 100644 --- a/internal/backend/kernel/namespace.go +++ b/internal/backend/kernel/namespace.go @@ -10,8 +10,8 @@ import "strings" // quoteIdent renders name as a backtick-quoted Databricks SQL identifier, doubling // any embedded backtick. The kernel C ABI exposes no catalog/schema config setter, // so the initial namespace is applied post-connect by running USE CATALOG / USE -// SCHEMA (the same workaround the OSS ODBC driver uses); quoting makes those -// statements injection-safe for arbitrary identifier text. +// SCHEMA; quoting makes those statements injection-safe for arbitrary identifier +// text. func quoteIdent(name string) string { return "`" + strings.ReplaceAll(name, "`", "``") + "`" } diff --git a/internal/config/config.go b/internal/config/config.go index 2113a64f..f689f414 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -189,7 +189,6 @@ type UserConfig struct { // connector copies it into Config.ArrowConfig.UseArrowNativeDecimal when it is // assembled. The name intentionally differs from ArrowConfig's field so the // promoted selector Config.UseArrowNativeDecimal stays unambiguous. - // See databricks/databricks-sql-go#274. UseArrowNativeDecimalDSN bool CloudFetchConfig // UseKernel selects the SEA-via-kernel backend instead of Thrift. See the diff --git a/kernel_config.go b/kernel_config.go index cf5fa0d0..adc97e98 100644 --- a/kernel_config.go +++ b/kernel_config.go @@ -34,8 +34,8 @@ import ( func validateKernelConfig(cfg *config.Config) (kernel.Auth, error) { // 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 (the OSS ODBC driver's - // workaround). No per-backend handling needed here. + // selects it post-connect with USE CATALOG / USE SCHEMA. No per-backend handling + // needed here. // EnableMetricViewMetadata is forwarded, not rejected: config.EffectiveSessionParams // folds its server conf (spark.sql.thriftserver.metadata.metricview.enabled=true) // into SessionConf backend-neutrally, so the kernel path sends the identical conf @@ -90,9 +90,9 @@ func validateKernelConfig(cfg *config.Config) (kernel.Auth, error) { } // resolveKernelAuth picks the kernel auth form from the config. The kernel backend -// drives the kernel's own OAuth flow from raw credentials (mirroring pyo3/napi and -// the Node/Python kernel bindings) rather than reusing the Go authenticator's -// Authenticate method. It reads those credentials off cfg.Authenticator — the +// drives the kernel's own OAuth flow from raw credentials rather than reusing the Go +// authenticator's Authenticate method. It reads those credentials off +// cfg.Authenticator — the // single source of truth for auth, so the last WithX option applied wins for both // backends (matching Thrift's last-writer-wins on cfg.Authenticator). The M2M/U2M // authenticator types are unexported, so it asserts the small diff --git a/kernel_parity_test.go b/kernel_parity_test.go index b36d6550..75ae619f 100644 --- a/kernel_parity_test.go +++ b/kernel_parity_test.go @@ -78,7 +78,7 @@ type paramCase struct { // The kernel binds these via kernel_statement_bind_parameter (the driver's // backend.Param{Name, Type, Value}); Thrift binds via toSparkParameters. Covers // positional (?) and named (:n) markers, each scalar type, SQL NULL, multi-param, -// and a predicate — mirroring the C2 POC gate. +// and a predicate. var paramCases = []paramCase{ {"pos_int", "SELECT ? AS v", []any{int64(42)}}, {"pos_string", "SELECT ? AS v", []any{"hello"}},