From 1f29fd791cda54b2d3bb9caf6314a8f3b1c84ad0 Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Thu, 9 Jul 2026 05:52:49 +0000 Subject: [PATCH 01/32] feat(kernel): add SEA-via-kernel backend (scalar read path) Add a second execution backend that runs statements over the Statement Execution API via the Rust databricks-sql-kernel, reached through a cgo C ABI. It is opt-in behind WithUseKernel + the databricks_kernel build tag, so the default build stays pure-Go (CGO_ENABLED=0) and is unaffected. Pure-Go side (always compiled): - WithUseKernel / WithWarehouseID connector options; UseKernel / WarehouseID on config with DSN parsing (useKernel, warehouseId). - connector.Connect selects the backend via newKernelBackend, which in a build without the tag is a stub returning a clear "not compiled in" error rather than silently falling back to Thrift. Kernel side (//go:build cgo && databricks_kernel): - KernelBackend/kernelOp implement backend.Backend/Operation over the C ABI: PAT session open (warehouse id or http path), blocking execute with out-of-band context cancellation (a watcher goroutine drives the kernel's detached statement canceller), and result streaming. - kernelRows imports Arrow batches zero-copy via the Arrow C Data Interface and scans the scalar types (ints, floats, bool, string, binary, date, timestamp, and top-level decimal as an exact string per #274); unsupported types return an explicit error. KernelError maps to the driver's error surface with sqlstate, and to driver.ErrBadConn for session-unusable statuses. Tests: tagged unit tests for error/status mapping and scalar scanning; live e2e (exercised against a staging warehouse) for select, per-type scanning, CloudFetch, and context cancellation, plus a Thrift-parity check that both backends render scalars identically. The cgo link directives point at a locally built kernel; committing a prebuilt per-platform static lib and adding a tagged CI job are a follow-up, so the kernel path is not yet covered by CI. Co-authored-by: Isaac Signed-off-by: Mani Kaustubh Mathur --- connector.go | 45 ++++- doc.go | 36 ++++ internal/backend/kernel/backend.go | 134 +++++++++++++++ internal/backend/kernel/cgo.go | 179 ++++++++++++++++++++ internal/backend/kernel/kernel_test.go | 158 +++++++++++++++++ internal/backend/kernel/operation.go | 210 +++++++++++++++++++++++ internal/backend/kernel/rows.go | 224 +++++++++++++++++++++++++ internal/backend/kernel/scan.go | 40 +++++ internal/config/config.go | 20 +++ kernel_backend.go | 24 +++ kernel_e2e_test.go | 170 +++++++++++++++++++ kernel_parity_test.go | 94 +++++++++++ kernel_stub.go | 20 +++ 13 files changed, 1350 insertions(+), 4 deletions(-) create mode 100644 internal/backend/kernel/backend.go create mode 100644 internal/backend/kernel/cgo.go create mode 100644 internal/backend/kernel/kernel_test.go create mode 100644 internal/backend/kernel/operation.go create mode 100644 internal/backend/kernel/rows.go create mode 100644 internal/backend/kernel/scan.go create mode 100644 kernel_backend.go create mode 100644 kernel_e2e_test.go create mode 100644 kernel_parity_test.go create mode 100644 kernel_stub.go diff --git a/connector.go b/connector.go index e1701fcb..61d676ad 100644 --- a/connector.go +++ b/connector.go @@ -15,6 +15,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" + "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" "github.com/databricks/databricks-sql-go/internal/config" @@ -32,9 +33,18 @@ type connector struct { func (c *connector) Connect(ctx context.Context) (driver.Conn, error) { defer debuglog.Track(ctx, "connector.Connect", "host=%s", c.cfg.Host)() - // Build the execution backend. Currently always the Thrift backend; a second - // SEA-via-kernel backend will be selected here by a connect-time flag. - be, err := thrift.New(ctx, c.cfg, c.client) + // Build the execution backend. Thrift is the default; the SEA-via-kernel + // backend is selected when UseKernel is set. newKernelBackend is build-tag + // gated: in the default pure-Go build it returns a clear "not linked in" + // error, so the kernel path compiles and links only under -tags + // databricks_kernel + CGO_ENABLED=1. + var be backend.Backend + var err error + if c.cfg.UseKernel { + be, err = newKernelBackend(ctx, c.cfg) + } else { + be, err = thrift.New(ctx, c.cfg, c.client) + } if err != nil { return nil, err } @@ -77,7 +87,14 @@ func (c *connector) Connect(ctx context.Context) (driver.Conn, error) { conn.telemetry.RecordOperation(ctx, conn.id, "", telemetry.OperationTypeCreateSession, sessionLatencyMs, nil) } - log.Info().Msgf("connect: host=%s port=%d httpPath=%s serverProtocolVersion=0x%X", c.cfg.Host, c.cfg.Port, c.cfg.HTTPPath, be.ServerProtocolVersion()) + // ServerProtocolVersion is Thrift-specific (not on the neutral backend + // interface); the kernel backend has no negotiated Thrift protocol, so log it + // only when present. + if tb, ok := be.(*thrift.Backend); ok { + log.Info().Msgf("connect: host=%s port=%d httpPath=%s serverProtocolVersion=0x%X", c.cfg.Host, c.cfg.Port, c.cfg.HTTPPath, tb.ServerProtocolVersion()) + } else { + log.Info().Msgf("connect: host=%s port=%d httpPath=%s backend=kernel", c.cfg.Host, c.cfg.Port, c.cfg.HTTPPath) + } return conn, nil } @@ -301,6 +318,26 @@ func WithHTTPPath(path string) ConnOption { } } +// WithUseKernel selects the SEA-via-kernel backend instead of the default +// Thrift backend. It has effect only in a build compiled with +// `-tags databricks_kernel` and CGO_ENABLED=1; in the default pure-Go build a +// connection made with this option set returns a clear error at connect time +// (the kernel backend is not linked in). +func WithUseKernel(useKernel bool) ConnOption { + return func(c *config.Config) { + c.UseKernel = useKernel + } +} + +// WithWarehouseID sets the bare SQL warehouse id. The kernel backend addresses a +// warehouse by id; when set it is preferred over the http path. The Thrift +// backend ignores it and continues to route by http path. +func WithWarehouseID(id string) ConnOption { + return func(c *config.Config) { + c.WarehouseID = id + } +} + // WithMaxRows sets up the max rows fetched per request. Default is 10000 func WithMaxRows(n int) ConnOption { return func(c *config.Config) { diff --git a/doc.go b/doc.go index 6ae9616b..becab1f8 100644 --- a/doc.go +++ b/doc.go @@ -155,6 +155,42 @@ The result log may look like this: {"level":"debug","connId":"01ed6545-5669-1ec7-8c7e-6d8a1ea0ab16","corrId":"workflow-example","queryId":"01ed6545-57cc-188a-bfc5-d9c0eaf8e189","time":1668558402,"message":"Run Main elapsed time: 1.298712292s"} +# SEA-via-kernel backend (experimental) + +By default the driver uses the Thrift/HiveServer2 backend and is pure Go +(CGO_ENABLED=0, go-gettable, cross-compilable). An experimental second backend +runs statements over the Statement Execution API via the Rust +databricks-sql-kernel, reached through a cgo C ABI. It is opt-in and compiled in +only under a build tag, so the default build is unchanged. + +To use it, build with the databricks_kernel tag and CGO enabled, and select it +per connection with WithUseKernel: + + CGO_ENABLED=1 go build -tags databricks_kernel ./... + + connector, _ := dbsql.NewConnector( + dbsql.WithServerHostname(host), + dbsql.WithHTTPPath(httpPath), + dbsql.WithAccessToken(token), + dbsql.WithUseKernel(true), + ) + db := sql.OpenDB(connector) + +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, enable both on the +same stderr. The kernel filters on the target databricks::sql::kernel (note the +colons — it is not the crate module path), and the driver's step tracer is +enabled by the standard logging level: + + RUST_LOG=databricks::sql::kernel=debug DATABRICKS_LOG_LEVEL=debug ./your_app 2>&1 + +The kernel backend currently supports PAT authentication and reading +scalar-typed results (CloudFetch is handled transparently), with context +cancellation. Bound parameters, nested and complex types, and other +authentication and transport options are not yet supported on this backend. + # Programmatically Retrieving Connection and Query Id Use the driverctx package under driverctx/ctx.go to add callbacks to the query context to receive the connection id and query id. diff --git a/internal/backend/kernel/backend.go b/internal/backend/kernel/backend.go new file mode 100644 index 00000000..9e160252 --- /dev/null +++ b/internal/backend/kernel/backend.go @@ -0,0 +1,134 @@ +//go:build cgo && databricks_kernel + +package kernel + +/* +#include +#include "databricks_kernel.h" +*/ +import "C" + +import ( + "context" + "fmt" + + "github.com/databricks/databricks-sql-go/internal/backend" +) + +// Config is the connection config for the kernel backend: host, the warehouse +// (by bare id or http path), and a PAT. +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...) +} + +// KernelBackend implements backend.Backend over the kernel C ABI. One backend +// backs one conn, which database/sql serializes to a single goroutine at a time, +// so the kernel session inherits single-owner-ship and needs no locks; the only +// concurrency is the per-statement cancel watcher (see operation.go), which +// touches only the kernel's internal inflight-id slot. +type KernelBackend struct { + cfg Config + session *C.kernel_session_t + sessionID string + valid bool +} + +var _ backend.Backend = (*KernelBackend)(nil) + +// New builds a kernel backend without opening the session; the connector calls +// OpenSession immediately after, mirroring the Thrift backend's shape. +func New(cfg Config) *KernelBackend { + return &KernelBackend{cfg: cfg} +} + +// OpenSession builds a session config (warehouse/http-path + PAT), opens the +// session, and captures a per-conn id. Called once by the connector at connect +// time. The config handle is consumed by kernel_session_open on success and +// freed by us on any earlier failure. +func (k *KernelBackend) OpenSession(ctx context.Context) error { + klog("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 { + return fmt.Errorf("kernel: config_new: %w", toDriverError(err)) + } + consumed := false + defer func() { + if !consumed { + C.kernel_session_config_free(cfg) + } + }() + + // Warehouse addressing: bare id when provided, else the http path (which also + // carries ?o= org routing for shared hosts). + host := newCStr(k.cfg.Host) + defer host.free() + if k.cfg.WarehouseID != "" { + wh := newCStr(k.cfg.WarehouseID) + defer wh.free() + if err := call(func() C.KernelStatusCode { + return C.kernel_session_config_set_warehouse(cfg, host.c, wh.c) + }); err != nil { + return fmt.Errorf("kernel: set_warehouse: %w", toDriverError(err)) + } + } else { + path := newCStr(k.cfg.HTTPPath) + defer path.free() + if err := call(func() C.KernelStatusCode { + return C.kernel_session_config_set_http_path(cfg, host.c, path.c) + }); err != nil { + return fmt.Errorf("kernel: set_http_path: %w", toDriverError(err)) + } + } + + 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", toDriverError(err)) + } + + var sess *C.kernel_session_t + if err := call(func() C.KernelStatusCode { return C.kernel_session_open(cfg, &sess) }); err != nil { + return fmt.Errorf("kernel: session_open: %w", toDriverError(err)) + } + consumed = true // kernel_session_open took ownership of cfg + k.session = sess + k.valid = true + // The C ABI exposes no formatted session-id accessor; use the handle pointer + // as a stable per-conn id for logging / telemetry correlation. + k.sessionID = fmt.Sprintf("kernel-%p", sess) + klog("OpenSession OK session=%s", k.sessionID) + return nil +} + +// 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. +func (k *KernelBackend) CloseSession(ctx context.Context) error { + if k.session == nil { + return nil + } + klog("CloseSession session=%s", k.sessionID) + err := call(func() C.KernelStatusCode { return C.kernel_session_close(k.session) }) + k.session = nil + k.valid = false + return toDriverError(err) +} + +// SessionValid backs conn.IsValid → pool eviction. No I/O; inspects state +// captured at OpenSession. +func (k *KernelBackend) SessionValid() bool { return k.valid && k.session != nil } + +// SessionID is the per-conn id (conn.id). Valid after OpenSession. +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. +func (k *KernelBackend) Execute(ctx context.Context, req backend.ExecRequest) (backend.Operation, error) { + return k.execute(ctx, req) +} diff --git a/internal/backend/kernel/cgo.go b/internal/backend/kernel/cgo.go new file mode 100644 index 00000000..f967c2a0 --- /dev/null +++ b/internal/backend/kernel/cgo.go @@ -0,0 +1,179 @@ +//go:build cgo && databricks_kernel + +// Package kernel implements backend.Backend over the Databricks SQL kernel's C +// ABI (databricks-sql-kernel, src/c_abi). It is build-tag-gated behind +// `databricks_kernel` so the default driver build stays pure-Go (CGO_ENABLED=0, +// go-gettable, cross-compilable); only a build that opts in with +// `-tags databricks_kernel` and CGO_ENABLED=1 links the kernel static lib. +// +// This file holds the cgo plumbing: the C ABI include/link directives, the +// fallible-call helper that makes the kernel's thread-local last error readable, +// the error mapping to the driver's error surface, and the gated step logger. +// The backend, operation, and rows layers live in sibling files. +// +// The link directives below point at a locally built kernel +// (~/oss/databricks-sql-kernel/target/release + include), so this package builds +// only on a machine with that checkout. A shippable build instead links a +// committed per-platform prebuilt static lib via a ${SRCDIR}-relative path (the +// go-duckdb duckdb-go-bindings model); see cabi-distribution-references.md. +package kernel + +/* +#cgo CFLAGS: -I/home/mani.mathur/oss/databricks-sql-kernel/include +#cgo LDFLAGS: -L/home/mani.mathur/oss/databricks-sql-kernel/target/release -ldatabricks_sql_kernel -Wl,-rpath,/home/mani.mathur/oss/databricks-sql-kernel/target/release -ldl -lm +#include +#include "databricks_kernel.h" +*/ +import "C" + +import ( + "fmt" + "os" + "runtime" + "unsafe" + + dbsqlerrint "github.com/databricks/databricks-sql-go/internal/errors" +) + +// ─── 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 kernel's own +// Rust logs are enabled separately via kernel_init_logging / RUST_LOG and +// interleave on the same stderr. +var kdebug = os.Getenv("DBSQL_KERNEL_DEBUG") != "" + +func klog(format string, args ...any) { + if !kdebug { + return + } + fmt.Fprintf(os.Stderr, "[kernel] "+format+"\n", args...) +} + +// KernelDebugEnabled reports whether binding-level debug logging is on. Tests +// assert the flag wiring; benchmarks assert it is false before measuring. +func KernelDebugEnabled() bool { return kdebug } + +// 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. +// +// The kernel reports rich errors via a thread-local buffer read by a *second* +// call (kernel_get_last_error). Go's M:N scheduler can move a goroutine to a +// different OS thread between two cgo calls, so a naive call-then-read pair can +// observe the wrong thread's buffer. LockOSThread pins the goroutine to its OS +// thread across the call and its error read, closing that window. (The kernel's +// preferred long-term shape is a same-call out-param error, which would remove +// the need for this; until then LockOSThread is correct and call-site-local.) +func call(fn func() C.KernelStatusCode) error { + runtime.LockOSThread() + defer runtime.UnlockOSThread() + st := fn() + if st == C.KernelStatusCode_Success { + return nil + } + return lastError(st) +} + +// lastError reads the kernel's thread-local last error and copies its string +// fields out immediately — the C `char*` fields are valid only until the next +// FFI call on this thread. Must run on the same OS thread as the failing call; +// call guarantees that via LockOSThread. +func lastError(code C.KernelStatusCode) *KernelError { + var e C.KernelError + if !bool(C.kernel_get_last_error(&e)) { + return &KernelError{Code: int(code), Message: fmt.Sprintf("kernel status %d (no detail)", int(code))} + } + ke := &KernelError{ + Code: int(e.code), + Message: C.GoString(e.message), + VendorCode: int32(e.vendor_code), + HTTPStatus: uint16(e.http_status), + Retryable: bool(e.retryable), + } + if e.sql_state != nil { + ke.SQLState = C.GoString(e.sql_state) + } + if e.query_id != nil { + ke.QueryID = C.GoString(e.query_id) + } + klog("kernel error: code=%d sqlstate=%q vendor=%d http=%d retryable=%v msg=%q", + ke.Code, ke.SQLState, ke.VendorCode, ke.HTTPStatus, ke.Retryable, ke.Message) + return ke +} + +// KernelError is the Go-side structured error mapped from the kernel's +// KernelError struct. It carries the sqlstate so the backend's ExecutionError +// can attach it, matching the Thrift error surface. +type KernelError struct { + Code int + Message string + SQLState string + VendorCode int32 + HTTPStatus uint16 + Retryable bool + QueryID string +} + +func (e *KernelError) Error() string { + if e.SQLState != "" { + return fmt.Sprintf("kernel: %s (sqlstate=%s, code=%d)", e.Message, e.SQLState, e.Code) + } + return fmt.Sprintf("kernel: %s (code=%d)", e.Message, e.Code) +} + +// Status codes mirrored as Go ints so non-cgo code (tests, error mapping) can +// reference them without the C import. Kept in lockstep with the C enum. +const ( + statusInvalidArgument = int(C.KernelStatusCode_InvalidArgument) + statusUnauthenticated = int(C.KernelStatusCode_Unauthenticated) + statusUnavailable = int(C.KernelStatusCode_Unavailable) + statusTimeout = int(C.KernelStatusCode_Timeout) + statusNetworkError = int(C.KernelStatusCode_NetworkError) + statusSqlError = int(C.KernelStatusCode_SqlError) +) + +// isBadConnection reports whether a status code means the session is no longer +// usable, so the driver evicts the conn from the pool (via driver.ErrBadConn's +// surface) rather than reusing it. +func isBadConnection(code int) bool { + switch code { + case statusUnauthenticated, statusUnavailable, statusNetworkError: + return true + default: + return false + } +} + +// toDriverError classifies a kernel error for the pool: a status that means the +// session is unusable is wrapped as a bad-connection error (which identifies as +// driver.ErrBadConn, so database/sql evicts the conn), matching how the Thrift +// backend surfaces connection loss. Other kernel errors — and plain +// (non-KernelError) errors — are returned unchanged, carrying their sqlstate. +func toDriverError(err error) error { + if err == nil { + return nil + } + ke, ok := err.(*KernelError) + if !ok { + return err + } + if isBadConnection(ke.Code) { + return dbsqlerrint.NewBadConnectionError(ke) + } + return ke +} + +// cStr wraps C.CString with a guaranteed free. The kernel copies strings into +// owned Rust memory on receipt, so freeing immediately after the call is safe. +// Use: cs := newCStr(s); defer cs.free(); ...C.fn(cs.c)... +type cStr struct{ c *C.char } + +func newCStr(s string) cStr { 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 new file mode 100644 index 00000000..2efb5512 --- /dev/null +++ b/internal/backend/kernel/kernel_test.go @@ -0,0 +1,158 @@ +//go:build cgo && databricks_kernel + +package kernel + +import ( + "database/sql/driver" + "errors" + "testing" + + "github.com/apache/arrow/go/v12/arrow" + "github.com/apache/arrow/go/v12/arrow/array" + "github.com/apache/arrow/go/v12/arrow/decimal128" + "github.com/apache/arrow/go/v12/arrow/memory" +) + +// isBadConnection maps the session-unusable status codes so the pool evicts the +// conn; every other code stays a plain kernel error. +func TestIsBadConnection(t *testing.T) { + bad := []int{statusUnauthenticated, statusUnavailable, statusNetworkError} + for _, code := range bad { + if !isBadConnection(code) { + t.Errorf("code %d should be a bad connection", code) + } + } + notBad := []int{statusInvalidArgument, statusSqlError, statusTimeout} + for _, code := range notBad { + if isBadConnection(code) { + t.Errorf("code %d should not be a bad connection", code) + } + } +} + +// toDriverError wraps a session-unusable KernelError as driver.ErrBadConn (so +// database/sql evicts the conn) and leaves other errors, and their sqlstate, +// intact. +func TestToDriverError(t *testing.T) { + if toDriverError(nil) != nil { + t.Fatal("nil should map to nil") + } + + badConn := &KernelError{Code: statusUnavailable, Message: "gone"} + if !errors.Is(toDriverError(badConn), driver.ErrBadConn) { + t.Errorf("unavailable kernel error should identify as driver.ErrBadConn") + } + + sqlErr := &KernelError{Code: statusSqlError, Message: "boom", SQLState: "42703"} + got := toDriverError(sqlErr) + ke, ok := got.(*KernelError) + if !ok { + t.Fatalf("sql error should remain a *KernelError, got %T", got) + } + if ke.SQLState != "42703" { + t.Errorf("sqlstate lost: got %q", ke.SQLState) + } +} + +// scanCell renders the supported scalar types and rejects an unsupported type +// (rather than returning a silently wrong value). +func TestScanCellScalars(t *testing.T) { + pool := memory.NewGoAllocator() + + t.Run("int64", func(t *testing.T) { + b := array.NewInt64Builder(pool) + defer b.Release() + b.Append(42) + arr := b.NewArray() + defer arr.Release() + v, err := scanCell(arr, 0) + if err != nil { + t.Fatal(err) + } + if v.(int64) != 42 { + t.Errorf("got %v", v) + } + }) + + t.Run("string", func(t *testing.T) { + b := array.NewStringBuilder(pool) + defer b.Release() + b.Append("hi") + arr := b.NewArray() + defer arr.Release() + v, err := scanCell(arr, 0) + if err != nil { + t.Fatal(err) + } + if v.(string) != "hi" { + t.Errorf("got %v", v) + } + }) + + t.Run("null", func(t *testing.T) { + b := array.NewInt64Builder(pool) + defer b.Release() + b.AppendNull() + arr := b.NewArray() + defer arr.Release() + v, err := scanCell(arr, 0) + if err != nil { + t.Fatal(err) + } + if v != nil { + t.Errorf("null should scan to nil, got %v", v) + } + }) + + t.Run("decimal_exact_string", func(t *testing.T) { + // 12345 at scale 2 = "123.45", exact (not a float64). + dt := &arrow.Decimal128Type{Precision: 10, Scale: 2} + b := array.NewDecimal128Builder(pool, dt) + defer b.Release() + b.Append(decimal128.FromU64(12345)) + arr := b.NewArray() + defer arr.Release() + v, err := scanCell(arr, 0) + if err != nil { + t.Fatal(err) + } + if v.(string) != "123.45" { + t.Errorf("got %v, want 123.45", v) + } + }) + + t.Run("unsupported_type_errors", func(t *testing.T) { + // A List is unsupported: must error, not return a wrong value. + b := array.NewListBuilder(pool, arrow.PrimitiveTypes.Int64) + defer b.Release() + b.Append(true) + b.ValueBuilder().(*array.Int64Builder).Append(1) + arr := b.NewArray() + defer arr.Release() + if _, err := scanCell(arr, 0); err == nil { + t.Error("scanning a List should return an unsupported-type error") + } + }) +} + +// decimal128ToExactString applies scale by string placement, preserving digits a +// float64 would lose. +func TestDecimal128ToExactString(t *testing.T) { + cases := []struct { + unscaled uint64 + scale int32 + want string + }{ + {12345, 2, "123.45"}, + {5, 3, "0.005"}, + {100, 0, "100"}, + } + for _, c := range cases { + got := decimal128ToExactString(decimal128.FromU64(c.unscaled), c.scale) + if got != c.want { + t.Errorf("unscaled=%d scale=%d: got %q want %q", c.unscaled, c.scale, got, c.want) + } + } +} + +var _ driver.Value // keep database/sql/driver imported for the scanCell signature diff --git a/internal/backend/kernel/operation.go b/internal/backend/kernel/operation.go new file mode 100644 index 00000000..568c2624 --- /dev/null +++ b/internal/backend/kernel/operation.go @@ -0,0 +1,210 @@ +//go:build cgo && databricks_kernel + +package kernel + +/* +#include +#include "databricks_kernel.h" +*/ +import "C" + +import ( + "context" + "database/sql/driver" + "fmt" + "sync" + "time" + + "github.com/databricks/databricks-sql-go/internal/backend" + dbsqlrows "github.com/databricks/databricks-sql-go/internal/rows" +) + +// execute runs one statement to a terminal state on the blocking-execute path, +// with out-of-band cancellation from a watcher goroutine: +// 1. new_statement + set_sql +// 2. canceller_new BEFORE execute, so it can observe the server statement id +// 3. a watcher goroutine that fires the canceller on ctx.Done() +// 4. the single blocking kernel_statement_execute (inline/CloudFetch and +// long-query polling all happen inside the kernel, invisibly) +// 5. drain the watcher before returning, so a late cancel cannot land on a +// statement that reuses this handle +// +// Executes SQL text only; req.Params (bound parameters) are not yet wired. +func (k *KernelBackend) execute(ctx context.Context, req backend.ExecRequest) (backend.Operation, error) { + klog("Execute sql=%q", truncate(req.Query, 120)) + + var stmt *C.kernel_statement_t + if err := call(func() C.KernelStatusCode { + return C.kernel_session_new_statement(k.session, &stmt) + }); err != nil { + return &kernelOp{}, fmt.Errorf("kernel: new_statement: %w", toDriverError(err)) + } + + sql := newCStr(req.Query) + if err := call(func() C.KernelStatusCode { + return C.kernel_statement_set_sql(stmt, sql.c) + }); err != nil { + sql.free() + C.kernel_statement_close(stmt) + return &kernelOp{}, fmt.Errorf("kernel: set_sql: %w", toDriverError(err)) + } + sql.free() + + // 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. + var canceller *C.kernel_statement_canceller_t + 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) + canceller = nil + } + + // Watcher goroutine (only when there is both a cancellable ctx and a + // canceller). The server publishes the statement id to the canceller's + // inflight slot only when the initial POST returns — held up to the server's + // inline wait (~10s) even for a long query — so a cancel fired before that is + // a no-op. Re-fire every 250ms after ctx.Done until execute returns, so the + // cancel takes effect once the id appears, with no kernel change. + done := make(chan struct{}) + var watcherWg sync.WaitGroup + if canceller != nil && ctx.Done() != nil { + watcherWg.Add(1) + go func() { + defer watcherWg.Done() + select { + case <-ctx.Done(): + case <-done: + return + } + klog("ctx.Done (%v) → firing canceller (with retry)", ctx.Err()) + ticker := time.NewTicker(250 * time.Millisecond) + defer ticker.Stop() + C.kernel_statement_canceller_cancel(canceller) + for { + select { + case <-done: + return + case <-ticker.C: + C.kernel_statement_canceller_cancel(canceller) + } + } + }() + } + + // The one blocking call. inline vs CloudFetch and long-query polling are all + // resolved inside the kernel; Go just waits here. + var exec *C.kernel_executed_statement_t + execErr := call(func() C.KernelStatusCode { + return C.kernel_statement_execute(stmt, &exec) + }) + + // Drain the watcher before returning so a late canceller fire cannot land on + // a subsequent statement reusing this handle. + close(done) + watcherWg.Wait() + if canceller != nil { + C.kernel_statement_canceller_free(canceller) + } + + op := &kernelOp{stmt: stmt} + if execErr != nil { + // Prefer the caller's ctx error when the ctx was cancelled (database/sql + // convention), keeping the kernel error as the cause. + 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()) + } + klog("Execute failed: %v", execErr) + op.close() + return op, fmt.Errorf("kernel: execute: %w", toDriverError(execErr)) + } + op.exec = exec + klog("Execute OK stmt=%p exec=%p", stmt, exec) + return op, nil +} + +// kernelOp implements backend.Operation over a sync executed statement. +type kernelOp struct { + stmt *C.kernel_statement_t + exec *C.kernel_executed_statement_t + closed bool +} + +var _ backend.Operation = (*kernelOp)(nil) + +// StatementID returns "": the C ABI exposes no server statement id accessor. The +// id is used only for logging/telemetry correlation, which falls back to the +// session id. +func (o *kernelOp) StatementID() string { return "" } + +// AffectedRows is the modified-row count for ExecContext. +func (o *kernelOp) AffectedRows() int64 { + if o.exec == nil { + return 0 + } + return int64(C.kernel_executed_statement_num_modified_rows(o.exec)) +} + +// Results builds the driver.Rows over the executed statement's result stream. +// On the query path the returned Rows owns closing the server-side operation. +func (o *kernelOp) Results(ctx context.Context, callbacks *dbsqlrows.TelemetryCallbacks) (driver.Rows, error) { + if o.exec == nil { + return nil, fmt.Errorf("kernel: no executed statement") + } + var stream *C.kernel_result_stream_t + if err := call(func() C.KernelStatusCode { + return C.kernel_executed_statement_get_result_stream(o.exec, &stream) + }); err != nil { + return nil, fmt.Errorf("kernel: get_result_stream: %w", toDriverError(err)) + } + return newKernelRows(ctx, o, stream, callbacks) +} + +// IsStaging reports whether this is a staging (PUT/GET/REMOVE) operation. The +// kernel backend does not support staging, so this is always false. +func (o *kernelOp) IsStaging(ctx context.Context) (bool, error) { return false, nil } + +// Close best-effort closes the executed statement and its statement handle. It is +// idempotent (a second call, or a call after Rows.Close already tore the +// operation down, is a no-op) per the backend.Operation contract. closed reports +// whether a close was actually issued. +func (o *kernelOp) Close(ctx context.Context) (bool, error) { + return o.close(), nil +} + +// close is the shared idempotent teardown used by both Operation.Close and +// Rows.Close. Closing the executed handle first, then the statement, matches the +// C ABI teardown order (result stream is closed by Rows before this runs). +func (o *kernelOp) close() bool { + if o.closed { + return false + } + o.closed = true + if o.exec != nil { + C.kernel_executed_statement_close(o.exec) + o.exec = nil + } + if o.stmt != nil { + C.kernel_statement_close(o.stmt) + o.stmt = nil + } + klog("kernelOp closed") + return true +} + +// ExecutionError wraps cause as the driver's execution error. The kernel error +// already carries the sqlstate (see KernelError), so this returns cause as-is +// (nil when cause is nil), matching the neutral contract. +func (o *kernelOp) ExecutionError(ctx context.Context, cause error) error { + return toDriverError(cause) +} + +func truncate(s string, n int) string { + if len(s) <= n { + return s + } + return s[:n] + "…" +} diff --git a/internal/backend/kernel/rows.go b/internal/backend/kernel/rows.go new file mode 100644 index 00000000..c0c86eb1 --- /dev/null +++ b/internal/backend/kernel/rows.go @@ -0,0 +1,224 @@ +//go:build cgo && databricks_kernel + +package kernel + +/* +#include +#include "databricks_kernel.h" +// Forward-declare the Arrow C Data Interface structs so cgo can take their +// addresses; arrow-go's cdata package reinterprets these via unsafe.Pointer. +struct ArrowSchema; +struct ArrowArray; +*/ +import "C" + +import ( + "context" + "database/sql/driver" + "fmt" + "io" + "unsafe" + + "github.com/apache/arrow/go/v12/arrow" + "github.com/apache/arrow/go/v12/arrow/array" + "github.com/apache/arrow/go/v12/arrow/cdata" + dbsqlrows "github.com/databricks/databricks-sql-go/internal/rows" +) + +var _ driver.Rows = (*kernelRows)(nil) + +// kernelRows implements driver.Rows over the kernel result stream. It pulls one +// Arrow RecordBatch at a time via kernel_result_stream_next_batch (inline and +// CloudFetch are transparent below the C ABI), imports it zero-copy through the +// Arrow C Data Interface, and scans rows out on demand. +// +// The batch/row split mirrors the kernel's own ResultStream: next_batch does the +// per-batch network/decode work once; row reads walk the already-imported +// arrow.Record with O(1) indexing. +type kernelRows struct { + ctx context.Context + op *kernelOp + stream *C.kernel_result_stream_t + callbacks *dbsqlrows.TelemetryCallbacks + + cols []string + cur arrow.Record // current batch (nil until first Next) + rowInCur int // next row index within cur + closed bool + eof bool +} + +// 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} + + var csch C.struct_ArrowSchema + if err := call(func() C.KernelStatusCode { + return C.kernel_result_stream_get_schema(stream, &csch) + }); err != nil { + r.Close() + return nil, fmt.Errorf("kernel: get_schema: %w", toDriverError(err)) + } + sch, err := cdata.ImportCArrowSchema((*cdata.CArrowSchema)(unsafe.Pointer(&csch))) + if err != nil { + r.Close() + return nil, fmt.Errorf("kernel: import schema: %w", err) + } + fields := sch.Fields() + r.cols = make([]string, len(fields)) + for i, f := range fields { + r.cols[i] = f.Name + } + klog("newKernelRows: %d columns", len(r.cols)) + return r, nil +} + +// Columns returns the result-set column names. +func (r *kernelRows) Columns() []string { return r.cols } + +// Close releases the current batch, the kernel result stream, and (query-path +// ownership) the server operation. Idempotent. +func (r *kernelRows) Close() error { + if r.closed { + return nil + } + r.closed = true + if r.cur != nil { + r.cur.Release() + r.cur = nil + } + if r.stream != nil { + C.kernel_result_stream_close(r.stream) + r.stream = nil + } + if r.op != nil { + r.op.close() + } + klog("kernelRows closed") + return nil +} + +// 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 { + if r.closed { + return io.EOF + } + for r.cur == nil || r.rowInCur >= int(r.cur.NumRows()) { + if r.eof { + return io.EOF + } + if err := r.nextBatch(); err != nil { + return err + } + } + rec := r.cur + for c := 0; c < len(dest); c++ { + v, err := scanCell(rec.Column(c), r.rowInCur) + if err != nil { + return fmt.Errorf("kernel: scan col %d (%s): %w", c, r.cols[c], err) + } + dest[c] = v + } + r.rowInCur++ + return nil +} + +// nextBatch pulls the next Arrow batch. A released array (release==NULL) is the +// kernel's end-of-stream sentinel. +func (r *kernelRows) nextBatch() error { + if r.cur != nil { + r.cur.Release() + r.cur = nil + } + var carr C.struct_ArrowArray + var csch C.struct_ArrowSchema + if err := call(func() C.KernelStatusCode { + return C.kernel_result_stream_next_batch(r.stream, &carr, &csch) + }); err != nil { + return fmt.Errorf("kernel: next_batch: %w", toDriverError(err)) + } + if carr.release == nil { + r.eof = true + klog("nextBatch: EOF") + return io.EOF + } + // Zero-copy import. The kernel exports self-contained batches (Rust to_ffi + // moves the Arc-owned buffers in), so the arrow.Record safely outlives the + // stream; we still Release each batch explicitly as we advance. + rec, err := cdata.ImportCRecordBatch( + (*cdata.CArrowArray)(unsafe.Pointer(&carr)), + (*cdata.CArrowSchema)(unsafe.Pointer(&csch))) + if err != nil { + return fmt.Errorf("kernel: import batch: %w", err) + } + r.cur = rec + r.rowInCur = 0 + if r.callbacks != nil && r.callbacks.OnChunkFetched != nil { + r.callbacks.OnChunkFetched(1, 0, 0, 0, 0) + } + klog("nextBatch: %d rows", rec.NumRows()) + return nil +} + +// scanCell extracts one cell as a driver.Value. It handles the scalar types: +// bool, all int/uint widths, float, 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). +// NULLs map to nil. An unsupported type (nested List/Map/Struct, Variant, +// Geometry, interval/duration) returns an error rather than a silently wrong +// value. +func scanCell(col arrow.Array, row int) (driver.Value, error) { + if col.IsNull(row) { + return nil, nil + } + switch c := col.(type) { + case *array.Null: + return nil, nil + case *array.Boolean: + return c.Value(row), nil + case *array.Int8: + return int64(c.Value(row)), nil + case *array.Int16: + return int64(c.Value(row)), nil + case *array.Int32: + return int64(c.Value(row)), nil + case *array.Int64: + return c.Value(row), nil + case *array.Uint8: + return int64(c.Value(row)), nil + case *array.Uint16: + return int64(c.Value(row)), nil + case *array.Uint32: + return int64(c.Value(row)), nil + case *array.Uint64: + return int64(c.Value(row)), nil + case *array.Float32: + return float64(c.Value(row)), nil + case *array.Float64: + return c.Value(row), nil + case *array.String: + return c.Value(row), nil + case *array.LargeString: + return c.Value(row), nil + case *array.Binary: + return c.Value(row), nil + case *array.Date32: + return c.Value(row).ToTime(), nil + case *array.Date64: + return c.Value(row).ToTime(), nil + case *array.Timestamp: + dt, ok := col.DataType().(*arrow.TimestampType) + if !ok { + return nil, fmt.Errorf("timestamp column has unexpected datatype %s", col.DataType()) + } + return c.Value(row).ToTime(dt.Unit), nil + case *array.Decimal128: + dt := col.DataType().(*arrow.Decimal128Type) + return decimal128ToExactString(c.Value(row), dt.Scale), nil + default: + return nil, fmt.Errorf("kernel: scanning arrow type %s is not supported "+ + "(nested types, variant, geometry, and intervals are not yet handled)", col.DataType()) + } +} diff --git a/internal/backend/kernel/scan.go b/internal/backend/kernel/scan.go new file mode 100644 index 00000000..ca924322 --- /dev/null +++ b/internal/backend/kernel/scan.go @@ -0,0 +1,40 @@ +//go:build cgo && databricks_kernel + +package kernel + +import ( + "math/big" + "strings" + + "github.com/apache/arrow/go/v12/arrow/decimal128" +) + +// decimal128ToExactString renders an Arrow decimal128 as an exact fixed-point +// string, applying scale by string placement rather than float conversion. This +// matches the Thrift path's default DECIMAL rendering and preserves precision a +// float64 would lose beyond ~17 significant digits (databricks-sql-go#274). +func decimal128ToExactString(n decimal128.Num, scale int32) string { + unscaled := n.BigInt() + neg := unscaled.Sign() < 0 + digits := new(big.Int).Abs(unscaled).String() + + var b strings.Builder + if neg { + b.WriteByte('-') + } + if scale <= 0 { + b.WriteString(digits) + return b.String() + } + s := int(scale) + if len(digits) <= s { + b.WriteString("0.") + b.WriteString(strings.Repeat("0", s-len(digits))) + b.WriteString(digits) + } else { + b.WriteString(digits[:len(digits)-s]) + b.WriteByte('.') + b.WriteString(digits[len(digits)-s:]) + } + return b.String() +} diff --git a/internal/config/config.go b/internal/config/config.go index fb7c6afb..12b59395 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -125,6 +125,13 @@ type UserConfig struct { // See databricks/databricks-sql-go#274. UseArrowNativeDecimalDSN bool CloudFetchConfig + // UseKernel selects the SEA-via-kernel backend instead of Thrift. See the + // WithUseKernel connector option for the build requirements. DSN: useKernel=true. + UseKernel bool + // WarehouseID is the bare SQL warehouse id, used by the kernel backend (which + // addresses a warehouse by id) in preference to HTTPPath. The Thrift backend + // ignores it and routes by HTTPPath. DSN: warehouseId=. + WarehouseID string } // DeepCopy returns a true deep copy of UserConfig @@ -171,6 +178,8 @@ func (ucfg UserConfig) DeepCopy() UserConfig { EnableTelemetry: ucfg.EnableTelemetry, TelemetryBatchSize: ucfg.TelemetryBatchSize, TelemetryFlushInterval: ucfg.TelemetryFlushInterval, + UseKernel: ucfg.UseKernel, + WarehouseID: ucfg.WarehouseID, } } @@ -320,6 +329,17 @@ func ParseDSN(dsn string) (UserConfig, error) { ucfg.UseArrowNativeDecimalDSN = useArrowNativeDecimal } + // Kernel backend parameters + if useKernel, ok, err := params.extractAsBool("useKernel"); ok { + if err != nil { + return UserConfig{}, err + } + ucfg.UseKernel = useKernel + } + if warehouseID, ok := params.extract("warehouseId"); ok { + ucfg.WarehouseID = warehouseID + } + // Telemetry parameters if enableTelemetry, ok, err := params.extractAsBool("enableTelemetry"); ok { if err != nil { diff --git a/kernel_backend.go b/kernel_backend.go new file mode 100644 index 00000000..d89349f2 --- /dev/null +++ b/kernel_backend.go @@ -0,0 +1,24 @@ +//go:build cgo && databricks_kernel + +package dbsql + +import ( + "context" + + "github.com/databricks/databricks-sql-go/internal/backend" + "github.com/databricks/databricks-sql-go/internal/backend/kernel" + "github.com/databricks/databricks-sql-go/internal/config" +) + +// newKernelBackend builds the SEA-via-kernel backend from the driver config; the +// connector opens the session right after, matching the Thrift path. It maps the +// config fields the kernel backend currently reads — host, warehouse/http path, +// and PAT. +func newKernelBackend(_ context.Context, cfg *config.Config) (backend.Backend, error) { + return kernel.New(kernel.Config{ + Host: cfg.Host, + HTTPPath: cfg.HTTPPath, + WarehouseID: cfg.WarehouseID, + Token: cfg.AccessToken, + }), nil +} diff --git a/kernel_e2e_test.go b/kernel_e2e_test.go new file mode 100644 index 00000000..2a3ea393 --- /dev/null +++ b/kernel_e2e_test.go @@ -0,0 +1,170 @@ +//go:build cgo && databricks_kernel + +package dbsql + +import ( + "bytes" + "context" + "database/sql" + "database/sql/driver" + "os" + "testing" + "time" +) + +// 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. +func kernelTestDB(t *testing.T) *sql.DB { + t.Helper() + host := os.Getenv("DATABRICKS_HOST") + httpPath := os.Getenv("DATABRICKS_HTTP_PATH") + token := os.Getenv("DATABRICKS_TOKEN") + if host == "" || httpPath == "" || token == "" { + t.Skip("set DATABRICKS_HOST / DATABRICKS_HTTP_PATH / DATABRICKS_TOKEN for the kernel e2e") + } + connector, err := NewConnector( + WithServerHostname(host), + WithHTTPPath(httpPath), + WithAccessToken(token), + WithUseKernel(true), + ) + if err != nil { + t.Fatalf("NewConnector: %v", err) + } + return sql.OpenDB(connector) +} + +// TestKernelE2ESelect1 is the smallest end-to-end proof: PAT session over the +// kernel, execute, scan one scalar row. +func TestKernelE2ESelect1(t *testing.T) { + db := kernelTestDB(t) + defer db.Close() + + var got int64 + if err := db.QueryRowContext(context.Background(), "SELECT 1").Scan(&got); err != nil { + t.Fatalf("query: %v", err) + } + if got != 1 { + t.Errorf("SELECT 1 = %d, want 1", got) + } +} + +// TestKernelE2EDataTypes scans each supported scalar type in its own subtest, so +// a failure names the exact type rather than being masked by others in a shared +// row. Each case selects a single value and compares the scanned result. NULL is +// covered as its own case. +func TestKernelE2EDataTypes(t *testing.T) { + db := kernelTestDB(t) + defer db.Close() + + cases := []struct { + name string + expr string // the single SELECT expression + want driver.Value // expected scanned value (nil for SQL NULL) + }{ + {"bigint", "CAST(42 AS BIGINT)", int64(42)}, + {"int", "CAST(7 AS INT)", int64(7)}, + {"smallint", "CAST(3 AS SMALLINT)", int64(3)}, + {"tinyint", "CAST(1 AS TINYINT)", int64(1)}, + {"double", "CAST(3.5 AS DOUBLE)", float64(3.5)}, + {"float", "CAST(1.5 AS FLOAT)", float64(1.5)}, + {"boolean", "true", true}, + {"string", "'hi'", "hi"}, + {"binary", "CAST('abc' AS BINARY)", []byte("abc")}, + {"decimal_exact", "CAST(1.25 AS DECIMAL(5,2))", "1.25"}, + {"date", "CAST('2026-07-09' AS DATE)", time.Date(2026, time.July, 9, 0, 0, 0, 0, time.UTC)}, + {"timestamp", "CAST('2026-07-09 12:34:56' AS TIMESTAMP)", time.Date(2026, time.July, 9, 12, 34, 56, 0, time.UTC)}, + {"null", "CAST(NULL AS STRING)", nil}, + } + + for _, c := range cases { + t.Run(c.name, func(t *testing.T) { + var got any + err := db.QueryRowContext(context.Background(), "SELECT "+c.expr).Scan(&got) + if err != nil { + t.Fatalf("scan %s: %v", c.expr, err) + } + if !dataTypeEqual(got, c.want) { + t.Errorf("%s = %#v (%T), want %#v (%T)", c.expr, got, got, c.want, c.want) + } + }) + } +} + +// dataTypeEqual compares scanned values, handling the two non-comparable cases: +// []byte (bytes.Equal) and time.Time (Equal, which is instant-based and ignores +// the location the value was materialized in). +func dataTypeEqual(got, want driver.Value) bool { + switch w := want.(type) { + case nil: + return got == nil + case []byte: + g, ok := got.([]byte) + return ok && bytes.Equal(g, w) + case time.Time: + g, ok := got.(time.Time) + return ok && g.Equal(w) + default: + return got == want + } +} + +// 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. +func TestKernelE2ECloudFetch(t *testing.T) { + db := kernelTestDB(t) + defer db.Close() + + const want = 1_000_000 + rows, err := db.QueryContext(context.Background(), + "SELECT id FROM range(0, 1000000)") + if err != nil { + t.Fatalf("query: %v", err) + } + defer rows.Close() + + var count, last int64 + for rows.Next() { + var id int64 + if err := rows.Scan(&id); err != nil { + t.Fatalf("scan at row %d: %v", count, err) + } + count++ + last = id + } + if err := rows.Err(); err != nil { + t.Fatalf("iteration: %v", err) + } + if count != want { + t.Errorf("row count = %d, want %d", count, want) + } + if last != want-1 { + t.Errorf("last id = %d, want %d", last, want-1) + } +} + +// TestKernelE2ECancellation cancels a long-running query via ctx and asserts it +// returns well before its uncancelled runtime. +func TestKernelE2ECancellation(t *testing.T) { + db := kernelTestDB(t) + defer db.Close() + + ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + defer cancel() + + start := time.Now() + // A query that would run far longer than the 3s deadline. + _, err := db.QueryContext(ctx, "SELECT count(*) FROM range(0, 100000000000) WHERE id % 7 = 0") + elapsed := time.Since(start) + + if err == nil { + t.Fatal("expected a cancellation error, got nil") + } + if elapsed > 30*time.Second { + t.Errorf("cancellation took %v; expected it to abandon well before the query's natural runtime", elapsed) + } + t.Logf("cancelled after %v with err=%v", elapsed, err) +} diff --git a/kernel_parity_test.go b/kernel_parity_test.go new file mode 100644 index 00000000..b687a5c2 --- /dev/null +++ b/kernel_parity_test.go @@ -0,0 +1,94 @@ +//go:build cgo && databricks_kernel + +package dbsql + +import ( + "context" + "database/sql" + "os" + "testing" +) + +// thriftTestDB opens the same warehouse over the default Thrift backend, for +// 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") + if host == "" || httpPath == "" || token == "" { + t.Skip("set DATABRICKS_HOST / DATABRICKS_HTTP_PATH / DATABRICKS_TOKEN for the parity test") + } + connector, err := NewConnector( + WithServerHostname(host), + WithHTTPPath(httpPath), + WithAccessToken(token), + ) + if err != nil { + t.Fatalf("NewConnector (thrift): %v", err) + } + return sql.OpenDB(connector) +} + +// TestKernelThriftParity runs the same scalar query through both backends and +// asserts identical row output. The two scanners are intentionally separate, so +// this golden comparison is the guarantee that they render the scalar types +// equivalently. +func TestKernelThriftParity(t *testing.T) { + const query = "SELECT CAST(7 AS BIGINT), CAST(2.5 AS DOUBLE), 'parity', " + + "CAST(NULL AS STRING), CAST(9.99 AS DECIMAL(5,2)), true" + + kernelDB := kernelTestDB(t) + defer kernelDB.Close() + thriftDB := thriftTestDB(t) + defer thriftDB.Close() + + kernelRow := scanOneRowAsStrings(t, kernelDB, query) + thriftRow := scanOneRowAsStrings(t, thriftDB, query) + + 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", i, kernelRow[i], thriftRow[i]) + } + } +} + +// 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 { + t.Helper() + rows, err := db.QueryContext(context.Background(), query) + if err != nil { + t.Fatalf("query: %v", err) + } + defer rows.Close() + + cols, err := rows.Columns() + if err != nil { + t.Fatalf("columns: %v", err) + } + if !rows.Next() { + t.Fatalf("no row: %v", rows.Err()) + } + raw := make([]sql.RawBytes, len(cols)) + dest := make([]any, len(cols)) + for i := range raw { + dest[i] = &raw[i] + } + if err := rows.Scan(dest...); err != nil { + t.Fatalf("scan: %v", err) + } + out := make([]string, len(cols)) + for i, b := range raw { + if b == nil { + out[i] = "" + } else { + out[i] = string(b) + } + } + return out +} diff --git a/kernel_stub.go b/kernel_stub.go new file mode 100644 index 00000000..4774278b --- /dev/null +++ b/kernel_stub.go @@ -0,0 +1,20 @@ +//go:build !cgo || !databricks_kernel + +package dbsql + +import ( + "context" + "errors" + + "github.com/databricks/databricks-sql-go/internal/backend" + "github.com/databricks/databricks-sql-go/internal/config" +) + +// newKernelBackend is the stub compiled when the kernel backend is not built in. +// It fails loudly rather than silently falling back to Thrift, so a mismatch +// between WithUseKernel and the build tags surfaces at connect time. The real +// implementation is in kernel_backend.go. +func newKernelBackend(_ context.Context, _ *config.Config) (backend.Backend, error) { + return nil, errors.New("databricks: the SEA-via-kernel backend is not compiled into this binary; " + + "rebuild with -tags databricks_kernel and CGO_ENABLED=1, or unset WithUseKernel") +} From 8bc7f17143836490b61959e50e8388b62fccb8a0 Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Thu, 9 Jul 2026 06:02:24 +0000 Subject: [PATCH 02/32] feat(kernel): render nested and complex types MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Extend the kernel backend's scanner to the complex data types, so the SEA path covers the full scalar-plus-nested type set. Nested Arrow values — list, map, struct, and VARIANT (which arrives as a nested value) — render to a JSON string that is byte-identical to the Thrift arrow path, so a query's result is the same across backends. GEOMETRY arrives as a WKT string and is read by the existing string arm. The renderer (scan_nested.go) recurses into child arrays and mirrors the Thrift marshal() rules: time.Time as a quoted .String(), and nested decimals as float64 (the exact-string decimal applies only to a top-level decimal column, #274). scanCell delegates nested columns here; genuinely unhandled types (interval/ duration) still return an explicit error. Tests: unit tests for list/map/struct/nested-null rendering; the live e2e data types table gains array/map/struct/variant/geometry cases; the Thrift-parity query gains the same, asserting byte-identical output across backends. Signed-off-by: Mani Kaustubh Mathur --- internal/backend/kernel/kernel_test.go | 88 +++++++++++- internal/backend/kernel/rows.go | 16 ++- internal/backend/kernel/scan_nested.go | 183 +++++++++++++++++++++++++ kernel_e2e_test.go | 7 + kernel_parity_test.go | 4 +- 5 files changed, 288 insertions(+), 10 deletions(-) create mode 100644 internal/backend/kernel/scan_nested.go diff --git a/internal/backend/kernel/kernel_test.go b/internal/backend/kernel/kernel_test.go index 2efb5512..cf313d69 100644 --- a/internal/backend/kernel/kernel_test.go +++ b/internal/backend/kernel/kernel_test.go @@ -122,15 +122,95 @@ func TestScanCellScalars(t *testing.T) { }) t.Run("unsupported_type_errors", func(t *testing.T) { - // A List is unsupported: must error, not return a wrong value. + // A duration (INTERVAL) is not yet handled: must error, not return a + // wrong value. + b := array.NewDurationBuilder(pool, &arrow.DurationType{Unit: arrow.Microsecond}) + defer b.Release() + b.Append(1000) + arr := b.NewArray() + defer arr.Release() + if _, err := scanCell(arr, 0); err == nil { + t.Error("scanning a Duration should return an unsupported-type error") + } + }) +} + +// scanCell renders nested types (list/struct/map) to a JSON string matching the +// Thrift path. +func TestScanCellNested(t *testing.T) { + pool := memory.NewGoAllocator() + + t.Run("list", func(t *testing.T) { b := array.NewListBuilder(pool, arrow.PrimitiveTypes.Int64) defer b.Release() + vb := b.ValueBuilder().(*array.Int64Builder) b.Append(true) - b.ValueBuilder().(*array.Int64Builder).Append(1) + vb.Append(1) + vb.Append(2) + vb.Append(3) arr := b.NewArray() defer arr.Release() - if _, err := scanCell(arr, 0); err == nil { - t.Error("scanning a List should return an unsupported-type error") + v, err := scanCell(arr, 0) + if err != nil { + t.Fatal(err) + } + if v.(string) != "[1,2,3]" { + t.Errorf("got %q, want [1,2,3]", v) + } + }) + + t.Run("struct", func(t *testing.T) { + dt := arrow.StructOf( + arrow.Field{Name: "a", Type: arrow.PrimitiveTypes.Int64}, + arrow.Field{Name: "b", Type: arrow.BinaryTypes.String}, + ) + b := array.NewStructBuilder(pool, dt) + defer b.Release() + b.Append(true) + b.FieldBuilder(0).(*array.Int64Builder).Append(1) + b.FieldBuilder(1).(*array.StringBuilder).Append("x") + arr := b.NewArray() + defer arr.Release() + v, err := scanCell(arr, 0) + if err != nil { + t.Fatal(err) + } + if v.(string) != `{"a":1,"b":"x"}` { + t.Errorf("got %q, want {\"a\":1,\"b\":\"x\"}", v) + } + }) + + t.Run("map", func(t *testing.T) { + b := array.NewMapBuilder(pool, arrow.BinaryTypes.String, arrow.PrimitiveTypes.Int64, false) + defer b.Release() + kb := b.KeyBuilder().(*array.StringBuilder) + ib := b.ItemBuilder().(*array.Int64Builder) + b.Append(true) + kb.Append("k") + ib.Append(9) + arr := b.NewArray() + defer arr.Release() + v, err := scanCell(arr, 0) + if err != nil { + t.Fatal(err) + } + if v.(string) != `{"k":9}` { + t.Errorf("got %q, want {\"k\":9}", v) + } + }) + + t.Run("nested_null", func(t *testing.T) { + b := array.NewListBuilder(pool, arrow.PrimitiveTypes.Int64) + defer b.Release() + b.AppendNull() + arr := b.NewArray() + defer arr.Release() + v, err := scanCell(arr, 0) + if err != nil { + t.Fatal(err) + } + if v != nil { + t.Errorf("null list should scan to nil, got %v", v) } }) } diff --git a/internal/backend/kernel/rows.go b/internal/backend/kernel/rows.go index c0c86eb1..c0258753 100644 --- a/internal/backend/kernel/rows.go +++ b/internal/backend/kernel/rows.go @@ -162,13 +162,15 @@ func (r *kernelRows) nextBatch() error { return nil } -// scanCell extracts one cell as a driver.Value. It handles the scalar types: +// scanCell extracts one cell as a driver.Value. Scalars map to their Go value: // bool, all int/uint widths, float, 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). -// NULLs map to nil. An unsupported type (nested List/Map/Struct, Variant, -// Geometry, interval/duration) returns an error rather than a silently wrong -// value. +// Nested types (List/Map/Struct, and VARIANT which arrives nested) render to a +// JSON string byte-identical to the Thrift path (see scan_nested.go); 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. func scanCell(col arrow.Array, row int) (driver.Value, error) { if col.IsNull(row) { return nil, nil @@ -217,8 +219,12 @@ func scanCell(col arrow.Array, row int) (driver.Value, error) { case *array.Decimal128: dt := col.DataType().(*arrow.Decimal128Type) return decimal128ToExactString(c.Value(row), dt.Scale), 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) default: return nil, fmt.Errorf("kernel: scanning arrow type %s is not supported "+ - "(nested types, variant, geometry, and intervals are not yet handled)", col.DataType()) + "(intervals are not yet handled)", col.DataType()) } } diff --git a/internal/backend/kernel/scan_nested.go b/internal/backend/kernel/scan_nested.go new file mode 100644 index 00000000..75138c26 --- /dev/null +++ b/internal/backend/kernel/scan_nested.go @@ -0,0 +1,183 @@ +//go:build cgo && databricks_kernel + +package kernel + +// Recursive rendering of nested Arrow values (List/Map/Struct) to a JSON string, +// byte-compatible with the Thrift arrow path +// (internal/rows/arrowbased/columnValues.go). scanCell delegates here for nested +// columns; database/sql consumers then get the same JSON shape from either +// backend: +// - list → [v0,v1,...] +// - map → {"k0":v0,"k1":v1,...} (keys stringified) +// - struct → {"field0":v0,...} +// - nested NULL → null +// - time.Time → quoted .String() (matches the Thrift marshal() special-case) +// - nested decimal → float64 (lossy, matching Thrift's in-JSON rendering; the +// exact-string path applies only to top-level decimal columns, #274) +// +// VARIANT arrives as a nested value and renders through this path; GEOMETRY +// arrives as a WKB/WKT string and is handled by scanCell's scalar arm. +// +// Rendering to JSON (not a Go map/slice) is deliberate: it is what the Thrift +// path returns, so a query's result is identical across backends — the property +// the Thrift-parity test asserts. + +import ( + "database/sql/driver" + "encoding/json" + "fmt" + "strings" + "time" + + "github.com/apache/arrow/go/v12/arrow" + "github.com/apache/arrow/go/v12/arrow/array" +) + +// renderJSONString renders one nested cell as a JSON string (nil for a NULL cell). +func renderJSONString(col arrow.Array, row int) (driver.Value, error) { + if col.IsNull(row) { + return nil, nil + } + var b strings.Builder + if err := writeJSON(&b, col, row); err != nil { + return nil, err + } + return b.String(), nil +} + +// writeJSON writes the JSON form of col[row] into b, recursing for nested types. +func writeJSON(b *strings.Builder, col arrow.Array, row int) error { + if col.IsNull(row) { + b.WriteString("null") + return nil + } + switch c := col.(type) { + case *array.List: + return writeListJSON(b, c.ListValues(), int(c.Offsets()[row]), int(c.Offsets()[row+1])) + case *array.LargeList: + return writeListJSON(b, c.ListValues(), int(c.Offsets()[row]), int(c.Offsets()[row+1])) + case *array.FixedSizeList: + n := int(c.DataType().(*arrow.FixedSizeListType).Len()) + return writeListJSON(b, c.ListValues(), row*n, row*n+n) + case *array.Map: + return writeMapJSON(b, c, row) + case *array.Struct: + return writeStructJSON(b, c, row) + default: + v, err := scalarForJSON(col, row) + if err != nil { + return err + } + return writeScalarJSON(b, v) + } +} + +func writeListJSON(b *strings.Builder, values arrow.Array, start, end int) error { + b.WriteByte('[') + for i := start; i < end; i++ { + if i > start { + b.WriteByte(',') + } + if err := writeJSON(b, values, i); err != nil { + return err + } + } + b.WriteByte(']') + return nil +} + +func writeMapJSON(b *strings.Builder, m *array.Map, row int) error { + start, end := int(m.Offsets()[row]), int(m.Offsets()[row+1]) + keys := m.Keys() + items := m.Items() + b.WriteByte('{') + for i := start; i < end; i++ { + if i > start { + b.WriteByte(',') + } + // JSON object keys must be strings; stringify the key value. + kv, err := scalarForJSON(keys, i) + if err != nil { + return err + } + writeJSONKey(b, kv) + b.WriteByte(':') + if err := writeJSON(b, items, i); err != nil { + return err + } + } + b.WriteByte('}') + return nil +} + +func writeStructJSON(b *strings.Builder, s *array.Struct, row int) error { + st := s.DataType().(*arrow.StructType) + b.WriteByte('{') + for f := 0; f < s.NumField(); f++ { + if f > 0 { + b.WriteByte(',') + } + keyBytes, _ := json.Marshal(st.Field(f).Name) + b.Write(keyBytes) + b.WriteByte(':') + if err := writeJSON(b, s.Field(f), row); err != nil { + return err + } + } + b.WriteByte('}') + return nil +} + +// writeJSONKey writes a value as a JSON object key (always a quoted string). +func writeJSONKey(b *strings.Builder, v any) { + switch k := v.(type) { + case string: + kb, _ := json.Marshal(k) + b.Write(kb) + default: + kb, _ := json.Marshal(fmt.Sprintf("%v", k)) + b.Write(kb) + } +} + +// writeScalarJSON writes a scalar leaf, mirroring the Thrift marshal(): a +// time.Time becomes a quoted .String(); everything else uses json.Marshal. +func writeScalarJSON(b *strings.Builder, v any) error { + if v == nil { + b.WriteString("null") + return nil + } + if t, ok := v.(time.Time); ok { + b.WriteByte('"') + b.WriteString(t.String()) + b.WriteByte('"') + return nil + } + vb, err := json.Marshal(v) + if err != nil { + return err + } + b.Write(vb) + return nil +} + +// scalarForJSON returns the Go value used inside JSON for a leaf cell. It differs +// from scanCell in one way: a nested decimal renders as float64 (lossy), matching +// the Thrift path's in-JSON decimal rendering (#274); the exact-string decimal +// applies only to a top-level decimal column. +func scalarForJSON(col arrow.Array, row int) (any, error) { + if col.IsNull(row) { + return nil, nil + } + switch c := col.(type) { + case *array.Decimal128: + dt := col.DataType().(*arrow.Decimal128Type) + return c.Value(row).ToFloat64(dt.Scale), nil + case *array.Decimal256: + dt := col.DataType().(*arrow.Decimal256Type) + return c.Value(row).ToFloat64(dt.Scale), nil + default: + // Reuse scanCell's scalar arm for every non-nested, non-decimal leaf. + return scanCell(col, row) + } +} diff --git a/kernel_e2e_test.go b/kernel_e2e_test.go index 2a3ea393..f77c7ed1 100644 --- a/kernel_e2e_test.go +++ b/kernel_e2e_test.go @@ -77,6 +77,13 @@ func TestKernelE2EDataTypes(t *testing.T) { {"date", "CAST('2026-07-09' AS DATE)", time.Date(2026, time.July, 9, 0, 0, 0, 0, time.UTC)}, {"timestamp", "CAST('2026-07-09 12:34:56' AS TIMESTAMP)", time.Date(2026, time.July, 9, 12, 34, 56, 0, time.UTC)}, {"null", "CAST(NULL AS STRING)", nil}, + // Nested types render to a JSON string; VARIANT arrives nested, GEOMETRY + // as a WKT/WKB string. + {"array", "array(1, 2, 3)", "[1,2,3]"}, + {"map", "map('k', 9)", `{"k":9}`}, + {"struct", "named_struct('a', 1, 'b', 'x')", `{"a":1,"b":"x"}`}, + {"variant", `parse_json('{"a":1,"b":[2,3]}')`, `{"a":1,"b":[2,3]}`}, + {"geometry", "st_point(1, 2)", "POINT(1 2)"}, } for _, c := range cases { diff --git a/kernel_parity_test.go b/kernel_parity_test.go index b687a5c2..22d8f4d1 100644 --- a/kernel_parity_test.go +++ b/kernel_parity_test.go @@ -36,7 +36,9 @@ func thriftTestDB(t *testing.T) *sql.DB { // equivalently. func TestKernelThriftParity(t *testing.T) { const query = "SELECT CAST(7 AS BIGINT), CAST(2.5 AS DOUBLE), 'parity', " + - "CAST(NULL AS STRING), CAST(9.99 AS DECIMAL(5,2)), true" + "CAST(NULL AS STRING), CAST(9.99 AS DECIMAL(5,2)), true, " + + "array(1, 2, 3), map('k', 9), named_struct('a', 1, 'b', 'x'), " + + `parse_json('{"a":1,"b":[2,3]}'), st_point(1, 2)` kernelDB := kernelTestDB(t) defer kernelDB.Close() From 20950c9a24ccb1d6bc6b86d75ec2c82afadcf95e Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Thu, 9 Jul 2026 06:33:22 +0000 Subject: [PATCH 03/32] feat(kernel): map TLS, proxy, and session-conf connection options MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Route the driver's existing connection options through to the kernel session, so a kernel-backed connection honors the same knobs as Thrift with no change to the user-facing API — only WithUseKernel selects the backend. The connector reads the same config the Thrift backend reads and translates it to the kernel's flat connection config: - Session confs (STATEMENT_TIMEOUT, QUERY_TAGS, TIMEZONE, …): the same SessionParams map Thrift forwards, applied one key at a time via set_session_conf. This covers Query Tags and server statement timeout. - TLS: TLSConfig.InsecureSkipVerify (WithSkipTLSHostVerify) maps to the kernel's skip-hostname-verification setter — the one TLS knob the driver actually honors today. - HTTP proxy: resolved via http.ProxyFromEnvironment for the connection's endpoint, the same cached HTTP(S)_PROXY / NO_PROXY decision the Thrift transport makes, then passed to set_proxy. No new dependency or option. - SPOG org routing continues to ride in the http path's ?o= (parsed kernel-side). M2M/OAuth is deliberately not included here: its credentials are held inside the authenticator rather than on config, so wiring it needs a small config change that is better done on its own. Tests: a proxy-resolution unit test; live e2e that reads query tags and statement timeout back from the server via SET (proving they were applied, not just accepted) and that a TLS skip-verify connection still succeeds. Signed-off-by: Mani Kaustubh Mathur --- internal/backend/kernel/backend.go | 58 +++++++++++++++++++++++- kernel_backend.go | 48 +++++++++++++++++--- kernel_backend_test.go | 33 ++++++++++++++ kernel_e2e_test.go | 71 ++++++++++++++++++++++++++++++ 4 files changed, 203 insertions(+), 7 deletions(-) create mode 100644 kernel_backend_test.go diff --git a/internal/backend/kernel/backend.go b/internal/backend/kernel/backend.go index 9e160252..b785bebe 100644 --- a/internal/backend/kernel/backend.go +++ b/internal/backend/kernel/backend.go @@ -15,13 +15,29 @@ import ( "github.com/databricks/databricks-sql-go/internal/backend" ) -// Config is the connection config for the kernel backend: host, the warehouse -// (by bare id or http path), and a PAT. +// 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. 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...) + + // SessionConf carries server-bound session confs verbatim — the same map the + // Thrift backend forwards (STATEMENT_TIMEOUT, QUERY_TAGS, TIMEZONE, …). + SessionConf map[string]string + + // TLSSkipVerify disables server-cert hostname verification (maps the driver's + // WithSkipTLSHostVerify / TLSConfig.InsecureSkipVerify). + TLSSkipVerify 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 + // connection. + ProxyURL string } // KernelBackend implements backend.Backend over the kernel C ABI. One backend @@ -92,6 +108,44 @@ func (k *KernelBackend) OpenSession(ctx context.Context) error { return fmt.Errorf("kernel: set_auth_pat: %w", toDriverError(err)) } + // TLS: skip server-cert hostname verification when the driver requested it. + if k.cfg.TLSSkipVerify { + 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", toDriverError(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), + // so username/password are NULL. + if k.cfg.ProxyURL != "" { + url := newCStr(k.cfg.ProxyURL) + defer url.free() + if err := call(func() C.KernelStatusCode { + return C.kernel_session_config_set_proxy(cfg, url.c, nil, nil, nil) + }); err != nil { + return fmt.Errorf("kernel: set_proxy: %w", toDriverError(err)) + } + } + + // Session confs (STATEMENT_TIMEOUT, QUERY_TAGS, TIMEZONE, …) — the same map + // the Thrift backend forwards, applied one key at a time. + for key, val := range k.cfg.SessionConf { + ck := newCStr(key) + cv := newCStr(val) + errSet := call(func() C.KernelStatusCode { + return C.kernel_session_config_set_session_conf(cfg, ck.c, cv.c) + }) + ck.free() + cv.free() + if errSet != nil { + return fmt.Errorf("kernel: set_session_conf[%s]: %w", key, toDriverError(errSet)) + } + } + var sess *C.kernel_session_t if err := call(func() C.KernelStatusCode { return C.kernel_session_open(cfg, &sess) }); err != nil { return fmt.Errorf("kernel: session_open: %w", toDriverError(err)) diff --git a/kernel_backend.go b/kernel_backend.go index d89349f2..2b823cda 100644 --- a/kernel_backend.go +++ b/kernel_backend.go @@ -4,6 +4,7 @@ package dbsql import ( "context" + "net/http" "github.com/databricks/databricks-sql-go/internal/backend" "github.com/databricks/databricks-sql-go/internal/backend/kernel" @@ -11,14 +12,51 @@ import ( ) // newKernelBackend builds the SEA-via-kernel backend from the driver config; the -// connector opens the session right after, matching the Thrift path. It maps the -// config fields the kernel backend currently reads — host, warehouse/http path, -// and PAT. +// connector opens the session right after, matching the Thrift path. It reads the +// same config fields Thrift does and translates them to the kernel's flat +// connection config, so the user-facing options are unchanged — only the routing +// differs. The public API adds nothing beyond WithUseKernel. func newKernelBackend(_ context.Context, cfg *config.Config) (backend.Backend, error) { - return kernel.New(kernel.Config{ + kc := kernel.Config{ Host: cfg.Host, HTTPPath: cfg.HTTPPath, WarehouseID: cfg.WarehouseID, Token: cfg.AccessToken, - }), nil + // Session confs (STATEMENT_TIMEOUT, QUERY_TAGS, TIMEZONE, …) are the same + // map the Thrift backend forwards; SPOG org routing rides in HTTPPath's ?o= + // and is parsed kernel-side, matching Thrift's URL-based routing. + SessionConf: cfg.SessionParams, + } + // TLS: the driver honors TLSConfig only for InsecureSkipVerify (see + // internal/client), so map exactly that knob to the kernel. + if cfg.TLSConfig != nil && cfg.TLSConfig.InsecureSkipVerify { + kc.TLSSkipVerify = true + } + // 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) + return kernel.New(kc), nil +} + +// proxyForEndpoint resolves the proxy the Thrift path would use for this +// connection, via the same http.ProxyFromEnvironment (HTTP(S)_PROXY / NO_PROXY) +// the Thrift transport applies at request time. Building the endpoint request +// lets ProxyFromEnvironment apply the NO_PROXY rules for this exact host, so the +// kernel sees the same effective proxy decision — returning "" (direct) when +// NO_PROXY excludes the host or no proxy is set. No extra dependency: this is the +// stdlib function the driver already relies on. +func proxyForEndpoint(cfg *config.Config) string { + endpoint, err := cfg.ToEndpointURL() + if err != nil { + return "" + } + req, err := http.NewRequest(http.MethodPost, endpoint, nil) + if err != nil { + return "" + } + proxyURL, err := http.ProxyFromEnvironment(req) + if err != nil || proxyURL == nil { + return "" + } + return proxyURL.String() } diff --git a/kernel_backend_test.go b/kernel_backend_test.go new file mode 100644 index 00000000..36748926 --- /dev/null +++ b/kernel_backend_test.go @@ -0,0 +1,33 @@ +//go:build cgo && databricks_kernel + +package dbsql + +import ( + "testing" + + "github.com/databricks/databricks-sql-go/internal/config" +) + +// proxyForEndpoint returns a valid config's endpoint proxy without error. Its +// value comes from http.ProxyFromEnvironment, which snapshots the proxy env once +// per process (a sync.Once) — the same cached decision the Thrift transport +// makes — so the resolved value can't be re-driven by setting env vars mid-test. +// This asserts the invariant that matters here: a well-formed config never makes +// the resolver error out (it returns "" for direct), and a malformed config +// (missing host) is handled gracefully rather than panicking. +func TestProxyForEndpoint(t *testing.T) { + valid := config.WithDefaults() + valid.Host = "my-workspace.databricks.com" + valid.Port = 443 + valid.HTTPPath = "/sql/1.0/warehouses/abc" + // Must not panic; with no proxy env in the test environment this is "". + _ = proxyForEndpoint(valid) + + // A config whose endpoint URL can't be built (no host) resolves to direct + // rather than erroring. + bad := config.WithDefaults() + bad.Host = "" + if got := proxyForEndpoint(bad); got != "" { + t.Errorf("unbuildable endpoint should resolve to direct, got %q", got) + } +} diff --git a/kernel_e2e_test.go b/kernel_e2e_test.go index f77c7ed1..48a0cdb8 100644 --- a/kernel_e2e_test.go +++ b/kernel_e2e_test.go @@ -51,6 +51,77 @@ func TestKernelE2ESelect1(t *testing.T) { } } +// kernelTestDBWith opens a kernel-backed *sql.DB with extra connector options on +// top of the base host/path/PAT, or skips when creds are unset. It is the config +// 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") + if host == "" || httpPath == "" || token == "" { + t.Skip("set DATABRICKS_HOST / DATABRICKS_HTTP_PATH / DATABRICKS_TOKEN for the kernel e2e") + } + opts := append([]ConnOption{ + WithServerHostname(host), + WithHTTPPath(httpPath), + WithAccessToken(token), + WithUseKernel(true), + }, extra...) + connector, err := NewConnector(opts...) + if err != nil { + t.Fatalf("NewConnector: %v", err) + } + return sql.OpenDB(connector) +} + +// TestKernelE2EQueryTags proves session confs reach the server: WithQueryTags +// (the same option the Thrift path uses) is routed to the kernel and read back +// via SET, which echoes each tag by key. +func TestKernelE2EQueryTags(t *testing.T) { + db := kernelTestDBWith(t, WithQueryTags(map[string]string{"team": "peco"})) + defer db.Close() + + var key, val string + if err := db.QueryRowContext(context.Background(), "SET query_tags").Scan(&key, &val); err != nil { + t.Fatalf("SET query_tags: %v", err) + } + if key != "team" || val != "peco" { + t.Errorf("query tag read back as %q=%q, want team=peco", key, val) + } +} + +// TestKernelE2EStatementTimeout proves a STATEMENT_TIMEOUT session param (via +// WithSessionParams) is applied on the kernel session and read back via SET. +func TestKernelE2EStatementTimeout(t *testing.T) { + db := kernelTestDBWith(t, WithSessionParams(map[string]string{"STATEMENT_TIMEOUT": "300"})) + defer db.Close() + + var key, val string + if err := db.QueryRowContext(context.Background(), "SET statement_timeout").Scan(&key, &val); err != nil { + t.Fatalf("SET statement_timeout: %v", err) + } + if val != "300" { + t.Errorf("statement_timeout read back as %q, want 300", val) + } +} + +// TestKernelE2ETLSSkipVerify checks that WithSkipTLSHostVerify (a relaxation +// knob) is accepted on the kernel path; the connection must still succeed +// against the warehouse's valid certificate. +func TestKernelE2ETLSSkipVerify(t *testing.T) { + db := kernelTestDBWith(t, WithSkipTLSHostVerify()) + defer db.Close() + + var got int64 + if err := db.QueryRowContext(context.Background(), "SELECT 1").Scan(&got); err != nil { + t.Fatalf("query with TLS skip-verify: %v", err) + } + if got != 1 { + t.Errorf("SELECT 1 = %d, want 1", got) + } +} + // TestKernelE2EDataTypes scans each supported scalar type in its own subtest, so // a failure names the exact type rather than being masked by others in a shared // row. Each case selects a single value and compares the scanned result. NULL is From 0c112018523ba3f2451090766a6e586935abc4b1 Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Thu, 9 Jul 2026 07:17:26 +0000 Subject: [PATCH 04/32] fix(kernel): honor session timezone; reject unsupported options; fix leak MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Address three review findings on the kernel backend. Session timezone: DATE/TIMESTAMP values were rendered in UTC, ignoring the configured location, so a connection with a timezone returned times in a different zone than the Thrift backend. Forward cfg.Location into the kernel config and apply it (.In(loc)) when scanning dates/timestamps, including inside nested/JSON values — matching the Thrift path. nil location keeps UTC. Silently-dropped options: newKernelBackend ignored Catalog/Schema and EnableMetricViewMetadata. Neither is wired for the kernel yet — Catalog/Schema have no kernel C-ABI setter, and metric-view maps to a server session conf we want to route backend-neutrally rather than duplicate the Thrift literal here — so the backend now returns a clear error at connect time when either is set, instead of running with different behavior than Thrift. Both are follow-ups. Handle leak: kernelOp.Results returned without closing the operation when get_result_stream failed — and on the query path nothing else closes it, since the (absent) Rows was to own teardown. Close the operation on that error path. Tests: unit tests for timezone rendering (location applied vs nil=UTC) and the unsupported-option rejections; a live timezone e2e asserting the scanned timestamp carries the configured location. Signed-off-by: Mani Kaustubh Mathur --- internal/backend/kernel/backend.go | 5 +++ internal/backend/kernel/kernel_test.go | 61 ++++++++++++++++++++++---- internal/backend/kernel/operation.go | 9 +++- internal/backend/kernel/rows.go | 22 +++++++--- internal/backend/kernel/scan_nested.go | 40 +++++++++-------- kernel_backend.go | 25 +++++++++-- kernel_backend_test.go | 47 ++++++++++++++++++++ kernel_e2e_test.go | 18 ++++++++ 8 files changed, 189 insertions(+), 38 deletions(-) diff --git a/internal/backend/kernel/backend.go b/internal/backend/kernel/backend.go index b785bebe..7f92b2be 100644 --- a/internal/backend/kernel/backend.go +++ b/internal/backend/kernel/backend.go @@ -11,6 +11,7 @@ import "C" import ( "context" "fmt" + "time" "github.com/databricks/databricks-sql-go/internal/backend" ) @@ -38,6 +39,10 @@ type Config struct { // is applied during resolution). Empty leaves the kernel on a direct // connection. ProxyURL string + + // 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 } // KernelBackend implements backend.Backend over the kernel C ABI. One backend diff --git a/internal/backend/kernel/kernel_test.go b/internal/backend/kernel/kernel_test.go index cf313d69..de6bf4f5 100644 --- a/internal/backend/kernel/kernel_test.go +++ b/internal/backend/kernel/kernel_test.go @@ -6,6 +6,7 @@ import ( "database/sql/driver" "errors" "testing" + "time" "github.com/apache/arrow/go/v12/arrow" "github.com/apache/arrow/go/v12/arrow/array" @@ -65,7 +66,7 @@ func TestScanCellScalars(t *testing.T) { b.Append(42) arr := b.NewArray() defer arr.Release() - v, err := scanCell(arr, 0) + v, err := scanCell(arr, 0, nil) if err != nil { t.Fatal(err) } @@ -80,7 +81,7 @@ func TestScanCellScalars(t *testing.T) { b.Append("hi") arr := b.NewArray() defer arr.Release() - v, err := scanCell(arr, 0) + v, err := scanCell(arr, 0, nil) if err != nil { t.Fatal(err) } @@ -95,7 +96,7 @@ func TestScanCellScalars(t *testing.T) { b.AppendNull() arr := b.NewArray() defer arr.Release() - v, err := scanCell(arr, 0) + v, err := scanCell(arr, 0, nil) if err != nil { t.Fatal(err) } @@ -112,7 +113,7 @@ func TestScanCellScalars(t *testing.T) { b.Append(decimal128.FromU64(12345)) arr := b.NewArray() defer arr.Release() - v, err := scanCell(arr, 0) + v, err := scanCell(arr, 0, nil) if err != nil { t.Fatal(err) } @@ -129,12 +130,54 @@ func TestScanCellScalars(t *testing.T) { b.Append(1000) arr := b.NewArray() defer arr.Release() - if _, err := scanCell(arr, 0); err == nil { + if _, err := scanCell(arr, 0, nil); err == nil { t.Error("scanning a Duration should return an unsupported-type error") } }) } +// 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) { + pool := memory.NewGoAllocator() + loc, err := time.LoadLocation("America/New_York") + if err != nil { + t.Skipf("tz database unavailable: %v", err) + } + + // 2026-07-09T12:00:00Z as microseconds since epoch. + utcTS := time.Date(2026, time.July, 9, 12, 0, 0, 0, time.UTC) + b := array.NewTimestampBuilder(pool, &arrow.TimestampType{Unit: arrow.Microsecond}) + defer b.Release() + b.Append(arrow.Timestamp(utcTS.UnixMicro())) + arr := b.NewArray() + defer arr.Release() + + t.Run("location applied", func(t *testing.T) { + v, err := scanCell(arr, 0, loc) + if err != nil { + t.Fatal(err) + } + got := v.(time.Time) + if got.Location() != loc { + t.Errorf("location = %v, want %v", got.Location(), loc) + } + if !got.Equal(utcTS) { + t.Errorf("instant changed: got %v, want %v", got, utcTS) + } + }) + + t.Run("nil location is UTC", func(t *testing.T) { + v, err := scanCell(arr, 0, nil) + if err != nil { + t.Fatal(err) + } + if v.(time.Time).Location() != time.UTC { + t.Errorf("nil location should render UTC, got %v", v.(time.Time).Location()) + } + }) +} + // scanCell renders nested types (list/struct/map) to a JSON string matching the // Thrift path. func TestScanCellNested(t *testing.T) { @@ -150,7 +193,7 @@ func TestScanCellNested(t *testing.T) { vb.Append(3) arr := b.NewArray() defer arr.Release() - v, err := scanCell(arr, 0) + v, err := scanCell(arr, 0, nil) if err != nil { t.Fatal(err) } @@ -171,7 +214,7 @@ func TestScanCellNested(t *testing.T) { b.FieldBuilder(1).(*array.StringBuilder).Append("x") arr := b.NewArray() defer arr.Release() - v, err := scanCell(arr, 0) + v, err := scanCell(arr, 0, nil) if err != nil { t.Fatal(err) } @@ -190,7 +233,7 @@ func TestScanCellNested(t *testing.T) { ib.Append(9) arr := b.NewArray() defer arr.Release() - v, err := scanCell(arr, 0) + v, err := scanCell(arr, 0, nil) if err != nil { t.Fatal(err) } @@ -205,7 +248,7 @@ func TestScanCellNested(t *testing.T) { b.AppendNull() arr := b.NewArray() defer arr.Release() - v, err := scanCell(arr, 0) + v, err := scanCell(arr, 0, nil) if err != nil { t.Fatal(err) } diff --git a/internal/backend/kernel/operation.go b/internal/backend/kernel/operation.go index 568c2624..730a4ddd 100644 --- a/internal/backend/kernel/operation.go +++ b/internal/backend/kernel/operation.go @@ -108,7 +108,7 @@ func (k *KernelBackend) execute(ctx context.Context, req backend.ExecRequest) (b C.kernel_statement_canceller_free(canceller) } - op := &kernelOp{stmt: stmt} + op := &kernelOp{stmt: stmt, location: k.cfg.Location} if execErr != nil { // Prefer the caller's ctx error when the ctx was cancelled (database/sql // convention), keeping the kernel error as the cause. @@ -131,6 +131,9 @@ type kernelOp struct { stmt *C.kernel_statement_t exec *C.kernel_executed_statement_t closed bool + // 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 } var _ backend.Operation = (*kernelOp)(nil) @@ -158,6 +161,10 @@ func (o *kernelOp) Results(ctx context.Context, callbacks *dbsqlrows.TelemetryCa if err := call(func() C.KernelStatusCode { return C.kernel_executed_statement_get_result_stream(o.exec, &stream) }); err != nil { + // No Rows is returned to own teardown, and the query path does not call + // Operation.Close on a Results error — so close the handles here to avoid + // leaking the statement / executed handle (and its server operation). + o.close() return nil, fmt.Errorf("kernel: get_result_stream: %w", toDriverError(err)) } return newKernelRows(ctx, o, stream, callbacks) diff --git a/internal/backend/kernel/rows.go b/internal/backend/kernel/rows.go index c0258753..107048b5 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" @@ -115,7 +116,7 @@ func (r *kernelRows) Next(dest []driver.Value) error { } rec := r.cur for c := 0; c < len(dest); c++ { - v, err := scanCell(rec.Column(c), r.rowInCur) + v, err := scanCell(rec.Column(c), r.rowInCur, r.op.location) if err != nil { return fmt.Errorf("kernel: scan col %d (%s): %w", c, r.cols[c], err) } @@ -171,7 +172,7 @@ func (r *kernelRows) nextBatch() error { // 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. -func scanCell(col arrow.Array, row int) (driver.Value, error) { +func scanCell(col arrow.Array, row int, loc *time.Location) (driver.Value, error) { if col.IsNull(row) { return nil, nil } @@ -207,24 +208,33 @@ func scanCell(col arrow.Array, row int) (driver.Value, error) { case *array.Binary: return c.Value(row), nil case *array.Date32: - return c.Value(row).ToTime(), nil + return inLocation(c.Value(row).ToTime(), loc), nil case *array.Date64: - return c.Value(row).ToTime(), nil + return inLocation(c.Value(row).ToTime(), loc), nil case *array.Timestamp: dt, ok := col.DataType().(*arrow.TimestampType) if !ok { return nil, fmt.Errorf("timestamp column has unexpected datatype %s", col.DataType()) } - return c.Value(row).ToTime(dt.Unit), nil + return inLocation(c.Value(row).ToTime(dt.Unit), loc), nil case *array.Decimal128: dt := col.DataType().(*arrow.Decimal128Type) return decimal128ToExactString(c.Value(row), dt.Scale), 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) + return renderJSONString(col, row, loc) default: return nil, fmt.Errorf("kernel: scanning arrow type %s is not supported "+ "(intervals are not yet handled)", col.DataType()) } } + +// 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 { + if loc == nil { + return t + } + return t.In(loc) +} diff --git a/internal/backend/kernel/scan_nested.go b/internal/backend/kernel/scan_nested.go index 75138c26..013342b1 100644 --- a/internal/backend/kernel/scan_nested.go +++ b/internal/backend/kernel/scan_nested.go @@ -33,38 +33,40 @@ import ( "github.com/apache/arrow/go/v12/arrow/array" ) -// renderJSONString renders one nested cell as a JSON string (nil for a NULL cell). -func renderJSONString(col arrow.Array, row int) (driver.Value, error) { +// renderJSONString renders one nested cell as a JSON string (nil for a NULL +// cell). loc is applied to any timestamp/date leaves, matching the top-level +// scan and the Thrift path. +func renderJSONString(col arrow.Array, row int, loc *time.Location) (driver.Value, error) { if col.IsNull(row) { return nil, nil } var b strings.Builder - if err := writeJSON(&b, col, row); err != nil { + if err := writeJSON(&b, col, row, loc); err != nil { return nil, err } return b.String(), nil } // writeJSON writes the JSON form of col[row] into b, recursing for nested types. -func writeJSON(b *strings.Builder, col arrow.Array, row int) error { +func writeJSON(b *strings.Builder, col arrow.Array, row int, loc *time.Location) error { if col.IsNull(row) { b.WriteString("null") return nil } switch c := col.(type) { case *array.List: - return writeListJSON(b, c.ListValues(), int(c.Offsets()[row]), int(c.Offsets()[row+1])) + return writeListJSON(b, c.ListValues(), int(c.Offsets()[row]), int(c.Offsets()[row+1]), loc) case *array.LargeList: - return writeListJSON(b, c.ListValues(), int(c.Offsets()[row]), int(c.Offsets()[row+1])) + return writeListJSON(b, c.ListValues(), int(c.Offsets()[row]), int(c.Offsets()[row+1]), loc) case *array.FixedSizeList: n := int(c.DataType().(*arrow.FixedSizeListType).Len()) - return writeListJSON(b, c.ListValues(), row*n, row*n+n) + return writeListJSON(b, c.ListValues(), row*n, row*n+n, loc) case *array.Map: - return writeMapJSON(b, c, row) + return writeMapJSON(b, c, row, loc) case *array.Struct: - return writeStructJSON(b, c, row) + return writeStructJSON(b, c, row, loc) default: - v, err := scalarForJSON(col, row) + v, err := scalarForJSON(col, row, loc) if err != nil { return err } @@ -72,13 +74,13 @@ func writeJSON(b *strings.Builder, col arrow.Array, row int) error { } } -func writeListJSON(b *strings.Builder, values arrow.Array, start, end int) error { +func writeListJSON(b *strings.Builder, values arrow.Array, start, end int, loc *time.Location) error { b.WriteByte('[') for i := start; i < end; i++ { if i > start { b.WriteByte(',') } - if err := writeJSON(b, values, i); err != nil { + if err := writeJSON(b, values, i, loc); err != nil { return err } } @@ -86,7 +88,7 @@ func writeListJSON(b *strings.Builder, values arrow.Array, start, end int) error return nil } -func writeMapJSON(b *strings.Builder, m *array.Map, row int) error { +func writeMapJSON(b *strings.Builder, m *array.Map, row int, loc *time.Location) error { start, end := int(m.Offsets()[row]), int(m.Offsets()[row+1]) keys := m.Keys() items := m.Items() @@ -96,13 +98,13 @@ func writeMapJSON(b *strings.Builder, m *array.Map, row int) error { b.WriteByte(',') } // JSON object keys must be strings; stringify the key value. - kv, err := scalarForJSON(keys, i) + kv, err := scalarForJSON(keys, i, loc) if err != nil { return err } writeJSONKey(b, kv) b.WriteByte(':') - if err := writeJSON(b, items, i); err != nil { + if err := writeJSON(b, items, i, loc); err != nil { return err } } @@ -110,7 +112,7 @@ func writeMapJSON(b *strings.Builder, m *array.Map, row int) error { return nil } -func writeStructJSON(b *strings.Builder, s *array.Struct, row int) error { +func writeStructJSON(b *strings.Builder, s *array.Struct, row int, loc *time.Location) error { st := s.DataType().(*arrow.StructType) b.WriteByte('{') for f := 0; f < s.NumField(); f++ { @@ -120,7 +122,7 @@ func writeStructJSON(b *strings.Builder, s *array.Struct, row int) error { keyBytes, _ := json.Marshal(st.Field(f).Name) b.Write(keyBytes) b.WriteByte(':') - if err := writeJSON(b, s.Field(f), row); err != nil { + if err := writeJSON(b, s.Field(f), row, loc); err != nil { return err } } @@ -165,7 +167,7 @@ func writeScalarJSON(b *strings.Builder, v any) error { // from scanCell in one way: a nested decimal renders as float64 (lossy), matching // the Thrift path's in-JSON decimal rendering (#274); the exact-string decimal // applies only to a top-level decimal column. -func scalarForJSON(col arrow.Array, row int) (any, error) { +func scalarForJSON(col arrow.Array, row int, loc *time.Location) (any, error) { if col.IsNull(row) { return nil, nil } @@ -178,6 +180,6 @@ func scalarForJSON(col arrow.Array, row int) (any, error) { return c.Value(row).ToFloat64(dt.Scale), nil default: // Reuse scanCell's scalar arm for every non-nested, non-decimal leaf. - return scanCell(col, row) + return scanCell(col, row, loc) } } diff --git a/kernel_backend.go b/kernel_backend.go index 2b823cda..779859f2 100644 --- a/kernel_backend.go +++ b/kernel_backend.go @@ -4,6 +4,7 @@ package dbsql import ( "context" + "errors" "net/http" "github.com/databricks/databricks-sql-go/internal/backend" @@ -17,14 +18,32 @@ 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) { + // A few options aren't wired for the kernel backend yet. Fail loudly rather + // than silently ignore them (which would behave differently than Thrift): + // - Catalog/Schema (WithInitialNamespace): no kernel C-ABI setter yet, so + // the session would run in the default namespace and unqualified names + // would resolve differently. + // - EnableMetricViewMetadata: deferred — it maps to a server session conf, + // which we want to route backend-neutrally rather than duplicate here. + if cfg.Catalog != "" || cfg.Schema != "" { + return nil, errors.New("databricks: WithInitialNamespace (catalog/schema) is not yet supported by the kernel backend; " + + "omit it or use the default (Thrift) backend") + } + if cfg.EnableMetricViewMetadata { + return nil, errors.New("databricks: WithEnableMetricViewMetadata is not yet supported by the kernel backend; " + + "omit it or use the default (Thrift) backend") + } + kc := kernel.Config{ Host: cfg.Host, HTTPPath: cfg.HTTPPath, WarehouseID: cfg.WarehouseID, Token: cfg.AccessToken, - // Session confs (STATEMENT_TIMEOUT, QUERY_TAGS, TIMEZONE, …) are the same - // map the Thrift backend forwards; SPOG org routing rides in HTTPPath's ?o= - // and is parsed kernel-side, matching Thrift's URL-based routing. + 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 + // rides in HTTPPath's ?o= and is parsed kernel-side. SessionConf: cfg.SessionParams, } // TLS: the driver honors TLSConfig only for InsecureSkipVerify (see diff --git a/kernel_backend_test.go b/kernel_backend_test.go index 36748926..ce5356d7 100644 --- a/kernel_backend_test.go +++ b/kernel_backend_test.go @@ -3,11 +3,58 @@ package dbsql import ( + "context" "testing" "github.com/databricks/databricks-sql-go/internal/config" ) +// newKernelBackend rejects options it can't yet honor (initial namespace, +// metric-view metadata) loudly, rather than silently ignoring them — which would +// behave differently than the Thrift backend. +func TestNewKernelBackendRejectsUnsupportedOptions(t *testing.T) { + base := func() *config.Config { + c := config.WithDefaults() + c.Host = "h.databricks.com" + c.Port = 443 + c.HTTPPath = "/sql/1.0/warehouses/abc" + c.AccessToken = "dapi-x" + return c + } + + t.Run("catalog rejected", func(t *testing.T) { + c := base() + c.Catalog = "main" + if _, err := newKernelBackend(context.Background(), c); err == nil { + t.Error("expected an error when a catalog is set on the kernel backend") + } + }) + + t.Run("schema rejected", func(t *testing.T) { + c := base() + c.Schema = "default" + if _, err := newKernelBackend(context.Background(), c); err == nil { + t.Error("expected an error when a schema is set on the kernel backend") + } + }) + + t.Run("metric view rejected", func(t *testing.T) { + c := base() + c.EnableMetricViewMetadata = true + if _, err := newKernelBackend(context.Background(), c); err == nil { + t.Error("expected an error when metric-view metadata is enabled on the kernel backend") + } + }) + + t.Run("supported options ok", func(t *testing.T) { + c := base() + c.SessionParams = map[string]string{"QUERY_TAGS": "a:1"} + if _, err := newKernelBackend(context.Background(), c); err != nil { + t.Errorf("a supported config should build cleanly, got %v", err) + } + }) +} + // proxyForEndpoint returns a valid config's endpoint proxy without error. Its // value comes from http.ProxyFromEnvironment, which snapshots the proxy env once // per process (a sync.Once) — the same cached decision the Thrift transport diff --git a/kernel_e2e_test.go b/kernel_e2e_test.go index 48a0cdb8..1d2b3c00 100644 --- a/kernel_e2e_test.go +++ b/kernel_e2e_test.go @@ -106,6 +106,24 @@ func TestKernelE2EStatementTimeout(t *testing.T) { } } +// TestKernelE2ETimeZone proves the session time zone (WithSessionParams +// timezone) is applied to scanned TIMESTAMP values, matching the Thrift path — +// the returned time.Time carries the configured location, not UTC. +func TestKernelE2ETimeZone(t *testing.T) { + const tz = "America/New_York" + db := kernelTestDBWith(t, WithSessionParams(map[string]string{"timezone": tz})) + defer db.Close() + + var ts time.Time + if err := db.QueryRowContext(context.Background(), + "SELECT CAST('2026-07-09 12:00:00' AS TIMESTAMP)").Scan(&ts); err != nil { + t.Fatalf("query: %v", err) + } + if ts.Location().String() != tz { + t.Errorf("timestamp location = %q, want %q", ts.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 031c7834c76c19bdfa2fe8090898e6f75b4c7cf5 Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Thu, 9 Jul 2026 08:55:46 +0000 Subject: [PATCH 05/32] fix(kernel): render nested decimals as exact JSON numbers MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A DECIMAL inside a nested value (struct field, list element, map value/key) was rendered via ToFloat64, so it emitted a lossy float64 — e.g. a struct decimal 19.99 came out as {"d":19.990000000000002}, diverging from the Thrift path's {"d":19.99}. The top-level decimal column was already exact (#274); only the nested path was lossy. writeJSON now renders a nested Decimal128 with the same exact-string helper the top-level scan uses, emitted as a raw JSON number literal — mirroring the Thrift arrow path's marshalScalar → ValueString (databricks-sql-go#253/#274). Map keys go through scalarForJSON, which now defers to the scalar scan (exact string) as well, so a decimal key is exact too. The earlier nested tests missed this because they used float64-exact values (1.5, 2.5). Add a struct-decimal unit case (19.99) and a nested-decimal column to the live Thrift-parity query as regression guards. Signed-off-by: Mani Kaustubh Mathur --- internal/backend/kernel/kernel_test.go | 21 ++++++++++++++++++++ internal/backend/kernel/scan_nested.go | 27 ++++++++++++-------------- kernel_parity_test.go | 5 ++++- 3 files changed, 37 insertions(+), 16 deletions(-) diff --git a/internal/backend/kernel/kernel_test.go b/internal/backend/kernel/kernel_test.go index de6bf4f5..86217f13 100644 --- a/internal/backend/kernel/kernel_test.go +++ b/internal/backend/kernel/kernel_test.go @@ -242,6 +242,27 @@ func TestScanCellNested(t *testing.T) { } }) + t.Run("nested_decimal_exact", func(t *testing.T) { + // A decimal inside a struct must render as an exact JSON number, not a + // lossy float64 (19.99, not 19.990000000000002) — matching Thrift's + // marshalScalar. Regression guard for a bug the POC missed by only testing + // float64-exact values. + dt := arrow.StructOf(arrow.Field{Name: "d", Type: &arrow.Decimal128Type{Precision: 5, Scale: 2}}) + b := array.NewStructBuilder(pool, dt) + defer b.Release() + b.Append(true) + b.FieldBuilder(0).(*array.Decimal128Builder).Append(decimal128.FromU64(1999)) + arr := b.NewArray() + defer arr.Release() + v, err := scanCell(arr, 0, nil) + if err != nil { + t.Fatal(err) + } + if v.(string) != `{"d":19.99}` { + t.Errorf("got %q, want {\"d\":19.99}", v) + } + }) + t.Run("nested_null", func(t *testing.T) { b := array.NewListBuilder(pool, arrow.PrimitiveTypes.Int64) defer b.Release() diff --git a/internal/backend/kernel/scan_nested.go b/internal/backend/kernel/scan_nested.go index 013342b1..f42213a5 100644 --- a/internal/backend/kernel/scan_nested.go +++ b/internal/backend/kernel/scan_nested.go @@ -65,6 +65,13 @@ func writeJSON(b *strings.Builder, col arrow.Array, row int, loc *time.Location) return writeMapJSON(b, c, row, loc) case *array.Struct: return writeStructJSON(b, c, row, loc) + case *array.Decimal128: + // 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). + b.WriteString(decimal128ToExactString(c.Value(row), col.DataType().(*arrow.Decimal128Type).Scale)) + return nil default: v, err := scalarForJSON(col, row, loc) if err != nil { @@ -163,23 +170,13 @@ func writeScalarJSON(b *strings.Builder, v any) error { return nil } -// scalarForJSON returns the Go value used inside JSON for a leaf cell. It differs -// from scanCell in one way: a nested decimal renders as float64 (lossy), matching -// the Thrift path's in-JSON decimal rendering (#274); the exact-string decimal -// applies only to a top-level decimal column. +// scalarForJSON returns the Go value used for a nested leaf that is not itself a +// container — today only a map key (values are written directly by writeJSON). +// It reuses scanCell's scalar arm, so a decimal key renders via the exact-string +// path (writeJSONKey then quotes it), never a lossy float64. func scalarForJSON(col arrow.Array, row int, loc *time.Location) (any, error) { if col.IsNull(row) { return nil, nil } - switch c := col.(type) { - case *array.Decimal128: - dt := col.DataType().(*arrow.Decimal128Type) - return c.Value(row).ToFloat64(dt.Scale), nil - case *array.Decimal256: - dt := col.DataType().(*arrow.Decimal256Type) - return c.Value(row).ToFloat64(dt.Scale), nil - default: - // Reuse scanCell's scalar arm for every non-nested, non-decimal leaf. - return scanCell(col, row, loc) - } + return scanCell(col, row, loc) } diff --git a/kernel_parity_test.go b/kernel_parity_test.go index 22d8f4d1..223a0081 100644 --- a/kernel_parity_test.go +++ b/kernel_parity_test.go @@ -38,7 +38,10 @@ func TestKernelThriftParity(t *testing.T) { const query = "SELECT CAST(7 AS BIGINT), CAST(2.5 AS DOUBLE), 'parity', " + "CAST(NULL AS STRING), CAST(9.99 AS DECIMAL(5,2)), true, " + "array(1, 2, 3), map('k', 9), named_struct('a', 1, 'b', 'x'), " + - `parse_json('{"a":1,"b":[2,3]}'), st_point(1, 2)` + `parse_json('{"a":1,"b":[2,3]}'), st_point(1, 2), ` + + // A decimal inside a struct: exercises exact-string nested-decimal + // rendering (19.99, not a lossy 19.990000000000002). + "named_struct('d', CAST(19.99 AS DECIMAL(5,2)))" kernelDB := kernelTestDB(t) defer kernelDB.Close() From 8c99345394c2ca8fe60ee75e12722983ace60431 Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Fri, 10 Jul 2026 10:17:04 +0000 Subject: [PATCH 06/32] fix(kernel): sync to merged kernel C ABI (cancel arity, logging, double-free) Adapt the SEA-via-kernel backend to the kernel C ABI as merged (canceller #163 + logging/U2M #162), and fix a latent crash the sync surfaced. - Cancellation: kernel_statement_canceller_cancel now takes a bool* dispatched out-param. Route it through a fireCancel helper and stop the 250ms re-fire loop the moment a cancel actually dispatches (server id observed, RPC sent), instead of firing blindly until execute returns. - Logging: wire kernel_init_logging under a sync.Once, gated on the same DBSQL_KERNEL_DEBUG flag as the binding tracer, so one switch interleaves Go and kernel logs on stderr and both stay off by default / during benchmarks. level=NULL honors RUST_LOG; file=NULL uses stderr. Filter on target databricks::sql::kernel (colons), not the underscore module path. - Double-free: OpenSession freed the session config on the kernel_session_open failure path, but the kernel consumes the config on every path. Set consumed before checking the error so a failed open (e.g. HTTP 403) returns a clean error instead of aborting the process. - doc.go: correct the debug-logging env vars and refresh the supported-features summary (nested/complex types, TLS/proxy/session-conf now supported). Co-authored-by: Isaac Signed-off-by: Mani Kaustubh Mathur --- doc.go | 31 ++++++++++++------- internal/backend/kernel/backend.go | 14 +++++++-- internal/backend/kernel/cgo.go | 46 ++++++++++++++++++++++++++-- internal/backend/kernel/operation.go | 19 +++++++++--- 4 files changed, 89 insertions(+), 21 deletions(-) diff --git a/doc.go b/doc.go index becab1f8..f5e48063 100644 --- a/doc.go +++ b/doc.go @@ -179,17 +179,26 @@ per connection with WithUseKernel: 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, enable both on the -same stderr. The kernel filters on the target databricks::sql::kernel (note the -colons — it is not the crate module path), and the driver's step tracer is -enabled by the standard logging level: - - RUST_LOG=databricks::sql::kernel=debug DATABRICKS_LOG_LEVEL=debug ./your_app 2>&1 - -The kernel backend currently supports PAT authentication and reading -scalar-typed results (CloudFetch is handled transparently), with context -cancellation. Bound parameters, nested and complex types, and other -authentication and transport options are not yet supported on this backend. +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: + 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 + +The kernel backend currently supports PAT authentication; reading scalar, nested, +and complex-typed results (CloudFetch is handled transparently); context +cancellation; and the TLS, proxy, and session-conf (query tags, statement timeout, +time zone) connection options. Bound parameters, OAuth (M2M/U2M), metadata +commands, and initial catalog/schema are not yet supported on this backend and +return a clear error at connect time rather than being silently ignored. # Programmatically Retrieving Connection and Query Id diff --git a/internal/backend/kernel/backend.go b/internal/backend/kernel/backend.go index 7f92b2be..4f528f86 100644 --- a/internal/backend/kernel/backend.go +++ b/internal/backend/kernel/backend.go @@ -70,12 +70,18 @@ func New(cfg Config) *KernelBackend { // time. The config handle is consumed by kernel_session_open on success and // freed by us on any earlier failure. func (k *KernelBackend) OpenSession(ctx context.Context) error { + initKernelLogging() klog("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 { return fmt.Errorf("kernel: config_new: %w", toDriverError(err)) } + // kernel_session_open consumes the config on EVERY path — success and + // failure alike (it reclaims the box up front). So we free the config + // ourselves only when we bail out BEFORE reaching kernel_session_open (a + // setter error below); once that call is made, ownership has transferred and + // a free here would double-free. consumed := false defer func() { if !consumed { @@ -151,11 +157,15 @@ func (k *KernelBackend) OpenSession(ctx context.Context) error { } } + // kernel_session_open takes ownership of cfg here, on both success and + // failure — mark consumed before checking the error so the deferred free + // never runs on the already-consumed config. var sess *C.kernel_session_t - if err := call(func() C.KernelStatusCode { return C.kernel_session_open(cfg, &sess) }); err != nil { + err := call(func() C.KernelStatusCode { return C.kernel_session_open(cfg, &sess) }) + consumed = true + if err != nil { return fmt.Errorf("kernel: session_open: %w", toDriverError(err)) } - consumed = true // kernel_session_open took ownership of cfg k.session = sess k.valid = true // The C ABI exposes no formatted session-id accessor; use the handle pointer diff --git a/internal/backend/kernel/cgo.go b/internal/backend/kernel/cgo.go index f967c2a0..80a5b4c2 100644 --- a/internal/backend/kernel/cgo.go +++ b/internal/backend/kernel/cgo.go @@ -30,6 +30,7 @@ import ( "fmt" "os" "runtime" + "sync" "unsafe" dbsqlerrint "github.com/databricks/databricks-sql-go/internal/errors" @@ -40,9 +41,9 @@ import ( // 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 kernel's own -// Rust logs are enabled separately via kernel_init_logging / RUST_LOG and -// interleave on the same stderr. +// 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") != "" func klog(format string, args ...any) { @@ -56,6 +57,32 @@ func klog(format string, args ...any) { // assert the flag wiring; benchmarks assert it is false before measuring. func KernelDebugEnabled() bool { return kdebug } +// initLoggingOnce guards kernel_init_logging, which is process-wide and +// first-call-wins in the kernel. We install the kernel subscriber lazily on the +// first session open rather than in init(), so a process that never opens a +// 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. +func initKernelLogging() { + if !kdebug { + return + } + initLoggingOnce.Do(func() { + if err := call(func() C.KernelStatusCode { + return C.kernel_init_logging(nil, nil) + }); err != nil { + klog("kernel_init_logging: %v (kernel logs unavailable; proceeding)", err) + } + }) +} + // 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. // @@ -76,6 +103,19 @@ func call(fn func() C.KernelStatusCode) error { return lastError(st) } +// fireCancel dispatches a server-side cancel and reports whether the RPC was +// actually sent. The kernel returns dispatched=false while no server statement +// id has been observed yet (the F4 window before the execute POST returns), and +// true once the cancel RPC goes out — which is the watcher's cue to stop +// re-firing. Best-effort: a non-Success status is ignored (dispatched stays +// false, so the watcher keeps retrying), matching the "proceed without cancel" +// stance rather than surfacing a cancel-path error. +func fireCancel(canceller *C.kernel_statement_canceller_t) bool { + var dispatched C.bool + C.kernel_statement_canceller_cancel(canceller, &dispatched) + return bool(dispatched) +} + // lastError reads the kernel's thread-local last error and copies its string // fields out immediately — the C `char*` fields are valid only until the next // FFI call on this thread. Must run on the same OS thread as the failing call; diff --git a/internal/backend/kernel/operation.go b/internal/backend/kernel/operation.go index 730a4ddd..566394dc 100644 --- a/internal/backend/kernel/operation.go +++ b/internal/backend/kernel/operation.go @@ -65,8 +65,11 @@ func (k *KernelBackend) execute(ctx context.Context, req backend.ExecRequest) (b // canceller). The server publishes the statement id to the canceller's // inflight slot only when the initial POST returns — held up to the server's // inline wait (~10s) even for a long query — so a cancel fired before that is - // a no-op. Re-fire every 250ms after ctx.Done until execute returns, so the - // cancel takes effect once the id appears, with no kernel change. + // a no-op. Re-fire every 250ms after ctx.Done until the kernel reports the + // cancel RPC was actually dispatched (dispatched=true, i.e. the id appeared + // and the RPC went out), then stop: one real cancel is enough, and further + // fires would only hammer the server. Falls back to firing until execute + // returns (done) if the RPC never dispatches. done := make(chan struct{}) var watcherWg sync.WaitGroup if canceller != nil && ctx.Done() != nil { @@ -78,16 +81,22 @@ func (k *KernelBackend) execute(ctx context.Context, req backend.ExecRequest) (b case <-done: return } - klog("ctx.Done (%v) → firing canceller (with retry)", ctx.Err()) + klog("ctx.Done (%v) → firing canceller (with retry until dispatched)", ctx.Err()) ticker := time.NewTicker(250 * time.Millisecond) defer ticker.Stop() - C.kernel_statement_canceller_cancel(canceller) + if fireCancel(canceller) { + klog("cancel dispatched on first fire") + return + } for { select { case <-done: return case <-ticker.C: - C.kernel_statement_canceller_cancel(canceller) + if fireCancel(canceller) { + klog("cancel dispatched, watcher stopping") + return + } } } }() From 45aaf8490632c0fcf03c591b0f733f802bbf2ca7 Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Fri, 10 Jul 2026 14:01:06 +0000 Subject: [PATCH 07/32] fix(kernel): map InsecureSkipVerify to both kernel TLS relaxations MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit crypto/tls's InsecureSkipVerify accepts any server cert — it disables both chain validation and the hostname check. The kernel path mapped it to only kernel_session_config_set_tls_skip_hostname_verification, leaving the kernel stricter than the Thrift path it mirrors: a self-signed cert + skip-verify succeeds on Thrift but was rejected by the kernel at chain validation. Also call kernel_session_config_set_tls_allow_self_signed under the same flag so both relax together, matching crypto/tls semantics and the pyo3/napi mapping. Also trim two over-detailed doc comments (call, fireCancel): drop the kernel-out-param aside and the internal "F4 window" name. Co-authored-by: Isaac Signed-off-by: Mani Kaustubh Mathur --- internal/backend/kernel/backend.go | 16 +++++++++++++--- internal/backend/kernel/cgo.go | 10 +++------- 2 files changed, 16 insertions(+), 10 deletions(-) diff --git a/internal/backend/kernel/backend.go b/internal/backend/kernel/backend.go index 4f528f86..bbd058ad 100644 --- a/internal/backend/kernel/backend.go +++ b/internal/backend/kernel/backend.go @@ -30,8 +30,10 @@ type Config struct { // Thrift backend forwards (STATEMENT_TIMEOUT, QUERY_TAGS, TIMEZONE, …). SessionConf map[string]string - // TLSSkipVerify disables server-cert hostname verification (maps the driver's - // WithSkipTLSHostVerify / TLSConfig.InsecureSkipVerify). + // TLSSkipVerify accepts any server cert (maps the driver's + // WithSkipTLSHostVerify / TLSConfig.InsecureSkipVerify). crypto/tls's + // InsecureSkipVerify disables both chain validation and the hostname check, + // so the kernel path relaxes both to match. TLSSkipVerify bool // ProxyURL configures an HTTP proxy, already resolved for this endpoint from @@ -119,8 +121,16 @@ func (k *KernelBackend) OpenSession(ctx context.Context) error { return fmt.Errorf("kernel: set_auth_pat: %w", toDriverError(err)) } - // TLS: skip server-cert hostname verification when the driver requested it. + // TLS: crypto/tls's InsecureSkipVerify accepts any server cert, so relax both + // chain validation and the hostname check — mapping only one would leave the + // kernel path stricter than the Thrift path it mirrors (a self-signed cert + // would still be rejected). if k.cfg.TLSSkipVerify { + if err := call(func() C.KernelStatusCode { + return C.kernel_session_config_set_tls_allow_self_signed(cfg, C.bool(true)) + }); err != nil { + return fmt.Errorf("kernel: set_tls_allow_self_signed: %w", toDriverError(err)) + } if err := call(func() C.KernelStatusCode { return C.kernel_session_config_set_tls_skip_hostname_verification(cfg, C.bool(true)) }); err != nil { diff --git a/internal/backend/kernel/cgo.go b/internal/backend/kernel/cgo.go index 80a5b4c2..58294c6f 100644 --- a/internal/backend/kernel/cgo.go +++ b/internal/backend/kernel/cgo.go @@ -90,9 +90,7 @@ func initKernelLogging() { // call (kernel_get_last_error). Go's M:N scheduler can move a goroutine to a // different OS thread between two cgo calls, so a naive call-then-read pair can // observe the wrong thread's buffer. LockOSThread pins the goroutine to its OS -// thread across the call and its error read, closing that window. (The kernel's -// preferred long-term shape is a same-call out-param error, which would remove -// the need for this; until then LockOSThread is correct and call-site-local.) +// thread across the call and its error read, closing that window. func call(fn func() C.KernelStatusCode) error { runtime.LockOSThread() defer runtime.UnlockOSThread() @@ -105,11 +103,9 @@ func call(fn func() C.KernelStatusCode) error { // fireCancel dispatches a server-side cancel and reports whether the RPC was // actually sent. The kernel returns dispatched=false while no server statement -// id has been observed yet (the F4 window before the execute POST returns), and +// id has been observed yet (the window before the execute POST returns), and // true once the cancel RPC goes out — which is the watcher's cue to stop -// re-firing. Best-effort: a non-Success status is ignored (dispatched stays -// false, so the watcher keeps retrying), matching the "proceed without cancel" -// stance rather than surfacing a cancel-path error. +// re-firing. func fireCancel(canceller *C.kernel_statement_canceller_t) bool { var dispatched C.bool C.kernel_statement_canceller_cancel(canceller, &dispatched) From 220d0d2c3bd1bb93920efd25edb9f1795e52503c Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Fri, 10 Jul 2026 14:07:06 +0000 Subject: [PATCH 08/32] chore(kernel): remove hardcoded local kernel paths from cgo directives The #cgo CFLAGS/LDFLAGS hardcoded an absolute path to a developer's local kernel checkout, so the package built only on that one machine and leaked a local filesystem path into the repo. Drop the search paths from the directives (keep only the library name) and document supplying the header/lib locations at build time via the standard CGO_CFLAGS / CGO_LDFLAGS env vars. Also drop a dangling reference to an internal-only design doc. Committing a per-platform prebuilt static lib at a ${SRCDIR}-relative path + a tagged CI job is a separate distribution follow-up. Co-authored-by: Isaac Signed-off-by: Mani Kaustubh Mathur --- internal/backend/kernel/cgo.go | 19 ++++++++++++------- 1 file changed, 12 insertions(+), 7 deletions(-) diff --git a/internal/backend/kernel/cgo.go b/internal/backend/kernel/cgo.go index 58294c6f..3ee07b7d 100644 --- a/internal/backend/kernel/cgo.go +++ b/internal/backend/kernel/cgo.go @@ -11,16 +11,21 @@ // the error mapping to the driver's error surface, and the gated step logger. // The backend, operation, and rows layers live in sibling files. // -// The link directives below point at a locally built kernel -// (~/oss/databricks-sql-kernel/target/release + include), so this package builds -// only on a machine with that checkout. A shippable build instead links a -// committed per-platform prebuilt static lib via a ${SRCDIR}-relative path (the -// go-duckdb duckdb-go-bindings model); see cabi-distribution-references.md. +// The directives below name the kernel library to link but carry no search +// paths, so the header and static lib locations are supplied at build time via +// the standard CGO_CFLAGS / CGO_LDFLAGS environment variables, e.g.: +// +// CGO_CFLAGS="-I/include" \ +// CGO_LDFLAGS="-L/target/release -Wl,-rpath,/target/release" \ +// go build -tags databricks_kernel ./... +// +// A shippable build instead links a committed per-platform prebuilt static lib +// via a ${SRCDIR}-relative path (the go-duckdb duckdb-go-bindings model); wiring +// that + a tagged CI job is a distribution follow-up. package kernel /* -#cgo CFLAGS: -I/home/mani.mathur/oss/databricks-sql-kernel/include -#cgo LDFLAGS: -L/home/mani.mathur/oss/databricks-sql-kernel/target/release -ldatabricks_sql_kernel -Wl,-rpath,/home/mani.mathur/oss/databricks-sql-kernel/target/release -ldl -lm +#cgo LDFLAGS: -ldatabricks_sql_kernel -ldl -lm #include #include "databricks_kernel.h" */ From c7a19f0c039a8f142ac903e99ddd2146d8d86ac6 Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Fri, 10 Jul 2026 20:48:06 +0000 Subject: [PATCH 09/32] =?UTF-8?q?fix(kernel):=20address=20code=20review=20?= =?UTF-8?q?=E2=80=94=20correctness,=20doc-contract,=20parity,=20tests?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Close the GA-readiness defects from the #393 review (verified against source). Correctness / doc-contract (High): - AffectedRows: cache the modified-row count at execute time. conn.ExecContext closes the op (nulling exec) before reading AffectedRows, so the previous live read returned 0 for every INSERT/UPDATE/DELETE/MERGE. - OAuth: reject non-PAT authenticators in newKernelBackend. OAuth/token-provider options leave AccessToken empty, so an empty PAT reached the kernel and failed with an opaque Unauthenticated error instead of the documented clear error. - Bound parameters: reject len(req.Params) > 0 in Execute with a clear error (non-nil Operation per the contract); they were silently dropped. doc.go now says params error at execute time, OAuth/metadata/namespace at connect time. - Context: fail fast on an already-cancelled ctx in nextBatch and OpenSession before the blocking C calls (database/sql's Rows.Close watcher still tears down an in-flight fetch). Parity / quality (Medium/Low): - Decimal: hoist the exact fixed-point formatter into internal/decimalfmt, shared by the Thrift and kernel paths so a #274-class fix lands in both. Its unit test runs in the default (untagged) build. - Nested FLOAT: marshal the native float32, not a widened float64, so ARRAY/MAP/STRUCT match Thrift byte-for-byte (3.14, not 3.140000104904175). Added a nested-float parity test. - Observability: OnChunkFetched now reports a cumulative chunk count; kernel errors log at the driver's default (Warn) level (no SQL/PII) so a failure is visible without DBSQL_KERNEL_DEBUG. - klog logs sql.len, not raw SQL text (PII/secret safety), matching debuglog. - Precompute struct field-name JSON keys once per type instead of per row. - Fix the scan_nested header comment (nested decimals are exact, not lossy) and drop a POC-history comment. Tests: - config: DSN useKernel/warehouseId parse + malformed-useKernel error + both fields in the DeepCopy all-values case. - stub: default build asserts WithUseKernel(true) → Connect errors with a clear not-compiled-in message. - e2e cancellation: assert context.DeadlineExceeded and tighten timing. Co-authored-by: Isaac Signed-off-by: Mani Kaustubh Mathur --- doc.go | 8 ++-- internal/backend/kernel/backend.go | 15 ++++++++ internal/backend/kernel/cgo.go | 7 ++++ internal/backend/kernel/kernel_test.go | 45 +++++++++++----------- internal/backend/kernel/operation.go | 32 ++++++++-------- internal/backend/kernel/rows.go | 33 ++++++++++++---- internal/backend/kernel/scan.go | 40 -------------------- internal/backend/kernel/scan_nested.go | 41 ++++++++++++++++---- internal/config/config_test.go | 30 +++++++++++++++ internal/decimalfmt/decimalfmt.go | 47 +++++++++++++++++++++++ internal/decimalfmt/decimalfmt_test.go | 43 +++++++++++++++++++++ internal/rows/arrowbased/columnValues.go | 48 ++---------------------- kernel_backend.go | 16 ++++++++ kernel_e2e_test.go | 14 ++++++- kernel_stub_test.go | 33 ++++++++++++++++ 15 files changed, 312 insertions(+), 140 deletions(-) delete mode 100644 internal/backend/kernel/scan.go create mode 100644 internal/decimalfmt/decimalfmt.go create mode 100644 internal/decimalfmt/decimalfmt_test.go create mode 100644 kernel_stub_test.go diff --git a/doc.go b/doc.go index f5e48063..60dd1c37 100644 --- a/doc.go +++ b/doc.go @@ -196,9 +196,11 @@ path databricks_sql_kernel, which would match nothing): The kernel backend currently supports PAT authentication; reading scalar, nested, and complex-typed results (CloudFetch is handled transparently); context cancellation; and the TLS, proxy, and session-conf (query tags, statement timeout, -time zone) connection options. Bound parameters, OAuth (M2M/U2M), metadata -commands, and initial catalog/schema are not yet supported on this backend and -return a clear error at connect time rather than being silently ignored. +time zone) connection options. OAuth (M2M/U2M), metadata commands, and initial +catalog/schema are not yet supported and return a clear error at connect time +rather than being silently ignored. Bound query parameters are likewise not yet +supported and return a clear error at execute time (they arrive per-query, not at +connect). None of these is silently ignored. # Programmatically Retrieving Connection and Query Id diff --git a/internal/backend/kernel/backend.go b/internal/backend/kernel/backend.go index bbd058ad..e0933f09 100644 --- a/internal/backend/kernel/backend.go +++ b/internal/backend/kernel/backend.go @@ -10,6 +10,7 @@ import "C" import ( "context" + "errors" "fmt" "time" @@ -72,6 +73,11 @@ func New(cfg Config) *KernelBackend { // time. The config handle is consumed by kernel_session_open on success and // freed by us on any earlier failure. func (k *KernelBackend) OpenSession(ctx context.Context) error { + // Fail fast on an already-cancelled context before the blocking kernel_session + // _open (which the C ABI does not let us interrupt mid-call). + if err := ctx.Err(); err != nil { + return err + } initKernelLogging() klog("OpenSession host=%s httpPath=%s warehouse=%s", k.cfg.Host, k.cfg.HTTPPath, k.cfg.WarehouseID) @@ -209,5 +215,14 @@ func (k *KernelBackend) SessionID() string { return k.sessionID } // Per the Backend contract it returns a non-nil Operation even on error so the // caller can read StatementID / wrap the error / Close uniformly. 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{}, errors.New("databricks: query parameters are not yet supported by the kernel backend; " + + "inline the values or use the default (Thrift) backend") + } return k.execute(ctx, req) } diff --git a/internal/backend/kernel/cgo.go b/internal/backend/kernel/cgo.go index 3ee07b7d..d38445df 100644 --- a/internal/backend/kernel/cgo.go +++ b/internal/backend/kernel/cgo.go @@ -39,6 +39,7 @@ import ( "unsafe" dbsqlerrint "github.com/databricks/databricks-sql-go/internal/errors" + "github.com/databricks/databricks-sql-go/logger" ) // ─── Debug logging ─────────────────────────────────────────────────────────── @@ -141,6 +142,12 @@ func lastError(code C.KernelStatusCode) *KernelError { } klog("kernel error: code=%d sqlstate=%q vendor=%d http=%d retryable=%v msg=%q", ke.Code, ke.SQLState, ke.VendorCode, ke.HTTPStatus, ke.Retryable, ke.Message) + // Also emit at the driver's default (Warn) level — no SQL text or PII, just + // the status/sqlstate/http fields — so a kernel-path failure is visible + // without DBSQL_KERNEL_DEBUG. This is the error path only (never the hot + // per-row/per-batch path), so it does not perturb benchmarks. + logger.Logger.Warn().Msgf("databricks: kernel call failed: code=%d sqlstate=%q vendor=%d http=%d retryable=%v", + ke.Code, ke.SQLState, ke.VendorCode, ke.HTTPStatus, ke.Retryable) return ke } diff --git a/internal/backend/kernel/kernel_test.go b/internal/backend/kernel/kernel_test.go index 86217f13..33bd8d3f 100644 --- a/internal/backend/kernel/kernel_test.go +++ b/internal/backend/kernel/kernel_test.go @@ -245,8 +245,7 @@ func TestScanCellNested(t *testing.T) { t.Run("nested_decimal_exact", func(t *testing.T) { // A decimal inside a struct must render as an exact JSON number, not a // lossy float64 (19.99, not 19.990000000000002) — matching Thrift's - // marshalScalar. Regression guard for a bug the POC missed by only testing - // float64-exact values. + // marshalScalar. dt := arrow.StructOf(arrow.Field{Name: "d", Type: &arrow.Decimal128Type{Precision: 5, Scale: 2}}) b := array.NewStructBuilder(pool, dt) defer b.Release() @@ -263,6 +262,26 @@ func TestScanCellNested(t *testing.T) { } }) + t.Run("nested_float32_exact", func(t *testing.T) { + // A float32 inside a struct must marshal as the native float32 (3.14), not + // a widened float64 (3.140000104904175) — matching Thrift's nested path, + // which marshals the native float32. See databricks-sql-go#393 review M2. + dt := arrow.StructOf(arrow.Field{Name: "f", Type: arrow.PrimitiveTypes.Float32}) + b := array.NewStructBuilder(pool, dt) + defer b.Release() + b.Append(true) + b.FieldBuilder(0).(*array.Float32Builder).Append(3.14) + arr := b.NewArray() + defer arr.Release() + v, err := scanCell(arr, 0, nil) + if err != nil { + t.Fatal(err) + } + if v.(string) != `{"f":3.14}` { + t.Errorf("got %q, want {\"f\":3.14}", v) + } + }) + t.Run("nested_null", func(t *testing.T) { b := array.NewListBuilder(pool, arrow.PrimitiveTypes.Int64) defer b.Release() @@ -279,24 +298,8 @@ func TestScanCellNested(t *testing.T) { }) } -// decimal128ToExactString applies scale by string placement, preserving digits a -// float64 would lose. -func TestDecimal128ToExactString(t *testing.T) { - cases := []struct { - unscaled uint64 - scale int32 - want string - }{ - {12345, 2, "123.45"}, - {5, 3, "0.005"}, - {100, 0, "100"}, - } - for _, c := range cases { - got := decimal128ToExactString(decimal128.FromU64(c.unscaled), c.scale) - if got != c.want { - t.Errorf("unscaled=%d scale=%d: got %q want %q", c.unscaled, c.scale, got, c.want) - } - } -} +// The exact decimal formatter now lives in internal/decimalfmt (shared with the +// Thrift path); its unit test is decimalfmt.TestExactString. The nested/scalar +// tests above still assert the kernel path renders decimals exactly. var _ driver.Value // keep database/sql/driver imported for the scanCell signature diff --git a/internal/backend/kernel/operation.go b/internal/backend/kernel/operation.go index 566394dc..57f3f1a9 100644 --- a/internal/backend/kernel/operation.go +++ b/internal/backend/kernel/operation.go @@ -29,9 +29,12 @@ import ( // 5. drain the watcher before returning, so a late cancel cannot land on a // statement that reuses this handle // -// Executes SQL text only; req.Params (bound parameters) are not yet wired. +// Executes SQL text only; bound parameters are rejected up front by Execute. func (k *KernelBackend) execute(ctx context.Context, req backend.ExecRequest) (backend.Operation, error) { - klog("Execute sql=%q", truncate(req.Query, 120)) + // 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)) var stmt *C.kernel_statement_t if err := call(func() C.KernelStatusCode { @@ -131,7 +134,10 @@ func (k *KernelBackend) execute(ctx context.Context, req backend.ExecRequest) (b return op, fmt.Errorf("kernel: execute: %w", toDriverError(execErr)) } op.exec = exec - klog("Execute OK stmt=%p exec=%p", stmt, 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. + op.affectedRows = int64(C.kernel_executed_statement_num_modified_rows(exec)) + klog("Execute OK stmt=%p exec=%p affectedRows=%d", stmt, exec, op.affectedRows) return op, nil } @@ -140,6 +146,10 @@ type kernelOp struct { stmt *C.kernel_statement_t exec *C.kernel_executed_statement_t closed bool + // affectedRows is the modified-row count captured at execute time. It is + // cached (not read live from exec) because the caller closes the operation — + // which nulls exec — before reading AffectedRows (see conn.ExecContext). + affectedRows int64 // 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 @@ -152,12 +162,11 @@ var _ backend.Operation = (*kernelOp)(nil) // session id. func (o *kernelOp) StatementID() string { return "" } -// AffectedRows is the modified-row count for ExecContext. +// 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 +// (the ExecContext path closes the op before reading this). func (o *kernelOp) AffectedRows() int64 { - if o.exec == nil { - return 0 - } - return int64(C.kernel_executed_statement_num_modified_rows(o.exec)) + return o.affectedRows } // Results builds the driver.Rows over the executed statement's result stream. @@ -217,10 +226,3 @@ func (o *kernelOp) close() bool { func (o *kernelOp) ExecutionError(ctx context.Context, cause error) error { return toDriverError(cause) } - -func truncate(s string, n int) string { - if len(s) <= n { - return s - } - return s[:n] + "…" -} diff --git a/internal/backend/kernel/rows.go b/internal/backend/kernel/rows.go index 107048b5..aa7a540e 100644 --- a/internal/backend/kernel/rows.go +++ b/internal/backend/kernel/rows.go @@ -23,6 +23,7 @@ import ( "github.com/apache/arrow/go/v12/arrow" "github.com/apache/arrow/go/v12/arrow/array" "github.com/apache/arrow/go/v12/arrow/cdata" + "github.com/databricks/databricks-sql-go/internal/decimalfmt" dbsqlrows "github.com/databricks/databricks-sql-go/internal/rows" ) @@ -42,11 +43,12 @@ type kernelRows struct { stream *C.kernel_result_stream_t callbacks *dbsqlrows.TelemetryCallbacks - cols []string - cur arrow.Record // current batch (nil until first Next) - rowInCur int // next row index within cur - closed bool - eof bool + cols []string + cur arrow.Record // current batch (nil until first Next) + rowInCur int // next row index within cur + chunkCount int // cumulative batches fetched, for OnChunkFetched + closed bool + eof bool } // newKernelRows fetches the schema up front (for Columns()) and returns the row @@ -129,6 +131,15 @@ func (r *kernelRows) Next(dest []driver.Value) error { // nextBatch pulls the next Arrow batch. A released array (release==NULL) is the // kernel's end-of-stream sentinel. func (r *kernelRows) nextBatch() error { + // Fail fast if the caller's context is already done rather than entering the + // blocking C fetch (which cannot itself observe ctx). database/sql also runs + // its own watcher that calls Rows.Close on cancellation, so an in-flight fetch + // is still torn down; this just avoids starting a new one under a dead ctx. + if r.ctx != nil { + if err := r.ctx.Err(); err != nil { + return err + } + } if r.cur != nil { r.cur.Release() r.cur = nil @@ -156,10 +167,16 @@ func (r *kernelRows) nextBatch() error { } r.cur = rec r.rowInCur = 0 + r.chunkCount++ if r.callbacks != nil && r.callbacks.OnChunkFetched != nil { - r.callbacks.OnChunkFetched(1, 0, 0, 0, 0) + // chunkCount is cumulative (per the callback contract). bytesDownloaded, + // chunkIndex, and latency are left 0: the kernel does CloudFetch/decompress + // internally and hands back ready Arrow batches, so the Go side never sees + // the compressed wire bytes or per-chunk fetch latency the Thrift path + // reports — a fabricated number would be worse than a truthful zero. + r.callbacks.OnChunkFetched(r.chunkCount, 0, 0, 0, 0) } - klog("nextBatch: %d rows", rec.NumRows()) + klog("nextBatch: %d rows (chunk %d)", rec.NumRows(), r.chunkCount) return nil } @@ -219,7 +236,7 @@ func scanCell(col arrow.Array, row int, loc *time.Location) (driver.Value, error return inLocation(c.Value(row).ToTime(dt.Unit), loc), nil case *array.Decimal128: dt := col.DataType().(*arrow.Decimal128Type) - return decimal128ToExactString(c.Value(row), dt.Scale), nil + return decimalfmt.ExactString(c.Value(row), dt.Scale), 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. diff --git a/internal/backend/kernel/scan.go b/internal/backend/kernel/scan.go deleted file mode 100644 index ca924322..00000000 --- a/internal/backend/kernel/scan.go +++ /dev/null @@ -1,40 +0,0 @@ -//go:build cgo && databricks_kernel - -package kernel - -import ( - "math/big" - "strings" - - "github.com/apache/arrow/go/v12/arrow/decimal128" -) - -// decimal128ToExactString renders an Arrow decimal128 as an exact fixed-point -// string, applying scale by string placement rather than float conversion. This -// matches the Thrift path's default DECIMAL rendering and preserves precision a -// float64 would lose beyond ~17 significant digits (databricks-sql-go#274). -func decimal128ToExactString(n decimal128.Num, scale int32) string { - unscaled := n.BigInt() - neg := unscaled.Sign() < 0 - digits := new(big.Int).Abs(unscaled).String() - - var b strings.Builder - if neg { - b.WriteByte('-') - } - if scale <= 0 { - b.WriteString(digits) - return b.String() - } - s := int(scale) - if len(digits) <= s { - b.WriteString("0.") - b.WriteString(strings.Repeat("0", s-len(digits))) - b.WriteString(digits) - } else { - b.WriteString(digits[:len(digits)-s]) - b.WriteByte('.') - b.WriteString(digits[len(digits)-s:]) - } - return b.String() -} diff --git a/internal/backend/kernel/scan_nested.go b/internal/backend/kernel/scan_nested.go index f42213a5..f27ed42b 100644 --- a/internal/backend/kernel/scan_nested.go +++ b/internal/backend/kernel/scan_nested.go @@ -12,8 +12,8 @@ package kernel // - struct → {"field0":v0,...} // - nested NULL → null // - time.Time → quoted .String() (matches the Thrift marshal() special-case) -// - nested decimal → float64 (lossy, matching Thrift's in-JSON rendering; the -// exact-string path applies only to top-level decimal columns, #274) +// - nested decimal → exact scale-applied JSON number literal (never a lossy +// float64), matching Thrift's marshalScalar → ValueString (#253/#274) // // VARIANT arrives as a nested value and renders through this path; GEOMETRY // arrives as a WKB/WKT string and is handled by scanCell's scalar arm. @@ -27,12 +27,35 @@ import ( "encoding/json" "fmt" "strings" + "sync" "time" "github.com/apache/arrow/go/v12/arrow" "github.com/apache/arrow/go/v12/arrow/array" + "github.com/databricks/databricks-sql-go/internal/decimalfmt" ) +// structKeyPrefixCache memoizes the JSON-escaped, quote-wrapped, colon-terminated +// key prefixes for each struct type (e.g. `"field":`). Field names are invariant +// across rows, but writeStructJSON runs per row; arrow reuses the same +// *arrow.StructType pointer for a column across its rows, so this is a hit after +// the first row. Keyed by that pointer; safe for concurrent Rows via sync.Map. +var structKeyPrefixCache sync.Map // map[*arrow.StructType][]string + +func structKeyPrefixes(st *arrow.StructType) []string { + if v, ok := structKeyPrefixCache.Load(st); ok { + return v.([]string) + } + fields := st.Fields() + prefixes := make([]string, len(fields)) + for f := range fields { + name, _ := json.Marshal(fields[f].Name) // JSON-escapes the field name + prefixes[f] = string(name) + ":" + } + structKeyPrefixCache.Store(st, prefixes) + return prefixes +} + // renderJSONString renders one nested cell as a JSON string (nil for a NULL // cell). loc is applied to any timestamp/date leaves, matching the top-level // scan and the Thrift path. @@ -70,8 +93,14 @@ func writeJSON(b *strings.Builder, col arrow.Array, row int, loc *time.Location) // 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). - b.WriteString(decimal128ToExactString(c.Value(row), col.DataType().(*arrow.Decimal128Type).Scale)) + b.WriteString(decimalfmt.ExactString(c.Value(row), col.DataType().(*arrow.Decimal128Type).Scale)) return nil + case *array.Float32: + // Marshal the native float32, NOT a widened float64: json.Marshal(float32 + // (3.14)) is "3.14" but json.Marshal(float64(float32(3.14))) is + // "3.140000104904175". The Thrift nested path marshals the native float32, + // so widening here would break byte-parity for ARRAY/MAP/STRUCT<…FLOAT…>. + return writeScalarJSON(b, c.Value(row)) default: v, err := scalarForJSON(col, row, loc) if err != nil { @@ -120,15 +149,13 @@ func writeMapJSON(b *strings.Builder, m *array.Map, row int, loc *time.Location) } func writeStructJSON(b *strings.Builder, s *array.Struct, row int, loc *time.Location) error { - st := s.DataType().(*arrow.StructType) + prefixes := structKeyPrefixes(s.DataType().(*arrow.StructType)) b.WriteByte('{') for f := 0; f < s.NumField(); f++ { if f > 0 { b.WriteByte(',') } - keyBytes, _ := json.Marshal(st.Field(f).Name) - b.Write(keyBytes) - b.WriteByte(':') + b.WriteString(prefixes[f]) // pre-escaped `"name":` if err := writeJSON(b, s.Field(f), row, loc); err != nil { return err } diff --git a/internal/config/config_test.go b/internal/config/config_test.go index d3574fd1..4c70e993 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -586,6 +586,34 @@ func TestParseConfig(t *testing.T) { wantCfg: UserConfig{}, wantErr: true, }, + { + name: "kernel backend params useKernel and warehouseId", + args: args{dsn: "token:supersecret@example.cloud.databricks.com:443/sql/1.0/endpoints/12346a5b5b0e123a?useKernel=true&warehouseId=abc123"}, + wantCfg: UserConfig{ + Protocol: "https", + Host: "example.cloud.databricks.com", + Port: 443, + MaxRows: defaultMaxRows, + Authenticator: &pat.PATAuth{AccessToken: "supersecret"}, + AccessToken: "supersecret", + HTTPPath: "/sql/1.0/endpoints/12346a5b5b0e123a", + SessionParams: make(map[string]string), + RetryMax: 4, + RetryWaitMin: 1 * time.Second, + RetryWaitMax: 30 * time.Second, + CloudFetchConfig: defCloudConfig, + UseKernel: true, + WarehouseID: "abc123", + }, + wantURL: "https://example.cloud.databricks.com:443/sql/1.0/endpoints/12346a5b5b0e123a", + wantErr: false, + }, + { + name: "malformed useKernel is an error", + args: args{dsn: "token:supersecret@example.cloud.databricks.com:443/sql/1.0/endpoints/12346a5b5b0e123a?useKernel=notabool"}, + wantCfg: UserConfig{}, + wantErr: true, + }, } for i, tt := range tests { fmt.Println(i) @@ -649,6 +677,8 @@ func TestUserConfig_DeepCopy(t *testing.T) { UserAgentEntry: "test", Location: location, SessionParams: map[string]string{"a": "32", "b": "4"}, + UseKernel: true, + WarehouseID: "wh-abc123", } cfg_copy := cfg.DeepCopy() diff --git a/internal/decimalfmt/decimalfmt.go b/internal/decimalfmt/decimalfmt.go new file mode 100644 index 00000000..7d21ebf5 --- /dev/null +++ b/internal/decimalfmt/decimalfmt.go @@ -0,0 +1,47 @@ +// Package decimalfmt renders Arrow decimal128 values as exact fixed-point +// strings, shared by the Thrift (arrowbased) and kernel result paths so a +// DECIMAL renders identically regardless of backend. Keeping one implementation +// means a precision fix (e.g. databricks-sql-go#274) lands in both at once. +package decimalfmt + +import ( + "math/big" + "strings" + + "github.com/apache/arrow/go/v12/arrow/decimal128" +) + +// ExactString renders an Arrow decimal128 as an exact fixed-point string, +// applying scale by string placement rather than float conversion. This +// preserves precision a float64 would lose beyond ~17 significant digits +// (databricks-sql-go#274). A negative scale is not produced by the server and is +// treated as scale 0 rather than panicking. +func ExactString(n decimal128.Num, scale int32) string { + unscaled := n.BigInt() // exact signed unscaled integer + neg := unscaled.Sign() < 0 + digits := new(big.Int).Abs(unscaled).String() + + var b strings.Builder + if neg { + b.WriteByte('-') + } + + if scale <= 0 { + b.WriteString(digits) + return b.String() + } + + s := int(scale) + if len(digits) <= s { + // Pad with leading zeros so there are exactly `scale` fractional digits + // and a single leading integer zero, e.g. 5 with scale 3 -> "0.005". + b.WriteString("0.") + b.WriteString(strings.Repeat("0", s-len(digits))) + b.WriteString(digits) + } else { + b.WriteString(digits[:len(digits)-s]) + b.WriteByte('.') + b.WriteString(digits[len(digits)-s:]) + } + return b.String() +} diff --git a/internal/decimalfmt/decimalfmt_test.go b/internal/decimalfmt/decimalfmt_test.go new file mode 100644 index 00000000..7765c7c6 --- /dev/null +++ b/internal/decimalfmt/decimalfmt_test.go @@ -0,0 +1,43 @@ +package decimalfmt + +import ( + "math/big" + "testing" + + "github.com/apache/arrow/go/v12/arrow/decimal128" +) + +// ExactString applies scale by string placement, preserving digits a float64 +// would lose beyond ~17 significant figures. +func TestExactString(t *testing.T) { + cases := []struct { + unscaled int64 + scale int32 + want string + }{ + {12345, 2, "123.45"}, + {5, 3, "0.005"}, + {100, 0, "100"}, + {1999, 2, "19.99"}, + {-12345, 2, "-123.45"}, + {-5, 3, "-0.005"}, + {0, 2, "0.00"}, + } + for _, c := range cases { + got := ExactString(decimal128.FromI64(c.unscaled), c.scale) + if got != c.want { + t.Errorf("unscaled=%d scale=%d: got %q want %q", c.unscaled, c.scale, got, c.want) + } + } +} + +// A value beyond float64's exact integer range must survive intact — the whole +// point of formatting from the 128-bit unscaled integer rather than a float. +func TestExactStringHighPrecision(t *testing.T) { + // 12345678901234567890 (20 digits) — past float64's 2^53 exact-integer limit. + big20, _ := new(big.Int).SetString("12345678901234567890", 10) + got := ExactString(decimal128.FromBigInt(big20), 4) + if want := "1234567890123456.7890"; got != want { + t.Errorf("got %q want %q", got, want) + } +} diff --git a/internal/rows/arrowbased/columnValues.go b/internal/rows/arrowbased/columnValues.go index 72e7f5f5..53bf8551 100644 --- a/internal/rows/arrowbased/columnValues.go +++ b/internal/rows/arrowbased/columnValues.go @@ -2,13 +2,12 @@ package arrowbased import ( "encoding/json" - "math/big" "strings" "time" "github.com/apache/arrow/go/v12/arrow" "github.com/apache/arrow/go/v12/arrow/array" - "github.com/apache/arrow/go/v12/arrow/decimal128" + "github.com/databricks/databricks-sql-go/internal/decimalfmt" "github.com/databricks/databricks-sql-go/internal/rows/rowscanner" dbsqllog "github.com/databricks/databricks-sql-go/logger" "github.com/pkg/errors" @@ -550,48 +549,9 @@ func (tvc *decimal128Container) Value(i int) (any, error) { // for backwards-compatible JSON rendering inside complex types. // See databricks/databricks-sql-go#274. func (tvc *decimal128Container) ValueString(i int) string { - return decimal128ToString(tvc.decimalArray.Value(i), tvc.scale) -} - -// decimal128ToString renders a decimal128 value exactly as a fixed-point -// string with `scale` fractional digits. It formats from the value's exact -// 128-bit unscaled integer (BigInt) and inserts the decimal point by string -// manipulation, so — unlike arrow's decimal128.Num.ToString, which divides via -// big.Float and rounds beyond ~17 significant digits — it never loses -// precision. See databricks/databricks-sql-go#274. -func decimal128ToString(n decimal128.Num, scale int32) string { - unscaled := n.BigInt() // exact signed unscaled integer - - neg := unscaled.Sign() < 0 - digits := new(big.Int).Abs(unscaled).String() - - var b strings.Builder - if neg { - b.WriteByte('-') - } - - if scale <= 0 { - // No fractional part (negative scale is not produced by the server, - // but treat it as scale 0 rather than panicking). - b.WriteString(digits) - return b.String() - } - - s := int(scale) - if len(digits) <= s { - // Pad with leading zeros so there are exactly `scale` fractional digits - // and a single leading integer zero, e.g. 5 with scale 3 -> "0.005". - b.WriteString("0.") - b.WriteString(strings.Repeat("0", s-len(digits))) - b.WriteString(digits) - } else { - intPart := digits[:len(digits)-s] - fracPart := digits[len(digits)-s:] - b.WriteString(intPart) - b.WriteByte('.') - b.WriteString(fracPart) - } - return b.String() + // Exact fixed-point rendering lives in internal/decimalfmt, shared with the + // kernel backend so both paths render DECIMAL identically. See #274. + return decimalfmt.ExactString(tvc.decimalArray.Value(i), tvc.scale) } func (tvc *decimal128Container) IsNull(i int) bool { diff --git a/kernel_backend.go b/kernel_backend.go index 779859f2..4cea76c6 100644 --- a/kernel_backend.go +++ b/kernel_backend.go @@ -7,6 +7,8 @@ import ( "errors" "net/http" + "github.com/databricks/databricks-sql-go/auth/noop" + "github.com/databricks/databricks-sql-go/auth/pat" "github.com/databricks/databricks-sql-go/internal/backend" "github.com/databricks/databricks-sql-go/internal/backend/kernel" "github.com/databricks/databricks-sql-go/internal/config" @@ -33,6 +35,20 @@ func newKernelBackend(_ context.Context, cfg *config.Config) (backend.Backend, e return nil, errors.New("databricks: WithEnableMetricViewMetadata is not yet supported by the kernel backend; " + "omit it or use the default (Thrift) backend") } + // Auth: the kernel backend authenticates with a PAT only (kc.Token below). + // Any other authenticator — OAuth M2M/U2M, a token provider, external/static + // token, federated — 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, per + // the doc.go contract. nil / NoopAuth / PATAuth are the PAT-or-none cases. + switch cfg.Authenticator.(type) { + case nil, *noop.NoopAuth, *pat.PATAuth: + // PAT (or no explicit authenticator) — supported. + default: + return nil, errors.New("databricks: only personal access token (WithAccessToken) auth is supported by the kernel backend; " + + "OAuth (M2M/U2M), token-provider, external/static, and federated authenticators are not yet supported — " + + "use PAT or the default (Thrift) backend") + } kc := kernel.Config{ Host: cfg.Host, diff --git a/kernel_e2e_test.go b/kernel_e2e_test.go index 1d2b3c00..b6d2fdd0 100644 --- a/kernel_e2e_test.go +++ b/kernel_e2e_test.go @@ -7,6 +7,7 @@ import ( "context" "database/sql" "database/sql/driver" + "errors" "os" "testing" "time" @@ -259,8 +260,17 @@ func TestKernelE2ECancellation(t *testing.T) { if err == nil { t.Fatal("expected a cancellation error, got nil") } - if elapsed > 30*time.Second { - t.Errorf("cancellation took %v; expected it to abandon well before the query's natural runtime", elapsed) + // It must be the deadline firing, not an unrelated failure (syntax error, + // transient network, server-side timeout) — otherwise the test would pass + // without proving cancellation works. + if !errors.Is(err, context.DeadlineExceeded) { + t.Errorf("expected context.DeadlineExceeded, got %v", err) + } + // Well under the query's natural runtime and not far past the 3s deadline; + // a no-op cancel that only returned when the query finished would blow past + // this. + if elapsed > 10*time.Second { + t.Errorf("cancellation took %v; expected it to abandon close to the 3s deadline", elapsed) } t.Logf("cancelled after %v with err=%v", elapsed, err) } diff --git a/kernel_stub_test.go b/kernel_stub_test.go new file mode 100644 index 00000000..9f8e8cde --- /dev/null +++ b/kernel_stub_test.go @@ -0,0 +1,33 @@ +//go:build !cgo || !databricks_kernel + +package dbsql + +import ( + "context" + "strings" + "testing" +) + +// In the default pure-Go build (no databricks_kernel tag) the kernel backend is +// not compiled in. Connecting with WithUseKernel(true) must fail loudly with a +// clear "not compiled in" error rather than silently falling back to Thrift or +// panicking on a nil backend. This guards the stub that fails closed. +func TestKernelBackendNotCompiledIn(t *testing.T) { + connector, err := NewConnector( + WithServerHostname("example.cloud.databricks.com"), + WithPort(443), + WithHTTPPath("/sql/1.0/endpoints/12346a5b5b0e123a"), + WithAccessToken("supersecret"), + WithUseKernel(true), + ) + if err != nil { + t.Fatalf("NewConnector: %v", err) + } + _, err = connector.Connect(context.Background()) + if err == nil { + t.Fatal("Connect with WithUseKernel(true) in a non-kernel build should error, got nil") + } + if !strings.Contains(err.Error(), "databricks_kernel") && !strings.Contains(err.Error(), "not compiled") { + t.Errorf("error should mention the missing build tag / not-compiled-in; got: %v", err) + } +} From 5852480760aa3236aace35f80cd08c7cdcf10ca4 Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Fri, 10 Jul 2026 20:54:44 +0000 Subject: [PATCH 10/32] ci(kernel): document deferral of the tagged kernel-path test job (H5) The default CI matrix builds CGO_ENABLED=0, so the SEA-via-kernel backend (//go:build cgo && databricks_kernel) is not compiled or tested here. A dedicated CGO_ENABLED=1 -tags databricks_kernel job needs the kernel static library linked in CI, which is the kernel distribution work (pinned-source build or published .a) tracked as a #393 follow-up. Add a comment in go.yml so the gap is explicit rather than silent. The shared pure-Go decimal formatter (internal/decimalfmt) is already covered by the default matrix via its own untagged test. Co-authored-by: Isaac Signed-off-by: Mani Kaustubh Mathur --- .github/workflows/go.yml | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/.github/workflows/go.yml b/.github/workflows/go.yml index beff4157..6c2e29b3 100644 --- a/.github/workflows/go.yml +++ b/.github/workflows/go.yml @@ -90,6 +90,14 @@ jobs: fi go get -v -t -d ./... + # This matrix builds the default pure-Go driver (CGO_ENABLED=0), which does + # NOT compile the SEA-via-kernel backend (//go:build cgo && databricks_kernel). + # A dedicated CGO_ENABLED=1 `-tags databricks_kernel` test job needs the + # kernel static library linked in CI, which arrives with the kernel + # distribution work (a pinned-revision source build or a published .a) — see + # the #393 follow-ups. Until then the kernel path's cgo files are not + # exercised here; the pure-Go decimal formatter it shares + # (internal/decimalfmt) IS covered by this matrix via its own untagged test. - name: Test run: make test env: From a0d1b47b9776d4b336dac6f1a94329f79892c576 Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Fri, 10 Jul 2026 21:08:11 +0000 Subject: [PATCH 11/32] ci(kernel): trim the kernel-path CI deferral comment Drop the decimalfmt sentence from the go.yml comment; the deferral note only needs to explain why the tagged kernel-path job is absent. Co-authored-by: Isaac Signed-off-by: Mani Kaustubh Mathur --- .github/workflows/go.yml | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/.github/workflows/go.yml b/.github/workflows/go.yml index 6c2e29b3..03a0e4f9 100644 --- a/.github/workflows/go.yml +++ b/.github/workflows/go.yml @@ -96,8 +96,7 @@ jobs: # kernel static library linked in CI, which arrives with the kernel # distribution work (a pinned-revision source build or a published .a) — see # the #393 follow-ups. Until then the kernel path's cgo files are not - # exercised here; the pure-Go decimal formatter it shares - # (internal/decimalfmt) IS covered by this matrix via its own untagged test. + # exercised here. - name: Test run: make test env: From 37196b94a54506ef3aa4321f2bcdfc8ae93e5e2f Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Fri, 10 Jul 2026 22:43:04 +0000 Subject: [PATCH 12/32] =?UTF-8?q?docs(kernel):=20round-2=20review=20?= =?UTF-8?q?=E2=80=94=20doc/comment/test-helper=20cleanups?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit No behavior change. Addresses the doc/maintainability findings: - doc.go: drop "metadata commands" from the connect-time-error clause — there is no such guard and metadata runs as ordinary SQL (SHOW/DESCRIBE/ information_schema); note that explicitly (N8). - Remove rotting review/PR labels from comments (kernel_test.go, go.yml) while keeping the substantive rationale (N10). - kernelTestDB delegates to kernelTestDBWith instead of duplicating it (N11). Co-authored-by: Isaac Signed-off-by: Mani Kaustubh Mathur --- .github/workflows/go.yml | 5 ++--- doc.go | 11 ++++++----- internal/backend/kernel/kernel_test.go | 2 +- kernel_e2e_test.go | 17 +---------------- 4 files changed, 10 insertions(+), 25 deletions(-) diff --git a/.github/workflows/go.yml b/.github/workflows/go.yml index 03a0e4f9..9164ced8 100644 --- a/.github/workflows/go.yml +++ b/.github/workflows/go.yml @@ -94,9 +94,8 @@ jobs: # NOT compile the SEA-via-kernel backend (//go:build cgo && databricks_kernel). # A dedicated CGO_ENABLED=1 `-tags databricks_kernel` test job needs the # kernel static library linked in CI, which arrives with the kernel - # distribution work (a pinned-revision source build or a published .a) — see - # the #393 follow-ups. Until then the kernel path's cgo files are not - # exercised here. + # distribution work (a pinned-revision source build or a published .a). + # Until then the kernel path's cgo files are not exercised here. - name: Test run: make test env: diff --git a/doc.go b/doc.go index 60dd1c37..83d7d821 100644 --- a/doc.go +++ b/doc.go @@ -196,11 +196,12 @@ path databricks_sql_kernel, which would match nothing): The kernel backend currently supports PAT authentication; reading scalar, nested, and complex-typed results (CloudFetch is handled transparently); context cancellation; and the TLS, proxy, and session-conf (query tags, statement timeout, -time zone) connection options. OAuth (M2M/U2M), metadata commands, and initial -catalog/schema are not yet supported and return a clear error at connect time -rather than being silently ignored. Bound query parameters are likewise not yet -supported and return a clear error at execute time (they arrive per-query, not at -connect). None of these is silently ignored. +time zone) connection options. OAuth (M2M/U2M) and initial catalog/schema are not +yet supported and return a clear error at connect time rather than being silently +ignored. Bound query parameters are likewise not yet supported and return a clear +error at execute time (they arrive per-query, not at connect). 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.) # Programmatically Retrieving Connection and Query Id diff --git a/internal/backend/kernel/kernel_test.go b/internal/backend/kernel/kernel_test.go index 33bd8d3f..6b84f042 100644 --- a/internal/backend/kernel/kernel_test.go +++ b/internal/backend/kernel/kernel_test.go @@ -265,7 +265,7 @@ func TestScanCellNested(t *testing.T) { t.Run("nested_float32_exact", func(t *testing.T) { // A float32 inside a struct must marshal as the native float32 (3.14), not // a widened float64 (3.140000104904175) — matching Thrift's nested path, - // which marshals the native float32. See databricks-sql-go#393 review M2. + // which marshals the native float32. dt := arrow.StructOf(arrow.Field{Name: "f", Type: arrow.PrimitiveTypes.Float32}) b := array.NewStructBuilder(pool, dt) defer b.Release() diff --git a/kernel_e2e_test.go b/kernel_e2e_test.go index b6d2fdd0..f3695a87 100644 --- a/kernel_e2e_test.go +++ b/kernel_e2e_test.go @@ -19,22 +19,7 @@ import ( // consumer uses — not a kernel-only connector. func kernelTestDB(t *testing.T) *sql.DB { t.Helper() - host := os.Getenv("DATABRICKS_HOST") - httpPath := os.Getenv("DATABRICKS_HTTP_PATH") - token := os.Getenv("DATABRICKS_TOKEN") - if host == "" || httpPath == "" || token == "" { - t.Skip("set DATABRICKS_HOST / DATABRICKS_HTTP_PATH / DATABRICKS_TOKEN for the kernel e2e") - } - connector, err := NewConnector( - WithServerHostname(host), - WithHTTPPath(httpPath), - WithAccessToken(token), - WithUseKernel(true), - ) - if err != nil { - t.Fatalf("NewConnector: %v", err) - } - return sql.OpenDB(connector) + return kernelTestDBWith(t) } // TestKernelE2ESelect1 is the smallest end-to-end proof: PAT session over the From 7d063b25b9482348227cd768158197ee13038cec Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Fri, 10 Jul 2026 22:52:37 +0000 Subject: [PATCH 13/32] =?UTF-8?q?fix(kernel):=20round-2=20review=20?= =?UTF-8?q?=E2=80=94=20float32=20parity=20+=20close()=20contract?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Top-level FLOAT columns now scan to a native float32, not a widened float64 (N3). Thrift returns float32 for a bare FLOAT, and database/sql's asString formats at bit-size 32, so widening rendered CAST(0.1 AS FLOAT) as "0.10000000149011612" vs Thrift's "0.1". The round-1 M2 fix covered only the nested path; this closes the top-level scalar arm. Parity test gains a top-level FLOAT case at a non-exactly-representable value, and the e2e float case moves from 1.5 (exactly representable, masked the bug) to 0.1. - close() reports closed=false for a handle-less op (N9). The bound-params error path returns &kernelOp{} with no handle; conn.ExecContext closes it unconditionally, and the previous unconditional true recorded a phantom CLOSE_STATEMENT for a statement that never reached the server. Now matches the backend.Operation contract (closed=false when there was no handle). Tests: add float32_native scalar case, top-level FLOAT to the parity query, and TestExecuteRejectsParams (asserts non-nil op, closed=false, AffectedRows 0). Co-authored-by: Isaac Signed-off-by: Mani Kaustubh Mathur --- internal/backend/kernel/kernel_test.go | 54 ++++++++++++++++++++++++++ internal/backend/kernel/operation.go | 11 +++++- internal/backend/kernel/rows.go | 6 ++- kernel_e2e_test.go | 5 ++- kernel_parity_test.go | 3 ++ 5 files changed, 75 insertions(+), 4 deletions(-) diff --git a/internal/backend/kernel/kernel_test.go b/internal/backend/kernel/kernel_test.go index 6b84f042..6a00d532 100644 --- a/internal/backend/kernel/kernel_test.go +++ b/internal/backend/kernel/kernel_test.go @@ -3,6 +3,7 @@ package kernel import ( + "context" "database/sql/driver" "errors" "testing" @@ -12,6 +13,8 @@ import ( "github.com/apache/arrow/go/v12/arrow/array" "github.com/apache/arrow/go/v12/arrow/decimal128" "github.com/apache/arrow/go/v12/arrow/memory" + + "github.com/databricks/databricks-sql-go/internal/backend" ) // isBadConnection maps the session-unusable status codes so the pool evicts the @@ -90,6 +93,28 @@ func TestScanCellScalars(t *testing.T) { } }) + t.Run("float32_native", func(t *testing.T) { + // A top-level FLOAT column must scan to a native float32, not a widened + // float64 — matching Thrift, so database/sql's asString renders + // CAST(0.1 AS FLOAT) as "0.1", not "0.10000000149011612". + b := array.NewFloat32Builder(pool) + defer b.Release() + b.Append(0.1) + arr := b.NewArray() + defer arr.Release() + v, err := scanCell(arr, 0, nil) + if err != nil { + t.Fatal(err) + } + got, ok := v.(float32) + if !ok { + t.Fatalf("want float32, got %T", v) + } + if got != float32(0.1) { + t.Errorf("got %v, want 0.1", got) + } + }) + t.Run("null", func(t *testing.T) { b := array.NewInt64Builder(pool) defer b.Release() @@ -298,6 +323,35 @@ func TestScanCellNested(t *testing.T) { }) } +// Bound parameters are rejected up front by Execute with a clear error, before +// any session/C work — so this runs on a zero-value backend. The returned +// 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{} + op, err := k.Execute(context.Background(), backend.ExecRequest{ + Query: "SELECT ?", + Params: []backend.Param{{Name: "x"}}, + }) + if err == nil { + t.Fatal("expected an error for bound parameters, got nil") + } + if op == nil { + t.Fatal("Execute must return a non-nil Operation per the Backend contract") + } + closed, closeErr := op.Close(context.Background()) + if closeErr != nil { + t.Errorf("Close error = %v, want nil", closeErr) + } + if closed { + t.Error("Close on a handle-less op must report closed=false (no CLOSE_STATEMENT)") + } + if got := op.AffectedRows(); got != 0 { + t.Errorf("AffectedRows on a handle-less op = %d, want 0", got) + } +} + // The exact decimal formatter now lives in internal/decimalfmt (shared with the // Thrift path); its unit test is decimalfmt.TestExactString. The nested/scalar // tests above still assert the kernel path renders decimals exactly. diff --git a/internal/backend/kernel/operation.go b/internal/backend/kernel/operation.go index 57f3f1a9..6fbee539 100644 --- a/internal/backend/kernel/operation.go +++ b/internal/backend/kernel/operation.go @@ -208,6 +208,13 @@ func (o *kernelOp) close() bool { return false } o.closed = true + // Report closed=true only when there was actually a handle to tear down. A + // handle-less op (e.g. the &kernelOp{} returned on the bound-params error + // path) reaches here via conn.ExecContext's unconditional Close; returning + // true would record a phantom CLOSE_STATEMENT for a statement that never hit + // the server. Matches the backend.Operation contract (closed=false when the + // operation had no handle) and the Thrift backend's !hasHandle() behavior. + didClose := o.exec != nil || o.stmt != nil if o.exec != nil { C.kernel_executed_statement_close(o.exec) o.exec = nil @@ -216,8 +223,8 @@ func (o *kernelOp) close() bool { C.kernel_statement_close(o.stmt) o.stmt = nil } - klog("kernelOp closed") - return true + klog("kernelOp closed (didClose=%v)", didClose) + return didClose } // ExecutionError wraps cause as the driver's execution error. The kernel error diff --git a/internal/backend/kernel/rows.go b/internal/backend/kernel/rows.go index aa7a540e..7477ddec 100644 --- a/internal/backend/kernel/rows.go +++ b/internal/backend/kernel/rows.go @@ -215,7 +215,11 @@ func scanCell(col arrow.Array, row int, loc *time.Location) (driver.Value, error case *array.Uint64: return int64(c.Value(row)), nil case *array.Float32: - return float64(c.Value(row)), nil + // Return the native float32, NOT a widened float64: the Thrift path returns + // a float32 driver.Value for a bare FLOAT column, and database/sql's + // asString formats it at bit-size 32 — so widening here would render + // CAST(0.1 AS FLOAT) as "0.10000000149011612" vs Thrift's "0.1". + return c.Value(row), nil case *array.Float64: return c.Value(row), nil case *array.String: diff --git a/kernel_e2e_test.go b/kernel_e2e_test.go index f3695a87..4b150b1d 100644 --- a/kernel_e2e_test.go +++ b/kernel_e2e_test.go @@ -144,7 +144,10 @@ func TestKernelE2EDataTypes(t *testing.T) { {"smallint", "CAST(3 AS SMALLINT)", int64(3)}, {"tinyint", "CAST(1 AS TINYINT)", int64(1)}, {"double", "CAST(3.5 AS DOUBLE)", float64(3.5)}, - {"float", "CAST(1.5 AS FLOAT)", float64(1.5)}, + // 0.1 is NOT exactly representable, so a float32-vs-float64 mismatch would + // surface here (float64(float32(0.1)) != float32(0.1)); a bare FLOAT must + // scan to a native float32, matching Thrift. + {"float", "CAST(0.1 AS FLOAT)", float32(0.1)}, {"boolean", "true", true}, {"string", "'hi'", "hi"}, {"binary", "CAST('abc' AS BINARY)", []byte("abc")}, diff --git a/kernel_parity_test.go b/kernel_parity_test.go index 223a0081..8f570302 100644 --- a/kernel_parity_test.go +++ b/kernel_parity_test.go @@ -37,6 +37,9 @@ func thriftTestDB(t *testing.T) *sql.DB { func TestKernelThriftParity(t *testing.T) { const query = "SELECT CAST(7 AS BIGINT), CAST(2.5 AS DOUBLE), 'parity', " + "CAST(NULL AS STRING), CAST(9.99 AS DECIMAL(5,2)), true, " + + // A bare FLOAT at a non-exactly-representable value: catches the + // float32-vs-widened-float64 divergence (0.1 vs 0.10000000149011612). + "CAST(0.1 AS FLOAT), " + "array(1, 2, 3), map('k', 9), named_struct('a', 1, 'b', 'x'), " + `parse_json('{"a":1,"b":[2,3]}'), st_point(1, 2), ` + // A decimal inside a struct: exercises exact-string nested-decimal From 0a4ab9b50a8c6cd87511cbe8209bdef9bc1ded1f Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Fri, 10 Jul 2026 22:54:30 +0000 Subject: [PATCH 14/32] =?UTF-8?q?docs(kernel):=20round-2=20review=20?= =?UTF-8?q?=E2=80=94=20correct=20cancellation=20+=20telemetry=20docs?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit No behavior change; the code is honest about two real limitations now. - nextBatch comment (N2): the previous comment claimed database/sql's cancel watcher tears down an in-flight fetch. It does not — Rows.Close takes closemu.Lock() which blocks until the in-progress Next (holding the RLock) returns, so the stream close waits for the blocking C call to finish. Read-path cancellation is honored only at batch boundaries; a single hung CloudFetch batch is uninterruptible (no per-download timeout). doc.go's "context cancellation" claim is scoped accordingly (real server cancel on execute; batch-boundary on read). - StatementID comment (N4): the kernel C ABI has no success-path statement/query id accessor (query_id is error-path only), so StatementID() is "" and per-statement telemetry (gated on != "") does not fire — there is no session-id fallback, contrary to the old comment. Documented; the accessor is a kernel follow-up. Co-authored-by: Isaac Signed-off-by: Mani Kaustubh Mathur --- doc.go | 6 ++++-- internal/backend/kernel/operation.go | 10 +++++++--- internal/backend/kernel/rows.go | 12 ++++++++---- 3 files changed, 19 insertions(+), 9 deletions(-) diff --git a/doc.go b/doc.go index 83d7d821..b5586992 100644 --- a/doc.go +++ b/doc.go @@ -195,8 +195,10 @@ path databricks_sql_kernel, which would match nothing): The kernel backend currently supports PAT authentication; reading scalar, nested, and complex-typed results (CloudFetch is handled transparently); context -cancellation; and the TLS, proxy, and session-conf (query tags, statement timeout, -time zone) connection options. OAuth (M2M/U2M) and initial catalog/schema are not +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) and initial catalog/schema are not yet supported and return a clear error at connect time rather than being silently ignored. Bound query parameters are likewise not yet supported and return a clear error at execute time (they arrive per-query, not at connect). None of these is diff --git a/internal/backend/kernel/operation.go b/internal/backend/kernel/operation.go index 6fbee539..8531859d 100644 --- a/internal/backend/kernel/operation.go +++ b/internal/backend/kernel/operation.go @@ -157,9 +157,13 @@ type kernelOp struct { var _ backend.Operation = (*kernelOp)(nil) -// StatementID returns "": the C ABI exposes no server statement id accessor. The -// id is used only for logging/telemetry correlation, which falls back to the -// session id. +// 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 "" } // AffectedRows is the modified-row count for ExecContext. It returns the value diff --git a/internal/backend/kernel/rows.go b/internal/backend/kernel/rows.go index 7477ddec..ee3e9112 100644 --- a/internal/backend/kernel/rows.go +++ b/internal/backend/kernel/rows.go @@ -131,10 +131,14 @@ func (r *kernelRows) Next(dest []driver.Value) error { // nextBatch pulls the next Arrow batch. A released array (release==NULL) is the // kernel's end-of-stream sentinel. func (r *kernelRows) nextBatch() error { - // Fail fast if the caller's context is already done rather than entering the - // blocking C fetch (which cannot itself observe ctx). database/sql also runs - // its own watcher that calls Rows.Close on cancellation, so an in-flight fetch - // is still torn down; this just avoids starting a new one under a dead ctx. + // Honor cancellation at batch boundaries: check ctx before entering the + // blocking C fetch (which cannot itself observe ctx). This does NOT interrupt + // a fetch already in flight — and database/sql's own cancel watcher can't + // either: its Rows.Close takes rs.closemu.Lock(), which blocks until the + // in-progress Next (holding the RLock) returns, so the stream close waits for + // the C call to finish on its own. A single hung CloudFetch batch is therefore + // uninterruptible (the kernel exposes no per-download timeout); mid-fetch + // cancellation would need the execute path's watcher/canceller applied here. if r.ctx != nil { if err := r.ctx.Err(); err != nil { return err From 490c97d96e5e045ecb742b20897b4eb4015cfe83 Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Fri, 10 Jul 2026 22:56:26 +0000 Subject: [PATCH 15/32] test(kernel): strengthen proxy test via a resolver seam (N7) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The old TestProxyForEndpoint discarded the valid-config result and only asserted "" for the bad-config arm — which every no-proxy outcome returns, so a `return ""` stub would have passed. http.ProxyFromEnvironment snapshots the proxy env once per process (sync.Once), so env-based cases can't be driven mid-test. Extract the core into proxyForEndpointFunc(cfg, resolve), keeping proxyForEndpoint(cfg) as the thin production wrapper over http.ProxyFromEnvironment. The test now injects deterministic resolvers to assert all branches: proxy-set→URL, NO_PROXY/nil→direct, resolver-error→direct, and unbuildable-endpoint→direct (resolver never consulted). Co-authored-by: Isaac Signed-off-by: Mani Kaustubh Mathur --- kernel_backend.go | 12 +++++- kernel_backend_test.go | 84 ++++++++++++++++++++++++++++++++---------- 2 files changed, 76 insertions(+), 20 deletions(-) diff --git a/kernel_backend.go b/kernel_backend.go index 4cea76c6..d203d5b1 100644 --- a/kernel_backend.go +++ b/kernel_backend.go @@ -6,6 +6,7 @@ import ( "context" "errors" "net/http" + "net/url" "github.com/databricks/databricks-sql-go/auth/noop" "github.com/databricks/databricks-sql-go/auth/pat" @@ -81,6 +82,15 @@ func newKernelBackend(_ context.Context, cfg *config.Config) (backend.Backend, e // NO_PROXY excludes the host or no proxy is set. No extra dependency: this is the // stdlib function the driver already relies on. func proxyForEndpoint(cfg *config.Config) string { + return proxyForEndpointFunc(cfg, http.ProxyFromEnvironment) +} + +// proxyForEndpointFunc is the testable core: it builds the endpoint request and +// asks resolve for the proxy, returning "" (direct) on any error, no proxy, or an +// unbuildable endpoint. resolve is http.ProxyFromEnvironment in production; tests +// inject a deterministic resolver to exercise proxy-set / NO_PROXY / direct +// without depending on http.ProxyFromEnvironment's process-wide env caching. +func proxyForEndpointFunc(cfg *config.Config, resolve func(*http.Request) (*url.URL, error)) string { endpoint, err := cfg.ToEndpointURL() if err != nil { return "" @@ -89,7 +99,7 @@ func proxyForEndpoint(cfg *config.Config) string { if err != nil { return "" } - proxyURL, err := http.ProxyFromEnvironment(req) + proxyURL, err := resolve(req) if err != nil || proxyURL == nil { return "" } diff --git a/kernel_backend_test.go b/kernel_backend_test.go index ce5356d7..20361022 100644 --- a/kernel_backend_test.go +++ b/kernel_backend_test.go @@ -4,6 +4,9 @@ package dbsql import ( "context" + "errors" + "net/http" + "net/url" "testing" "github.com/databricks/databricks-sql-go/internal/config" @@ -55,26 +58,69 @@ func TestNewKernelBackendRejectsUnsupportedOptions(t *testing.T) { }) } -// proxyForEndpoint returns a valid config's endpoint proxy without error. Its -// value comes from http.ProxyFromEnvironment, which snapshots the proxy env once -// per process (a sync.Once) — the same cached decision the Thrift transport -// makes — so the resolved value can't be re-driven by setting env vars mid-test. -// This asserts the invariant that matters here: a well-formed config never makes -// the resolver error out (it returns "" for direct), and a malformed config -// (missing host) is handled gracefully rather than panicking. +// proxyForEndpointFunc maps the injected resolver's decision to a proxy URL +// string (or "" for direct), and returns "" for an unbuildable endpoint. The +// production path uses http.ProxyFromEnvironment, whose env is snapshotted once +// per process (sync.Once) and so can't be re-driven mid-test; the resolver seam +// lets us assert every branch deterministically. Each case mirrors an +// httpproxy-style outcome: proxy set for this host, host excluded by NO_PROXY, +// no proxy configured, and an unbuildable endpoint. func TestProxyForEndpoint(t *testing.T) { - valid := config.WithDefaults() - valid.Host = "my-workspace.databricks.com" - valid.Port = 443 - valid.HTTPPath = "/sql/1.0/warehouses/abc" - // Must not panic; with no proxy env in the test environment this is "". - _ = proxyForEndpoint(valid) + validCfg := func() *config.Config { + c := config.WithDefaults() + c.Host = "my-workspace.databricks.com" + c.Port = 443 + c.HTTPPath = "/sql/1.0/warehouses/abc" + return c + } + proxyURL, _ := url.Parse("http://corp-proxy:3128") - // A config whose endpoint URL can't be built (no host) resolves to direct - // rather than erroring. - bad := config.WithDefaults() - bad.Host = "" - if got := proxyForEndpoint(bad); got != "" { - t.Errorf("unbuildable endpoint should resolve to direct, got %q", got) + cases := []struct { + name string + cfg *config.Config + resolve func(*http.Request) (*url.URL, error) + want string + }{ + { + name: "proxy set for host", + cfg: validCfg(), + resolve: func(*http.Request) (*url.URL, error) { return proxyURL, nil }, + want: "http://corp-proxy:3128", + }, + { + name: "host excluded by NO_PROXY -> direct", + cfg: validCfg(), + resolve: func(*http.Request) (*url.URL, error) { return nil, nil }, + want: "", + }, + { + name: "resolver error -> direct", + cfg: validCfg(), + resolve: func(*http.Request) (*url.URL, error) { return nil, errors.New("bad proxy url") }, + want: "", + }, + { + name: "unbuildable endpoint -> direct (resolver never consulted)", + cfg: func() *config.Config { + c := config.WithDefaults() + c.Host = "" + return c + }(), + resolve: func(*http.Request) (*url.URL, error) { + t.Error("resolver must not be called for an unbuildable endpoint") + return proxyURL, nil + }, + want: "", + }, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + if got := proxyForEndpointFunc(tc.cfg, tc.resolve); got != tc.want { + t.Errorf("proxyForEndpointFunc = %q, want %q", got, tc.want) + } + }) } + + // The production wrapper wires http.ProxyFromEnvironment and must not panic. + _ = proxyForEndpoint(validCfg()) } From d7263ea7a2772402ea5526af88ee0a2fe7a0e3b6 Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Fri, 10 Jul 2026 22:58:29 +0000 Subject: [PATCH 16/32] =?UTF-8?q?fix(kernel):=20remove=20struct-key=20pref?= =?UTF-8?q?ix=20cache=20=E2=80=94=20unbounded=20memory=20leak=20(N1)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The round-1 L4 change memoized struct field-name JSON prefixes in a process- global sync.Map keyed by *arrow.StructType. That leaks without bound: cdata.ImportCRecordBatch → importSchema → arrow.StructOf allocates a FRESH *StructType every batch (no interning), so the key never repeats across batches and the map grows one never-evicted entry per batch — a monotonic leak for a struct/nested-struct column over a large multi-batch CloudFetch result. The intended cross-row win only ever existed within a single batch anyway. Drop the cache and marshal the field name inline in writeStructJSON (the proven pre-L4 form). The per-row json.Marshal is cheap next to the once-per-batch cgo crossing; a correct per-batch precompute can return with the nested-renderer extraction if profiling warrants it. Co-authored-by: Isaac Signed-off-by: Mani Kaustubh Mathur --- internal/backend/kernel/scan_nested.go | 33 +++++++------------------- 1 file changed, 9 insertions(+), 24 deletions(-) diff --git a/internal/backend/kernel/scan_nested.go b/internal/backend/kernel/scan_nested.go index f27ed42b..66a1ef54 100644 --- a/internal/backend/kernel/scan_nested.go +++ b/internal/backend/kernel/scan_nested.go @@ -27,7 +27,6 @@ import ( "encoding/json" "fmt" "strings" - "sync" "time" "github.com/apache/arrow/go/v12/arrow" @@ -35,27 +34,6 @@ import ( "github.com/databricks/databricks-sql-go/internal/decimalfmt" ) -// structKeyPrefixCache memoizes the JSON-escaped, quote-wrapped, colon-terminated -// key prefixes for each struct type (e.g. `"field":`). Field names are invariant -// across rows, but writeStructJSON runs per row; arrow reuses the same -// *arrow.StructType pointer for a column across its rows, so this is a hit after -// the first row. Keyed by that pointer; safe for concurrent Rows via sync.Map. -var structKeyPrefixCache sync.Map // map[*arrow.StructType][]string - -func structKeyPrefixes(st *arrow.StructType) []string { - if v, ok := structKeyPrefixCache.Load(st); ok { - return v.([]string) - } - fields := st.Fields() - prefixes := make([]string, len(fields)) - for f := range fields { - name, _ := json.Marshal(fields[f].Name) // JSON-escapes the field name - prefixes[f] = string(name) + ":" - } - structKeyPrefixCache.Store(st, prefixes) - return prefixes -} - // renderJSONString renders one nested cell as a JSON string (nil for a NULL // cell). loc is applied to any timestamp/date leaves, matching the top-level // scan and the Thrift path. @@ -149,13 +127,20 @@ func writeMapJSON(b *strings.Builder, m *array.Map, row int, loc *time.Location) } func writeStructJSON(b *strings.Builder, s *array.Struct, row int, loc *time.Location) error { - prefixes := structKeyPrefixes(s.DataType().(*arrow.StructType)) + st := s.DataType().(*arrow.StructType) b.WriteByte('{') for f := 0; f < s.NumField(); f++ { if f > 0 { b.WriteByte(',') } - b.WriteString(prefixes[f]) // pre-escaped `"name":` + // Marshal the field name to JSON-escape it. Not memoized per struct type: a + // process-global cache keyed by *arrow.StructType leaks, because + // cdata.ImportCRecordBatch allocates a fresh *StructType every batch (no + // interning), so the key never repeats across batches. Recomputing here is + // cheap relative to the once-per-batch cgo crossing. + keyBytes, _ := json.Marshal(st.Field(f).Name) + b.Write(keyBytes) + b.WriteByte(':') if err := writeJSON(b, s.Field(f), row, loc); err != nil { return err } From 13a997f32575eb6d975724ac698cb3d1f7d25c2f Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Fri, 10 Jul 2026 23:05:52 +0000 Subject: [PATCH 17/32] refactor(kernel): extract Arrow cell/JSON rendering to untagged arrowscan (N5/N6) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Move ScanCell (scalar + nested→JSON grammar: native float32, exact decimals, time.Time formatting, list/map/struct) out of the cgo-tagged kernel package into a new pure-Go internal/arrowscan package. rows.go delegates via arrowscan.ScanCell. The renderer imports no C, so this is a pure move. Why: the rendering rules are the parity-critical contract both backends must agree on, but every kernel-package test is //go:build cgo && databricks_kernel and dead in CI. Relocating the tests to arrowscan makes them run in the default CGO_ENABLED=0 matrix (N6) — TestScanCellScalars/Nested/TimestampLocation now guard the native-float32 (N3/M2) and exact-decimal rendering on every build. The grammar is single-sourced for the kernel side (N5). Kernel-specific tests (error mapping, bad-connection, bound-params rejection) stay in the kernel package. Verified: default suite + arrowscan race pass; kernel unit tests pass; live Thrift-parity + all data types still byte-identical. Note: the Thrift arrowbased path still has its own nested renderer (built on its columnValues container abstraction, not raw arrow.Array). Having it delegate to arrowscan too — the remaining half of N5 — is a separate, higher-risk refactor of the primary production path; deferred. Co-authored-by: Isaac Signed-off-by: Mani Kaustubh Mathur --- .../scan_nested.go => arrowscan/arrowscan.go} | 120 ++++++-- internal/arrowscan/arrowscan_test.go | 276 +++++++++++++++++ internal/backend/kernel/kernel_test.go | 281 +----------------- internal/backend/kernel/rows.go | 86 +----- 4 files changed, 382 insertions(+), 381 deletions(-) rename internal/{backend/kernel/scan_nested.go => arrowscan/arrowscan.go} (57%) create mode 100644 internal/arrowscan/arrowscan_test.go diff --git a/internal/backend/kernel/scan_nested.go b/internal/arrowscan/arrowscan.go similarity index 57% rename from internal/backend/kernel/scan_nested.go rename to internal/arrowscan/arrowscan.go index 66a1ef54..f13125b0 100644 --- a/internal/backend/kernel/scan_nested.go +++ b/internal/arrowscan/arrowscan.go @@ -1,12 +1,13 @@ -//go:build cgo && databricks_kernel - -package kernel - -// Recursive rendering of nested Arrow values (List/Map/Struct) to a JSON string, -// byte-compatible with the Thrift arrow path -// (internal/rows/arrowbased/columnValues.go). scanCell delegates here for nested -// columns; database/sql consumers then get the same JSON shape from either -// backend: +// Package arrowscan converts Arrow array cells to database/sql driver.Values, +// with nested types (List/Map/Struct, and VARIANT which arrives nested) rendered +// to a JSON string byte-identical to the Thrift arrow path +// (internal/rows/arrowbased). It is pure Go (no cgo), so it is shared by the +// kernel backend and testable in the default CGO_ENABLED=0 build — the tests here +// are the regression guard for the exact rendering rules (native float32, exact +// decimals, time.Time formatting, JSON grammar) both backends must agree on. +// +// Rendering to JSON (not a Go map/slice) is deliberate: it is what the Thrift +// path returns, so a query's result is identical across backends. // - list → [v0,v1,...] // - map → {"k0":v0,"k1":v1,...} (keys stringified) // - struct → {"field0":v0,...} @@ -14,13 +15,9 @@ package kernel // - 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) -// -// VARIANT arrives as a nested value and renders through this path; GEOMETRY -// arrives as a WKB/WKT string and is handled by scanCell's scalar arm. -// -// Rendering to JSON (not a Go map/slice) is deliberate: it is what the Thrift -// path returns, so a query's result is identical across backends — the property -// the Thrift-parity test asserts. +// - float32 → native float32 (not widened to float64), so JSON renders +// 3.14, not 3.140000104904175 +package arrowscan import ( "database/sql/driver" @@ -34,6 +31,87 @@ import ( "github.com/databricks/databricks-sql-go/internal/decimalfmt" ) +// 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 +// 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). +func ScanCell(col arrow.Array, row int, loc *time.Location) (driver.Value, error) { + if col.IsNull(row) { + return nil, nil + } + switch c := col.(type) { + case *array.Null: + return nil, nil + case *array.Boolean: + return c.Value(row), nil + case *array.Int8: + return int64(c.Value(row)), nil + case *array.Int16: + return int64(c.Value(row)), nil + case *array.Int32: + return int64(c.Value(row)), nil + case *array.Int64: + return c.Value(row), nil + case *array.Uint8: + return int64(c.Value(row)), nil + case *array.Uint16: + return int64(c.Value(row)), nil + case *array.Uint32: + return int64(c.Value(row)), nil + case *array.Uint64: + return int64(c.Value(row)), nil + case *array.Float32: + // Return the native float32, NOT a widened float64: the Thrift path returns + // a float32 driver.Value for a bare FLOAT column, and database/sql's + // asString formats it at bit-size 32 — so widening here would render + // CAST(0.1 AS FLOAT) as "0.10000000149011612" vs Thrift's "0.1". + return c.Value(row), nil + case *array.Float64: + return c.Value(row), nil + case *array.String: + return c.Value(row), nil + case *array.LargeString: + return c.Value(row), nil + case *array.Binary: + return c.Value(row), nil + case *array.Date32: + return inLocation(c.Value(row).ToTime(), loc), nil + case *array.Date64: + return inLocation(c.Value(row).ToTime(), loc), nil + case *array.Timestamp: + dt, ok := col.DataType().(*arrow.TimestampType) + if !ok { + return nil, fmt.Errorf("timestamp column has unexpected datatype %s", col.DataType()) + } + return inLocation(c.Value(row).ToTime(dt.Unit), loc), nil + case *array.Decimal128: + dt := col.DataType().(*arrow.Decimal128Type) + return decimalfmt.ExactString(c.Value(row), dt.Scale), 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) + default: + return nil, fmt.Errorf("scanning arrow type %s is not supported "+ + "(intervals are not yet handled)", col.DataType()) + } +} + +// 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 { + if loc == nil { + return t + } + return t.In(loc) +} + // renderJSONString renders one nested cell as a JSON string (nil for a NULL // cell). loc is applied to any timestamp/date leaves, matching the top-level // scan and the Thrift path. @@ -134,10 +212,10 @@ func writeStructJSON(b *strings.Builder, s *array.Struct, row int, loc *time.Loc b.WriteByte(',') } // Marshal the field name to JSON-escape it. Not memoized per struct type: a - // process-global cache keyed by *arrow.StructType leaks, because - // cdata.ImportCRecordBatch allocates a fresh *StructType every batch (no + // process-global cache keyed by *arrow.StructType leaks, because the Arrow C + // Data Interface import allocates a fresh *StructType every batch (no // interning), so the key never repeats across batches. Recomputing here is - // cheap relative to the once-per-batch cgo crossing. + // cheap relative to the once-per-batch cost. keyBytes, _ := json.Marshal(st.Field(f).Name) b.Write(keyBytes) b.WriteByte(':') @@ -184,11 +262,11 @@ func writeScalarJSON(b *strings.Builder, v any) error { // scalarForJSON returns the Go value used for a nested leaf that is not itself a // container — today only a map key (values are written directly by writeJSON). -// It reuses scanCell's scalar arm, so a decimal key renders via the exact-string +// It reuses ScanCell's scalar arm, so a decimal key renders via the exact-string // path (writeJSONKey then quotes it), never a lossy float64. func scalarForJSON(col arrow.Array, row int, loc *time.Location) (any, error) { if col.IsNull(row) { return nil, nil } - return scanCell(col, row, loc) + return ScanCell(col, row, loc) } diff --git a/internal/arrowscan/arrowscan_test.go b/internal/arrowscan/arrowscan_test.go new file mode 100644 index 00000000..2da3adb4 --- /dev/null +++ b/internal/arrowscan/arrowscan_test.go @@ -0,0 +1,276 @@ +package arrowscan + +import ( + "testing" + "time" + + "github.com/apache/arrow/go/v12/arrow" + "github.com/apache/arrow/go/v12/arrow/array" + "github.com/apache/arrow/go/v12/arrow/decimal128" + "github.com/apache/arrow/go/v12/arrow/memory" +) + +// ScanCell renders the supported scalar types and rejects an unsupported type +// (rather than returning a silently wrong value). +func TestScanCellScalars(t *testing.T) { + pool := memory.NewGoAllocator() + + t.Run("int64", func(t *testing.T) { + b := array.NewInt64Builder(pool) + defer b.Release() + b.Append(42) + arr := b.NewArray() + defer arr.Release() + v, err := ScanCell(arr, 0, nil) + if err != nil { + t.Fatal(err) + } + if v.(int64) != 42 { + t.Errorf("got %v", v) + } + }) + + t.Run("string", func(t *testing.T) { + b := array.NewStringBuilder(pool) + defer b.Release() + b.Append("hi") + arr := b.NewArray() + defer arr.Release() + v, err := ScanCell(arr, 0, nil) + if err != nil { + t.Fatal(err) + } + if v.(string) != "hi" { + t.Errorf("got %v", v) + } + }) + + t.Run("float32_native", func(t *testing.T) { + // A top-level FLOAT column must scan to a native float32, not a widened + // float64 — matching Thrift, so database/sql's asString renders + // CAST(0.1 AS FLOAT) as "0.1", not "0.10000000149011612". + b := array.NewFloat32Builder(pool) + defer b.Release() + b.Append(0.1) + arr := b.NewArray() + defer arr.Release() + v, err := ScanCell(arr, 0, nil) + if err != nil { + t.Fatal(err) + } + got, ok := v.(float32) + if !ok { + t.Fatalf("want float32, got %T", v) + } + if got != float32(0.1) { + t.Errorf("got %v, want 0.1", got) + } + }) + + t.Run("null", func(t *testing.T) { + b := array.NewInt64Builder(pool) + defer b.Release() + b.AppendNull() + arr := b.NewArray() + defer arr.Release() + v, err := ScanCell(arr, 0, nil) + if err != nil { + t.Fatal(err) + } + if v != nil { + t.Errorf("null should scan to nil, got %v", v) + } + }) + + t.Run("decimal_exact_string", func(t *testing.T) { + // 12345 at scale 2 = "123.45", exact (not a float64). + dt := &arrow.Decimal128Type{Precision: 10, Scale: 2} + b := array.NewDecimal128Builder(pool, dt) + defer b.Release() + b.Append(decimal128.FromU64(12345)) + arr := b.NewArray() + defer arr.Release() + v, err := ScanCell(arr, 0, nil) + if err != nil { + t.Fatal(err) + } + if v.(string) != "123.45" { + t.Errorf("got %v, want 123.45", v) + } + }) + + 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}) + 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") + } + }) +} + +// 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) { + pool := memory.NewGoAllocator() + loc, err := time.LoadLocation("America/New_York") + if err != nil { + t.Skipf("tz database unavailable: %v", err) + } + + // 2026-07-09T12:00:00Z as microseconds since epoch. + utcTS := time.Date(2026, time.July, 9, 12, 0, 0, 0, time.UTC) + b := array.NewTimestampBuilder(pool, &arrow.TimestampType{Unit: arrow.Microsecond}) + defer b.Release() + b.Append(arrow.Timestamp(utcTS.UnixMicro())) + arr := b.NewArray() + defer arr.Release() + + t.Run("location applied", func(t *testing.T) { + v, err := ScanCell(arr, 0, loc) + if err != nil { + t.Fatal(err) + } + got := v.(time.Time) + if got.Location() != loc { + t.Errorf("location = %v, want %v", got.Location(), loc) + } + if !got.Equal(utcTS) { + t.Errorf("instant changed: got %v, want %v", got, utcTS) + } + }) + + t.Run("nil location is UTC", func(t *testing.T) { + v, err := ScanCell(arr, 0, nil) + if err != nil { + t.Fatal(err) + } + if v.(time.Time).Location() != time.UTC { + t.Errorf("nil location should render UTC, got %v", v.(time.Time).Location()) + } + }) +} + +// ScanCell renders nested types (list/struct/map) to a JSON string matching the +// Thrift path. +func TestScanCellNested(t *testing.T) { + pool := memory.NewGoAllocator() + + t.Run("list", func(t *testing.T) { + b := array.NewListBuilder(pool, arrow.PrimitiveTypes.Int64) + defer b.Release() + vb := b.ValueBuilder().(*array.Int64Builder) + b.Append(true) + vb.Append(1) + vb.Append(2) + vb.Append(3) + arr := b.NewArray() + defer arr.Release() + v, err := ScanCell(arr, 0, nil) + if err != nil { + t.Fatal(err) + } + if v.(string) != "[1,2,3]" { + t.Errorf("got %q, want [1,2,3]", v) + } + }) + + t.Run("struct", func(t *testing.T) { + dt := arrow.StructOf( + arrow.Field{Name: "a", Type: arrow.PrimitiveTypes.Int64}, + arrow.Field{Name: "b", Type: arrow.BinaryTypes.String}, + ) + b := array.NewStructBuilder(pool, dt) + defer b.Release() + b.Append(true) + b.FieldBuilder(0).(*array.Int64Builder).Append(1) + b.FieldBuilder(1).(*array.StringBuilder).Append("x") + arr := b.NewArray() + defer arr.Release() + v, err := ScanCell(arr, 0, nil) + if err != nil { + t.Fatal(err) + } + if v.(string) != `{"a":1,"b":"x"}` { + t.Errorf("got %q, want {\"a\":1,\"b\":\"x\"}", v) + } + }) + + t.Run("map", func(t *testing.T) { + b := array.NewMapBuilder(pool, arrow.BinaryTypes.String, arrow.PrimitiveTypes.Int64, false) + defer b.Release() + kb := b.KeyBuilder().(*array.StringBuilder) + ib := b.ItemBuilder().(*array.Int64Builder) + b.Append(true) + kb.Append("k") + ib.Append(9) + arr := b.NewArray() + defer arr.Release() + v, err := ScanCell(arr, 0, nil) + if err != nil { + t.Fatal(err) + } + if v.(string) != `{"k":9}` { + t.Errorf("got %q, want {\"k\":9}", v) + } + }) + + t.Run("nested_decimal_exact", func(t *testing.T) { + // A decimal inside a struct must render as an exact JSON number, not a + // lossy float64 (19.99, not 19.990000000000002) — matching Thrift's + // marshalScalar. + dt := arrow.StructOf(arrow.Field{Name: "d", Type: &arrow.Decimal128Type{Precision: 5, Scale: 2}}) + b := array.NewStructBuilder(pool, dt) + defer b.Release() + b.Append(true) + b.FieldBuilder(0).(*array.Decimal128Builder).Append(decimal128.FromU64(1999)) + arr := b.NewArray() + defer arr.Release() + v, err := ScanCell(arr, 0, nil) + if err != nil { + t.Fatal(err) + } + if v.(string) != `{"d":19.99}` { + t.Errorf("got %q, want {\"d\":19.99}", v) + } + }) + + t.Run("nested_float32_exact", func(t *testing.T) { + // A float32 inside a struct must marshal as the native float32 (3.14), not + // a widened float64 (3.140000104904175) — matching Thrift's nested path, + // which marshals the native float32. + dt := arrow.StructOf(arrow.Field{Name: "f", Type: arrow.PrimitiveTypes.Float32}) + b := array.NewStructBuilder(pool, dt) + defer b.Release() + b.Append(true) + b.FieldBuilder(0).(*array.Float32Builder).Append(3.14) + arr := b.NewArray() + defer arr.Release() + v, err := ScanCell(arr, 0, nil) + if err != nil { + t.Fatal(err) + } + if v.(string) != `{"f":3.14}` { + t.Errorf("got %q, want {\"f\":3.14}", v) + } + }) + + t.Run("nested_null", func(t *testing.T) { + b := array.NewListBuilder(pool, arrow.PrimitiveTypes.Int64) + defer b.Release() + b.AppendNull() + arr := b.NewArray() + defer arr.Release() + v, err := ScanCell(arr, 0, nil) + if err != nil { + t.Fatal(err) + } + if v != nil { + t.Errorf("null list should scan to nil, got %v", v) + } + }) +} diff --git a/internal/backend/kernel/kernel_test.go b/internal/backend/kernel/kernel_test.go index 6a00d532..6aae1256 100644 --- a/internal/backend/kernel/kernel_test.go +++ b/internal/backend/kernel/kernel_test.go @@ -7,12 +7,6 @@ import ( "database/sql/driver" "errors" "testing" - "time" - - "github.com/apache/arrow/go/v12/arrow" - "github.com/apache/arrow/go/v12/arrow/array" - "github.com/apache/arrow/go/v12/arrow/decimal128" - "github.com/apache/arrow/go/v12/arrow/memory" "github.com/databricks/databricks-sql-go/internal/backend" ) @@ -58,271 +52,6 @@ func TestToDriverError(t *testing.T) { } } -// scanCell renders the supported scalar types and rejects an unsupported type -// (rather than returning a silently wrong value). -func TestScanCellScalars(t *testing.T) { - pool := memory.NewGoAllocator() - - t.Run("int64", func(t *testing.T) { - b := array.NewInt64Builder(pool) - defer b.Release() - b.Append(42) - arr := b.NewArray() - defer arr.Release() - v, err := scanCell(arr, 0, nil) - if err != nil { - t.Fatal(err) - } - if v.(int64) != 42 { - t.Errorf("got %v", v) - } - }) - - t.Run("string", func(t *testing.T) { - b := array.NewStringBuilder(pool) - defer b.Release() - b.Append("hi") - arr := b.NewArray() - defer arr.Release() - v, err := scanCell(arr, 0, nil) - if err != nil { - t.Fatal(err) - } - if v.(string) != "hi" { - t.Errorf("got %v", v) - } - }) - - t.Run("float32_native", func(t *testing.T) { - // A top-level FLOAT column must scan to a native float32, not a widened - // float64 — matching Thrift, so database/sql's asString renders - // CAST(0.1 AS FLOAT) as "0.1", not "0.10000000149011612". - b := array.NewFloat32Builder(pool) - defer b.Release() - b.Append(0.1) - arr := b.NewArray() - defer arr.Release() - v, err := scanCell(arr, 0, nil) - if err != nil { - t.Fatal(err) - } - got, ok := v.(float32) - if !ok { - t.Fatalf("want float32, got %T", v) - } - if got != float32(0.1) { - t.Errorf("got %v, want 0.1", got) - } - }) - - t.Run("null", func(t *testing.T) { - b := array.NewInt64Builder(pool) - defer b.Release() - b.AppendNull() - arr := b.NewArray() - defer arr.Release() - v, err := scanCell(arr, 0, nil) - if err != nil { - t.Fatal(err) - } - if v != nil { - t.Errorf("null should scan to nil, got %v", v) - } - }) - - t.Run("decimal_exact_string", func(t *testing.T) { - // 12345 at scale 2 = "123.45", exact (not a float64). - dt := &arrow.Decimal128Type{Precision: 10, Scale: 2} - b := array.NewDecimal128Builder(pool, dt) - defer b.Release() - b.Append(decimal128.FromU64(12345)) - arr := b.NewArray() - defer arr.Release() - v, err := scanCell(arr, 0, nil) - if err != nil { - t.Fatal(err) - } - if v.(string) != "123.45" { - t.Errorf("got %v, want 123.45", v) - } - }) - - 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}) - 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") - } - }) -} - -// 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) { - pool := memory.NewGoAllocator() - loc, err := time.LoadLocation("America/New_York") - if err != nil { - t.Skipf("tz database unavailable: %v", err) - } - - // 2026-07-09T12:00:00Z as microseconds since epoch. - utcTS := time.Date(2026, time.July, 9, 12, 0, 0, 0, time.UTC) - b := array.NewTimestampBuilder(pool, &arrow.TimestampType{Unit: arrow.Microsecond}) - defer b.Release() - b.Append(arrow.Timestamp(utcTS.UnixMicro())) - arr := b.NewArray() - defer arr.Release() - - t.Run("location applied", func(t *testing.T) { - v, err := scanCell(arr, 0, loc) - if err != nil { - t.Fatal(err) - } - got := v.(time.Time) - if got.Location() != loc { - t.Errorf("location = %v, want %v", got.Location(), loc) - } - if !got.Equal(utcTS) { - t.Errorf("instant changed: got %v, want %v", got, utcTS) - } - }) - - t.Run("nil location is UTC", func(t *testing.T) { - v, err := scanCell(arr, 0, nil) - if err != nil { - t.Fatal(err) - } - if v.(time.Time).Location() != time.UTC { - t.Errorf("nil location should render UTC, got %v", v.(time.Time).Location()) - } - }) -} - -// scanCell renders nested types (list/struct/map) to a JSON string matching the -// Thrift path. -func TestScanCellNested(t *testing.T) { - pool := memory.NewGoAllocator() - - t.Run("list", func(t *testing.T) { - b := array.NewListBuilder(pool, arrow.PrimitiveTypes.Int64) - defer b.Release() - vb := b.ValueBuilder().(*array.Int64Builder) - b.Append(true) - vb.Append(1) - vb.Append(2) - vb.Append(3) - arr := b.NewArray() - defer arr.Release() - v, err := scanCell(arr, 0, nil) - if err != nil { - t.Fatal(err) - } - if v.(string) != "[1,2,3]" { - t.Errorf("got %q, want [1,2,3]", v) - } - }) - - t.Run("struct", func(t *testing.T) { - dt := arrow.StructOf( - arrow.Field{Name: "a", Type: arrow.PrimitiveTypes.Int64}, - arrow.Field{Name: "b", Type: arrow.BinaryTypes.String}, - ) - b := array.NewStructBuilder(pool, dt) - defer b.Release() - b.Append(true) - b.FieldBuilder(0).(*array.Int64Builder).Append(1) - b.FieldBuilder(1).(*array.StringBuilder).Append("x") - arr := b.NewArray() - defer arr.Release() - v, err := scanCell(arr, 0, nil) - if err != nil { - t.Fatal(err) - } - if v.(string) != `{"a":1,"b":"x"}` { - t.Errorf("got %q, want {\"a\":1,\"b\":\"x\"}", v) - } - }) - - t.Run("map", func(t *testing.T) { - b := array.NewMapBuilder(pool, arrow.BinaryTypes.String, arrow.PrimitiveTypes.Int64, false) - defer b.Release() - kb := b.KeyBuilder().(*array.StringBuilder) - ib := b.ItemBuilder().(*array.Int64Builder) - b.Append(true) - kb.Append("k") - ib.Append(9) - arr := b.NewArray() - defer arr.Release() - v, err := scanCell(arr, 0, nil) - if err != nil { - t.Fatal(err) - } - if v.(string) != `{"k":9}` { - t.Errorf("got %q, want {\"k\":9}", v) - } - }) - - t.Run("nested_decimal_exact", func(t *testing.T) { - // A decimal inside a struct must render as an exact JSON number, not a - // lossy float64 (19.99, not 19.990000000000002) — matching Thrift's - // marshalScalar. - dt := arrow.StructOf(arrow.Field{Name: "d", Type: &arrow.Decimal128Type{Precision: 5, Scale: 2}}) - b := array.NewStructBuilder(pool, dt) - defer b.Release() - b.Append(true) - b.FieldBuilder(0).(*array.Decimal128Builder).Append(decimal128.FromU64(1999)) - arr := b.NewArray() - defer arr.Release() - v, err := scanCell(arr, 0, nil) - if err != nil { - t.Fatal(err) - } - if v.(string) != `{"d":19.99}` { - t.Errorf("got %q, want {\"d\":19.99}", v) - } - }) - - t.Run("nested_float32_exact", func(t *testing.T) { - // A float32 inside a struct must marshal as the native float32 (3.14), not - // a widened float64 (3.140000104904175) — matching Thrift's nested path, - // which marshals the native float32. - dt := arrow.StructOf(arrow.Field{Name: "f", Type: arrow.PrimitiveTypes.Float32}) - b := array.NewStructBuilder(pool, dt) - defer b.Release() - b.Append(true) - b.FieldBuilder(0).(*array.Float32Builder).Append(3.14) - arr := b.NewArray() - defer arr.Release() - v, err := scanCell(arr, 0, nil) - if err != nil { - t.Fatal(err) - } - if v.(string) != `{"f":3.14}` { - t.Errorf("got %q, want {\"f\":3.14}", v) - } - }) - - t.Run("nested_null", func(t *testing.T) { - b := array.NewListBuilder(pool, arrow.PrimitiveTypes.Int64) - defer b.Release() - b.AppendNull() - arr := b.NewArray() - defer arr.Release() - v, err := scanCell(arr, 0, nil) - if err != nil { - t.Fatal(err) - } - if v != nil { - t.Errorf("null list should scan to nil, got %v", v) - } - }) -} - // Bound parameters are rejected up front by Execute with a clear error, before // any session/C work — so this runs on a zero-value backend. The returned // Operation must be non-nil (Backend contract) and its Close must report @@ -352,8 +81,8 @@ func TestExecuteRejectsParams(t *testing.T) { } } -// The exact decimal formatter now lives in internal/decimalfmt (shared with the -// Thrift path); its unit test is decimalfmt.TestExactString. The nested/scalar -// tests above still assert the kernel path renders decimals exactly. - -var _ driver.Value // keep database/sql/driver imported for the scanCell signature +// 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 +// internal/decimalfmt. This file keeps the kernel-specific tests: error mapping, +// bad-connection classification, and the bound-params rejection. diff --git a/internal/backend/kernel/rows.go b/internal/backend/kernel/rows.go index ee3e9112..70f05326 100644 --- a/internal/backend/kernel/rows.go +++ b/internal/backend/kernel/rows.go @@ -17,13 +17,11 @@ import ( "database/sql/driver" "fmt" "io" - "time" "unsafe" "github.com/apache/arrow/go/v12/arrow" - "github.com/apache/arrow/go/v12/arrow/array" "github.com/apache/arrow/go/v12/arrow/cdata" - "github.com/databricks/databricks-sql-go/internal/decimalfmt" + "github.com/databricks/databricks-sql-go/internal/arrowscan" dbsqlrows "github.com/databricks/databricks-sql-go/internal/rows" ) @@ -118,7 +116,7 @@ func (r *kernelRows) Next(dest []driver.Value) error { } rec := r.cur for c := 0; c < len(dest); c++ { - v, err := scanCell(rec.Column(c), r.rowInCur, r.op.location) + v, err := arrowscan.ScanCell(rec.Column(c), r.rowInCur, r.op.location) if err != nil { return fmt.Errorf("kernel: scan col %d (%s): %w", c, r.cols[c], err) } @@ -183,83 +181,3 @@ func (r *kernelRows) nextBatch() error { klog("nextBatch: %d rows (chunk %d)", rec.NumRows(), r.chunkCount) return nil } - -// scanCell extracts one cell as a driver.Value. Scalars map to their Go value: -// bool, all int/uint widths, float, 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 arrives nested) render to a -// JSON string byte-identical to the Thrift path (see scan_nested.go); 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. -func scanCell(col arrow.Array, row int, loc *time.Location) (driver.Value, error) { - if col.IsNull(row) { - return nil, nil - } - switch c := col.(type) { - case *array.Null: - return nil, nil - case *array.Boolean: - return c.Value(row), nil - case *array.Int8: - return int64(c.Value(row)), nil - case *array.Int16: - return int64(c.Value(row)), nil - case *array.Int32: - return int64(c.Value(row)), nil - case *array.Int64: - return c.Value(row), nil - case *array.Uint8: - return int64(c.Value(row)), nil - case *array.Uint16: - return int64(c.Value(row)), nil - case *array.Uint32: - return int64(c.Value(row)), nil - case *array.Uint64: - return int64(c.Value(row)), nil - case *array.Float32: - // Return the native float32, NOT a widened float64: the Thrift path returns - // a float32 driver.Value for a bare FLOAT column, and database/sql's - // asString formats it at bit-size 32 — so widening here would render - // CAST(0.1 AS FLOAT) as "0.10000000149011612" vs Thrift's "0.1". - return c.Value(row), nil - case *array.Float64: - return c.Value(row), nil - case *array.String: - return c.Value(row), nil - case *array.LargeString: - return c.Value(row), nil - case *array.Binary: - return c.Value(row), nil - case *array.Date32: - return inLocation(c.Value(row).ToTime(), loc), nil - case *array.Date64: - return inLocation(c.Value(row).ToTime(), loc), nil - case *array.Timestamp: - dt, ok := col.DataType().(*arrow.TimestampType) - if !ok { - return nil, fmt.Errorf("timestamp column has unexpected datatype %s", col.DataType()) - } - return inLocation(c.Value(row).ToTime(dt.Unit), loc), nil - case *array.Decimal128: - dt := col.DataType().(*arrow.Decimal128Type) - return decimalfmt.ExactString(c.Value(row), dt.Scale), 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) - default: - return nil, fmt.Errorf("kernel: scanning arrow type %s is not supported "+ - "(intervals are not yet handled)", col.DataType()) - } -} - -// 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 { - if loc == nil { - return t - } - return t.In(loc) -} From 44e180fc267a9e12a765c384ee92f12ea6b051da Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Sat, 11 Jul 2026 06:30:47 +0000 Subject: [PATCH 18/32] fix(arrowscan): suppress G115 on the defensive Uint64 arm Moving ScanCell out of the cgo-tagged kernel package into the untagged internal/arrowscan (N5/N6) put it under gosec in the default-build lint, which flags the uint64->int64 conversion (G115). Databricks SQL has no unsigned types, so a Uint64 column never occurs; driver.Value has no uint64 and the driver convention is int64, so the conversion is correct for every reachable value. Annotate the intent and suppress G115 on this unreachable arm, matching the repo's existing #nosec convention (connector.go G402). Verified with golangci-lint v2.12.2 (the CI version): 0 issues repo-wide. Co-authored-by: Isaac Signed-off-by: Mani Kaustubh Mathur --- internal/arrowscan/arrowscan.go | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/internal/arrowscan/arrowscan.go b/internal/arrowscan/arrowscan.go index f13125b0..6c481754 100644 --- a/internal/arrowscan/arrowscan.go +++ b/internal/arrowscan/arrowscan.go @@ -65,7 +65,11 @@ func ScanCell(col arrow.Array, row int, loc *time.Location) (driver.Value, error case *array.Uint32: return int64(c.Value(row)), nil case *array.Uint64: - return int64(c.Value(row)), nil + // Databricks SQL has no unsigned types, so a Uint64 column does not occur + // in practice; this arm is defensive. driver.Value has no uint64 and the + // driver convention is int64 for integers, so a value above MaxInt64 would + // wrap — acceptable for an unreachable path. + return int64(c.Value(row)), nil // #nosec G115 -- see above; unreachable for Databricks types case *array.Float32: // Return the native float32, NOT a widened float64: the Thrift path returns // a float32 driver.Value for a bare FLOAT column, and database/sql's From 0955e00974017cf994317c632fa19f76ae26e704 Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Sat, 11 Jul 2026 08:35:57 +0000 Subject: [PATCH 19/32] fix(kernel): surface server queryId in kernel error log and message (M5) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit lastError already parsed the kernel's query_id into KernelError.QueryID, but the always-on Warn log and Error() both omitted it — so a kernel-path failure gave on-call one log line with no queryId and (because StatementID() is "") no metric, leaving no way to pivot to server-side query history. Include the queryId in the Warn log and append it to Error() when non-empty. A query id is a correlation token, not PII; error-path only, so no benchmark impact. Co-authored-by: Isaac Signed-off-by: Mani Kaustubh Mathur --- internal/backend/kernel/cgo.go | 22 +++++++++++++++------- 1 file changed, 15 insertions(+), 7 deletions(-) diff --git a/internal/backend/kernel/cgo.go b/internal/backend/kernel/cgo.go index d38445df..ca83080b 100644 --- a/internal/backend/kernel/cgo.go +++ b/internal/backend/kernel/cgo.go @@ -143,11 +143,13 @@ func lastError(code C.KernelStatusCode) *KernelError { klog("kernel error: code=%d sqlstate=%q vendor=%d http=%d retryable=%v msg=%q", ke.Code, ke.SQLState, ke.VendorCode, ke.HTTPStatus, ke.Retryable, ke.Message) // Also emit at the driver's default (Warn) level — no SQL text or PII, just - // the status/sqlstate/http fields — so a kernel-path failure is visible - // without DBSQL_KERNEL_DEBUG. This is the error path only (never the hot - // per-row/per-batch path), so it does not perturb benchmarks. - logger.Logger.Warn().Msgf("databricks: kernel call failed: code=%d sqlstate=%q vendor=%d http=%d retryable=%v", - ke.Code, ke.SQLState, ke.VendorCode, ke.HTTPStatus, ke.Retryable) + // the status/sqlstate/http fields plus the server query id (a correlation + // token, not PII) so on-call can pivot to server-side query history — so a + // kernel-path failure is visible without DBSQL_KERNEL_DEBUG. This is the error + // path only (never the hot per-row/per-batch path), so it does not perturb + // benchmarks. + logger.Logger.Warn().Msgf("databricks: kernel call failed: code=%d sqlstate=%q vendor=%d http=%d retryable=%v queryId=%q", + ke.Code, ke.SQLState, ke.VendorCode, ke.HTTPStatus, ke.Retryable, ke.QueryID) return ke } @@ -165,10 +167,16 @@ 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. + q := "" + if e.QueryID != "" { + q = fmt.Sprintf(", queryId=%s", e.QueryID) + } if e.SQLState != "" { - return fmt.Sprintf("kernel: %s (sqlstate=%s, code=%d)", e.Message, e.SQLState, e.Code) + return fmt.Sprintf("kernel: %s (sqlstate=%s, code=%d%s)", e.Message, e.SQLState, e.Code, q) } - return fmt.Sprintf("kernel: %s (code=%d)", e.Message, e.Code) + return fmt.Sprintf("kernel: %s (code=%d%s)", e.Message, e.Code, q) } // Status codes mirrored as Go ints so non-cgo code (tests, error mapping) can From 9a5f62d052f3f6ab51059fcc6f5e218c6211cde3 Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Sat, 11 Jul 2026 08:40:38 +0000 Subject: [PATCH 20/32] fix(kernel): close fail-loud contract holes for auth/timeout/retries (M1/M2/M3) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The kernel backend promises "nothing silently ignored", but three options leaked: - PAT via WithAuthenticator (M1): the guard admits *pat.PATAuth, but the token was read from cfg.AccessToken — which WithAuthenticator(&pat.PATAuth{AccessToken:...}) leaves empty (only WithAccessToken sets both). So a valid, Thrift-supported PAT config reached the kernel with an empty token → opaque Unauthenticated. Resolve the token from the authenticator when cfg.AccessToken is empty, and reject an empty resolved token loudly. - WithTimeout (M2): cfg.QueryTimeout maps to a per-statement server timeout on Thrift (TExecuteStatementReq.QueryTimeout); the kernel C ABI has no equivalent setter (verified against the header), so reject QueryTimeout > 0 rather than run with no server-side timeout. - WithRetries(-1) (M3): explicitly disables retries, but the kernel retries internally with no toggle — reject the disable request. Positive retry tuning and WithMaxRows can't be distinguished from defaults and are managed kernel-side, so they're documented in doc.go as accepted-but-not-applied rather than rejected. Tests cover all three rejects, the PAT-via-authenticator success path, and the accepted positive-tuning path. Co-authored-by: Isaac Signed-off-by: Mani Kaustubh Mathur --- doc.go | 14 ++++++--- kernel_backend.go | 38 ++++++++++++++++++++++--- kernel_backend_test.go | 64 ++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 108 insertions(+), 8 deletions(-) diff --git a/doc.go b/doc.go index b5586992..e3c2abed 100644 --- a/doc.go +++ b/doc.go @@ -200,10 +200,16 @@ 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) and initial catalog/schema are not yet supported and return a clear error at connect time rather than being silently -ignored. Bound query parameters are likewise not yet supported and return a clear -error at execute time (they arrive per-query, not at connect). 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.) +ignored; 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 are likewise not yet supported and return a clear error at +execute time (they arrive per-query, not at connect). 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 +internally, below the C ABI, with no user-facing knob. # Programmatically Retrieving Connection and Query Id diff --git a/kernel_backend.go b/kernel_backend.go index d203d5b1..d2b16bd5 100644 --- a/kernel_backend.go +++ b/kernel_backend.go @@ -42,20 +42,50 @@ func newKernelBackend(_ context.Context, cfg *config.Config) (backend.Backend, e // so an empty PAT would reach the kernel and fail with an opaque // Unauthenticated error. Reject it here so the failure names the cause, per // the doc.go contract. nil / NoopAuth / PATAuth are the PAT-or-none cases. - switch cfg.Authenticator.(type) { - case nil, *noop.NoopAuth, *pat.PATAuth: - // PAT (or no explicit authenticator) — supported. + 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 nil, errors.New("databricks: only personal access token (WithAccessToken) auth is supported by the kernel backend; " + "OAuth (M2M/U2M), token-provider, external/static, and federated authenticators are not yet supported — " + "use PAT or the default (Thrift) backend") } + 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)") + } + // WithTimeout maps to a per-statement server timeout on Thrift + // (TExecuteStatementReq.QueryTimeout); the kernel C ABI exposes no equivalent, + // so honor the "nothing silently ignored" contract by rejecting it rather than + // running the query with no server-side timeout. (Default 0 = unset.) + if cfg.QueryTimeout > 0 { + return nil, errors.New("databricks: WithTimeout (server query timeout) is not yet supported by the kernel backend; " + + "omit it or use the default (Thrift) backend") + } + // WithRetries(-1) explicitly disables retries, but the kernel retries + // internally below the C ABI with no user-facing toggle — so a disable request + // 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, errors.New("databricks: disabling retries via WithRetries is not supported by the kernel backend " + + "(the kernel retries internally); omit it or use the default (Thrift) backend") + } kc := kernel.Config{ Host: cfg.Host, HTTPPath: cfg.HTTPPath, WarehouseID: cfg.WarehouseID, - Token: cfg.AccessToken, + Token: token, Location: cfg.Location, // Session confs (STATEMENT_TIMEOUT, QUERY_TAGS, TIMEZONE, …) — the same // SessionParams map the Thrift backend forwards, so they flow to the diff --git a/kernel_backend_test.go b/kernel_backend_test.go index 20361022..7be6f44e 100644 --- a/kernel_backend_test.go +++ b/kernel_backend_test.go @@ -8,10 +8,18 @@ import ( "net/http" "net/url" "testing" + "time" + "github.com/databricks/databricks-sql-go/auth/pat" "github.com/databricks/databricks-sql-go/internal/config" ) +// nonPATAuth is a stand-in for any non-PAT authenticator (OAuth / token-provider / +// external / federated) — the kernel backend must reject it at connect. +type nonPATAuth struct{} + +func (nonPATAuth) Authenticate(*http.Request) error { return nil } + // newKernelBackend rejects options it can't yet honor (initial namespace, // metric-view metadata) loudly, rather than silently ignoring them — which would // behave differently than the Thrift backend. @@ -56,6 +64,62 @@ func TestNewKernelBackendRejectsUnsupportedOptions(t *testing.T) { t.Errorf("a supported config should build cleanly, got %v", err) } }) + + t.Run("PAT via WithAuthenticator ok", func(t *testing.T) { + // WithAuthenticator(&pat.PATAuth{...}) sets only cfg.Authenticator, leaving + // cfg.AccessToken empty. The backend must still build (token sourced from + // the authenticator), not fail with an empty PAT at connect. + c := base() + c.AccessToken = "" + c.Authenticator = &pat.PATAuth{AccessToken: "dapi-y"} + if _, err := newKernelBackend(context.Background(), c); err != nil { + t.Errorf("PAT supplied via WithAuthenticator should build cleanly, got %v", err) + } + }) + + t.Run("empty token rejected", func(t *testing.T) { + c := base() + c.AccessToken = "" + c.Authenticator = &pat.PATAuth{AccessToken: ""} + if _, err := newKernelBackend(context.Background(), c); err == nil { + t.Error("expected an error when the resolved PAT is empty") + } + }) + + t.Run("non-PAT authenticator rejected", func(t *testing.T) { + c := base() + c.Authenticator = nonPATAuth{} + if _, err := newKernelBackend(context.Background(), c); err == nil { + t.Error("expected an error for a non-PAT authenticator") + } + }) + + t.Run("query timeout rejected", func(t *testing.T) { + c := base() + c.QueryTimeout = 30 * time.Second + if _, err := newKernelBackend(context.Background(), c); err == nil { + t.Error("expected an error when WithTimeout (query timeout) is set on the kernel backend") + } + }) + + t.Run("disabling retries rejected", func(t *testing.T) { + c := base() + c.RetryMax = -1 + if _, err := newKernelBackend(context.Background(), c); err == nil { + t.Error("expected an error when retries are disabled (WithRetries(-1)) on the kernel backend") + } + }) + + t.Run("positive retry tuning + maxrows accepted", func(t *testing.T) { + // These are accepted (not applied) per doc.go — the kernel manages fetch / + // retries internally, so they must not error. + c := base() + c.RetryMax = 8 + c.MaxRows = 5000 + if _, err := newKernelBackend(context.Background(), c); err != nil { + t.Errorf("positive retry/maxrows tuning should build cleanly, got %v", err) + } + }) } // proxyForEndpointFunc maps the injected resolver's decision to a proxy URL From d2721571a39e4c75a859d45294e46679d59f78bf Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Sat, 11 Jul 2026 08:44:08 +0000 Subject: [PATCH 21/32] refactor(kernel): move proxy resolution to an untagged file so its test runs in CI (M4) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit proxyForEndpoint / proxyForEndpointFunc are pure Go (only config + net/http/url, no kernel C symbol), but lived in the cgo-tagged kernel_backend.go — so the now-meaningful TestProxyForEndpoint (four asserted branches) was inert under the only CI job (CGO_ENABLED=0). Split them into an untagged kernel_proxy.go + kernel_proxy_test.go, mirroring the arrowscan/decimalfmt move; newKernelBackend stays gated and calls proxyForEndpoint. The proxy test now runs in the default matrix without a kernel lib. Co-authored-by: Isaac Signed-off-by: Mani Kaustubh Mathur --- kernel_backend.go | 35 ++----------------- kernel_backend_test.go | 69 ------------------------------------- kernel_proxy.go | 45 ++++++++++++++++++++++++ kernel_proxy_test.go | 78 ++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 125 insertions(+), 102 deletions(-) create mode 100644 kernel_proxy.go create mode 100644 kernel_proxy_test.go diff --git a/kernel_backend.go b/kernel_backend.go index d2b16bd5..500a3f4b 100644 --- a/kernel_backend.go +++ b/kernel_backend.go @@ -5,8 +5,6 @@ package dbsql import ( "context" "errors" - "net/http" - "net/url" "github.com/databricks/databricks-sql-go/auth/noop" "github.com/databricks/databricks-sql-go/auth/pat" @@ -104,34 +102,5 @@ func newKernelBackend(_ context.Context, cfg *config.Config) (backend.Backend, e return kernel.New(kc), nil } -// proxyForEndpoint resolves the proxy the Thrift path would use for this -// connection, via the same http.ProxyFromEnvironment (HTTP(S)_PROXY / NO_PROXY) -// the Thrift transport applies at request time. Building the endpoint request -// lets ProxyFromEnvironment apply the NO_PROXY rules for this exact host, so the -// kernel sees the same effective proxy decision — returning "" (direct) when -// NO_PROXY excludes the host or no proxy is set. No extra dependency: this is the -// stdlib function the driver already relies on. -func proxyForEndpoint(cfg *config.Config) string { - return proxyForEndpointFunc(cfg, http.ProxyFromEnvironment) -} - -// proxyForEndpointFunc is the testable core: it builds the endpoint request and -// asks resolve for the proxy, returning "" (direct) on any error, no proxy, or an -// unbuildable endpoint. resolve is http.ProxyFromEnvironment in production; tests -// inject a deterministic resolver to exercise proxy-set / NO_PROXY / direct -// without depending on http.ProxyFromEnvironment's process-wide env caching. -func proxyForEndpointFunc(cfg *config.Config, resolve func(*http.Request) (*url.URL, error)) string { - endpoint, err := cfg.ToEndpointURL() - if err != nil { - return "" - } - req, err := http.NewRequest(http.MethodPost, endpoint, nil) - if err != nil { - return "" - } - proxyURL, err := resolve(req) - if err != nil || proxyURL == nil { - return "" - } - return proxyURL.String() -} +// 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 7be6f44e..6fffa44e 100644 --- a/kernel_backend_test.go +++ b/kernel_backend_test.go @@ -4,9 +4,7 @@ package dbsql import ( "context" - "errors" "net/http" - "net/url" "testing" "time" @@ -121,70 +119,3 @@ func TestNewKernelBackendRejectsUnsupportedOptions(t *testing.T) { } }) } - -// proxyForEndpointFunc maps the injected resolver's decision to a proxy URL -// string (or "" for direct), and returns "" for an unbuildable endpoint. The -// production path uses http.ProxyFromEnvironment, whose env is snapshotted once -// per process (sync.Once) and so can't be re-driven mid-test; the resolver seam -// lets us assert every branch deterministically. Each case mirrors an -// httpproxy-style outcome: proxy set for this host, host excluded by NO_PROXY, -// no proxy configured, and an unbuildable endpoint. -func TestProxyForEndpoint(t *testing.T) { - validCfg := func() *config.Config { - c := config.WithDefaults() - c.Host = "my-workspace.databricks.com" - c.Port = 443 - c.HTTPPath = "/sql/1.0/warehouses/abc" - return c - } - proxyURL, _ := url.Parse("http://corp-proxy:3128") - - cases := []struct { - name string - cfg *config.Config - resolve func(*http.Request) (*url.URL, error) - want string - }{ - { - name: "proxy set for host", - cfg: validCfg(), - resolve: func(*http.Request) (*url.URL, error) { return proxyURL, nil }, - want: "http://corp-proxy:3128", - }, - { - name: "host excluded by NO_PROXY -> direct", - cfg: validCfg(), - resolve: func(*http.Request) (*url.URL, error) { return nil, nil }, - want: "", - }, - { - name: "resolver error -> direct", - cfg: validCfg(), - resolve: func(*http.Request) (*url.URL, error) { return nil, errors.New("bad proxy url") }, - want: "", - }, - { - name: "unbuildable endpoint -> direct (resolver never consulted)", - cfg: func() *config.Config { - c := config.WithDefaults() - c.Host = "" - return c - }(), - resolve: func(*http.Request) (*url.URL, error) { - t.Error("resolver must not be called for an unbuildable endpoint") - return proxyURL, nil - }, - want: "", - }, - } - for _, tc := range cases { - t.Run(tc.name, func(t *testing.T) { - if got := proxyForEndpointFunc(tc.cfg, tc.resolve); got != tc.want { - t.Errorf("proxyForEndpointFunc = %q, want %q", got, tc.want) - } - }) - } - - // The production wrapper wires http.ProxyFromEnvironment and must not panic. - _ = proxyForEndpoint(validCfg()) -} diff --git a/kernel_proxy.go b/kernel_proxy.go new file mode 100644 index 00000000..0e89377b --- /dev/null +++ b/kernel_proxy.go @@ -0,0 +1,45 @@ +package dbsql + +import ( + "net/http" + "net/url" + + "github.com/databricks/databricks-sql-go/internal/config" +) + +// This file is intentionally NOT behind the `cgo && databricks_kernel` build tag: +// proxy resolution is pure Go (no kernel C symbol), so keeping it in the default +// build lets its test run under CGO_ENABLED=0 rather than being dead behind the +// tag. newKernelBackend (tagged) calls proxyForEndpoint. + +// proxyForEndpoint resolves the proxy the Thrift path would use for this +// connection, via the same http.ProxyFromEnvironment (HTTP(S)_PROXY / NO_PROXY) +// the Thrift transport applies at request time. Building the endpoint request +// lets ProxyFromEnvironment apply the NO_PROXY rules for this exact host, so the +// kernel sees the same effective proxy decision — returning "" (direct) when +// NO_PROXY excludes the host or no proxy is set. No extra dependency: this is the +// stdlib function the driver already relies on. +func proxyForEndpoint(cfg *config.Config) string { + return proxyForEndpointFunc(cfg, http.ProxyFromEnvironment) +} + +// proxyForEndpointFunc is the testable core: it builds the endpoint request and +// asks resolve for the proxy, returning "" (direct) on any error, no proxy, or an +// unbuildable endpoint. resolve is http.ProxyFromEnvironment in production; tests +// inject a deterministic resolver to exercise proxy-set / NO_PROXY / direct +// without depending on http.ProxyFromEnvironment's process-wide env caching. +func proxyForEndpointFunc(cfg *config.Config, resolve func(*http.Request) (*url.URL, error)) string { + endpoint, err := cfg.ToEndpointURL() + if err != nil { + return "" + } + req, err := http.NewRequest(http.MethodPost, endpoint, nil) + if err != nil { + return "" + } + proxyURL, err := resolve(req) + if err != nil || proxyURL == nil { + return "" + } + return proxyURL.String() +} diff --git a/kernel_proxy_test.go b/kernel_proxy_test.go new file mode 100644 index 00000000..c7f73b13 --- /dev/null +++ b/kernel_proxy_test.go @@ -0,0 +1,78 @@ +package dbsql + +import ( + "errors" + "net/http" + "net/url" + "testing" + + "github.com/databricks/databricks-sql-go/internal/config" +) + +// proxyForEndpointFunc maps the injected resolver's decision to a proxy URL +// string (or "" for direct), and returns "" for an unbuildable endpoint. The +// production path uses http.ProxyFromEnvironment, whose env is snapshotted once +// per process (sync.Once) and so can't be re-driven mid-test; the resolver seam +// lets us assert every branch deterministically. Each case mirrors an +// httpproxy-style outcome: proxy set for this host, host excluded by NO_PROXY, +// no proxy configured, and an unbuildable endpoint. Pure Go, so it runs in the +// default CGO_ENABLED=0 build (no kernel lib required). +func TestProxyForEndpoint(t *testing.T) { + validCfg := func() *config.Config { + c := config.WithDefaults() + c.Host = "my-workspace.databricks.com" + c.Port = 443 + c.HTTPPath = "/sql/1.0/warehouses/abc" + return c + } + proxyURL, _ := url.Parse("http://corp-proxy:3128") + + cases := []struct { + name string + cfg *config.Config + resolve func(*http.Request) (*url.URL, error) + want string + }{ + { + name: "proxy set for host", + cfg: validCfg(), + resolve: func(*http.Request) (*url.URL, error) { return proxyURL, nil }, + want: "http://corp-proxy:3128", + }, + { + name: "host excluded by NO_PROXY -> direct", + cfg: validCfg(), + resolve: func(*http.Request) (*url.URL, error) { return nil, nil }, + want: "", + }, + { + name: "resolver error -> direct", + cfg: validCfg(), + resolve: func(*http.Request) (*url.URL, error) { return nil, errors.New("bad proxy url") }, + want: "", + }, + { + name: "unbuildable endpoint -> direct (resolver never consulted)", + cfg: func() *config.Config { + c := config.WithDefaults() + c.Host = "" + return c + }(), + resolve: func(*http.Request) (*url.URL, error) { + t.Error("resolver must not be called for an unbuildable endpoint") + return proxyURL, nil + }, + want: "", + }, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + if got := proxyForEndpointFunc(tc.cfg, tc.resolve); got != tc.want { + t.Errorf("proxyForEndpointFunc = %q, want %q", got, tc.want) + } + }) + } + + // The production wrapper wires http.ProxyFromEnvironment and must not panic. + _ = proxyForEndpoint(validCfg()) +} From bf3b2f13fbfdfbe280368e0b81b7d12187c1718c Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Sat, 11 Jul 2026 08:58:51 +0000 Subject: [PATCH 22/32] fix(arrowbased): JSON-escape struct field-name keys + guard cross-backend parity (H1) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The nested-JSON grammar is rendered independently by the kernel path (internal/arrowscan) and the Thrift path (arrowbased), and the two had diverged on struct-key escaping: arrowscan does json.Marshal(fieldName) (escaped) while arrowbased did a raw `"` + name + `"` concat — invalid JSON for a field name containing a quote/backslash/control char (e.g. `a"b` → `{"a"b":1}`). Thrift was the buggy one; fix it to json.Marshal the key so both agree on valid escaped JSON. The only prior guard, TestKernelThriftParity, is build-tagged AND needs a live warehouse, so it never runs in CI and the divergence shipped unnoticed. Add an untagged cross-backend parity test (in package arrowbased, calling the private container factory and importing arrowscan as a pure leaf) that feeds the same arrow.Array through both renderers and asserts byte-equality — runs in the default CGO_ENABLED=0 matrix, no kernel lib. Covers special-char struct keys, string/ special/int map keys, float32, and decimal; verified it fails on the pre-fix concat and passes after. Map-key rendering was confirmed already-identical (different code, same bytes), so left unchanged. Full arrowbased single-sourcing (Thrift delegating to arrowscan) remains a separate, higher-risk follow-up; this closes the correctness + CI-guard gap. Co-authored-by: Isaac Signed-off-by: Mani Kaustubh Mathur --- .../rows/arrowbased/arrowscan_parity_test.go | 132 ++++++++++++++++++ internal/rows/arrowbased/columnValues.go | 6 +- 2 files changed, 137 insertions(+), 1 deletion(-) create mode 100644 internal/rows/arrowbased/arrowscan_parity_test.go diff --git a/internal/rows/arrowbased/arrowscan_parity_test.go b/internal/rows/arrowbased/arrowscan_parity_test.go new file mode 100644 index 00000000..431a6b8a --- /dev/null +++ b/internal/rows/arrowbased/arrowscan_parity_test.go @@ -0,0 +1,132 @@ +package arrowbased + +import ( + "testing" + "time" + + "github.com/apache/arrow/go/v12/arrow" + "github.com/apache/arrow/go/v12/arrow/array" + "github.com/apache/arrow/go/v12/arrow/decimal128" + "github.com/apache/arrow/go/v12/arrow/memory" + + "github.com/databricks/databricks-sql-go/internal/arrowscan" +) + +// The kernel backend (internal/arrowscan) and the Thrift backend (this package) +// each render Arrow cells to driver.Values — nested types to a JSON string — and +// the "identical results across backends" contract requires byte-for-byte +// agreement. That contract used to be guarded only by TestKernelThriftParity, +// which is build-tagged AND needs a live warehouse, so it never runs in CI. This +// test feeds the same arrow.Array through both renderers with no cgo and no +// warehouse, so a divergence (e.g. struct-key JSON escaping) fails a default +// CGO_ENABLED=0 run. +// +// It lives in package arrowbased because the container factory +// (makeColumnValueContainer) is package-private; arrowscan is a pure leaf import +// (no cycle). +func renderViaArrowbased(t *testing.T, arr arrow.Array, row int) any { + t.Helper() + maker := &arrowValueContainerMaker{} + holder, err := maker.makeColumnValueContainer(arr.DataType(), time.UTC, func(ts arrow.Timestamp) time.Time { + return ts.ToTime(arrow.Microsecond) + }, nil) + if err != nil { + t.Fatalf("makeColumnValueContainer(%s): %v", arr.DataType(), err) + } + if err := holder.SetValueArray(arr.Data()); err != nil { + t.Fatalf("SetValueArray(%s): %v", arr.DataType(), err) + } + if holder.IsNull(row) { + return nil + } + v, err := holder.Value(row) + if err != nil { + t.Fatalf("Value(%s): %v", arr.DataType(), err) + } + return v +} + +func TestArrowbasedKernelRenderParity(t *testing.T) { + pool := memory.NewGoAllocator() + + // structWithKey builds a single-row STRUCT whose field is `name`. + structWithKey := func(name string) arrow.Array { + dt := arrow.StructOf(arrow.Field{Name: name, Type: arrow.PrimitiveTypes.Int64}) + b := array.NewStructBuilder(pool, dt) + b.Append(true) + b.FieldBuilder(0).(*array.Int64Builder).Append(1) + return b.NewArray() + } + + cases := []struct { + name string + build func() arrow.Array + }{ + {"list_int", func() arrow.Array { + b := array.NewListBuilder(pool, arrow.PrimitiveTypes.Int64) + vb := b.ValueBuilder().(*array.Int64Builder) + b.Append(true) + vb.Append(1) + vb.Append(2) + return b.NewArray() + }}, + {"struct_simple", func() arrow.Array { return structWithKey("a") }}, + // The escaping divergence: a field name with a quote must render as valid, + // identically-escaped JSON on both backends. + {"struct_key_with_quote", func() arrow.Array { return structWithKey(`a"b`) }}, + {"struct_key_with_backslash", func() arrow.Array { return structWithKey(`a\b`) }}, + {"struct_key_with_newline", func() arrow.Array { return structWithKey("a\nb") }}, + {"map_string_key", func() arrow.Array { + b := array.NewMapBuilder(pool, arrow.BinaryTypes.String, arrow.PrimitiveTypes.Int64, false) + b.Append(true) // open row 0's map + b.KeyBuilder().(*array.StringBuilder).Append("k") + b.ItemBuilder().(*array.Int64Builder).Append(9) + return b.NewArray() + }}, + {"map_special_string_key", func() arrow.Array { + b := array.NewMapBuilder(pool, arrow.BinaryTypes.String, arrow.PrimitiveTypes.Int64, false) + b.Append(true) + b.KeyBuilder().(*array.StringBuilder).Append(`k"x`) + b.ItemBuilder().(*array.Int64Builder).Append(9) + return b.NewArray() + }}, + {"map_int_key", func() arrow.Array { + b := array.NewMapBuilder(pool, arrow.PrimitiveTypes.Int64, arrow.PrimitiveTypes.Int64, false) + b.Append(true) + b.KeyBuilder().(*array.Int64Builder).Append(7) + b.ItemBuilder().(*array.Int64Builder).Append(9) + return b.NewArray() + }}, + {"nested_float32", func() arrow.Array { + dt := arrow.StructOf(arrow.Field{Name: "f", Type: arrow.PrimitiveTypes.Float32}) + b := array.NewStructBuilder(pool, dt) + b.Append(true) + b.FieldBuilder(0).(*array.Float32Builder).Append(0.1) + return b.NewArray() + }}, + {"nested_decimal", func() arrow.Array { + dt := arrow.StructOf(arrow.Field{Name: "d", Type: &arrow.Decimal128Type{Precision: 5, Scale: 2}}) + b := array.NewStructBuilder(pool, dt) + b.Append(true) + b.FieldBuilder(0).(*array.Decimal128Builder).Append(decimal128.FromU64(1999)) + return b.NewArray() + }}, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + arr := tc.build() + defer arr.Release() + + kernel, err := arrowscan.ScanCell(arr, 0, time.UTC) + if err != nil { + t.Fatalf("arrowscan.ScanCell: %v", err) + } + thrift := renderViaArrowbased(t, arr, 0) + + if kernel != thrift { + t.Errorf("backend divergence for %s:\n kernel = %#v\n thrift = %#v", tc.name, kernel, thrift) + } + }) + } +} diff --git a/internal/rows/arrowbased/columnValues.go b/internal/rows/arrowbased/columnValues.go index 53bf8551..14e9786b 100644 --- a/internal/rows/arrowbased/columnValues.go +++ b/internal/rows/arrowbased/columnValues.go @@ -374,7 +374,11 @@ func (svc *structValueContainer) Value(i int) (any, error) { if i < svc.structArray.Len() { r := "{" for j := range svc.fieldValues { - r = r + "\"" + svc.fieldNames[j] + "\":" + // JSON-escape the field name — a raw `"` + name + `"` concat produces + // invalid JSON for a name containing a quote/backslash/control char + // (e.g. `a"b` → `{"a"b":...}`). json.Marshal a string never errors. + keyBytes, _ := json.Marshal(svc.fieldNames[j]) + r = r + string(keyBytes) + ":" if svc.fieldValues[j].IsNull(int(i)) { r = r + "null" From f90c230278450e2b0b456b5bdb53ea895be3ed32 Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Sat, 11 Jul 2026 11:20:12 +0000 Subject: [PATCH 23/32] =?UTF-8?q?docs(kernel):=20round-4=20review=20?= =?UTF-8?q?=E2=80=94=20doc=20caveats=20+=20Thrift=20grammar=20back-pointer?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit No behavior change. - doc.go: add WithEnableMetricViewMetadata to the connect-time reject list (L2); note INTERVAL is not handled by the kernel scanner and returns a scan error (L1); caveat that the kernel backend surfaces no per-statement query id, so QueryIdCallback fires with "" and no EXECUTE_STATEMENT metric is emitted (M6). - columnValues.go: add a back-pointer above structValueContainer.Value noting the grammar is mirrored by internal/arrowscan and guarded by the parity test (L3), so a maintainer editing the Thrift renderer gets the signal. Co-authored-by: Isaac Signed-off-by: Mani Kaustubh Mathur --- doc.go | 24 ++++++++++++++++-------- internal/rows/arrowbased/columnValues.go | 3 +++ 2 files changed, 19 insertions(+), 8 deletions(-) diff --git a/doc.go b/doc.go index e3c2abed..feab8247 100644 --- a/doc.go +++ b/doc.go @@ -198,19 +198,27 @@ 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) and initial catalog/schema are not -yet supported and return a clear error at connect time rather than being silently -ignored; 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 are likewise not yet supported and return a clear error at -execute time (they arrive per-query, not at connect). 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.) +connection options. OAuth (M2M/U2M), initial catalog/schema, and +WithEnableMetricViewMetadata are not yet supported and return a clear error at +connect time rather than being silently ignored; 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 are likewise not yet +supported and return a clear error at execute time (they arrive per-query, not at +connect). 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 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, so a +QueryIdCallback (see below) fires with "" and no EXECUTE_STATEMENT telemetry is +emitted for kernel queries. + # Programmatically Retrieving Connection and Query Id Use the driverctx package under driverctx/ctx.go to add callbacks to the query context to receive the connection id and query id. diff --git a/internal/rows/arrowbased/columnValues.go b/internal/rows/arrowbased/columnValues.go index 14e9786b..44d81550 100644 --- a/internal/rows/arrowbased/columnValues.go +++ b/internal/rows/arrowbased/columnValues.go @@ -370,6 +370,9 @@ type structValueContainer struct { var _ columnValues = (*structValueContainer)(nil) +// Value renders the struct at row i as a JSON object string. This grammar is +// mirrored by internal/arrowscan (the kernel backend); the two must render +// byte-identically — internal/rows/arrowbased/arrowscan_parity_test.go guards it. func (svc *structValueContainer) Value(i int) (any, error) { if i < svc.structArray.Len() { r := "{" From 04e12c50a8ef517e9c205e1d1c307ec0a781e4d2 Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Sat, 11 Jul 2026 11:21:56 +0000 Subject: [PATCH 24/32] perf(arrowbased): precompute escaped struct keys once, not per row (M1) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Round 3's struct-key escaping fix (json.Marshal the field name) was correct but ran in structValueContainer.Value, the per-row hot path of the DEFAULT Thrift backend every customer uses — re-marshaling the same constant field names N_rows × N_fields times (~5M marshals + allocs on a 5-field × 1M-row result). Field names are invariant across rows and the container is built once per result (fieldNames is batch-stable, unlike the kernel's per-batch-fresh *StructType), so precompute the JSON-escaped `"name":` prefixes into svc.fieldKeys at construction and concat them in Value. Byte-identical output — the cross-backend parity test still passes — with zero per-row marshal. Co-authored-by: Isaac Signed-off-by: Mani Kaustubh Mathur --- internal/rows/arrowbased/arrowRows.go | 8 ++++++++ internal/rows/arrowbased/columnValues.go | 12 +++++++----- 2 files changed, 15 insertions(+), 5 deletions(-) diff --git a/internal/rows/arrowbased/arrowRows.go b/internal/rows/arrowbased/arrowRows.go index b7a03fb5..f54e27c8 100644 --- a/internal/rows/arrowbased/arrowRows.go +++ b/internal/rows/arrowbased/arrowRows.go @@ -4,6 +4,7 @@ import ( "bytes" "context" "database/sql/driver" + "encoding/json" "io" "time" @@ -687,10 +688,17 @@ func (vcm *arrowValueContainerMaker) makeColumnValueContainer(t arrow.DataType, case *arrow.StructType: svc := &structValueContainer{structArrayType: t} svc.fieldNames = make([]string, len(t.Fields())) + svc.fieldKeys = make([]string, len(t.Fields())) svc.fieldValues = make([]columnValues, len(t.Fields())) svc.complexValue = make([]bool, len(t.Fields())) for i, f := range t.Fields() { svc.fieldNames[i] = f.Name + // Precompute the JSON-escaped `"name":` key once — field names are + // invariant across rows, so escaping them per row (structValueContainer + // .Value runs per row) would re-marshal constant strings. json.Marshal a + // string never errors. + keyBytes, _ := json.Marshal(f.Name) + svc.fieldKeys[i] = string(keyBytes) + ":" c, err := vcm.makeColumnValueContainer(f.Type, location, toTimestampFn, nil) if err != nil { return nil, err diff --git a/internal/rows/arrowbased/columnValues.go b/internal/rows/arrowbased/columnValues.go index 44d81550..0e47d105 100644 --- a/internal/rows/arrowbased/columnValues.go +++ b/internal/rows/arrowbased/columnValues.go @@ -363,6 +363,7 @@ func (mvc *mapValueContainer) SetValueArray(colData arrow.ArrayData) error { type structValueContainer struct { structArray *array.Struct fieldNames []string + fieldKeys []string // precomputed JSON-escaped `"name":` prefixes, one per field complexValue []bool fieldValues []columnValues structArrayType *arrow.StructType @@ -377,11 +378,12 @@ func (svc *structValueContainer) Value(i int) (any, error) { if i < svc.structArray.Len() { r := "{" for j := range svc.fieldValues { - // JSON-escape the field name — a raw `"` + name + `"` concat produces - // invalid JSON for a name containing a quote/backslash/control char - // (e.g. `a"b` → `{"a"b":...}`). json.Marshal a string never errors. - keyBytes, _ := json.Marshal(svc.fieldNames[j]) - r = r + string(keyBytes) + ":" + // fieldKeys[j] is the precomputed JSON-escaped `"name":` prefix (built + // once at construction). Escaping matters — a raw `"` + name + `"` concat + // produces invalid JSON for a name with a quote/backslash/control char + // (`a"b` → `{"a"b":...}`) — but doing it per row would re-marshal the same + // constant strings N_rows × N_fields times on this hot path. + r = r + svc.fieldKeys[j] if svc.fieldValues[j].IsNull(int(i)) { r = r + "null" From 4707b32c142c03fcbcf44def5f77df52ab3de659 Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Sat, 11 Jul 2026 11:23:43 +0000 Subject: [PATCH 25/32] fix(arrowscan): render map keys like Thrift (base64 []byte, not fmt %v) (M4) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit writeJSONKey's default arm did json.Marshal(fmt.Sprintf("%v", k)), so a []byte (BINARY) map key rendered as "[97 98 99]" while the Thrift path (marshalScalar → json.Marshal) renders base64 "YWJj" — a byte-parity violation for MAP (and large-magnitude float keys via %v vs marshal). Marshal the key value directly (matching Thrift), quote-wrapping only when the result isn't already a JSON string; keep the time.Time → quoted .String() special-case. Parity test gains MAP and MAP cases; verified map_binary_key fails on the old %v path ({"[97 98 99]":9} vs {"YWJj":9}) and passes after. Co-authored-by: Isaac Signed-off-by: Mani Kaustubh Mathur --- internal/arrowscan/arrowscan.go | 29 ++++++++++++++----- .../rows/arrowbased/arrowscan_parity_test.go | 15 ++++++++++ 2 files changed, 37 insertions(+), 7 deletions(-) diff --git a/internal/arrowscan/arrowscan.go b/internal/arrowscan/arrowscan.go index 6c481754..23207491 100644 --- a/internal/arrowscan/arrowscan.go +++ b/internal/arrowscan/arrowscan.go @@ -231,16 +231,31 @@ func writeStructJSON(b *strings.Builder, s *array.Struct, row int, loc *time.Loc return nil } -// writeJSONKey writes a value as a JSON object key (always a quoted string). +// writeJSONKey writes a value as a JSON object key (always a quoted string), +// matching the Thrift path's mapValueContainer: marshal the value the same way a +// leaf is marshaled (marshalScalar → json.Marshal, so a []byte key renders as +// base64 "YWJj", NOT fmt "%v" [97 98 99]), then quote-wrap the result if it is +// not already a JSON string (numbers/bools become "7"/"true" keys). func writeJSONKey(b *strings.Builder, v any) { - switch k := v.(type) { - case string: - kb, _ := json.Marshal(k) - b.Write(kb) - default: - kb, _ := json.Marshal(fmt.Sprintf("%v", k)) + // A time.Time key renders as its quoted .String() (same special-case as + // writeScalarJSON / the Thrift marshal()), not RFC3339. + if t, ok := v.(time.Time); ok { + kb, _ := json.Marshal(t.String()) b.Write(kb) + return + } + kb, err := json.Marshal(v) + if err != nil { + // Unmarshalable key value — fall back to its stringified form, quoted. + kb, _ = json.Marshal(fmt.Sprintf("%v", v)) + } + if len(kb) > 0 && kb[0] == '"' { + b.Write(kb) // already a JSON string (string / []byte via marshal) + return } + b.WriteByte('"') + b.Write(kb) + b.WriteByte('"') } // writeScalarJSON writes a scalar leaf, mirroring the Thrift marshal(): a diff --git a/internal/rows/arrowbased/arrowscan_parity_test.go b/internal/rows/arrowbased/arrowscan_parity_test.go index 431a6b8a..e6116ec3 100644 --- a/internal/rows/arrowbased/arrowscan_parity_test.go +++ b/internal/rows/arrowbased/arrowscan_parity_test.go @@ -97,6 +97,21 @@ func TestArrowbasedKernelRenderParity(t *testing.T) { b.ItemBuilder().(*array.Int64Builder).Append(9) return b.NewArray() }}, + {"map_binary_key", func() arrow.Array { + // []byte key: json.Marshal → base64 "YWJj", NOT fmt "%v" [97 98 99]. + b := array.NewMapBuilder(pool, arrow.BinaryTypes.Binary, arrow.PrimitiveTypes.Int64, false) + b.Append(true) + b.KeyBuilder().(*array.BinaryBuilder).Append([]byte("abc")) + b.ItemBuilder().(*array.Int64Builder).Append(9) + return b.NewArray() + }}, + {"map_date_key", func() arrow.Array { + b := array.NewMapBuilder(pool, arrow.FixedWidthTypes.Date32, arrow.PrimitiveTypes.Int64, false) + b.Append(true) + b.KeyBuilder().(*array.Date32Builder).Append(arrow.Date32FromTime(time.Date(2026, time.July, 9, 0, 0, 0, 0, time.UTC))) + b.ItemBuilder().(*array.Int64Builder).Append(9) + return b.NewArray() + }}, {"nested_float32", func() arrow.Array { dt := arrow.StructOf(arrow.Field{Name: "f", Type: arrow.PrimitiveTypes.Float32}) b := array.NewStructBuilder(pool, dt) From f276304bc6c206957b3265d6794492eb32835872 Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Sat, 11 Jul 2026 11:25:18 +0000 Subject: [PATCH 26/32] test(arrowbased): extend cross-backend parity to high-drift nested shapes (M5) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The parity test covered the flat surface + escaping/binary-key cases but omitted the shapes where two hand-maintained grammars are most likely to drift. Add recursive nesting (ARRAY, STRUCT, MAP<_,STRUCT>), a nested TIMESTAMP leaf (the time.Time → quoted .String() special-case, previously only tested top-level), and a null leaf inside a struct (vs a null container). Render through both backends under a non-UTC location so the timestamp .In(loc) path is actually exercised rather than a UTC no-op. All agree. Co-authored-by: Isaac Signed-off-by: Mani Kaustubh Mathur --- .../rows/arrowbased/arrowscan_parity_test.go | 69 +++++++++++++++++-- 1 file changed, 65 insertions(+), 4 deletions(-) diff --git a/internal/rows/arrowbased/arrowscan_parity_test.go b/internal/rows/arrowbased/arrowscan_parity_test.go index e6116ec3..1268e7f4 100644 --- a/internal/rows/arrowbased/arrowscan_parity_test.go +++ b/internal/rows/arrowbased/arrowscan_parity_test.go @@ -24,10 +24,10 @@ import ( // It lives in package arrowbased because the container factory // (makeColumnValueContainer) is package-private; arrowscan is a pure leaf import // (no cycle). -func renderViaArrowbased(t *testing.T, arr arrow.Array, row int) any { +func renderViaArrowbased(t *testing.T, arr arrow.Array, row int, loc *time.Location) any { t.Helper() maker := &arrowValueContainerMaker{} - holder, err := maker.makeColumnValueContainer(arr.DataType(), time.UTC, func(ts arrow.Timestamp) time.Time { + holder, err := maker.makeColumnValueContainer(arr.DataType(), loc, func(ts arrow.Timestamp) time.Time { return ts.ToTime(arrow.Microsecond) }, nil) if err != nil { @@ -126,6 +126,67 @@ func TestArrowbasedKernelRenderParity(t *testing.T) { b.FieldBuilder(0).(*array.Decimal128Builder).Append(decimal128.FromU64(1999)) return b.NewArray() }}, + // Highest-drift shapes: recursive nesting, a nested timestamp leaf (the + // time.Time → quoted .String() special-case), and a null leaf inside a + // container (vs a null container, which is already covered). + {"array_of_struct", func() arrow.Array { + elem := arrow.StructOf(arrow.Field{Name: "a", Type: arrow.PrimitiveTypes.Int64}) + b := array.NewListBuilder(pool, elem) + sb := b.ValueBuilder().(*array.StructBuilder) + b.Append(true) + sb.Append(true) + sb.FieldBuilder(0).(*array.Int64Builder).Append(1) + sb.Append(true) + sb.FieldBuilder(0).(*array.Int64Builder).Append(2) + return b.NewArray() + }}, + {"struct_of_list", func() arrow.Array { + dt := arrow.StructOf(arrow.Field{Name: "xs", Type: arrow.ListOf(arrow.PrimitiveTypes.Int64)}) + b := array.NewStructBuilder(pool, dt) + b.Append(true) + lb := b.FieldBuilder(0).(*array.ListBuilder) + lb.Append(true) + vb := lb.ValueBuilder().(*array.Int64Builder) + vb.Append(1) + vb.Append(2) + return b.NewArray() + }}, + {"map_of_struct_value", func() arrow.Array { + valT := arrow.StructOf(arrow.Field{Name: "a", Type: arrow.PrimitiveTypes.Int64}) + b := array.NewMapBuilder(pool, arrow.BinaryTypes.String, valT, false) + b.Append(true) + b.KeyBuilder().(*array.StringBuilder).Append("k") + sb := b.ItemBuilder().(*array.StructBuilder) + sb.Append(true) + sb.FieldBuilder(0).(*array.Int64Builder).Append(5) + return b.NewArray() + }}, + {"nested_timestamp", func() arrow.Array { + dt := arrow.StructOf(arrow.Field{Name: "ts", Type: &arrow.TimestampType{Unit: arrow.Microsecond}}) + b := array.NewStructBuilder(pool, dt) + b.Append(true) + ts := time.Date(2026, time.July, 9, 12, 34, 56, 0, time.UTC) + b.FieldBuilder(0).(*array.TimestampBuilder).Append(arrow.Timestamp(ts.UnixMicro())) + return b.NewArray() + }}, + {"null_leaf_in_struct", func() arrow.Array { + dt := arrow.StructOf( + arrow.Field{Name: "a", Type: arrow.PrimitiveTypes.Int64}, + arrow.Field{Name: "b", Type: arrow.PrimitiveTypes.Int64}, + ) + b := array.NewStructBuilder(pool, dt) + b.Append(true) + b.FieldBuilder(0).(*array.Int64Builder).Append(1) + b.FieldBuilder(1).(*array.Int64Builder).AppendNull() + return b.NewArray() + }}, + } + + // A non-UTC location, so the timestamp/date rendering path (both backends + // apply .In(loc)) is actually exercised rather than a UTC no-op. + loc, err := time.LoadLocation("America/New_York") + if err != nil { + t.Skipf("tz database unavailable: %v", err) } for _, tc := range cases { @@ -133,11 +194,11 @@ func TestArrowbasedKernelRenderParity(t *testing.T) { arr := tc.build() defer arr.Release() - kernel, err := arrowscan.ScanCell(arr, 0, time.UTC) + kernel, err := arrowscan.ScanCell(arr, 0, loc) if err != nil { t.Fatalf("arrowscan.ScanCell: %v", err) } - thrift := renderViaArrowbased(t, arr, 0) + thrift := renderViaArrowbased(t, arr, 0, loc) if kernel != thrift { t.Errorf("backend divergence for %s:\n kernel = %#v\n thrift = %#v", tc.name, kernel, thrift) From 08cbcfe1b45117b09c474c7ade260206b8bebc64 Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Sat, 11 Jul 2026 11:28:25 +0000 Subject: [PATCH 27/32] fix(kernel): never return ErrBadConn on the execute/read path (H1) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit toDriverError mapped a NetworkError/Unavailable status to driver.ErrBadConn on every path, including execute and result-read. On the statement path that is unsafe: a network failure surfaced *after* the SEA statement was sent may have committed server-side, and driver.ErrBadConn makes database/sql transparently re-run the whole statement — a silent duplicate write for a non-idempotent INSERT/UPDATE/MERGE. This violates Go's own ErrBadConn rule ("never return ErrBadConn if the server might have performed the operation"). Split classification by path: - toConnError (session lifecycle: open/close/config — nothing executed) keeps the bad-conn mapping so the pool evicts a dead session. Safe: no statement ran. - toStatementError (execute + result read) never returns ErrBadConn; it returns the KernelError unchanged (sqlstate preserved). This restores the contract the rest of the stack already honors: the kernel itself classifies ExecuteStatement as NonIdempotent (retried only on connect-phase failures, per src/client/retry.rs), and the Thrift backend marks ExecuteStatement non-retryable. The kernel has already exhausted its safe internal retries by the time Go sees the error, so a second database/sql-level retry was pure hazard. Tests: TestToStatementErrorNeverBadConn asserts no ErrBadConn for network/unavailable/unauthenticated on the statement path; TestToConnError keeps the session-path eviction behavior. Co-authored-by: Isaac Signed-off-by: Mani Kaustubh Mathur --- internal/backend/kernel/backend.go | 20 +++++------ internal/backend/kernel/cgo.go | 35 ++++++++++++++++---- internal/backend/kernel/kernel_test.go | 46 ++++++++++++++++++++------ internal/backend/kernel/operation.go | 10 +++--- internal/backend/kernel/rows.go | 4 +-- 5 files changed, 82 insertions(+), 33 deletions(-) diff --git a/internal/backend/kernel/backend.go b/internal/backend/kernel/backend.go index e0933f09..4726279a 100644 --- a/internal/backend/kernel/backend.go +++ b/internal/backend/kernel/backend.go @@ -83,7 +83,7 @@ func (k *KernelBackend) OpenSession(ctx context.Context) error { var cfg *C.KernelSessionConfig if err := call(func() C.KernelStatusCode { return C.kernel_session_config_new(&cfg) }); err != nil { - return fmt.Errorf("kernel: config_new: %w", toDriverError(err)) + return fmt.Errorf("kernel: config_new: %w", toConnError(err)) } // kernel_session_open consumes the config on EVERY path — success and // failure alike (it reclaims the box up front). So we free the config @@ -107,7 +107,7 @@ func (k *KernelBackend) OpenSession(ctx context.Context) error { if err := call(func() C.KernelStatusCode { return C.kernel_session_config_set_warehouse(cfg, host.c, wh.c) }); err != nil { - return fmt.Errorf("kernel: set_warehouse: %w", toDriverError(err)) + return fmt.Errorf("kernel: set_warehouse: %w", toConnError(err)) } } else { path := newCStr(k.cfg.HTTPPath) @@ -115,7 +115,7 @@ func (k *KernelBackend) OpenSession(ctx context.Context) error { if err := call(func() C.KernelStatusCode { return C.kernel_session_config_set_http_path(cfg, host.c, path.c) }); err != nil { - return fmt.Errorf("kernel: set_http_path: %w", toDriverError(err)) + return fmt.Errorf("kernel: set_http_path: %w", toConnError(err)) } } @@ -124,7 +124,7 @@ func (k *KernelBackend) OpenSession(ctx context.Context) error { 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", toDriverError(err)) + return fmt.Errorf("kernel: set_auth_pat: %w", toConnError(err)) } // TLS: crypto/tls's InsecureSkipVerify accepts any server cert, so relax both @@ -135,12 +135,12 @@ func (k *KernelBackend) OpenSession(ctx context.Context) error { if err := call(func() C.KernelStatusCode { return C.kernel_session_config_set_tls_allow_self_signed(cfg, C.bool(true)) }); err != nil { - return fmt.Errorf("kernel: set_tls_allow_self_signed: %w", toDriverError(err)) + return fmt.Errorf("kernel: set_tls_allow_self_signed: %w", toConnError(err)) } 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", toDriverError(err)) + return fmt.Errorf("kernel: set_tls_skip_hostname_verification: %w", toConnError(err)) } } @@ -154,7 +154,7 @@ func (k *KernelBackend) OpenSession(ctx context.Context) error { if err := call(func() C.KernelStatusCode { return C.kernel_session_config_set_proxy(cfg, url.c, nil, nil, nil) }); err != nil { - return fmt.Errorf("kernel: set_proxy: %w", toDriverError(err)) + return fmt.Errorf("kernel: set_proxy: %w", toConnError(err)) } } @@ -169,7 +169,7 @@ func (k *KernelBackend) OpenSession(ctx context.Context) error { ck.free() cv.free() if errSet != nil { - return fmt.Errorf("kernel: set_session_conf[%s]: %w", key, toDriverError(errSet)) + return fmt.Errorf("kernel: set_session_conf[%s]: %w", key, toConnError(errSet)) } } @@ -180,7 +180,7 @@ func (k *KernelBackend) OpenSession(ctx context.Context) error { err := call(func() C.KernelStatusCode { return C.kernel_session_open(cfg, &sess) }) consumed = true if err != nil { - return fmt.Errorf("kernel: session_open: %w", toDriverError(err)) + return fmt.Errorf("kernel: session_open: %w", toConnError(err)) } k.session = sess k.valid = true @@ -201,7 +201,7 @@ func (k *KernelBackend) CloseSession(ctx context.Context) error { err := call(func() C.KernelStatusCode { return C.kernel_session_close(k.session) }) k.session = nil k.valid = false - return toDriverError(err) + return toConnError(err) } // SessionValid backs conn.IsValid → pool eviction. No I/O; inspects state diff --git a/internal/backend/kernel/cgo.go b/internal/backend/kernel/cgo.go index ca83080b..f5d5f4d6 100644 --- a/internal/backend/kernel/cgo.go +++ b/internal/backend/kernel/cgo.go @@ -202,12 +202,13 @@ func isBadConnection(code int) bool { } } -// toDriverError classifies a kernel error for the pool: a status that means the -// session is unusable is wrapped as a bad-connection error (which identifies as -// driver.ErrBadConn, so database/sql evicts the conn), matching how the Thrift -// backend surfaces connection loss. Other kernel errors — and plain -// (non-KernelError) errors — are returned unchanged, carrying their sqlstate. -func toDriverError(err error) error { +// toConnError classifies a kernel error on a SESSION-lifecycle path (open/close/ +// config, where nothing has executed): a status that means the session is unusable +// is wrapped as a bad-connection error, which identifies as driver.ErrBadConn so +// database/sql evicts the conn from the pool. Safe here because no statement ran, +// so there is nothing for database/sql to unsafely re-run. Other kernel errors — +// and plain (non-KernelError) errors — are returned unchanged, carrying sqlstate. +func toConnError(err error) error { if err == nil { return nil } @@ -221,6 +222,28 @@ func toDriverError(err error) error { return ke } +// toStatementError classifies a kernel error on the STATEMENT path (execute and +// result read). It NEVER returns driver.ErrBadConn: once a statement has been +// sent, a network/unavailable failure surfaced afterward may have committed +// server-side, and driver.ErrBadConn would make database/sql transparently +// re-run the statement — a silent duplicate write for a non-idempotent +// INSERT/UPDATE/MERGE. This mirrors the kernel's own retry contract +// (ExecuteStatement is NonIdempotent, retried only on connect-phase failures) and +// the Thrift backend (ExecuteStatement is non-retryable), and honors Go's +// driver.ErrBadConn rule ("never return ErrBadConn if the server might have +// performed the operation"). The kernel has already exhausted its safe internal +// retries by the time we see the error. Returns the KernelError (or plain error) +// unchanged, carrying sqlstate. +func toStatementError(err error) error { + if err == nil { + return nil + } + if ke, ok := err.(*KernelError); ok { + return ke + } + return err +} + // cStr wraps C.CString with a guaranteed free. The kernel copies strings into // owned Rust memory on receipt, so freeing immediately after the call is safe. // Use: cs := newCStr(s); defer cs.free(); ...C.fn(cs.c)... diff --git a/internal/backend/kernel/kernel_test.go b/internal/backend/kernel/kernel_test.go index 6aae1256..9a72fbe7 100644 --- a/internal/backend/kernel/kernel_test.go +++ b/internal/backend/kernel/kernel_test.go @@ -28,24 +28,50 @@ func TestIsBadConnection(t *testing.T) { } } -// toDriverError wraps a session-unusable KernelError as driver.ErrBadConn (so -// database/sql evicts the conn) and leaves other errors, and their sqlstate, -// intact. -func TestToDriverError(t *testing.T) { - if toDriverError(nil) != nil { +// toConnError (session-lifecycle path) wraps a session-unusable KernelError as +// driver.ErrBadConn so database/sql evicts the conn, and leaves other errors and +// their sqlstate intact. +func TestToConnError(t *testing.T) { + if toConnError(nil) != nil { t.Fatal("nil should map to nil") } badConn := &KernelError{Code: statusUnavailable, Message: "gone"} - if !errors.Is(toDriverError(badConn), driver.ErrBadConn) { - t.Errorf("unavailable kernel error should identify as driver.ErrBadConn") + if !errors.Is(toConnError(badConn), driver.ErrBadConn) { + t.Errorf("unavailable kernel error on the session path should identify as driver.ErrBadConn") } sqlErr := &KernelError{Code: statusSqlError, Message: "boom", SQLState: "42703"} - got := toDriverError(sqlErr) - ke, ok := got.(*KernelError) + ke, ok := toConnError(sqlErr).(*KernelError) if !ok { - t.Fatalf("sql error should remain a *KernelError, got %T", got) + t.Fatalf("sql error should remain a *KernelError, got %T", toConnError(sqlErr)) + } + if ke.SQLState != "42703" { + t.Errorf("sqlstate lost: got %q", ke.SQLState) + } +} + +// toStatementError (execute/read path) must NEVER return driver.ErrBadConn — even +// for a network/unavailable status — so database/sql cannot transparently re-run a +// statement that may have already executed server-side (silent duplicate write). +func TestToStatementErrorNeverBadConn(t *testing.T) { + if toStatementError(nil) != nil { + t.Fatal("nil should map to nil") + } + + for _, code := range []int{statusUnavailable, statusNetworkError, statusUnauthenticated} { + err := toStatementError(&KernelError{Code: code, Message: "post-execute failure"}) + if errors.Is(err, driver.ErrBadConn) { + t.Errorf("statement-path error (code=%d) must NOT identify as driver.ErrBadConn "+ + "(would let database/sql re-run a possibly-committed statement)", code) + } + } + + // sqlstate still preserved. + sqlErr := &KernelError{Code: statusSqlError, Message: "boom", SQLState: "42703"} + ke, ok := toStatementError(sqlErr).(*KernelError) + if !ok { + t.Fatalf("sql error should remain a *KernelError, got %T", toStatementError(sqlErr)) } if ke.SQLState != "42703" { t.Errorf("sqlstate lost: got %q", ke.SQLState) diff --git a/internal/backend/kernel/operation.go b/internal/backend/kernel/operation.go index 8531859d..b1a4d19c 100644 --- a/internal/backend/kernel/operation.go +++ b/internal/backend/kernel/operation.go @@ -40,7 +40,7 @@ func (k *KernelBackend) execute(ctx context.Context, req backend.ExecRequest) (b if err := call(func() C.KernelStatusCode { return C.kernel_session_new_statement(k.session, &stmt) }); err != nil { - return &kernelOp{}, fmt.Errorf("kernel: new_statement: %w", toDriverError(err)) + return &kernelOp{}, fmt.Errorf("kernel: new_statement: %w", toStatementError(err)) } sql := newCStr(req.Query) @@ -49,7 +49,7 @@ func (k *KernelBackend) execute(ctx context.Context, req backend.ExecRequest) (b }); err != nil { sql.free() C.kernel_statement_close(stmt) - return &kernelOp{}, fmt.Errorf("kernel: set_sql: %w", toDriverError(err)) + return &kernelOp{}, fmt.Errorf("kernel: set_sql: %w", toStatementError(err)) } sql.free() @@ -131,7 +131,7 @@ func (k *KernelBackend) execute(ctx context.Context, req backend.ExecRequest) (b } klog("Execute failed: %v", execErr) op.close() - return op, fmt.Errorf("kernel: execute: %w", toDriverError(execErr)) + 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 @@ -187,7 +187,7 @@ func (o *kernelOp) Results(ctx context.Context, callbacks *dbsqlrows.TelemetryCa // Operation.Close on a Results error — so close the handles here to avoid // leaking the statement / executed handle (and its server operation). o.close() - return nil, fmt.Errorf("kernel: get_result_stream: %w", toDriverError(err)) + return nil, fmt.Errorf("kernel: get_result_stream: %w", toStatementError(err)) } return newKernelRows(ctx, o, stream, callbacks) } @@ -235,5 +235,5 @@ func (o *kernelOp) close() bool { // already carries the sqlstate (see KernelError), so this returns cause as-is // (nil when cause is nil), matching the neutral contract. func (o *kernelOp) ExecutionError(ctx context.Context, cause error) error { - return toDriverError(cause) + return toStatementError(cause) } diff --git a/internal/backend/kernel/rows.go b/internal/backend/kernel/rows.go index 70f05326..973fb5d4 100644 --- a/internal/backend/kernel/rows.go +++ b/internal/backend/kernel/rows.go @@ -59,7 +59,7 @@ func newKernelRows(ctx context.Context, op *kernelOp, stream *C.kernel_result_st return C.kernel_result_stream_get_schema(stream, &csch) }); err != nil { r.Close() - return nil, fmt.Errorf("kernel: get_schema: %w", toDriverError(err)) + return nil, fmt.Errorf("kernel: get_schema: %w", toStatementError(err)) } sch, err := cdata.ImportCArrowSchema((*cdata.CArrowSchema)(unsafe.Pointer(&csch))) if err != nil { @@ -151,7 +151,7 @@ func (r *kernelRows) nextBatch() error { if err := call(func() C.KernelStatusCode { return C.kernel_result_stream_next_batch(r.stream, &carr, &csch) }); err != nil { - return fmt.Errorf("kernel: next_batch: %w", toDriverError(err)) + return fmt.Errorf("kernel: next_batch: %w", toStatementError(err)) } if carr.release == nil { r.eof = true From b2f3b575ec0b376961be80d6f16fbd7283954180 Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Sat, 11 Jul 2026 11:47:06 +0000 Subject: [PATCH 28/32] refactor(kernel): untag config validation + guard against dropped Config fields (M2/M3) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The round-3 auth/timeout/retry guards lived in newKernelBackend (cgo-tagged), and every test asserting them was tagged too — so CI (CGO_ENABLED=0) never compiled or ran them, and a config-field rename would silently delete the guards with make test still green (M2). Extract the pure-Go validation into an untagged kernel_config.go (validateKernelConfig: reject unsupported options + resolve the PAT). The tagged newKernelBackend now calls it, then assembles the cgo kernel.Config. The full reject/auth assertions move to an untagged kernel_config_test.go and run in the default build; the tagged test keeps only a smoke check that the wrapper wires validation + assembly. Add TestKernelConfigFieldsClassified (M3): a reflective test over config.UserConfig that fails if any field is unclassified (forwarded / rejected / inert), so a future WithFoo mapping to a wire setting can't be silently ignored on the kernel path — the divergence class rounds 2–4 closed. It already earned its keep by catching that TLSConfig lives on the outer Config, not UserConfig. Co-authored-by: Isaac Signed-off-by: Mani Kaustubh Mathur --- kernel_backend.go | 68 ++------------- kernel_backend_test.go | 103 +++-------------------- kernel_config.go | 81 ++++++++++++++++++ kernel_config_test.go | 186 +++++++++++++++++++++++++++++++++++++++++ 4 files changed, 285 insertions(+), 153 deletions(-) create mode 100644 kernel_config.go create mode 100644 kernel_config_test.go diff --git a/kernel_backend.go b/kernel_backend.go index 500a3f4b..412ffe78 100644 --- a/kernel_backend.go +++ b/kernel_backend.go @@ -4,10 +4,7 @@ package dbsql import ( "context" - "errors" - "github.com/databricks/databricks-sql-go/auth/noop" - "github.com/databricks/databricks-sql-go/auth/pat" "github.com/databricks/databricks-sql-go/internal/backend" "github.com/databricks/databricks-sql-go/internal/backend/kernel" "github.com/databricks/databricks-sql-go/internal/config" @@ -19,64 +16,13 @@ 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) { - // A few options aren't wired for the kernel backend yet. Fail loudly rather - // than silently ignore them (which would behave differently than Thrift): - // - Catalog/Schema (WithInitialNamespace): no kernel C-ABI setter yet, so - // the session would run in the default namespace and unqualified names - // would resolve differently. - // - EnableMetricViewMetadata: deferred — it maps to a server session conf, - // which we want to route backend-neutrally rather than duplicate here. - if cfg.Catalog != "" || cfg.Schema != "" { - return nil, errors.New("databricks: WithInitialNamespace (catalog/schema) is not yet supported by the kernel backend; " + - "omit it or use the default (Thrift) backend") - } - if cfg.EnableMetricViewMetadata { - return nil, errors.New("databricks: WithEnableMetricViewMetadata is not yet supported by the kernel backend; " + - "omit it or use the default (Thrift) backend") - } - // Auth: the kernel backend authenticates with a PAT only (kc.Token below). - // Any other authenticator — OAuth M2M/U2M, a token provider, external/static - // token, federated — 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, per - // the doc.go contract. nil / NoopAuth / PATAuth are the PAT-or-none cases. - 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 nil, errors.New("databricks: only personal access token (WithAccessToken) auth is supported by the kernel backend; " + - "OAuth (M2M/U2M), token-provider, external/static, and federated authenticators are not yet supported — " + - "use PAT or the default (Thrift) backend") - } - 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)") - } - // WithTimeout maps to a per-statement server timeout on Thrift - // (TExecuteStatementReq.QueryTimeout); the kernel C ABI exposes no equivalent, - // so honor the "nothing silently ignored" contract by rejecting it rather than - // running the query with no server-side timeout. (Default 0 = unset.) - if cfg.QueryTimeout > 0 { - return nil, errors.New("databricks: WithTimeout (server query timeout) is not yet supported by the kernel backend; " + - "omit it or use the default (Thrift) backend") - } - // WithRetries(-1) explicitly disables retries, but the kernel retries - // internally below the C ABI with no user-facing toggle — so a disable request - // 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, errors.New("databricks: disabling retries via WithRetries is not supported by the kernel backend " + - "(the kernel retries internally); omit it or use the default (Thrift) backend") + // Reject options the kernel path can't honor yet + resolve the PAT. 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) + if err != nil { + return nil, err } kc := kernel.Config{ diff --git a/kernel_backend_test.go b/kernel_backend_test.go index 6fffa44e..eb3244c4 100644 --- a/kernel_backend_test.go +++ b/kernel_backend_test.go @@ -4,24 +4,17 @@ package dbsql import ( "context" - "net/http" "testing" - "time" - "github.com/databricks/databricks-sql-go/auth/pat" "github.com/databricks/databricks-sql-go/internal/config" ) -// nonPATAuth is a stand-in for any non-PAT authenticator (OAuth / token-provider / -// external / federated) — the kernel backend must reject it at connect. -type nonPATAuth struct{} - -func (nonPATAuth) Authenticate(*http.Request) error { return nil } - -// newKernelBackend rejects options it can't yet honor (initial namespace, -// metric-view metadata) loudly, rather than silently ignoring them — which would -// behave differently than the Thrift backend. -func TestNewKernelBackendRejectsUnsupportedOptions(t *testing.T) { +// newKernelBackend builds a backend for a supported config and propagates a +// validation error otherwise. The exhaustive per-option reject/auth assertions +// live in the untagged TestValidateKernelConfig (kernel_config_test.go), which +// runs in the default CGO_ENABLED=0 build; this tagged smoke test just confirms +// the tagged wrapper wires validateKernelConfig + the kernel.Config assembly. +func TestNewKernelBackend(t *testing.T) { base := func() *config.Config { c := config.WithDefaults() c.Host = "h.databricks.com" @@ -31,91 +24,17 @@ func TestNewKernelBackendRejectsUnsupportedOptions(t *testing.T) { return c } - t.Run("catalog rejected", func(t *testing.T) { - c := base() - c.Catalog = "main" - if _, err := newKernelBackend(context.Background(), c); err == nil { - t.Error("expected an error when a catalog is set on the kernel backend") - } - }) - - t.Run("schema rejected", func(t *testing.T) { - c := base() - c.Schema = "default" - if _, err := newKernelBackend(context.Background(), c); err == nil { - t.Error("expected an error when a schema is set on the kernel backend") - } - }) - - t.Run("metric view rejected", func(t *testing.T) { - c := base() - c.EnableMetricViewMetadata = true - if _, err := newKernelBackend(context.Background(), c); err == nil { - t.Error("expected an error when metric-view metadata is enabled on the kernel backend") - } - }) - - t.Run("supported options ok", func(t *testing.T) { - c := base() - c.SessionParams = map[string]string{"QUERY_TAGS": "a:1"} - if _, err := newKernelBackend(context.Background(), c); err != nil { + t.Run("supported config builds", func(t *testing.T) { + if _, err := newKernelBackend(context.Background(), base()); err != nil { t.Errorf("a supported config should build cleanly, got %v", err) } }) - t.Run("PAT via WithAuthenticator ok", func(t *testing.T) { - // WithAuthenticator(&pat.PATAuth{...}) sets only cfg.Authenticator, leaving - // cfg.AccessToken empty. The backend must still build (token sourced from - // the authenticator), not fail with an empty PAT at connect. + t.Run("validation error propagates", func(t *testing.T) { c := base() - c.AccessToken = "" - c.Authenticator = &pat.PATAuth{AccessToken: "dapi-y"} - if _, err := newKernelBackend(context.Background(), c); err != nil { - t.Errorf("PAT supplied via WithAuthenticator should build cleanly, got %v", err) - } - }) - - t.Run("empty token rejected", func(t *testing.T) { - c := base() - c.AccessToken = "" - c.Authenticator = &pat.PATAuth{AccessToken: ""} - if _, err := newKernelBackend(context.Background(), c); err == nil { - t.Error("expected an error when the resolved PAT is empty") - } - }) - - t.Run("non-PAT authenticator rejected", func(t *testing.T) { - c := base() - c.Authenticator = nonPATAuth{} - if _, err := newKernelBackend(context.Background(), c); err == nil { - t.Error("expected an error for a non-PAT authenticator") - } - }) - - t.Run("query timeout rejected", func(t *testing.T) { - c := base() - c.QueryTimeout = 30 * time.Second - if _, err := newKernelBackend(context.Background(), c); err == nil { - t.Error("expected an error when WithTimeout (query timeout) is set on the kernel backend") - } - }) - - t.Run("disabling retries rejected", func(t *testing.T) { - c := base() - c.RetryMax = -1 + c.Catalog = "main" // rejected by validateKernelConfig if _, err := newKernelBackend(context.Background(), c); err == nil { - t.Error("expected an error when retries are disabled (WithRetries(-1)) on the kernel backend") - } - }) - - t.Run("positive retry tuning + maxrows accepted", func(t *testing.T) { - // These are accepted (not applied) per doc.go — the kernel manages fetch / - // retries internally, so they must not error. - c := base() - c.RetryMax = 8 - c.MaxRows = 5000 - if _, err := newKernelBackend(context.Background(), c); err != nil { - t.Errorf("positive retry/maxrows tuning should build cleanly, got %v", err) + t.Error("newKernelBackend should propagate the validation error") } }) } diff --git a/kernel_config.go b/kernel_config.go new file mode 100644 index 00000000..28cb35ad --- /dev/null +++ b/kernel_config.go @@ -0,0 +1,81 @@ +package dbsql + +import ( + "errors" + + "github.com/databricks/databricks-sql-go/auth/noop" + "github.com/databricks/databricks-sql-go/auth/pat" + "github.com/databricks/databricks-sql-go/internal/config" +) + +// This file is intentionally NOT behind the `cgo && databricks_kernel` build tag. +// The kernel backend's option-validation is pure Go (it reads config.Config and +// returns an error or a resolved PAT), so keeping it untagged lets its tests — +// including the reflective exhaustiveness check that guards against a future +// Config field being silently dropped — run under CGO_ENABLED=0. The tagged +// newKernelBackend calls validateKernelConfig, then assembles the cgo kernel.Config. + +// 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 +// intentionally accepted-but-inert (documented in doc.go and asserted by +// TestKernelConfigFieldsClassified). +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 "", errors.New("databricks: WithInitialNamespace (catalog/schema) is not yet supported by the kernel backend; " + + "omit it or use the default (Thrift) backend") + } + // EnableMetricViewMetadata: maps to a server session conf we want to route + // backend-neutrally rather than duplicate here. + if cfg.EnableMetricViewMetadata { + return "", errors.New("databricks: WithEnableMetricViewMetadata is not yet supported by the kernel backend; " + + "omit it or use the default (Thrift) backend") + } + // 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 "", errors.New("databricks: only personal access token (WithAccessToken) auth is supported by the kernel backend; " + + "OAuth (M2M/U2M), token-provider, external/static, and federated authenticators are not yet supported — " + + "use PAT or the default (Thrift) backend") + } + if token == "" { + return "", errors.New("databricks: the kernel backend requires a personal access token; " + + "set one with WithAccessToken (or a *pat.PATAuth via WithAuthenticator)") + } + // 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 "", errors.New("databricks: WithTimeout (server query timeout) is not yet supported by the kernel backend; " + + "omit it or use the default (Thrift) backend") + } + // WithRetries(-1) explicitly disables retries, but the kernel retries + // internally below the C ABI with no user-facing toggle — so a disable request + // 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 "", errors.New("databricks: disabling retries via WithRetries is not supported by the kernel backend " + + "(the kernel retries internally); omit it or use the default (Thrift) backend") + } + return token, nil +} diff --git a/kernel_config_test.go b/kernel_config_test.go new file mode 100644 index 00000000..729428b0 --- /dev/null +++ b/kernel_config_test.go @@ -0,0 +1,186 @@ +package dbsql + +import ( + "net/http" + "reflect" + "testing" + "time" + + "github.com/databricks/databricks-sql-go/auth/pat" + "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. +type nonPATAuth struct{} + +func (nonPATAuth) Authenticate(*http.Request) error { return nil } + +func baseKernelConfig() *config.Config { + c := config.WithDefaults() + c.Host = "h.databricks.com" + c.Port = 443 + c.HTTPPath = "/sql/1.0/warehouses/abc" + c.AccessToken = "dapi-x" + return c +} + +// validateKernelConfig enforces the kernel backend's "nothing silently ignored" +// contract: unsupported options are rejected loudly and the PAT is resolved. This +// is pure Go (no cgo), so these run in the default CGO_ENABLED=0 build — the +// round-4 fix, since the guards previously lived only in tag-gated code. +func TestValidateKernelConfig(t *testing.T) { + t.Run("supported config ok", func(t *testing.T) { + c := baseKernelConfig() + c.SessionParams = map[string]string{"QUERY_TAGS": "a:1"} + if _, err := validateKernelConfig(c); err != nil { + t.Errorf("a supported config should validate, got %v", err) + } + }) + + t.Run("catalog rejected", func(t *testing.T) { + c := baseKernelConfig() + c.Catalog = "main" + if _, err := validateKernelConfig(c); err == nil { + t.Error("expected an error when a catalog is set") + } + }) + + t.Run("schema rejected", func(t *testing.T) { + c := baseKernelConfig() + c.Schema = "sys" + if _, err := validateKernelConfig(c); err == nil { + t.Error("expected an error when a schema is set") + } + }) + + t.Run("metric view rejected", func(t *testing.T) { + c := baseKernelConfig() + c.EnableMetricViewMetadata = true + if _, err := validateKernelConfig(c); err == nil { + t.Error("expected an error when metric-view metadata is enabled") + } + }) + + t.Run("PAT via WithAuthenticator resolves the token", func(t *testing.T) { + c := baseKernelConfig() + c.AccessToken = "" + c.Authenticator = &pat.PATAuth{AccessToken: "dapi-y"} + tok, 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) + } + }) + + t.Run("empty token rejected", func(t *testing.T) { + c := baseKernelConfig() + c.AccessToken = "" + c.Authenticator = &pat.PATAuth{AccessToken: ""} + if _, err := validateKernelConfig(c); err == nil { + t.Error("expected an error when the resolved PAT is empty") + } + }) + + t.Run("non-PAT authenticator rejected", func(t *testing.T) { + c := baseKernelConfig() + c.Authenticator = nonPATAuth{} + if _, err := validateKernelConfig(c); err == nil { + t.Error("expected an error for a non-PAT authenticator") + } + }) + + t.Run("query timeout rejected", func(t *testing.T) { + c := baseKernelConfig() + c.QueryTimeout = 30 * time.Second + if _, err := validateKernelConfig(c); err == nil { + t.Error("expected an error when WithTimeout (query timeout) is set") + } + }) + + t.Run("disabling retries rejected", func(t *testing.T) { + c := baseKernelConfig() + c.RetryMax = -1 + if _, err := validateKernelConfig(c); err == nil { + t.Error("expected an error when retries are disabled (WithRetries(-1))") + } + }) + + t.Run("positive retry tuning + maxrows accepted", func(t *testing.T) { + c := baseKernelConfig() + c.RetryMax = 8 + c.MaxRows = 5000 + if _, err := validateKernelConfig(c); err != nil { + t.Errorf("positive retry/maxrows tuning should validate, got %v", err) + } + }) +} + +// kernelConfigFieldDisposition records, for every UserConfig field, how the kernel +// backend treats it. UserConfig is the user-facing option surface ("Only +// UserConfig are currently exposed to users"). Adding a field to config.UserConfig +// without classifying it here fails TestKernelConfigFieldsClassified — forcing a +// deliberate decision (forward it, reject it loudly, or accept it as inert) rather +// than silently dropping it on the kernel path, the exact divergence rounds 2–4 +// closed. (TLSConfig and ArrowConfig live on the outer config.Config, not +// UserConfig; newKernelBackend reads TLSConfig.InsecureSkipVerify explicitly, and +// the kernel renders decimals exactly regardless of ArrowConfig.) +var kernelConfigFieldDisposition = map[string]string{ + // Forwarded to kernel.Config (see newKernelBackend). + "Host": "forwarded", + "HTTPPath": "forwarded", + "WarehouseID": "forwarded", + "AccessToken": "forwarded", // as the resolved PAT (kc.Token) + "Authenticator": "forwarded", // PAT authenticator resolved to the token + "Location": "forwarded", + "SessionParams": "forwarded", + "Port": "forwarded", // part of the endpoint URL (host:port) + "Protocol": "forwarded", // part of the endpoint URL (scheme) + "UseKernel": "forwarded", // the routing flag itself + + // Rejected loudly by validateKernelConfig. + "Catalog": "rejected", + "Schema": "rejected", + "EnableMetricViewMetadata": "rejected", + "QueryTimeout": "rejected", // when > 0 (WithTimeout) + "RetryMax": "rejected", // when < 0 (disable retries) + + // 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. + "MaxRows": "inert", + "RetryWaitMin": "inert", + "RetryWaitMax": "inert", + "UseLz4Compression": "inert", // kernel negotiates compression internally + + // Not applicable to the kernel path (Thrift/HTTP-transport or telemetry knobs + // that don't reach the kernel binding). + "UserAgentEntry": "inert", // TODO: forward once the kernel exposes a UA setter + "Transport": "inert", // custom RoundTripper; kernel uses its own HTTP stack + "EnableTelemetry": "inert", + "TelemetryBatchSize": "inert", + "TelemetryFlushInterval": "inert", + "UseArrowNativeDecimalDSN": "inert", // DSN carrier; kernel renders decimals exactly regardless + "CloudFetchConfig": "inert", // kernel does CloudFetch internally +} + +func TestKernelConfigFieldsClassified(t *testing.T) { + ut := reflect.TypeOf(config.UserConfig{}) + for i := 0; i < ut.NumField(); i++ { + name := ut.Field(i).Name + if _, ok := kernelConfigFieldDisposition[name]; !ok { + t.Errorf("config.UserConfig field %q is not classified for the kernel backend. "+ + "Add it to kernelConfigFieldDisposition as forwarded/rejected/inert (and wire it in "+ + "validateKernelConfig / newKernelBackend if it must be honored) so it isn't silently "+ + "dropped on the kernel path.", name) + } + } + // Guard the reverse too: a disposition entry for a field that no longer exists + // is stale. + for name := range kernelConfigFieldDisposition { + if _, ok := ut.FieldByName(name); !ok { + t.Errorf("kernelConfigFieldDisposition has %q but config.UserConfig no longer does; remove it", name) + } + } +} From e39484843743c3bbfb94b3c81701e51318c25fdc Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Sat, 11 Jul 2026 16:18:35 +0000 Subject: [PATCH 29/32] test(kernel): flatten embedded structs in the config drop-guard (M2-r5) + comment cleanup (L2-r5) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit TestKernelConfigFieldsClassified used reflect NumField(), which reports the embedded CloudFetchConfig as a single field — so its six promoted fields (UseCloudFetch, MaxDownloadThreads, …) were invisible to the "no silent drop" guard. A future WithCloudFetchX would slip through with the test green, the exact divergence the guard exists to prevent. Recurse into anonymous struct fields so each promoted field is classified individually; classify the six CloudFetch fields as inert (the kernel does CloudFetch internally). Verified the guard now fails when any one is unclassified. Also drop review-round narration ("round-4", "rounds 2–4") from the file's comments per the comment-style ruleset, keeping the design rationale. Co-authored-by: Isaac Signed-off-by: Mani Kaustubh Mathur --- kernel_config_test.go | 54 ++++++++++++++++++++++++++++++++----------- 1 file changed, 40 insertions(+), 14 deletions(-) diff --git a/kernel_config_test.go b/kernel_config_test.go index 729428b0..9e1ff642 100644 --- a/kernel_config_test.go +++ b/kernel_config_test.go @@ -27,8 +27,7 @@ func baseKernelConfig() *config.Config { // validateKernelConfig enforces the kernel backend's "nothing silently ignored" // contract: unsupported options are rejected loudly and the PAT is resolved. This -// is pure Go (no cgo), so these run in the default CGO_ENABLED=0 build — the -// round-4 fix, since the guards previously lived only in tag-gated code. +// is pure Go (no cgo), so these run in the default CGO_ENABLED=0 build. func TestValidateKernelConfig(t *testing.T) { t.Run("supported config ok", func(t *testing.T) { c := baseKernelConfig() @@ -123,8 +122,8 @@ func TestValidateKernelConfig(t *testing.T) { // UserConfig are currently exposed to users"). Adding a field to config.UserConfig // without classifying it here fails TestKernelConfigFieldsClassified — forcing a // deliberate decision (forward it, reject it loudly, or accept it as inert) rather -// than silently dropping it on the kernel path, the exact divergence rounds 2–4 -// closed. (TLSConfig and ArrowConfig live on the outer config.Config, not +// than silently dropping it on the kernel path. (TLSConfig and ArrowConfig live on +// the outer config.Config, not // UserConfig; newKernelBackend reads TLSConfig.InsecureSkipVerify explicitly, and // the kernel renders decimals exactly regardless of ArrowConfig.) var kernelConfigFieldDisposition = map[string]string{ @@ -162,25 +161,52 @@ var kernelConfigFieldDisposition = map[string]string{ "TelemetryBatchSize": "inert", "TelemetryFlushInterval": "inert", "UseArrowNativeDecimalDSN": "inert", // DSN carrier; kernel renders decimals exactly regardless - "CloudFetchConfig": "inert", // kernel does CloudFetch internally + + // Fields promoted from the embedded CloudFetchConfig. The kernel does + // CloudFetch internally (below the C ABI), so none is forwarded — but each is + // classified individually so a new CloudFetch option can't slip the guard. + "UseCloudFetch": "inert", + "MaxDownloadThreads": "inert", + "MaxFilesInMemory": "inert", + "MinTimeToExpiry": "inert", + "CloudFetchSpeedThresholdMbps": "inert", + "HTTPClient": "inert", +} + +// kernelConfigClassifiedNames returns every UserConfig field name the kernel gate +// must account for, flattening anonymous embedded structs (e.g. CloudFetchConfig) +// into their promoted fields — reflect's NumField reports an embed as one field, +// which would hide new sub-fields from the drop guard below. +func kernelConfigClassifiedNames(t reflect.Type) []string { + var names []string + for i := 0; i < t.NumField(); i++ { + f := t.Field(i) + if f.Anonymous && f.Type.Kind() == reflect.Struct { + names = append(names, kernelConfigClassifiedNames(f.Type)...) + continue + } + names = append(names, f.Name) + } + return names } func TestKernelConfigFieldsClassified(t *testing.T) { - ut := reflect.TypeOf(config.UserConfig{}) - for i := 0; i < ut.NumField(); i++ { - name := ut.Field(i).Name + names := kernelConfigClassifiedNames(reflect.TypeOf(config.UserConfig{})) + classified := make(map[string]bool, len(names)) + for _, name := range names { + classified[name] = true if _, ok := kernelConfigFieldDisposition[name]; !ok { - t.Errorf("config.UserConfig field %q is not classified for the kernel backend. "+ - "Add it to kernelConfigFieldDisposition as forwarded/rejected/inert (and wire it in "+ - "validateKernelConfig / newKernelBackend if it must be honored) so it isn't silently "+ - "dropped on the kernel path.", name) + t.Errorf("config.UserConfig field %q (incl. promoted embed fields) is not classified "+ + "for the kernel backend. Add it to kernelConfigFieldDisposition as "+ + "forwarded/rejected/inert (and wire it in validateKernelConfig / newKernelBackend if "+ + "it must be honored) so it isn't silently dropped on the kernel path.", name) } } // Guard the reverse too: a disposition entry for a field that no longer exists // is stale. for name := range kernelConfigFieldDisposition { - if _, ok := ut.FieldByName(name); !ok { - t.Errorf("kernelConfigFieldDisposition has %q but config.UserConfig no longer does; remove it", name) + if !classified[name] { + t.Errorf("kernelConfigFieldDisposition has %q but config.UserConfig (incl. embeds) no longer does; remove it", name) } } } From 8c730da41ec2f129b3d2098a6dcec34cff08dd7b Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Sat, 11 Jul 2026 16:24:19 +0000 Subject: [PATCH 30/32] fix(kernel): evict a dead session on the statement path; stop retrying permanent auth (H1-r5, L1-r5) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit H1 (round 4) stopped the statement/read path from returning driver.ErrBadConn — which correctly killed silent DML replay — but ErrBadConn was also the ONLY thing that evicted a conn whose server session died mid-life (nothing on the statement/read path mutates k.valid, and kernelOp/kernelRows held no backend ref). So after a warehouse idle-stop/restart GC'd the session, the pooled conn stayed IsValid()==true and every subsequent query failed identically, with no self-healing (Thrift evicts via "Invalid SessionHandle" → ErrBadConn). Thread the H1 needle: give kernelOp a *KernelBackend back-reference (mirroring thriftOperation's *Backend) and, on a session-fatal status, call markSessionDead() so SessionValid()→conn.IsValid() returns false and database/sql discards the conn on return to the pool — eviction WITHOUT ErrBadConn, so the statement is never re-run. Wired on every statement/read error path (new_statement, set_sql, execute, get_result_stream, get_schema, next_batch). Split the two notions the old isBadConnection conflated: - isBadConnection = TRANSIENT (Unavailable/NetworkError) → produces ErrBadConn on the session-lifecycle path so connect is retried. Unauthenticated dropped from it (L1-r5): a wrong/expired PAT is permanent, so retrying connect 3× is wasteful and can worsen auth rate-limiting — matching Thrift, which only treats an invalid session handle (not a 401) as bad-conn. - isSessionFatal = the broader eviction set (adds Unauthenticated): a dead session from any of these is evicted, but via markSessionDead, never ErrBadConn. Tests: TestIsSessionFatal, TestEvictIfSessionFatal (evicts + asserts not ErrBadConn), and TestIsBadConnection updated to exclude Unauthenticated. Co-authored-by: Isaac Signed-off-by: Mani Kaustubh Mathur --- internal/backend/kernel/backend.go | 20 ++++++++- internal/backend/kernel/cgo.go | 28 +++++++++++-- internal/backend/kernel/kernel_test.go | 58 +++++++++++++++++++++++--- internal/backend/kernel/operation.go | 20 +++++++-- internal/backend/kernel/rows.go | 2 + 5 files changed, 114 insertions(+), 14 deletions(-) diff --git a/internal/backend/kernel/backend.go b/internal/backend/kernel/backend.go index 4726279a..8427ddeb 100644 --- a/internal/backend/kernel/backend.go +++ b/internal/backend/kernel/backend.go @@ -205,9 +205,27 @@ func (k *KernelBackend) CloseSession(ctx context.Context) error { } // SessionValid backs conn.IsValid → pool eviction. No I/O; inspects state -// captured at OpenSession. +// captured at OpenSession and updated by markSessionDead. func (k *KernelBackend) SessionValid() bool { return k.valid && k.session != nil } +// markSessionDead marks the session unusable so SessionValid() → conn.IsValid() +// returns false and database/sql discards this conn on return to the pool. Called +// from the statement/read path when a kernel error is session-fatal (isSessionFatal): +// it evicts the dead conn WITHOUT returning driver.ErrBadConn, so the statement is +// never transparently re-run (no duplicate write — the H1 constraint). The backend +// is single-owner per conn (only the cancel watcher shares state, and it touches +// only the canceller's inflight slot), so this write is race-free. +func (k *KernelBackend) markSessionDead() { k.valid = false } + +// evictIfSessionFatal marks the session dead when err is a session-fatal +// KernelError, so a conn whose server session died mid-life is evicted from the +// pool. A no-op for nil, non-KernelError, or non-fatal errors. +func (k *KernelBackend) evictIfSessionFatal(err error) { + if ke, ok := err.(*KernelError); ok && isSessionFatal(ke.Code) { + k.markSessionDead() + } +} + // SessionID is the per-conn id (conn.id). Valid after OpenSession. func (k *KernelBackend) SessionID() string { return k.sessionID } diff --git a/internal/backend/kernel/cgo.go b/internal/backend/kernel/cgo.go index f5d5f4d6..60310118 100644 --- a/internal/backend/kernel/cgo.go +++ b/internal/backend/kernel/cgo.go @@ -190,10 +190,32 @@ const ( statusSqlError = int(C.KernelStatusCode_SqlError) ) -// isBadConnection reports whether a status code means the session is no longer -// usable, so the driver evicts the conn from the pool (via driver.ErrBadConn's -// surface) rather than reusing it. +// isBadConnection reports whether a status code is a *transient* connection +// failure — one where retrying on a fresh connection could succeed. On the +// session-lifecycle path this is wrapped as driver.ErrBadConn so database/sql +// retries connect. Unauthenticated is deliberately excluded: a wrong/expired PAT +// is permanent, so retrying it just burns connect attempts (and can worsen +// server-side auth rate-limiting) and fails identically — matching Thrift, which +// only treats an invalid session handle (a liveness signal), not a 401, as +// bad-conn. (Auth failure still marks the session dead for pool eviction — see +// isSessionFatal.) func isBadConnection(code int) bool { + switch code { + case statusUnavailable, statusNetworkError: + return true + default: + return false + } +} + +// isSessionFatal reports whether a status code means the server-side session is no +// longer usable, so the conn must be evicted from the pool rather than reused for +// the next query. Broader than isBadConnection: it also covers Unauthenticated (an +// expired/revoked token kills the session) — but unlike isBadConnection it is NOT +// used to produce driver.ErrBadConn on the statement path, so eviction happens +// without database/sql replaying the statement (see toStatementError + the +// KernelBackend.markSessionDead call sites). +func isSessionFatal(code int) bool { switch code { case statusUnauthenticated, statusUnavailable, statusNetworkError: return true diff --git a/internal/backend/kernel/kernel_test.go b/internal/backend/kernel/kernel_test.go index 9a72fbe7..9ac4e0a4 100644 --- a/internal/backend/kernel/kernel_test.go +++ b/internal/backend/kernel/kernel_test.go @@ -11,16 +11,18 @@ import ( "github.com/databricks/databricks-sql-go/internal/backend" ) -// isBadConnection maps the session-unusable status codes so the pool evicts the -// conn; every other code stays a plain kernel error. +// isBadConnection is the TRANSIENT-failure set (retrying a fresh connect could +// succeed) — used to produce driver.ErrBadConn on the session-lifecycle path. +// Unauthenticated is excluded: a wrong/expired PAT is permanent, not retryable. func TestIsBadConnection(t *testing.T) { - bad := []int{statusUnauthenticated, statusUnavailable, statusNetworkError} - for _, code := range bad { + transient := []int{statusUnavailable, statusNetworkError} + for _, code := range transient { if !isBadConnection(code) { - t.Errorf("code %d should be a bad connection", code) + t.Errorf("code %d should be a (transient) bad connection", code) } } - notBad := []int{statusInvalidArgument, statusSqlError, statusTimeout} + // Unauthenticated is session-fatal but NOT retryable, so it is not bad-conn. + notBad := []int{statusUnauthenticated, statusInvalidArgument, statusSqlError, statusTimeout} for _, code := range notBad { if isBadConnection(code) { t.Errorf("code %d should not be a bad connection", code) @@ -28,6 +30,50 @@ func TestIsBadConnection(t *testing.T) { } } +// isSessionFatal is the broader "session is dead, evict the conn" set — it adds +// Unauthenticated (an expired token kills the session) on top of the transient set. +func TestIsSessionFatal(t *testing.T) { + fatal := []int{statusUnauthenticated, statusUnavailable, statusNetworkError} + for _, code := range fatal { + if !isSessionFatal(code) { + t.Errorf("code %d should be session-fatal", code) + } + } + notFatal := []int{statusInvalidArgument, statusSqlError, statusTimeout} + for _, code := range notFatal { + if isSessionFatal(code) { + t.Errorf("code %d should not be session-fatal", code) + } + } +} + +// evictIfSessionFatal flips SessionValid()→false on a session-fatal error (so the +// pool discards the conn) WITHOUT the error being driver.ErrBadConn (so the +// statement is never transparently re-run — the H1 constraint). +func TestEvictIfSessionFatal(t *testing.T) { + // valid tracks the session-dead flag SessionValid() gates on; the opaque + // session pointer is orthogonal here (can't construct the incomplete C type), + // so assert on k.valid directly. + k := &KernelBackend{valid: true} + + // Non-fatal (e.g. a SQL error) leaves the session valid. + k.evictIfSessionFatal(&KernelError{Code: statusSqlError}) + if !k.valid { + t.Error("a SQL error must not evict the session") + } + + // A session-fatal error marks the session dead, and the surfaced statement-path + // error is NOT driver.ErrBadConn (so database/sql won't re-run the statement). + fatal := &KernelError{Code: statusUnavailable, Message: "session gone"} + k.evictIfSessionFatal(fatal) + if k.valid { + t.Error("a session-fatal error must evict the session (valid=false)") + } + if errors.Is(toStatementError(fatal), driver.ErrBadConn) { + t.Error("the statement-path error must not be driver.ErrBadConn (no replay)") + } +} + // toConnError (session-lifecycle path) wraps a session-unusable KernelError as // driver.ErrBadConn so database/sql evicts the conn, and leaves other errors and // their sqlstate intact. diff --git a/internal/backend/kernel/operation.go b/internal/backend/kernel/operation.go index b1a4d19c..c551f56b 100644 --- a/internal/backend/kernel/operation.go +++ b/internal/backend/kernel/operation.go @@ -40,6 +40,7 @@ func (k *KernelBackend) execute(ctx context.Context, req backend.ExecRequest) (b if err := call(func() C.KernelStatusCode { return C.kernel_session_new_statement(k.session, &stmt) }); err != nil { + k.evictIfSessionFatal(err) return &kernelOp{}, fmt.Errorf("kernel: new_statement: %w", toStatementError(err)) } @@ -49,6 +50,7 @@ func (k *KernelBackend) execute(ctx context.Context, req backend.ExecRequest) (b }); err != nil { sql.free() C.kernel_statement_close(stmt) + k.evictIfSessionFatal(err) return &kernelOp{}, fmt.Errorf("kernel: set_sql: %w", toStatementError(err)) } sql.free() @@ -120,7 +122,7 @@ func (k *KernelBackend) execute(ctx context.Context, req backend.ExecRequest) (b C.kernel_statement_canceller_free(canceller) } - op := &kernelOp{stmt: stmt, location: k.cfg.Location} + op := &kernelOp{backend: k, stmt: stmt, location: k.cfg.Location} if execErr != nil { // Prefer the caller's ctx error when the ctx was cancelled (database/sql // convention), keeping the kernel error as the cause. @@ -130,6 +132,11 @@ func (k *KernelBackend) execute(ctx context.Context, req backend.ExecRequest) (b return op, fmt.Errorf("kernel: execute cancelled: %w", ctx.Err()) } 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 + // conn is discarded without database/sql re-running the statement (H1). + k.evictIfSessionFatal(execErr) op.close() return op, fmt.Errorf("kernel: execute: %w", toStatementError(execErr)) } @@ -143,9 +150,13 @@ func (k *KernelBackend) execute(ctx context.Context, req backend.ExecRequest) (b // kernelOp implements backend.Operation over a sync executed statement. type kernelOp struct { - stmt *C.kernel_statement_t - exec *C.kernel_executed_statement_t - closed bool + // backend is the owning connection's backend, held so a session-fatal error on + // the result-read path can evict the conn (markSessionDead), mirroring how the + // Thrift thriftOperation holds its *Backend. + backend *KernelBackend + stmt *C.kernel_statement_t + exec *C.kernel_executed_statement_t + closed bool // affectedRows is the modified-row count captured at execute time. It is // cached (not read live from exec) because the caller closes the operation — // which nulls exec — before reading AffectedRows (see conn.ExecContext). @@ -186,6 +197,7 @@ func (o *kernelOp) Results(ctx context.Context, callbacks *dbsqlrows.TelemetryCa // No Rows is returned to own teardown, and the query path does not call // Operation.Close on a Results error — so close the handles here to avoid // leaking the statement / executed handle (and its server operation). + o.backend.evictIfSessionFatal(err) o.close() return nil, fmt.Errorf("kernel: get_result_stream: %w", toStatementError(err)) } diff --git a/internal/backend/kernel/rows.go b/internal/backend/kernel/rows.go index 973fb5d4..f2f3c764 100644 --- a/internal/backend/kernel/rows.go +++ b/internal/backend/kernel/rows.go @@ -58,6 +58,7 @@ func newKernelRows(ctx context.Context, op *kernelOp, stream *C.kernel_result_st if err := call(func() C.KernelStatusCode { return C.kernel_result_stream_get_schema(stream, &csch) }); err != nil { + op.backend.evictIfSessionFatal(err) r.Close() return nil, fmt.Errorf("kernel: get_schema: %w", toStatementError(err)) } @@ -151,6 +152,7 @@ func (r *kernelRows) nextBatch() error { if err := call(func() C.KernelStatusCode { return C.kernel_result_stream_next_batch(r.stream, &carr, &csch) }); err != nil { + r.op.backend.evictIfSessionFatal(err) return fmt.Errorf("kernel: next_batch: %w", toStatementError(err)) } if carr.release == nil { From 5068f5a16108ef54db07d9b6e6db1c499739f066 Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Sat, 11 Jul 2026 16:29:04 +0000 Subject: [PATCH 31/32] perf(arrowscan): memoize struct-key escaping per result set on the kernel path (M3-r5) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit writeStructJSON did json.Marshal(fieldName) per field, reached per row via kernelRows.Next → ScanCell — the exact per-row-marshal antipattern M1 fixed on the Thrift path, on the backend whose reason to exist is throughput. (The old comment claiming it was "cheap relative to once-per-batch" was wrong: it's per row.) Add a caller-owned StructKeyCache and ScanCellCached; kernelRows creates one in newKernelRows and passes it down, so escaped `"name":` keys are computed once per result set. The cache is Rows-scoped and freed with it — deliberately NOT a process-global keyed by *arrow.StructType, which would leak (the C Data import allocates a fresh *StructType per batch — round-2 N1). ScanCell stays as the nil-cache one-shot, so existing callers/tests are unchanged; a nil cache just recomputes inline. TestScanCellCachedMatchesUncached asserts the cached path is byte-identical to the uncached one across rows. Co-authored-by: Isaac Signed-off-by: Mani Kaustubh Mathur --- internal/arrowscan/arrowscan.go | 90 ++++++++++++++++++++-------- internal/arrowscan/arrowscan_test.go | 35 +++++++++++ internal/backend/kernel/rows.go | 8 ++- 3 files changed, 105 insertions(+), 28 deletions(-) diff --git a/internal/arrowscan/arrowscan.go b/internal/arrowscan/arrowscan.go index 23207491..89b7c759 100644 --- a/internal/arrowscan/arrowscan.go +++ b/internal/arrowscan/arrowscan.go @@ -42,6 +42,48 @@ import ( // 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) +} + +// StructKeyCache memoizes the JSON-escaped `"name":` prefixes for a struct type, +// so writeStructJSON doesn't re-marshal constant field names on every row. It is +// caller-owned and must be scoped to a single result set (e.g. one driver.Rows) +// and discarded with it — NOT a process-global, which would leak because the Arrow +// C Data import allocates a fresh *StructType per batch (see databricks-sql-go +// round-2 N1). A nil cache is valid: rendering just recomputes the keys inline. +type StructKeyCache struct { + m map[*arrow.StructType][]string +} + +// NewStructKeyCache returns a cache ready to pass to ScanCellCached. +func NewStructKeyCache() *StructKeyCache { + return &StructKeyCache{m: make(map[*arrow.StructType][]string)} +} + +// keyPrefixes returns the escaped `"name":` prefix for each field of st, +// memoized. Safe on a nil receiver (computes without caching). +func (c *StructKeyCache) keyPrefixes(st *arrow.StructType) []string { + if c != nil { + if p, ok := c.m[st]; ok { + return p + } + } + fields := st.Fields() + prefixes := make([]string, len(fields)) + for i := range fields { + name, _ := json.Marshal(fields[i].Name) // JSON-escapes the field name + prefixes[i] = string(name) + ":" + } + if c != nil { + c.m[st] = prefixes + } + return prefixes +} + +// ScanCellCached is ScanCell with a caller-owned StructKeyCache (see +// StructKeyCache) so struct field-name keys are escaped once per result set +// rather than once per row. Pass nil for the un-memoized one-shot behavior. +func ScanCellCached(col arrow.Array, row int, loc *time.Location, keys *StructKeyCache) (driver.Value, error) { if col.IsNull(row) { return nil, nil } @@ -100,7 +142,7 @@ func ScanCell(col arrow.Array, row int, loc *time.Location) (driver.Value, error 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) + 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()) @@ -119,35 +161,35 @@ func inLocation(t time.Time, loc *time.Location) time.Time { // renderJSONString renders one nested cell as a JSON string (nil for a NULL // cell). loc is applied to any timestamp/date leaves, matching the top-level // scan and the Thrift path. -func renderJSONString(col arrow.Array, row int, loc *time.Location) (driver.Value, error) { +func renderJSONString(col arrow.Array, row int, loc *time.Location, keys *StructKeyCache) (driver.Value, error) { if col.IsNull(row) { return nil, nil } var b strings.Builder - if err := writeJSON(&b, col, row, loc); err != nil { + if err := writeJSON(&b, col, row, loc, keys); err != nil { return nil, err } return b.String(), nil } // writeJSON writes the JSON form of col[row] into b, recursing for nested types. -func writeJSON(b *strings.Builder, col arrow.Array, row int, loc *time.Location) error { +func writeJSON(b *strings.Builder, col arrow.Array, row int, loc *time.Location, keys *StructKeyCache) error { if col.IsNull(row) { b.WriteString("null") return nil } switch c := col.(type) { case *array.List: - return writeListJSON(b, c.ListValues(), int(c.Offsets()[row]), int(c.Offsets()[row+1]), loc) + return writeListJSON(b, c.ListValues(), int(c.Offsets()[row]), int(c.Offsets()[row+1]), loc, keys) case *array.LargeList: - return writeListJSON(b, c.ListValues(), int(c.Offsets()[row]), int(c.Offsets()[row+1]), loc) + return writeListJSON(b, c.ListValues(), int(c.Offsets()[row]), int(c.Offsets()[row+1]), loc, keys) case *array.FixedSizeList: n := int(c.DataType().(*arrow.FixedSizeListType).Len()) - return writeListJSON(b, c.ListValues(), row*n, row*n+n, loc) + return writeListJSON(b, c.ListValues(), row*n, row*n+n, loc, keys) case *array.Map: - return writeMapJSON(b, c, row, loc) + return writeMapJSON(b, c, row, loc, keys) case *array.Struct: - return writeStructJSON(b, c, row, loc) + return writeStructJSON(b, c, row, loc, keys) case *array.Decimal128: // 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 @@ -170,13 +212,13 @@ func writeJSON(b *strings.Builder, col arrow.Array, row int, loc *time.Location) } } -func writeListJSON(b *strings.Builder, values arrow.Array, start, end int, loc *time.Location) error { +func writeListJSON(b *strings.Builder, values arrow.Array, start, end int, loc *time.Location, keys *StructKeyCache) error { b.WriteByte('[') for i := start; i < end; i++ { if i > start { b.WriteByte(',') } - if err := writeJSON(b, values, i, loc); err != nil { + if err := writeJSON(b, values, i, loc, keys); err != nil { return err } } @@ -184,9 +226,9 @@ func writeListJSON(b *strings.Builder, values arrow.Array, start, end int, loc * return nil } -func writeMapJSON(b *strings.Builder, m *array.Map, row int, loc *time.Location) error { +func writeMapJSON(b *strings.Builder, m *array.Map, row int, loc *time.Location, keys *StructKeyCache) error { start, end := int(m.Offsets()[row]), int(m.Offsets()[row+1]) - keys := m.Keys() + mapKeys := m.Keys() items := m.Items() b.WriteByte('{') for i := start; i < end; i++ { @@ -194,13 +236,13 @@ func writeMapJSON(b *strings.Builder, m *array.Map, row int, loc *time.Location) b.WriteByte(',') } // JSON object keys must be strings; stringify the key value. - kv, err := scalarForJSON(keys, i, loc) + kv, err := scalarForJSON(mapKeys, i, loc) if err != nil { return err } writeJSONKey(b, kv) b.WriteByte(':') - if err := writeJSON(b, items, i, loc); err != nil { + if err := writeJSON(b, items, i, loc, keys); err != nil { return err } } @@ -208,22 +250,18 @@ func writeMapJSON(b *strings.Builder, m *array.Map, row int, loc *time.Location) return nil } -func writeStructJSON(b *strings.Builder, s *array.Struct, row int, loc *time.Location) error { - st := s.DataType().(*arrow.StructType) +func writeStructJSON(b *strings.Builder, s *array.Struct, row int, loc *time.Location, keys *StructKeyCache) error { + // Field-name keys are constant across rows; keys.keyPrefixes memoizes them per + // result set so this per-row/per-cell path doesn't re-marshal them (a nil cache + // recomputes inline — correct, just not memoized). + prefixes := keys.keyPrefixes(s.DataType().(*arrow.StructType)) b.WriteByte('{') for f := 0; f < s.NumField(); f++ { if f > 0 { b.WriteByte(',') } - // Marshal the field name to JSON-escape it. Not memoized per struct type: a - // process-global cache keyed by *arrow.StructType leaks, because the Arrow C - // Data Interface import allocates a fresh *StructType every batch (no - // interning), so the key never repeats across batches. Recomputing here is - // cheap relative to the once-per-batch cost. - keyBytes, _ := json.Marshal(st.Field(f).Name) - b.Write(keyBytes) - b.WriteByte(':') - if err := writeJSON(b, s.Field(f), row, loc); err != nil { + b.WriteString(prefixes[f]) // pre-escaped `"name":` + if err := writeJSON(b, s.Field(f), row, loc, keys); err != nil { return err } } diff --git a/internal/arrowscan/arrowscan_test.go b/internal/arrowscan/arrowscan_test.go index 2da3adb4..089cea73 100644 --- a/internal/arrowscan/arrowscan_test.go +++ b/internal/arrowscan/arrowscan_test.go @@ -274,3 +274,38 @@ func TestScanCellNested(t *testing.T) { } }) } + +// A StructKeyCache must produce byte-identical output to the nil (recompute-inline) +// path, and reused across rows of the same struct type — so the round-trip +// through the cache can never silently diverge from the one-shot rendering. +func TestScanCellCachedMatchesUncached(t *testing.T) { + pool := memory.NewGoAllocator() + dt := arrow.StructOf( + arrow.Field{Name: "a", Type: arrow.PrimitiveTypes.Int64}, + arrow.Field{Name: `q"x`, Type: arrow.BinaryTypes.String}, // needs escaping + ) + b := array.NewStructBuilder(pool, dt) + defer b.Release() + for i := 0; i < 3; i++ { // multiple rows: exercises the memoized second+ hit + b.Append(true) + b.FieldBuilder(0).(*array.Int64Builder).Append(int64(i)) + b.FieldBuilder(1).(*array.StringBuilder).Append("v") + } + arr := b.NewArray() + defer arr.Release() + + cache := NewStructKeyCache() + for row := 0; row < 3; row++ { + uncached, err := ScanCell(arr, row, nil) + if err != nil { + t.Fatal(err) + } + cached, err := ScanCellCached(arr, row, nil, cache) + if err != nil { + t.Fatal(err) + } + if cached != uncached { + t.Errorf("row %d: cached %q != uncached %q", row, cached, uncached) + } + } +} diff --git a/internal/backend/kernel/rows.go b/internal/backend/kernel/rows.go index f2f3c764..681e4077 100644 --- a/internal/backend/kernel/rows.go +++ b/internal/backend/kernel/rows.go @@ -47,12 +47,16 @@ type kernelRows struct { chunkCount int // cumulative batches fetched, for OnChunkFetched closed bool eof bool + // 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 (round-2 N1). + keyCache *arrowscan.StructKeyCache } // 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} + r := &kernelRows{ctx: ctx, op: op, stream: stream, callbacks: cb, keyCache: arrowscan.NewStructKeyCache()} var csch C.struct_ArrowSchema if err := call(func() C.KernelStatusCode { @@ -117,7 +121,7 @@ func (r *kernelRows) Next(dest []driver.Value) error { } rec := r.cur for c := 0; c < len(dest); c++ { - v, err := arrowscan.ScanCell(rec.Column(c), r.rowInCur, r.op.location) + v, err := arrowscan.ScanCellCached(rec.Column(c), r.rowInCur, r.op.location, r.keyCache) if err != nil { return fmt.Errorf("kernel: scan col %d (%s): %w", c, r.cols[c], err) } From 0eaf6007e75fa99e374da7273107fb203e5c3d50 Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Sat, 11 Jul 2026 16:32:08 +0000 Subject: [PATCH 32/32] test(arrowbased): add top-level scalar parity + document the DECIMAL rendering (M1-r5) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review r5 M1-r5 flagged a top-level DECIMAL divergence (kernel exact string vs Thrift lossy float64). Investigated + verified live against staging: it does NOT occur on any real configuration. In the DEFAULT config UseArrowNativeDecimal is false, so the execute request sets DecimalAsArrow=false and the server sends DECIMAL as a *string* column — no decimal128 arrives, so the container's ToFloat64 path is never reached. With the flag on, the top-level scan uses DecimalStringValue (exact). All three configs (thrift-default, thrift-nativeDecimal, kernel) return the exact string "1234567890123456.7890" (Go string) — confirmed by a live probe. The finding's premise (reasoned from the code path) missed the DecimalAsArrow wire behavior; no code change needed. The legitimate part — the parity test had no top-level scalar cases — is addressed: TestArrowbasedKernelTopLevelScalarParity covers int64/float32/float64/ string/timestamp/date32, and TestTopLevelDecimalRendering pins all three decimal behaviors (kernel exact string, Thrift DecimalStringValue exact, Thrift container Value lossy-float64) with the wire nuance documented, so the distinction can't silently drift into an actual result divergence. Co-authored-by: Isaac Signed-off-by: Mani Kaustubh Mathur --- .../rows/arrowbased/arrowscan_parity_test.go | 110 ++++++++++++++++++ 1 file changed, 110 insertions(+) diff --git a/internal/rows/arrowbased/arrowscan_parity_test.go b/internal/rows/arrowbased/arrowscan_parity_test.go index 1268e7f4..2882d05c 100644 --- a/internal/rows/arrowbased/arrowscan_parity_test.go +++ b/internal/rows/arrowbased/arrowscan_parity_test.go @@ -206,3 +206,113 @@ func TestArrowbasedKernelRenderParity(t *testing.T) { }) } } + +// Top-level (non-nested) scalars: assert the kernel ScanCell and the Thrift +// container render the same Go value. DECIMAL is deliberately excluded here — see +// TestTopLevelDecimalRendering for why the two *container-level* paths differ by +// arrow type while the actual driver results agree. +func TestArrowbasedKernelTopLevelScalarParity(t *testing.T) { + pool := memory.NewGoAllocator() + loc := time.UTC + + cases := []struct { + name string + build func() arrow.Array + }{ + {"int64", func() arrow.Array { + b := array.NewInt64Builder(pool) + b.Append(42) + return b.NewArray() + }}, + {"float32", func() arrow.Array { + b := array.NewFloat32Builder(pool) + b.Append(0.1) + return b.NewArray() + }}, + {"float64", func() arrow.Array { + b := array.NewFloat64Builder(pool) + b.Append(3.5) + return b.NewArray() + }}, + {"string", func() arrow.Array { + b := array.NewStringBuilder(pool) + b.Append("hi") + return b.NewArray() + }}, + {"timestamp", func() arrow.Array { + b := array.NewTimestampBuilder(pool, &arrow.TimestampType{Unit: arrow.Microsecond}) + b.Append(arrow.Timestamp(time.Date(2026, 7, 9, 12, 0, 0, 0, time.UTC).UnixMicro())) + return b.NewArray() + }}, + {"date32", func() arrow.Array { + b := array.NewDate32Builder(pool) + b.Append(arrow.Date32FromTime(time.Date(2026, 7, 9, 0, 0, 0, 0, time.UTC))) + return b.NewArray() + }}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + arr := tc.build() + 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("top-level %s divergence:\n kernel = %#v (%T)\n thrift = %#v (%T)", tc.name, kernel, kernel, thrift, thrift) + } + }) + } +} + +// TestTopLevelDecimalRendering documents the top-level DECIMAL story (review r5 +// M1-r5), which is subtler than "kernel string vs Thrift float64": +// +// - The kernel path always delivers DECIMAL as a native arrow decimal128 and +// renders it as an exact scale-applied string (arrowscan.ScanCell). +// - The Thrift *container* (decimal128Container.Value) converts a decimal128 to +// a lossy float64 — but that path is only reached when UseArrowNativeDecimal +// is on AND the value is read as a nested leaf. For a top-level column the +// scan uses DecimalStringValue (exact string) when the flag is on. +// - In the DEFAULT config the flag is off, so the server sends DECIMAL as a +// string column (DecimalAsArrow=false) — no decimal128 arrives at all, and the +// user gets the exact string. +// +// So across every real driver configuration a top-level DECIMAL comes back as the +// exact string on both backends (verified live). This test pins the kernel side +// (exact string) and the two Thrift container behaviors so the distinction can't +// silently drift into an actual result divergence. +func TestTopLevelDecimalRendering(t *testing.T) { + pool := memory.NewGoAllocator() + dt := &arrow.Decimal128Type{Precision: 38, Scale: 4} + b := array.NewDecimal128Builder(pool, dt) + defer b.Release() + n, _ := decimal128.FromString("1234567890123456.7890", dt.Precision, dt.Scale) + b.Append(n) + arr := b.NewArray() + defer arr.Release() + + // Kernel: exact string. + got, err := arrowscan.ScanCell(arr, 0, time.UTC) + if err != nil { + t.Fatal(err) + } + if got != "1234567890123456.7890" { + t.Errorf("kernel top-level decimal = %#v, want exact string", got) + } + + // Thrift container Value() is the lossy-float64 path (reached only with native + // decimal + nested read); DecimalStringValue is the exact top-level scan path. + // Assert both so a change to either is caught. + holder := &decimal128Container{scale: dt.Scale} + if err := holder.SetValueArray(arr.Data()); err != nil { + t.Fatal(err) + } + if s := holder.ValueString(0); s != "1234567890123456.7890" { + t.Errorf("Thrift DecimalStringValue = %q, want exact string", s) + } + if v, _ := holder.Value(0); v == "1234567890123456.7890" { + t.Errorf("Thrift container Value() unexpectedly exact; it is the documented lossy-float64 path") + } +}