From 1f29fd791cda54b2d3bb9caf6314a8f3b1c84ad0 Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Thu, 9 Jul 2026 05:52:49 +0000 Subject: [PATCH 1/5] 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 e1701fc..61d676a 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 6ae9616..becab1f 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 0000000..9e16025 --- /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 0000000..f967c2a --- /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 0000000..2efb551 --- /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 0000000..568c262 --- /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 0000000..c0c86eb --- /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 0000000..ca92432 --- /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 fb7c6af..12b5939 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 0000000..d89349f --- /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 0000000..2a3ea39 --- /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 0000000..b687a5c --- /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 0000000..4774278 --- /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 2/5] 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 2efb551..cf313d6 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 c0c86eb..c025875 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 0000000..75138c2 --- /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 2a3ea39..f77c7ed 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 b687a5c..22d8f4d 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 3/5] 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 9e16025..b785beb 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 d89349f..2b823cd 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 0000000..3674892 --- /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 f77c7ed..48a0cdb 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 4/5] 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 b785beb..7f92b2b 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 cf313d6..de6bf4f 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 568c262..730a4dd 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 c025875..107048b 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 75138c2..013342b 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 2b823cd..779859f 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 3674892..ce5356d 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 48a0cdb..1d2b3c0 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 5/5] 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 de6bf4f..86217f1 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 013342b..f42213a 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 22d8f4d..223a008 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()