diff --git a/.github/workflows/go.yml b/.github/workflows/go.yml index beff4157..9164ced8 100644 --- a/.github/workflows/go.yml +++ b/.github/workflows/go.yml @@ -90,6 +90,12 @@ jobs: fi go get -v -t -d ./... + # This matrix builds the default pure-Go driver (CGO_ENABLED=0), which does + # NOT compile the SEA-via-kernel backend (//go:build cgo && databricks_kernel). + # A dedicated CGO_ENABLED=1 `-tags databricks_kernel` test job needs the + # kernel static library linked in CI, which arrives with the kernel + # distribution work (a pinned-revision source build or a published .a). + # Until then the kernel path's cgo files are not exercised here. - name: Test run: make test env: diff --git a/connector.go b/connector.go index e1701fcb..61d676ad 100644 --- a/connector.go +++ b/connector.go @@ -15,6 +15,7 @@ import ( "github.com/databricks/databricks-sql-go/auth/pat" "github.com/databricks/databricks-sql-go/auth/tokenprovider" "github.com/databricks/databricks-sql-go/driverctx" + "github.com/databricks/databricks-sql-go/internal/backend" "github.com/databricks/databricks-sql-go/internal/backend/thrift" "github.com/databricks/databricks-sql-go/internal/client" "github.com/databricks/databricks-sql-go/internal/config" @@ -32,9 +33,18 @@ type connector struct { func (c *connector) Connect(ctx context.Context) (driver.Conn, error) { defer debuglog.Track(ctx, "connector.Connect", "host=%s", c.cfg.Host)() - // Build the execution backend. Currently always the Thrift backend; a second - // SEA-via-kernel backend will be selected here by a connect-time flag. - be, err := thrift.New(ctx, c.cfg, c.client) + // Build the execution backend. Thrift is the default; the SEA-via-kernel + // backend is selected when UseKernel is set. newKernelBackend is build-tag + // gated: in the default pure-Go build it returns a clear "not linked in" + // error, so the kernel path compiles and links only under -tags + // databricks_kernel + CGO_ENABLED=1. + var be backend.Backend + var err error + if c.cfg.UseKernel { + be, err = newKernelBackend(ctx, c.cfg) + } else { + be, err = thrift.New(ctx, c.cfg, c.client) + } if err != nil { return nil, err } @@ -77,7 +87,14 @@ func (c *connector) Connect(ctx context.Context) (driver.Conn, error) { conn.telemetry.RecordOperation(ctx, conn.id, "", telemetry.OperationTypeCreateSession, sessionLatencyMs, nil) } - log.Info().Msgf("connect: host=%s port=%d httpPath=%s serverProtocolVersion=0x%X", c.cfg.Host, c.cfg.Port, c.cfg.HTTPPath, be.ServerProtocolVersion()) + // ServerProtocolVersion is Thrift-specific (not on the neutral backend + // interface); the kernel backend has no negotiated Thrift protocol, so log it + // only when present. + if tb, ok := be.(*thrift.Backend); ok { + log.Info().Msgf("connect: host=%s port=%d httpPath=%s serverProtocolVersion=0x%X", c.cfg.Host, c.cfg.Port, c.cfg.HTTPPath, tb.ServerProtocolVersion()) + } else { + log.Info().Msgf("connect: host=%s port=%d httpPath=%s backend=kernel", c.cfg.Host, c.cfg.Port, c.cfg.HTTPPath) + } return conn, nil } @@ -301,6 +318,26 @@ func WithHTTPPath(path string) ConnOption { } } +// WithUseKernel selects the SEA-via-kernel backend instead of the default +// Thrift backend. It has effect only in a build compiled with +// `-tags databricks_kernel` and CGO_ENABLED=1; in the default pure-Go build a +// connection made with this option set returns a clear error at connect time +// (the kernel backend is not linked in). +func WithUseKernel(useKernel bool) ConnOption { + return func(c *config.Config) { + c.UseKernel = useKernel + } +} + +// WithWarehouseID sets the bare SQL warehouse id. The kernel backend addresses a +// warehouse by id; when set it is preferred over the http path. The Thrift +// backend ignores it and continues to route by http path. +func WithWarehouseID(id string) ConnOption { + return func(c *config.Config) { + c.WarehouseID = id + } +} + // WithMaxRows sets up the max rows fetched per request. Default is 10000 func WithMaxRows(n int) ConnOption { return func(c *config.Config) { diff --git a/doc.go b/doc.go index 6ae9616b..feab8247 100644 --- a/doc.go +++ b/doc.go @@ -155,6 +155,70 @@ 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, set DBSQL_KERNEL_DEBUG +to any non-empty value. That single flag turns on the driver's binding step tracer +AND installs the kernel's Rust log subscriber, so both write to the same stderr in +execution order. It is off by default (and must stay off during benchmarks — debug +logging perturbs latency). The kernel's verbosity is then controlled by RUST_LOG, +which the kernel honors directly; filter on the target databricks::sql::kernel +(note the colons — it is the kernel's explicit log target, NOT the crate module +path databricks_sql_kernel, which would match nothing): + + # kernel logs only: + DBSQL_KERNEL_DEBUG=1 RUST_LOG=databricks::sql::kernel=debug ./your_app 2>&1 + # kernel logs plus its HTTP stack (hyper/reqwest): + DBSQL_KERNEL_DEBUG=1 RUST_LOG=debug ./your_app 2>&1 + +The kernel backend currently supports PAT authentication; reading scalar, nested, +and complex-typed results (CloudFetch is handled transparently); context +cancellation during execute (a cancelled ctx fires a real server-side cancel; on +the read path cancellation is honored at result-batch boundaries, not mid-fetch); +and the TLS, proxy, and session-conf (query tags, statement timeout, time zone) +connection options. OAuth (M2M/U2M), initial catalog/schema, and +WithEnableMetricViewMetadata are not yet supported and return a clear error at +connect time rather than being silently ignored; likewise WithTimeout (a server +query timeout the kernel C ABI can't set) and WithRetries used to disable retries +(the kernel retries internally). Bound query parameters are likewise not yet +supported and return a clear error at execute time (they arrive per-query, not at +connect). None of these is silently ignored. (Metadata is issued as ordinary SQL +— SHOW/DESCRIBE/information_schema — and runs on this backend like any other +query.) + +WithMaxRows and retry tuning (WithRetries with a positive limit) are accepted but +not applied on the kernel path: the kernel manages result fetching and retries +internally, below the C ABI, with no user-facing knob. + +Two further kernel-backend caveats. The INTERVAL types (year-month / day-time) +listed in the type table below are not yet handled by the kernel scanner and +return a scan error; use the default (Thrift) backend for interval columns. And +the kernel backend does not yet surface a per-statement server query id, so a +QueryIdCallback (see below) fires with "" and no EXECUTE_STATEMENT telemetry is +emitted for kernel queries. + # Programmatically Retrieving Connection and Query Id Use the driverctx package under driverctx/ctx.go to add callbacks to the query context to receive the connection id and query id. diff --git a/internal/arrowscan/arrowscan.go b/internal/arrowscan/arrowscan.go new file mode 100644 index 00000000..89b7c759 --- /dev/null +++ b/internal/arrowscan/arrowscan.go @@ -0,0 +1,329 @@ +// Package arrowscan converts Arrow array cells to database/sql driver.Values, +// with nested types (List/Map/Struct, and VARIANT which arrives nested) rendered +// to a JSON string byte-identical to the Thrift arrow path +// (internal/rows/arrowbased). It is pure Go (no cgo), so it is shared by the +// kernel backend and testable in the default CGO_ENABLED=0 build — the tests here +// are the regression guard for the exact rendering rules (native float32, exact +// decimals, time.Time formatting, JSON grammar) both backends must agree on. +// +// Rendering to JSON (not a Go map/slice) is deliberate: it is what the Thrift +// path returns, so a query's result is identical across backends. +// - list → [v0,v1,...] +// - map → {"k0":v0,"k1":v1,...} (keys stringified) +// - struct → {"field0":v0,...} +// - nested NULL → null +// - time.Time → quoted .String() (matches the Thrift marshal() special-case) +// - nested decimal → exact scale-applied JSON number literal (never a lossy +// float64), matching Thrift's marshalScalar → ValueString (#253/#274) +// - float32 → native float32 (not widened to float64), so JSON renders +// 3.14, not 3.140000104904175 +package arrowscan + +import ( + "database/sql/driver" + "encoding/json" + "fmt" + "strings" + "time" + + "github.com/apache/arrow/go/v12/arrow" + "github.com/apache/arrow/go/v12/arrow/array" + "github.com/databricks/databricks-sql-go/internal/decimalfmt" +) + +// ScanCell extracts one cell as a driver.Value. Scalars map to their Go value: +// bool, all int/uint widths, float (native float32/float64), string, binary, +// date, timestamp, and top-level decimal (as an exact fixed-point string, +// matching the Thrift path — a float64 would lose precision beyond ~17 digits; +// see databricks-sql-go#274). Nested types (List/Map/Struct, and VARIANT which +// arrives nested) render to a JSON string byte-identical to the Thrift path; +// GEOMETRY arrives as a WKB/WKT string and is handled by the string arm. NULLs +// map to nil. A genuinely unhandled type (e.g. interval/duration) returns an +// error rather than a silently wrong value. loc renders DATE / TIMESTAMP in the +// session time zone (nil = UTC, arrow's ToTime default). +func ScanCell(col arrow.Array, row int, loc *time.Location) (driver.Value, error) { + return ScanCellCached(col, row, loc, nil) +} + +// StructKeyCache memoizes the JSON-escaped `"name":` prefixes for a struct type, +// so writeStructJSON doesn't re-marshal constant field names on every row. It is +// caller-owned and must be scoped to a single result set (e.g. one driver.Rows) +// and discarded with it — NOT a process-global, which would leak because the Arrow +// C Data import allocates a fresh *StructType per batch (see databricks-sql-go +// round-2 N1). A nil cache is valid: rendering just recomputes the keys inline. +type StructKeyCache struct { + m map[*arrow.StructType][]string +} + +// NewStructKeyCache returns a cache ready to pass to ScanCellCached. +func NewStructKeyCache() *StructKeyCache { + return &StructKeyCache{m: make(map[*arrow.StructType][]string)} +} + +// keyPrefixes returns the escaped `"name":` prefix for each field of st, +// memoized. Safe on a nil receiver (computes without caching). +func (c *StructKeyCache) keyPrefixes(st *arrow.StructType) []string { + if c != nil { + if p, ok := c.m[st]; ok { + return p + } + } + fields := st.Fields() + prefixes := make([]string, len(fields)) + for i := range fields { + name, _ := json.Marshal(fields[i].Name) // JSON-escapes the field name + prefixes[i] = string(name) + ":" + } + if c != nil { + c.m[st] = prefixes + } + return prefixes +} + +// ScanCellCached is ScanCell with a caller-owned StructKeyCache (see +// StructKeyCache) so struct field-name keys are escaped once per result set +// rather than once per row. Pass nil for the un-memoized one-shot behavior. +func ScanCellCached(col arrow.Array, row int, loc *time.Location, keys *StructKeyCache) (driver.Value, error) { + if col.IsNull(row) { + return nil, nil + } + 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: + // Databricks SQL has no unsigned types, so a Uint64 column does not occur + // in practice; this arm is defensive. driver.Value has no uint64 and the + // driver convention is int64 for integers, so a value above MaxInt64 would + // wrap — acceptable for an unreachable path. + return int64(c.Value(row)), nil // #nosec G115 -- see above; unreachable for Databricks types + case *array.Float32: + // Return the native float32, NOT a widened float64: the Thrift path returns + // a float32 driver.Value for a bare FLOAT column, and database/sql's + // asString formats it at bit-size 32 — so widening here would render + // CAST(0.1 AS FLOAT) as "0.10000000149011612" vs Thrift's "0.1". + return c.Value(row), nil + case *array.Float64: + return c.Value(row), nil + case *array.String: + return c.Value(row), nil + case *array.LargeString: + return c.Value(row), nil + case *array.Binary: + return c.Value(row), nil + case *array.Date32: + return inLocation(c.Value(row).ToTime(), loc), nil + case *array.Date64: + return inLocation(c.Value(row).ToTime(), loc), nil + case *array.Timestamp: + dt, ok := col.DataType().(*arrow.TimestampType) + if !ok { + return nil, fmt.Errorf("timestamp column has unexpected datatype %s", col.DataType()) + } + return inLocation(c.Value(row).ToTime(dt.Unit), loc), nil + case *array.Decimal128: + dt := col.DataType().(*arrow.Decimal128Type) + return decimalfmt.ExactString(c.Value(row), dt.Scale), nil + case *array.List, *array.LargeList, *array.FixedSizeList, *array.Map, *array.Struct: + // Nested types (and VARIANT, which arrives as a nested value) render to a + // JSON string matching the Thrift path. + return renderJSONString(col, row, loc, keys) + default: + return nil, fmt.Errorf("scanning arrow type %s is not supported "+ + "(intervals are not yet handled)", col.DataType()) + } +} + +// inLocation renders t in loc, matching the Thrift path's .In(location); a nil +// loc leaves the value in UTC (arrow's ToTime default). +func inLocation(t time.Time, loc *time.Location) time.Time { + if loc == nil { + return t + } + return t.In(loc) +} + +// renderJSONString renders one nested cell as a JSON string (nil for a NULL +// cell). loc is applied to any timestamp/date leaves, matching the top-level +// scan and the Thrift path. +func renderJSONString(col arrow.Array, row int, loc *time.Location, keys *StructKeyCache) (driver.Value, error) { + if col.IsNull(row) { + return nil, nil + } + var b strings.Builder + if err := writeJSON(&b, col, row, loc, keys); err != nil { + return nil, err + } + return b.String(), nil +} + +// writeJSON writes the JSON form of col[row] into b, recursing for nested types. +func writeJSON(b *strings.Builder, col arrow.Array, row int, loc *time.Location, keys *StructKeyCache) error { + if col.IsNull(row) { + b.WriteString("null") + return nil + } + switch c := col.(type) { + case *array.List: + return writeListJSON(b, c.ListValues(), int(c.Offsets()[row]), int(c.Offsets()[row+1]), loc, keys) + case *array.LargeList: + return writeListJSON(b, c.ListValues(), int(c.Offsets()[row]), int(c.Offsets()[row+1]), loc, keys) + case *array.FixedSizeList: + n := int(c.DataType().(*arrow.FixedSizeListType).Len()) + return writeListJSON(b, c.ListValues(), row*n, row*n+n, loc, keys) + case *array.Map: + return writeMapJSON(b, c, row, loc, keys) + case *array.Struct: + return writeStructJSON(b, c, row, loc, keys) + case *array.Decimal128: + // Emit the exact scale-applied decimal as a raw JSON number literal, not a + // float64 — a float64 would render DECIMAL(5,2) 19.99 as 19.990000000000002 + // and corrupt high-precision values. Matches the Thrift path's marshalScalar + // → ValueString (databricks-sql-go#253/#274). + b.WriteString(decimalfmt.ExactString(c.Value(row), col.DataType().(*arrow.Decimal128Type).Scale)) + return nil + case *array.Float32: + // Marshal the native float32, NOT a widened float64: json.Marshal(float32 + // (3.14)) is "3.14" but json.Marshal(float64(float32(3.14))) is + // "3.140000104904175". The Thrift nested path marshals the native float32, + // so widening here would break byte-parity for ARRAY/MAP/STRUCT<…FLOAT…>. + return writeScalarJSON(b, c.Value(row)) + default: + v, err := scalarForJSON(col, row, loc) + if err != nil { + return err + } + return writeScalarJSON(b, v) + } +} + +func writeListJSON(b *strings.Builder, values arrow.Array, start, end int, loc *time.Location, keys *StructKeyCache) error { + b.WriteByte('[') + for i := start; i < end; i++ { + if i > start { + b.WriteByte(',') + } + if err := writeJSON(b, values, i, loc, keys); err != nil { + return err + } + } + b.WriteByte(']') + return nil +} + +func writeMapJSON(b *strings.Builder, m *array.Map, row int, loc *time.Location, keys *StructKeyCache) error { + start, end := int(m.Offsets()[row]), int(m.Offsets()[row+1]) + mapKeys := 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(mapKeys, i, loc) + if err != nil { + return err + } + writeJSONKey(b, kv) + b.WriteByte(':') + if err := writeJSON(b, items, i, loc, keys); err != nil { + return err + } + } + b.WriteByte('}') + return nil +} + +func writeStructJSON(b *strings.Builder, s *array.Struct, row int, loc *time.Location, keys *StructKeyCache) error { + // Field-name keys are constant across rows; keys.keyPrefixes memoizes them per + // result set so this per-row/per-cell path doesn't re-marshal them (a nil cache + // recomputes inline — correct, just not memoized). + prefixes := keys.keyPrefixes(s.DataType().(*arrow.StructType)) + b.WriteByte('{') + for f := 0; f < s.NumField(); f++ { + if f > 0 { + b.WriteByte(',') + } + b.WriteString(prefixes[f]) // pre-escaped `"name":` + if err := writeJSON(b, s.Field(f), row, loc, keys); err != nil { + return err + } + } + b.WriteByte('}') + return nil +} + +// writeJSONKey writes a value as a JSON object key (always a quoted string), +// matching the Thrift path's mapValueContainer: marshal the value the same way a +// leaf is marshaled (marshalScalar → json.Marshal, so a []byte key renders as +// base64 "YWJj", NOT fmt "%v" [97 98 99]), then quote-wrap the result if it is +// not already a JSON string (numbers/bools become "7"/"true" keys). +func writeJSONKey(b *strings.Builder, v any) { + // A time.Time key renders as its quoted .String() (same special-case as + // writeScalarJSON / the Thrift marshal()), not RFC3339. + if t, ok := v.(time.Time); ok { + kb, _ := json.Marshal(t.String()) + b.Write(kb) + return + } + kb, err := json.Marshal(v) + if err != nil { + // Unmarshalable key value — fall back to its stringified form, quoted. + kb, _ = json.Marshal(fmt.Sprintf("%v", v)) + } + if len(kb) > 0 && kb[0] == '"' { + b.Write(kb) // already a JSON string (string / []byte via marshal) + return + } + b.WriteByte('"') + b.Write(kb) + b.WriteByte('"') +} + +// writeScalarJSON writes a scalar leaf, mirroring the Thrift marshal(): a +// 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 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 + } + return ScanCell(col, row, loc) +} diff --git a/internal/arrowscan/arrowscan_test.go b/internal/arrowscan/arrowscan_test.go new file mode 100644 index 00000000..089cea73 --- /dev/null +++ b/internal/arrowscan/arrowscan_test.go @@ -0,0 +1,311 @@ +package arrowscan + +import ( + "testing" + "time" + + "github.com/apache/arrow/go/v12/arrow" + "github.com/apache/arrow/go/v12/arrow/array" + "github.com/apache/arrow/go/v12/arrow/decimal128" + "github.com/apache/arrow/go/v12/arrow/memory" +) + +// ScanCell renders the supported scalar types and rejects an unsupported type +// (rather than returning a silently wrong value). +func TestScanCellScalars(t *testing.T) { + pool := memory.NewGoAllocator() + + t.Run("int64", func(t *testing.T) { + b := array.NewInt64Builder(pool) + defer b.Release() + b.Append(42) + arr := b.NewArray() + defer arr.Release() + v, err := ScanCell(arr, 0, nil) + if err != nil { + t.Fatal(err) + } + if v.(int64) != 42 { + t.Errorf("got %v", v) + } + }) + + t.Run("string", func(t *testing.T) { + b := array.NewStringBuilder(pool) + defer b.Release() + b.Append("hi") + arr := b.NewArray() + defer arr.Release() + v, err := ScanCell(arr, 0, nil) + if err != nil { + t.Fatal(err) + } + if v.(string) != "hi" { + t.Errorf("got %v", v) + } + }) + + t.Run("float32_native", func(t *testing.T) { + // A top-level FLOAT column must scan to a native float32, not a widened + // float64 — matching Thrift, so database/sql's asString renders + // CAST(0.1 AS FLOAT) as "0.1", not "0.10000000149011612". + b := array.NewFloat32Builder(pool) + defer b.Release() + b.Append(0.1) + arr := b.NewArray() + defer arr.Release() + v, err := ScanCell(arr, 0, nil) + if err != nil { + t.Fatal(err) + } + got, ok := v.(float32) + if !ok { + t.Fatalf("want float32, got %T", v) + } + if got != float32(0.1) { + t.Errorf("got %v, want 0.1", got) + } + }) + + t.Run("null", func(t *testing.T) { + b := array.NewInt64Builder(pool) + defer b.Release() + b.AppendNull() + arr := b.NewArray() + defer arr.Release() + v, err := ScanCell(arr, 0, nil) + if err != nil { + t.Fatal(err) + } + if v != nil { + t.Errorf("null should scan to nil, got %v", v) + } + }) + + t.Run("decimal_exact_string", func(t *testing.T) { + // 12345 at scale 2 = "123.45", exact (not a float64). + dt := &arrow.Decimal128Type{Precision: 10, Scale: 2} + b := array.NewDecimal128Builder(pool, dt) + defer b.Release() + b.Append(decimal128.FromU64(12345)) + arr := b.NewArray() + defer arr.Release() + v, err := ScanCell(arr, 0, nil) + if err != nil { + t.Fatal(err) + } + if v.(string) != "123.45" { + t.Errorf("got %v, want 123.45", v) + } + }) + + t.Run("unsupported_type_errors", func(t *testing.T) { + // A duration (INTERVAL) is not yet handled: must error, not return a + // wrong value. + b := array.NewDurationBuilder(pool, &arrow.DurationType{Unit: arrow.Microsecond}) + defer b.Release() + b.Append(1000) + arr := b.NewArray() + defer arr.Release() + if _, err := ScanCell(arr, 0, nil); err == nil { + t.Error("scanning a Duration should return an unsupported-type error") + } + }) +} + +// ScanCell renders DATE / TIMESTAMP in the requested location, matching the +// Thrift path's .In(location); a nil location leaves the value in UTC. +func TestScanCellTimestampLocation(t *testing.T) { + pool := memory.NewGoAllocator() + loc, err := time.LoadLocation("America/New_York") + if err != nil { + t.Skipf("tz database unavailable: %v", err) + } + + // 2026-07-09T12:00:00Z as microseconds since epoch. + utcTS := time.Date(2026, time.July, 9, 12, 0, 0, 0, time.UTC) + b := array.NewTimestampBuilder(pool, &arrow.TimestampType{Unit: arrow.Microsecond}) + defer b.Release() + b.Append(arrow.Timestamp(utcTS.UnixMicro())) + arr := b.NewArray() + defer arr.Release() + + t.Run("location applied", func(t *testing.T) { + v, err := ScanCell(arr, 0, loc) + if err != nil { + t.Fatal(err) + } + got := v.(time.Time) + if got.Location() != loc { + t.Errorf("location = %v, want %v", got.Location(), loc) + } + if !got.Equal(utcTS) { + t.Errorf("instant changed: got %v, want %v", got, utcTS) + } + }) + + t.Run("nil location is UTC", func(t *testing.T) { + v, err := ScanCell(arr, 0, nil) + if err != nil { + t.Fatal(err) + } + if v.(time.Time).Location() != time.UTC { + t.Errorf("nil location should render UTC, got %v", v.(time.Time).Location()) + } + }) +} + +// ScanCell renders nested types (list/struct/map) to a JSON string matching the +// Thrift path. +func TestScanCellNested(t *testing.T) { + pool := memory.NewGoAllocator() + + t.Run("list", func(t *testing.T) { + b := array.NewListBuilder(pool, arrow.PrimitiveTypes.Int64) + defer b.Release() + vb := b.ValueBuilder().(*array.Int64Builder) + b.Append(true) + vb.Append(1) + vb.Append(2) + vb.Append(3) + arr := b.NewArray() + defer arr.Release() + v, err := ScanCell(arr, 0, nil) + if err != nil { + t.Fatal(err) + } + if v.(string) != "[1,2,3]" { + t.Errorf("got %q, want [1,2,3]", v) + } + }) + + t.Run("struct", func(t *testing.T) { + dt := arrow.StructOf( + arrow.Field{Name: "a", Type: arrow.PrimitiveTypes.Int64}, + arrow.Field{Name: "b", Type: arrow.BinaryTypes.String}, + ) + b := array.NewStructBuilder(pool, dt) + defer b.Release() + b.Append(true) + b.FieldBuilder(0).(*array.Int64Builder).Append(1) + b.FieldBuilder(1).(*array.StringBuilder).Append("x") + arr := b.NewArray() + defer arr.Release() + v, err := ScanCell(arr, 0, nil) + if err != nil { + t.Fatal(err) + } + if v.(string) != `{"a":1,"b":"x"}` { + t.Errorf("got %q, want {\"a\":1,\"b\":\"x\"}", v) + } + }) + + t.Run("map", func(t *testing.T) { + b := array.NewMapBuilder(pool, arrow.BinaryTypes.String, arrow.PrimitiveTypes.Int64, false) + defer b.Release() + kb := b.KeyBuilder().(*array.StringBuilder) + ib := b.ItemBuilder().(*array.Int64Builder) + b.Append(true) + kb.Append("k") + ib.Append(9) + arr := b.NewArray() + defer arr.Release() + v, err := ScanCell(arr, 0, nil) + if err != nil { + t.Fatal(err) + } + if v.(string) != `{"k":9}` { + t.Errorf("got %q, want {\"k\":9}", v) + } + }) + + t.Run("nested_decimal_exact", func(t *testing.T) { + // A decimal inside a struct must render as an exact JSON number, not a + // lossy float64 (19.99, not 19.990000000000002) — matching Thrift's + // marshalScalar. + dt := arrow.StructOf(arrow.Field{Name: "d", Type: &arrow.Decimal128Type{Precision: 5, Scale: 2}}) + b := array.NewStructBuilder(pool, dt) + defer b.Release() + b.Append(true) + b.FieldBuilder(0).(*array.Decimal128Builder).Append(decimal128.FromU64(1999)) + arr := b.NewArray() + defer arr.Release() + v, err := ScanCell(arr, 0, nil) + if err != nil { + t.Fatal(err) + } + if v.(string) != `{"d":19.99}` { + t.Errorf("got %q, want {\"d\":19.99}", v) + } + }) + + t.Run("nested_float32_exact", func(t *testing.T) { + // A float32 inside a struct must marshal as the native float32 (3.14), not + // a widened float64 (3.140000104904175) — matching Thrift's nested path, + // which marshals the native float32. + dt := arrow.StructOf(arrow.Field{Name: "f", Type: arrow.PrimitiveTypes.Float32}) + b := array.NewStructBuilder(pool, dt) + defer b.Release() + b.Append(true) + b.FieldBuilder(0).(*array.Float32Builder).Append(3.14) + arr := b.NewArray() + defer arr.Release() + v, err := ScanCell(arr, 0, nil) + if err != nil { + t.Fatal(err) + } + if v.(string) != `{"f":3.14}` { + t.Errorf("got %q, want {\"f\":3.14}", v) + } + }) + + t.Run("nested_null", func(t *testing.T) { + b := array.NewListBuilder(pool, arrow.PrimitiveTypes.Int64) + defer b.Release() + b.AppendNull() + arr := b.NewArray() + defer arr.Release() + v, err := ScanCell(arr, 0, nil) + if err != nil { + t.Fatal(err) + } + if v != nil { + t.Errorf("null list should scan to nil, got %v", v) + } + }) +} + +// A StructKeyCache must produce byte-identical output to the nil (recompute-inline) +// path, and reused across rows of the same struct type — so the round-trip +// through the cache can never silently diverge from the one-shot rendering. +func TestScanCellCachedMatchesUncached(t *testing.T) { + pool := memory.NewGoAllocator() + dt := arrow.StructOf( + arrow.Field{Name: "a", Type: arrow.PrimitiveTypes.Int64}, + arrow.Field{Name: `q"x`, Type: arrow.BinaryTypes.String}, // needs escaping + ) + b := array.NewStructBuilder(pool, dt) + defer b.Release() + for i := 0; i < 3; i++ { // multiple rows: exercises the memoized second+ hit + b.Append(true) + b.FieldBuilder(0).(*array.Int64Builder).Append(int64(i)) + b.FieldBuilder(1).(*array.StringBuilder).Append("v") + } + arr := b.NewArray() + defer arr.Release() + + cache := NewStructKeyCache() + for row := 0; row < 3; row++ { + uncached, err := ScanCell(arr, row, nil) + if err != nil { + t.Fatal(err) + } + cached, err := ScanCellCached(arr, row, nil, cache) + if err != nil { + t.Fatal(err) + } + if cached != uncached { + t.Errorf("row %d: cached %q != uncached %q", row, cached, uncached) + } + } +} diff --git a/internal/backend/kernel/backend.go b/internal/backend/kernel/backend.go new file mode 100644 index 00000000..8427ddeb --- /dev/null +++ b/internal/backend/kernel/backend.go @@ -0,0 +1,246 @@ +//go:build cgo && databricks_kernel + +package kernel + +/* +#include +#include "databricks_kernel.h" +*/ +import "C" + +import ( + "context" + "errors" + "fmt" + "time" + + "github.com/databricks/databricks-sql-go/internal/backend" +) + +// 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 accepts any server cert (maps the driver's + // WithSkipTLSHostVerify / TLSConfig.InsecureSkipVerify). crypto/tls's + // InsecureSkipVerify disables both chain validation and the hostname check, + // so the kernel path relaxes both to match. + TLSSkipVerify bool + + // ProxyURL configures an HTTP proxy, already resolved for this endpoint from + // 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 + + // 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 +// 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 { + // Fail fast on an already-cancelled context before the blocking kernel_session + // _open (which the C ABI does not let us interrupt mid-call). + if err := ctx.Err(); err != nil { + return err + } + initKernelLogging() + klog("OpenSession host=%s httpPath=%s warehouse=%s", k.cfg.Host, k.cfg.HTTPPath, k.cfg.WarehouseID) + + 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", toConnError(err)) + } + // kernel_session_open consumes the config on EVERY path — success and + // failure alike (it reclaims the box up front). So we free the config + // ourselves only when we bail out BEFORE reaching kernel_session_open (a + // setter error below); once that call is made, ownership has transferred and + // a free here would double-free. + consumed := false + defer func() { + if !consumed { + 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", toConnError(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", toConnError(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", toConnError(err)) + } + + // TLS: crypto/tls's InsecureSkipVerify accepts any server cert, so relax both + // chain validation and the hostname check — mapping only one would leave the + // kernel path stricter than the Thrift path it mirrors (a self-signed cert + // would still be rejected). + if k.cfg.TLSSkipVerify { + if err := call(func() C.KernelStatusCode { + return C.kernel_session_config_set_tls_allow_self_signed(cfg, C.bool(true)) + }); err != nil { + return fmt.Errorf("kernel: set_tls_allow_self_signed: %w", toConnError(err)) + } + if err := call(func() C.KernelStatusCode { + return C.kernel_session_config_set_tls_skip_hostname_verification(cfg, C.bool(true)) + }); err != nil { + return fmt.Errorf("kernel: set_tls_skip_hostname_verification: %w", toConnError(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", toConnError(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, toConnError(errSet)) + } + } + + // kernel_session_open takes ownership of cfg here, on both success and + // failure — mark consumed before checking the error so the deferred free + // never runs on the already-consumed config. + var sess *C.kernel_session_t + err := call(func() C.KernelStatusCode { return C.kernel_session_open(cfg, &sess) }) + consumed = true + if err != nil { + return fmt.Errorf("kernel: session_open: %w", toConnError(err)) + } + 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 toConnError(err) +} + +// SessionValid backs conn.IsValid → pool eviction. No I/O; inspects state +// captured at OpenSession and updated by markSessionDead. +func (k *KernelBackend) SessionValid() bool { return k.valid && k.session != nil } + +// markSessionDead marks the session unusable so SessionValid() → conn.IsValid() +// returns false and database/sql discards this conn on return to the pool. Called +// from the statement/read path when a kernel error is session-fatal (isSessionFatal): +// it evicts the dead conn WITHOUT returning driver.ErrBadConn, so the statement is +// never transparently re-run (no duplicate write — the H1 constraint). The backend +// is single-owner per conn (only the cancel watcher shares state, and it touches +// only the canceller's inflight slot), so this write is race-free. +func (k *KernelBackend) markSessionDead() { k.valid = false } + +// evictIfSessionFatal marks the session dead when err is a session-fatal +// KernelError, so a conn whose server session died mid-life is evicted from the +// pool. A no-op for nil, non-KernelError, or non-fatal errors. +func (k *KernelBackend) evictIfSessionFatal(err error) { + if ke, ok := err.(*KernelError); ok && isSessionFatal(ke.Code) { + k.markSessionDead() + } +} + +// SessionID is the per-conn id (conn.id). Valid after OpenSession. +func (k *KernelBackend) SessionID() string { return k.sessionID } + +// 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) { + // Bound parameters are not yet wired for the kernel backend. Reject them with + // a clear error rather than silently shipping the query with unbound + // placeholders (which would behave differently than Thrift). Parameters arrive + // per-query, so this is an execute-time error, not a connect-time one. Return a + // non-nil Operation per the Backend contract. + if len(req.Params) > 0 { + return &kernelOp{}, errors.New("databricks: query parameters are not yet supported by the kernel backend; " + + "inline the values or use the default (Thrift) backend") + } + return k.execute(ctx, req) +} diff --git a/internal/backend/kernel/cgo.go b/internal/backend/kernel/cgo.go new file mode 100644 index 00000000..60310118 --- /dev/null +++ b/internal/backend/kernel/cgo.go @@ -0,0 +1,280 @@ +//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 directives below name the kernel library to link but carry no search +// paths, so the header and static lib locations are supplied at build time via +// the standard CGO_CFLAGS / CGO_LDFLAGS environment variables, e.g.: +// +// CGO_CFLAGS="-I/include" \ +// CGO_LDFLAGS="-L/target/release -Wl,-rpath,/target/release" \ +// go build -tags databricks_kernel ./... +// +// A shippable build instead links a committed per-platform prebuilt static lib +// via a ${SRCDIR}-relative path (the go-duckdb duckdb-go-bindings model); wiring +// that + a tagged CI job is a distribution follow-up. +package kernel + +/* +#cgo LDFLAGS: -ldatabricks_sql_kernel -ldl -lm +#include +#include "databricks_kernel.h" +*/ +import "C" + +import ( + "fmt" + "os" + "runtime" + "sync" + "unsafe" + + dbsqlerrint "github.com/databricks/databricks-sql-go/internal/errors" + "github.com/databricks/databricks-sql-go/logger" +) + +// ─── Debug logging ─────────────────────────────────────────────────────────── +// +// Gated on DBSQL_KERNEL_DEBUG so it is OFF by default and, in particular, OFF +// during benchmarks (debug logging perturbs latency). Every binding step logs +// through klog when enabled — entry, status codes, handle addresses, batch +// counts — which is what makes a failing e2e cheap to diagnose. The same flag +// also installs the kernel's own Rust (tracing) logs via initKernelLogging, so +// binding and kernel logs interleave on stderr as one stream. +var kdebug = os.Getenv("DBSQL_KERNEL_DEBUG") != "" + +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 } + +// initLoggingOnce guards kernel_init_logging, which is process-wide and +// first-call-wins in the kernel. We install the kernel subscriber lazily on the +// first session open rather than in init(), so a process that never opens a +// kernel session installs nothing. +var initLoggingOnce sync.Once + +// initKernelLogging turns on the kernel's own Rust (tracing) logs, gated on the +// same DBSQL_KERNEL_DEBUG flag as klog so both are OFF by default and, in +// particular, OFF during benchmarks — the subscriber is never installed in a +// benchmark run. level=NULL lets the kernel honor RUST_LOG (default warn); +// file_path=NULL sends kernel logs to stderr so they interleave with klog on one +// stream. Best-effort: Internal (e.g. the host already installed a global +// subscriber) is a documented, benign outcome — logged, never fatal to connect. +func initKernelLogging() { + if !kdebug { + return + } + initLoggingOnce.Do(func() { + if err := call(func() C.KernelStatusCode { + return C.kernel_init_logging(nil, nil) + }); err != nil { + klog("kernel_init_logging: %v (kernel logs unavailable; proceeding)", err) + } + }) +} + +// call runs a fallible kernel entry point and, on a non-Success status, reads +// the kernel's thread-local last error into a Go error. +// +// 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. +func call(fn func() C.KernelStatusCode) error { + runtime.LockOSThread() + defer runtime.UnlockOSThread() + st := fn() + if st == C.KernelStatusCode_Success { + return nil + } + return lastError(st) +} + +// fireCancel dispatches a server-side cancel and reports whether the RPC was +// actually sent. The kernel returns dispatched=false while no server statement +// id has been observed yet (the window before the execute POST returns), and +// true once the cancel RPC goes out — which is the watcher's cue to stop +// re-firing. +func fireCancel(canceller *C.kernel_statement_canceller_t) bool { + var dispatched C.bool + C.kernel_statement_canceller_cancel(canceller, &dispatched) + return bool(dispatched) +} + +// lastError reads the kernel's thread-local last error and copies its string +// fields out immediately — the C `char*` fields are valid only until the next +// FFI call on this thread. Must run on the same OS thread as the failing call; +// 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) + // Also emit at the driver's default (Warn) level — no SQL text or PII, just + // the status/sqlstate/http fields plus the server query id (a correlation + // token, not PII) so on-call can pivot to server-side query history — so a + // kernel-path failure is visible without DBSQL_KERNEL_DEBUG. This is the error + // path only (never the hot per-row/per-batch path), so it does not perturb + // benchmarks. + logger.Logger.Warn().Msgf("databricks: kernel call failed: code=%d sqlstate=%q vendor=%d http=%d retryable=%v queryId=%q", + ke.Code, ke.SQLState, ke.VendorCode, ke.HTTPStatus, ke.Retryable, ke.QueryID) + return ke +} + +// 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 { + // Append the server query id when present — it is the one correlation handle + // to server-side query history, and StatementID() is "" on this backend. + q := "" + if e.QueryID != "" { + q = fmt.Sprintf(", queryId=%s", e.QueryID) + } + if e.SQLState != "" { + return fmt.Sprintf("kernel: %s (sqlstate=%s, code=%d%s)", e.Message, e.SQLState, e.Code, q) + } + return fmt.Sprintf("kernel: %s (code=%d%s)", e.Message, e.Code, q) +} + +// Status codes mirrored as Go ints so non-cgo code (tests, error mapping) can +// 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 is a *transient* connection +// failure — one where retrying on a fresh connection could succeed. On the +// session-lifecycle path this is wrapped as driver.ErrBadConn so database/sql +// retries connect. Unauthenticated is deliberately excluded: a wrong/expired PAT +// is permanent, so retrying it just burns connect attempts (and can worsen +// server-side auth rate-limiting) and fails identically — matching Thrift, which +// only treats an invalid session handle (a liveness signal), not a 401, as +// bad-conn. (Auth failure still marks the session dead for pool eviction — see +// isSessionFatal.) +func isBadConnection(code int) bool { + switch code { + case statusUnavailable, statusNetworkError: + return true + default: + return false + } +} + +// isSessionFatal reports whether a status code means the server-side session is no +// longer usable, so the conn must be evicted from the pool rather than reused for +// the next query. Broader than isBadConnection: it also covers Unauthenticated (an +// expired/revoked token kills the session) — but unlike isBadConnection it is NOT +// used to produce driver.ErrBadConn on the statement path, so eviction happens +// without database/sql replaying the statement (see toStatementError + the +// KernelBackend.markSessionDead call sites). +func isSessionFatal(code int) bool { + switch code { + case statusUnauthenticated, statusUnavailable, statusNetworkError: + return true + default: + return false + } +} + +// toConnError classifies a kernel error on a SESSION-lifecycle path (open/close/ +// config, where nothing has executed): a status that means the session is unusable +// is wrapped as a bad-connection error, which identifies as driver.ErrBadConn so +// database/sql evicts the conn from the pool. Safe here because no statement ran, +// so there is nothing for database/sql to unsafely re-run. Other kernel errors — +// and plain (non-KernelError) errors — are returned unchanged, carrying sqlstate. +func toConnError(err error) error { + if err == nil { + return nil + } + ke, ok := err.(*KernelError) + if !ok { + return err + } + if isBadConnection(ke.Code) { + return dbsqlerrint.NewBadConnectionError(ke) + } + return ke +} + +// toStatementError classifies a kernel error on the STATEMENT path (execute and +// result read). It NEVER returns driver.ErrBadConn: once a statement has been +// sent, a network/unavailable failure surfaced afterward may have committed +// server-side, and driver.ErrBadConn would make database/sql transparently +// re-run the statement — a silent duplicate write for a non-idempotent +// INSERT/UPDATE/MERGE. This mirrors the kernel's own retry contract +// (ExecuteStatement is NonIdempotent, retried only on connect-phase failures) and +// the Thrift backend (ExecuteStatement is non-retryable), and honors Go's +// driver.ErrBadConn rule ("never return ErrBadConn if the server might have +// performed the operation"). The kernel has already exhausted its safe internal +// retries by the time we see the error. Returns the KernelError (or plain error) +// unchanged, carrying sqlstate. +func toStatementError(err error) error { + if err == nil { + return nil + } + if ke, ok := err.(*KernelError); ok { + return ke + } + return err +} + +// cStr wraps C.CString with a guaranteed free. The kernel copies strings into +// owned Rust memory on receipt, so freeing immediately after the call is safe. +// Use: cs := newCStr(s); defer cs.free(); ...C.fn(cs.c)... +type cStr struct{ c *C.char } + +func newCStr(s string) cStr { return cStr{c: C.CString(s)} } + +func (s cStr) free() { + if s.c != nil { + C.free(unsafe.Pointer(s.c)) + } +} diff --git a/internal/backend/kernel/kernel_test.go b/internal/backend/kernel/kernel_test.go new file mode 100644 index 00000000..9ac4e0a4 --- /dev/null +++ b/internal/backend/kernel/kernel_test.go @@ -0,0 +1,160 @@ +//go:build cgo && databricks_kernel + +package kernel + +import ( + "context" + "database/sql/driver" + "errors" + "testing" + + "github.com/databricks/databricks-sql-go/internal/backend" +) + +// isBadConnection is the TRANSIENT-failure set (retrying a fresh connect could +// succeed) — used to produce driver.ErrBadConn on the session-lifecycle path. +// Unauthenticated is excluded: a wrong/expired PAT is permanent, not retryable. +func TestIsBadConnection(t *testing.T) { + transient := []int{statusUnavailable, statusNetworkError} + for _, code := range transient { + if !isBadConnection(code) { + t.Errorf("code %d should be a (transient) bad connection", code) + } + } + // Unauthenticated is session-fatal but NOT retryable, so it is not bad-conn. + notBad := []int{statusUnauthenticated, statusInvalidArgument, statusSqlError, statusTimeout} + for _, code := range notBad { + if isBadConnection(code) { + t.Errorf("code %d should not be a bad connection", code) + } + } +} + +// isSessionFatal is the broader "session is dead, evict the conn" set — it adds +// Unauthenticated (an expired token kills the session) on top of the transient set. +func TestIsSessionFatal(t *testing.T) { + fatal := []int{statusUnauthenticated, statusUnavailable, statusNetworkError} + for _, code := range fatal { + if !isSessionFatal(code) { + t.Errorf("code %d should be session-fatal", code) + } + } + notFatal := []int{statusInvalidArgument, statusSqlError, statusTimeout} + for _, code := range notFatal { + if isSessionFatal(code) { + t.Errorf("code %d should not be session-fatal", code) + } + } +} + +// evictIfSessionFatal flips SessionValid()→false on a session-fatal error (so the +// pool discards the conn) WITHOUT the error being driver.ErrBadConn (so the +// statement is never transparently re-run — the H1 constraint). +func TestEvictIfSessionFatal(t *testing.T) { + // valid tracks the session-dead flag SessionValid() gates on; the opaque + // session pointer is orthogonal here (can't construct the incomplete C type), + // so assert on k.valid directly. + k := &KernelBackend{valid: true} + + // Non-fatal (e.g. a SQL error) leaves the session valid. + k.evictIfSessionFatal(&KernelError{Code: statusSqlError}) + if !k.valid { + t.Error("a SQL error must not evict the session") + } + + // A session-fatal error marks the session dead, and the surfaced statement-path + // error is NOT driver.ErrBadConn (so database/sql won't re-run the statement). + fatal := &KernelError{Code: statusUnavailable, Message: "session gone"} + k.evictIfSessionFatal(fatal) + if k.valid { + t.Error("a session-fatal error must evict the session (valid=false)") + } + if errors.Is(toStatementError(fatal), driver.ErrBadConn) { + t.Error("the statement-path error must not be driver.ErrBadConn (no replay)") + } +} + +// toConnError (session-lifecycle path) wraps a session-unusable KernelError as +// driver.ErrBadConn so database/sql evicts the conn, and leaves other errors and +// their sqlstate intact. +func TestToConnError(t *testing.T) { + if toConnError(nil) != nil { + t.Fatal("nil should map to nil") + } + + badConn := &KernelError{Code: statusUnavailable, Message: "gone"} + if !errors.Is(toConnError(badConn), driver.ErrBadConn) { + t.Errorf("unavailable kernel error on the session path should identify as driver.ErrBadConn") + } + + sqlErr := &KernelError{Code: statusSqlError, Message: "boom", SQLState: "42703"} + ke, ok := toConnError(sqlErr).(*KernelError) + if !ok { + t.Fatalf("sql error should remain a *KernelError, got %T", toConnError(sqlErr)) + } + if ke.SQLState != "42703" { + t.Errorf("sqlstate lost: got %q", ke.SQLState) + } +} + +// toStatementError (execute/read path) must NEVER return driver.ErrBadConn — even +// for a network/unavailable status — so database/sql cannot transparently re-run a +// statement that may have already executed server-side (silent duplicate write). +func TestToStatementErrorNeverBadConn(t *testing.T) { + if toStatementError(nil) != nil { + t.Fatal("nil should map to nil") + } + + for _, code := range []int{statusUnavailable, statusNetworkError, statusUnauthenticated} { + err := toStatementError(&KernelError{Code: code, Message: "post-execute failure"}) + if errors.Is(err, driver.ErrBadConn) { + t.Errorf("statement-path error (code=%d) must NOT identify as driver.ErrBadConn "+ + "(would let database/sql re-run a possibly-committed statement)", code) + } + } + + // sqlstate still preserved. + sqlErr := &KernelError{Code: statusSqlError, Message: "boom", SQLState: "42703"} + ke, ok := toStatementError(sqlErr).(*KernelError) + if !ok { + t.Fatalf("sql error should remain a *KernelError, got %T", toStatementError(sqlErr)) + } + if ke.SQLState != "42703" { + t.Errorf("sqlstate lost: got %q", ke.SQLState) + } +} + +// Bound parameters are rejected up front by Execute with a clear error, before +// any session/C work — so this runs on a zero-value backend. The returned +// Operation must be non-nil (Backend contract) and its Close must report +// closed=false, since no server statement was ever created (a phantom +// CLOSE_STATEMENT would otherwise be recorded for it). +func TestExecuteRejectsParams(t *testing.T) { + k := &KernelBackend{} + op, err := k.Execute(context.Background(), backend.ExecRequest{ + Query: "SELECT ?", + Params: []backend.Param{{Name: "x"}}, + }) + if err == nil { + t.Fatal("expected an error for bound parameters, got nil") + } + if op == nil { + t.Fatal("Execute must return a non-nil Operation per the Backend contract") + } + closed, closeErr := op.Close(context.Background()) + if closeErr != nil { + t.Errorf("Close error = %v, want nil", closeErr) + } + if closed { + t.Error("Close on a handle-less op must report closed=false (no CLOSE_STATEMENT)") + } + if got := op.AffectedRows(); got != 0 { + t.Errorf("AffectedRows on a handle-less op = %d, want 0", got) + } +} + +// The cell/nested rendering (ScanCell and the JSON grammar) now lives in the +// untagged internal/arrowscan package, where its tests run in the default +// CGO_ENABLED=0 build; see arrowscan_test.go. The decimal formatter lives in +// internal/decimalfmt. This file keeps the kernel-specific tests: error mapping, +// bad-connection classification, and the bound-params rejection. diff --git a/internal/backend/kernel/operation.go b/internal/backend/kernel/operation.go new file mode 100644 index 00000000..c551f56b --- /dev/null +++ b/internal/backend/kernel/operation.go @@ -0,0 +1,251 @@ +//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; bound parameters are rejected up front by Execute. +func (k *KernelBackend) execute(ctx context.Context, req backend.ExecRequest) (backend.Operation, error) { + // Log the SQL length, not the text: query bodies can carry PII/secrets in + // WHERE/INSERT/SET, and this goes to stderr. Matches the driver's own + // debuglog convention (conn.ExecContext logs sql.len=%d). + klog("Execute sql.len=%d", len(req.Query)) + + var stmt *C.kernel_statement_t + if err := call(func() C.KernelStatusCode { + return C.kernel_session_new_statement(k.session, &stmt) + }); err != nil { + k.evictIfSessionFatal(err) + return &kernelOp{}, fmt.Errorf("kernel: new_statement: %w", toStatementError(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) + k.evictIfSessionFatal(err) + return &kernelOp{}, fmt.Errorf("kernel: set_sql: %w", toStatementError(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 the kernel reports the + // cancel RPC was actually dispatched (dispatched=true, i.e. the id appeared + // and the RPC went out), then stop: one real cancel is enough, and further + // fires would only hammer the server. Falls back to firing until execute + // returns (done) if the RPC never dispatches. + done := make(chan struct{}) + var watcherWg sync.WaitGroup + if canceller != nil && ctx.Done() != nil { + watcherWg.Add(1) + go func() { + defer watcherWg.Done() + select { + case <-ctx.Done(): + case <-done: + return + } + klog("ctx.Done (%v) → firing canceller (with retry until dispatched)", ctx.Err()) + ticker := time.NewTicker(250 * time.Millisecond) + defer ticker.Stop() + if fireCancel(canceller) { + klog("cancel dispatched on first fire") + return + } + for { + select { + case <-done: + return + case <-ticker.C: + if fireCancel(canceller) { + klog("cancel dispatched, watcher stopping") + return + } + } + } + }() + } + + // 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{backend: k, stmt: stmt, location: k.cfg.Location} + if execErr != nil { + // Prefer the caller's ctx error when the ctx was cancelled (database/sql + // convention), keeping the kernel error as the cause. + 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) + // A session-fatal status (expired token, dropped/unavailable session) means + // this conn is unusable: evict it so the pool doesn't hand it out again. We + // still return the PLAIN error (toStatementError, never ErrBadConn), so the + // conn is discarded without database/sql re-running the statement (H1). + k.evictIfSessionFatal(execErr) + op.close() + return op, fmt.Errorf("kernel: execute: %w", toStatementError(execErr)) + } + op.exec = exec + // Capture the modified-row count now, while exec is live — the operation is + // closed (nulling exec) before AffectedRows is read on the ExecContext path. + op.affectedRows = int64(C.kernel_executed_statement_num_modified_rows(exec)) + klog("Execute OK stmt=%p exec=%p affectedRows=%d", stmt, exec, op.affectedRows) + return op, nil +} + +// kernelOp implements backend.Operation over a sync executed statement. +type kernelOp struct { + // backend is the owning connection's backend, held so a session-fatal error on + // the result-read path can evict the conn (markSessionDead), mirroring how the + // Thrift thriftOperation holds its *Backend. + backend *KernelBackend + stmt *C.kernel_statement_t + exec *C.kernel_executed_statement_t + closed bool + // affectedRows is the modified-row count captured at execute time. It is + // cached (not read live from exec) because the caller closes the operation — + // which nulls exec — before reading AffectedRows (see conn.ExecContext). + affectedRows int64 + // location renders DATE / TIMESTAMP values in the session time zone, matching + // the Thrift path; nil means UTC. Carried onto the rows built by Results. + location *time.Location +} + +var _ backend.Operation = (*kernelOp)(nil) + +// StatementID returns "": the kernel C ABI exposes no success-path statement/query +// id accessor (only the error path carries a query_id, on KernelError). Because +// conn.ExecContext/QueryContext gate per-statement telemetry on +// StatementID() != "", kernel queries currently emit no EXECUTE_STATEMENT metric +// and their query-id log field is empty — there is no session-id fallback on that +// gate. Surfacing an id needs a kernel accessor (e.g. +// kernel_executed_statement_query_id); tracked as a follow-up. +func (o *kernelOp) StatementID() string { return "" } + +// AffectedRows is the modified-row count for ExecContext. It returns the value +// cached at execute time, so it is correct even after the operation is closed +// (the ExecContext path closes the op before reading this). +func (o *kernelOp) AffectedRows() int64 { + return o.affectedRows +} + +// 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 { + // No Rows is returned to own teardown, and the query path does not call + // Operation.Close on a Results error — so close the handles here to avoid + // leaking the statement / executed handle (and its server operation). + o.backend.evictIfSessionFatal(err) + o.close() + return nil, fmt.Errorf("kernel: get_result_stream: %w", toStatementError(err)) + } + 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 + // Report closed=true only when there was actually a handle to tear down. A + // handle-less op (e.g. the &kernelOp{} returned on the bound-params error + // path) reaches here via conn.ExecContext's unconditional Close; returning + // true would record a phantom CLOSE_STATEMENT for a statement that never hit + // the server. Matches the backend.Operation contract (closed=false when the + // operation had no handle) and the Thrift backend's !hasHandle() behavior. + didClose := o.exec != nil || o.stmt != nil + if o.exec != nil { + C.kernel_executed_statement_close(o.exec) + o.exec = nil + } + if o.stmt != nil { + C.kernel_statement_close(o.stmt) + o.stmt = nil + } + klog("kernelOp closed (didClose=%v)", didClose) + return didClose +} + +// 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 toStatementError(cause) +} diff --git a/internal/backend/kernel/rows.go b/internal/backend/kernel/rows.go new file mode 100644 index 00000000..681e4077 --- /dev/null +++ b/internal/backend/kernel/rows.go @@ -0,0 +1,189 @@ +//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/cdata" + "github.com/databricks/databricks-sql-go/internal/arrowscan" + 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 + chunkCount int // cumulative batches fetched, for OnChunkFetched + closed bool + eof bool + // keyCache memoizes struct field-name JSON keys for this result set so + // per-row rendering doesn't re-marshal constant names. Scoped to this Rows + // (freed with it) — not a process-global, which would leak (round-2 N1). + keyCache *arrowscan.StructKeyCache +} + +// newKernelRows fetches the schema up front (for Columns()) and returns the row +// iterator; batches are pulled lazily on Next. +func newKernelRows(ctx context.Context, op *kernelOp, stream *C.kernel_result_stream_t, cb *dbsqlrows.TelemetryCallbacks) (driver.Rows, error) { + r := &kernelRows{ctx: ctx, op: op, stream: stream, callbacks: cb, keyCache: arrowscan.NewStructKeyCache()} + + var csch C.struct_ArrowSchema + if err := call(func() C.KernelStatusCode { + return C.kernel_result_stream_get_schema(stream, &csch) + }); err != nil { + op.backend.evictIfSessionFatal(err) + r.Close() + return nil, fmt.Errorf("kernel: get_schema: %w", toStatementError(err)) + } + 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 := arrowscan.ScanCellCached(rec.Column(c), r.rowInCur, r.op.location, r.keyCache) + if err != nil { + return fmt.Errorf("kernel: scan col %d (%s): %w", c, r.cols[c], err) + } + 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 { + // Honor cancellation at batch boundaries: check ctx before entering the + // blocking C fetch (which cannot itself observe ctx). This does NOT interrupt + // a fetch already in flight — and database/sql's own cancel watcher can't + // either: its Rows.Close takes rs.closemu.Lock(), which blocks until the + // in-progress Next (holding the RLock) returns, so the stream close waits for + // the C call to finish on its own. A single hung CloudFetch batch is therefore + // uninterruptible (the kernel exposes no per-download timeout); mid-fetch + // cancellation would need the execute path's watcher/canceller applied here. + if r.ctx != nil { + if err := r.ctx.Err(); err != nil { + return err + } + } + 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 { + r.op.backend.evictIfSessionFatal(err) + return fmt.Errorf("kernel: next_batch: %w", toStatementError(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 + r.chunkCount++ + if r.callbacks != nil && r.callbacks.OnChunkFetched != nil { + // chunkCount is cumulative (per the callback contract). bytesDownloaded, + // chunkIndex, and latency are left 0: the kernel does CloudFetch/decompress + // internally and hands back ready Arrow batches, so the Go side never sees + // the compressed wire bytes or per-chunk fetch latency the Thrift path + // reports — a fabricated number would be worse than a truthful zero. + r.callbacks.OnChunkFetched(r.chunkCount, 0, 0, 0, 0) + } + klog("nextBatch: %d rows (chunk %d)", rec.NumRows(), r.chunkCount) + return nil +} diff --git a/internal/config/config.go b/internal/config/config.go index fb7c6afb..12b59395 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -125,6 +125,13 @@ type UserConfig struct { // See databricks/databricks-sql-go#274. UseArrowNativeDecimalDSN bool CloudFetchConfig + // UseKernel selects the SEA-via-kernel backend instead of Thrift. See the + // WithUseKernel connector option for the build requirements. DSN: useKernel=true. + UseKernel bool + // WarehouseID is the bare SQL warehouse id, used by the kernel backend (which + // addresses a warehouse by id) in preference to HTTPPath. The Thrift backend + // ignores it and routes by HTTPPath. DSN: warehouseId=. + WarehouseID string } // DeepCopy returns a true deep copy of UserConfig @@ -171,6 +178,8 @@ func (ucfg UserConfig) DeepCopy() UserConfig { EnableTelemetry: ucfg.EnableTelemetry, TelemetryBatchSize: ucfg.TelemetryBatchSize, TelemetryFlushInterval: ucfg.TelemetryFlushInterval, + UseKernel: ucfg.UseKernel, + WarehouseID: ucfg.WarehouseID, } } @@ -320,6 +329,17 @@ func ParseDSN(dsn string) (UserConfig, error) { ucfg.UseArrowNativeDecimalDSN = useArrowNativeDecimal } + // Kernel backend parameters + if useKernel, ok, err := params.extractAsBool("useKernel"); ok { + if err != nil { + return UserConfig{}, err + } + ucfg.UseKernel = useKernel + } + if warehouseID, ok := params.extract("warehouseId"); ok { + ucfg.WarehouseID = warehouseID + } + // Telemetry parameters if enableTelemetry, ok, err := params.extractAsBool("enableTelemetry"); ok { if err != nil { diff --git a/internal/config/config_test.go b/internal/config/config_test.go index d3574fd1..4c70e993 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -586,6 +586,34 @@ func TestParseConfig(t *testing.T) { wantCfg: UserConfig{}, wantErr: true, }, + { + name: "kernel backend params useKernel and warehouseId", + args: args{dsn: "token:supersecret@example.cloud.databricks.com:443/sql/1.0/endpoints/12346a5b5b0e123a?useKernel=true&warehouseId=abc123"}, + wantCfg: UserConfig{ + Protocol: "https", + Host: "example.cloud.databricks.com", + Port: 443, + MaxRows: defaultMaxRows, + Authenticator: &pat.PATAuth{AccessToken: "supersecret"}, + AccessToken: "supersecret", + HTTPPath: "/sql/1.0/endpoints/12346a5b5b0e123a", + SessionParams: make(map[string]string), + RetryMax: 4, + RetryWaitMin: 1 * time.Second, + RetryWaitMax: 30 * time.Second, + CloudFetchConfig: defCloudConfig, + UseKernel: true, + WarehouseID: "abc123", + }, + wantURL: "https://example.cloud.databricks.com:443/sql/1.0/endpoints/12346a5b5b0e123a", + wantErr: false, + }, + { + name: "malformed useKernel is an error", + args: args{dsn: "token:supersecret@example.cloud.databricks.com:443/sql/1.0/endpoints/12346a5b5b0e123a?useKernel=notabool"}, + wantCfg: UserConfig{}, + wantErr: true, + }, } for i, tt := range tests { fmt.Println(i) @@ -649,6 +677,8 @@ func TestUserConfig_DeepCopy(t *testing.T) { UserAgentEntry: "test", Location: location, SessionParams: map[string]string{"a": "32", "b": "4"}, + UseKernel: true, + WarehouseID: "wh-abc123", } cfg_copy := cfg.DeepCopy() diff --git a/internal/decimalfmt/decimalfmt.go b/internal/decimalfmt/decimalfmt.go new file mode 100644 index 00000000..7d21ebf5 --- /dev/null +++ b/internal/decimalfmt/decimalfmt.go @@ -0,0 +1,47 @@ +// Package decimalfmt renders Arrow decimal128 values as exact fixed-point +// strings, shared by the Thrift (arrowbased) and kernel result paths so a +// DECIMAL renders identically regardless of backend. Keeping one implementation +// means a precision fix (e.g. databricks-sql-go#274) lands in both at once. +package decimalfmt + +import ( + "math/big" + "strings" + + "github.com/apache/arrow/go/v12/arrow/decimal128" +) + +// ExactString renders an Arrow decimal128 as an exact fixed-point string, +// applying scale by string placement rather than float conversion. This +// preserves precision a float64 would lose beyond ~17 significant digits +// (databricks-sql-go#274). A negative scale is not produced by the server and is +// treated as scale 0 rather than panicking. +func ExactString(n decimal128.Num, scale int32) string { + unscaled := n.BigInt() // exact signed unscaled integer + neg := unscaled.Sign() < 0 + digits := new(big.Int).Abs(unscaled).String() + + var b strings.Builder + if neg { + b.WriteByte('-') + } + + if scale <= 0 { + b.WriteString(digits) + return b.String() + } + + s := int(scale) + if len(digits) <= s { + // Pad with leading zeros so there are exactly `scale` fractional digits + // and a single leading integer zero, e.g. 5 with scale 3 -> "0.005". + b.WriteString("0.") + b.WriteString(strings.Repeat("0", s-len(digits))) + b.WriteString(digits) + } else { + b.WriteString(digits[:len(digits)-s]) + b.WriteByte('.') + b.WriteString(digits[len(digits)-s:]) + } + return b.String() +} diff --git a/internal/decimalfmt/decimalfmt_test.go b/internal/decimalfmt/decimalfmt_test.go new file mode 100644 index 00000000..7765c7c6 --- /dev/null +++ b/internal/decimalfmt/decimalfmt_test.go @@ -0,0 +1,43 @@ +package decimalfmt + +import ( + "math/big" + "testing" + + "github.com/apache/arrow/go/v12/arrow/decimal128" +) + +// ExactString applies scale by string placement, preserving digits a float64 +// would lose beyond ~17 significant figures. +func TestExactString(t *testing.T) { + cases := []struct { + unscaled int64 + scale int32 + want string + }{ + {12345, 2, "123.45"}, + {5, 3, "0.005"}, + {100, 0, "100"}, + {1999, 2, "19.99"}, + {-12345, 2, "-123.45"}, + {-5, 3, "-0.005"}, + {0, 2, "0.00"}, + } + for _, c := range cases { + got := ExactString(decimal128.FromI64(c.unscaled), c.scale) + if got != c.want { + t.Errorf("unscaled=%d scale=%d: got %q want %q", c.unscaled, c.scale, got, c.want) + } + } +} + +// A value beyond float64's exact integer range must survive intact — the whole +// point of formatting from the 128-bit unscaled integer rather than a float. +func TestExactStringHighPrecision(t *testing.T) { + // 12345678901234567890 (20 digits) — past float64's 2^53 exact-integer limit. + big20, _ := new(big.Int).SetString("12345678901234567890", 10) + got := ExactString(decimal128.FromBigInt(big20), 4) + if want := "1234567890123456.7890"; got != want { + t.Errorf("got %q want %q", got, want) + } +} diff --git a/internal/rows/arrowbased/arrowRows.go b/internal/rows/arrowbased/arrowRows.go index b7a03fb5..f54e27c8 100644 --- a/internal/rows/arrowbased/arrowRows.go +++ b/internal/rows/arrowbased/arrowRows.go @@ -4,6 +4,7 @@ import ( "bytes" "context" "database/sql/driver" + "encoding/json" "io" "time" @@ -687,10 +688,17 @@ func (vcm *arrowValueContainerMaker) makeColumnValueContainer(t arrow.DataType, case *arrow.StructType: svc := &structValueContainer{structArrayType: t} svc.fieldNames = make([]string, len(t.Fields())) + svc.fieldKeys = make([]string, len(t.Fields())) svc.fieldValues = make([]columnValues, len(t.Fields())) svc.complexValue = make([]bool, len(t.Fields())) for i, f := range t.Fields() { svc.fieldNames[i] = f.Name + // Precompute the JSON-escaped `"name":` key once — field names are + // invariant across rows, so escaping them per row (structValueContainer + // .Value runs per row) would re-marshal constant strings. json.Marshal a + // string never errors. + keyBytes, _ := json.Marshal(f.Name) + svc.fieldKeys[i] = string(keyBytes) + ":" c, err := vcm.makeColumnValueContainer(f.Type, location, toTimestampFn, nil) if err != nil { return nil, err diff --git a/internal/rows/arrowbased/arrowscan_parity_test.go b/internal/rows/arrowbased/arrowscan_parity_test.go new file mode 100644 index 00000000..2882d05c --- /dev/null +++ b/internal/rows/arrowbased/arrowscan_parity_test.go @@ -0,0 +1,318 @@ +package arrowbased + +import ( + "testing" + "time" + + "github.com/apache/arrow/go/v12/arrow" + "github.com/apache/arrow/go/v12/arrow/array" + "github.com/apache/arrow/go/v12/arrow/decimal128" + "github.com/apache/arrow/go/v12/arrow/memory" + + "github.com/databricks/databricks-sql-go/internal/arrowscan" +) + +// The kernel backend (internal/arrowscan) and the Thrift backend (this package) +// each render Arrow cells to driver.Values — nested types to a JSON string — and +// the "identical results across backends" contract requires byte-for-byte +// agreement. That contract used to be guarded only by TestKernelThriftParity, +// which is build-tagged AND needs a live warehouse, so it never runs in CI. This +// test feeds the same arrow.Array through both renderers with no cgo and no +// warehouse, so a divergence (e.g. struct-key JSON escaping) fails a default +// CGO_ENABLED=0 run. +// +// It lives in package arrowbased because the container factory +// (makeColumnValueContainer) is package-private; arrowscan is a pure leaf import +// (no cycle). +func renderViaArrowbased(t *testing.T, arr arrow.Array, row int, loc *time.Location) any { + t.Helper() + maker := &arrowValueContainerMaker{} + holder, err := maker.makeColumnValueContainer(arr.DataType(), loc, func(ts arrow.Timestamp) time.Time { + return ts.ToTime(arrow.Microsecond) + }, nil) + if err != nil { + t.Fatalf("makeColumnValueContainer(%s): %v", arr.DataType(), err) + } + if err := holder.SetValueArray(arr.Data()); err != nil { + t.Fatalf("SetValueArray(%s): %v", arr.DataType(), err) + } + if holder.IsNull(row) { + return nil + } + v, err := holder.Value(row) + if err != nil { + t.Fatalf("Value(%s): %v", arr.DataType(), err) + } + return v +} + +func TestArrowbasedKernelRenderParity(t *testing.T) { + pool := memory.NewGoAllocator() + + // structWithKey builds a single-row STRUCT whose field is `name`. + structWithKey := func(name string) arrow.Array { + dt := arrow.StructOf(arrow.Field{Name: name, Type: arrow.PrimitiveTypes.Int64}) + b := array.NewStructBuilder(pool, dt) + b.Append(true) + b.FieldBuilder(0).(*array.Int64Builder).Append(1) + return b.NewArray() + } + + cases := []struct { + name string + build func() arrow.Array + }{ + {"list_int", func() arrow.Array { + b := array.NewListBuilder(pool, arrow.PrimitiveTypes.Int64) + vb := b.ValueBuilder().(*array.Int64Builder) + b.Append(true) + vb.Append(1) + vb.Append(2) + return b.NewArray() + }}, + {"struct_simple", func() arrow.Array { return structWithKey("a") }}, + // The escaping divergence: a field name with a quote must render as valid, + // identically-escaped JSON on both backends. + {"struct_key_with_quote", func() arrow.Array { return structWithKey(`a"b`) }}, + {"struct_key_with_backslash", func() arrow.Array { return structWithKey(`a\b`) }}, + {"struct_key_with_newline", func() arrow.Array { return structWithKey("a\nb") }}, + {"map_string_key", func() arrow.Array { + b := array.NewMapBuilder(pool, arrow.BinaryTypes.String, arrow.PrimitiveTypes.Int64, false) + b.Append(true) // open row 0's map + b.KeyBuilder().(*array.StringBuilder).Append("k") + b.ItemBuilder().(*array.Int64Builder).Append(9) + return b.NewArray() + }}, + {"map_special_string_key", func() arrow.Array { + b := array.NewMapBuilder(pool, arrow.BinaryTypes.String, arrow.PrimitiveTypes.Int64, false) + b.Append(true) + b.KeyBuilder().(*array.StringBuilder).Append(`k"x`) + b.ItemBuilder().(*array.Int64Builder).Append(9) + return b.NewArray() + }}, + {"map_int_key", func() arrow.Array { + b := array.NewMapBuilder(pool, arrow.PrimitiveTypes.Int64, arrow.PrimitiveTypes.Int64, false) + b.Append(true) + b.KeyBuilder().(*array.Int64Builder).Append(7) + b.ItemBuilder().(*array.Int64Builder).Append(9) + return b.NewArray() + }}, + {"map_binary_key", func() arrow.Array { + // []byte key: json.Marshal → base64 "YWJj", NOT fmt "%v" [97 98 99]. + b := array.NewMapBuilder(pool, arrow.BinaryTypes.Binary, arrow.PrimitiveTypes.Int64, false) + b.Append(true) + b.KeyBuilder().(*array.BinaryBuilder).Append([]byte("abc")) + b.ItemBuilder().(*array.Int64Builder).Append(9) + return b.NewArray() + }}, + {"map_date_key", func() arrow.Array { + b := array.NewMapBuilder(pool, arrow.FixedWidthTypes.Date32, arrow.PrimitiveTypes.Int64, false) + b.Append(true) + b.KeyBuilder().(*array.Date32Builder).Append(arrow.Date32FromTime(time.Date(2026, time.July, 9, 0, 0, 0, 0, time.UTC))) + b.ItemBuilder().(*array.Int64Builder).Append(9) + return b.NewArray() + }}, + {"nested_float32", func() arrow.Array { + dt := arrow.StructOf(arrow.Field{Name: "f", Type: arrow.PrimitiveTypes.Float32}) + b := array.NewStructBuilder(pool, dt) + b.Append(true) + b.FieldBuilder(0).(*array.Float32Builder).Append(0.1) + return b.NewArray() + }}, + {"nested_decimal", func() arrow.Array { + dt := arrow.StructOf(arrow.Field{Name: "d", Type: &arrow.Decimal128Type{Precision: 5, Scale: 2}}) + b := array.NewStructBuilder(pool, dt) + b.Append(true) + b.FieldBuilder(0).(*array.Decimal128Builder).Append(decimal128.FromU64(1999)) + return b.NewArray() + }}, + // Highest-drift shapes: recursive nesting, a nested timestamp leaf (the + // time.Time → quoted .String() special-case), and a null leaf inside a + // container (vs a null container, which is already covered). + {"array_of_struct", func() arrow.Array { + elem := arrow.StructOf(arrow.Field{Name: "a", Type: arrow.PrimitiveTypes.Int64}) + b := array.NewListBuilder(pool, elem) + sb := b.ValueBuilder().(*array.StructBuilder) + b.Append(true) + sb.Append(true) + sb.FieldBuilder(0).(*array.Int64Builder).Append(1) + sb.Append(true) + sb.FieldBuilder(0).(*array.Int64Builder).Append(2) + return b.NewArray() + }}, + {"struct_of_list", func() arrow.Array { + dt := arrow.StructOf(arrow.Field{Name: "xs", Type: arrow.ListOf(arrow.PrimitiveTypes.Int64)}) + b := array.NewStructBuilder(pool, dt) + b.Append(true) + lb := b.FieldBuilder(0).(*array.ListBuilder) + lb.Append(true) + vb := lb.ValueBuilder().(*array.Int64Builder) + vb.Append(1) + vb.Append(2) + return b.NewArray() + }}, + {"map_of_struct_value", func() arrow.Array { + valT := arrow.StructOf(arrow.Field{Name: "a", Type: arrow.PrimitiveTypes.Int64}) + b := array.NewMapBuilder(pool, arrow.BinaryTypes.String, valT, false) + b.Append(true) + b.KeyBuilder().(*array.StringBuilder).Append("k") + sb := b.ItemBuilder().(*array.StructBuilder) + sb.Append(true) + sb.FieldBuilder(0).(*array.Int64Builder).Append(5) + return b.NewArray() + }}, + {"nested_timestamp", func() arrow.Array { + dt := arrow.StructOf(arrow.Field{Name: "ts", Type: &arrow.TimestampType{Unit: arrow.Microsecond}}) + b := array.NewStructBuilder(pool, dt) + b.Append(true) + ts := time.Date(2026, time.July, 9, 12, 34, 56, 0, time.UTC) + b.FieldBuilder(0).(*array.TimestampBuilder).Append(arrow.Timestamp(ts.UnixMicro())) + return b.NewArray() + }}, + {"null_leaf_in_struct", func() arrow.Array { + dt := arrow.StructOf( + arrow.Field{Name: "a", Type: arrow.PrimitiveTypes.Int64}, + arrow.Field{Name: "b", Type: arrow.PrimitiveTypes.Int64}, + ) + b := array.NewStructBuilder(pool, dt) + b.Append(true) + b.FieldBuilder(0).(*array.Int64Builder).Append(1) + b.FieldBuilder(1).(*array.Int64Builder).AppendNull() + return b.NewArray() + }}, + } + + // A non-UTC location, so the timestamp/date rendering path (both backends + // apply .In(loc)) is actually exercised rather than a UTC no-op. + loc, err := time.LoadLocation("America/New_York") + if err != nil { + t.Skipf("tz database unavailable: %v", err) + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + arr := tc.build() + defer arr.Release() + + kernel, err := arrowscan.ScanCell(arr, 0, loc) + if err != nil { + t.Fatalf("arrowscan.ScanCell: %v", err) + } + thrift := renderViaArrowbased(t, arr, 0, loc) + + if kernel != thrift { + t.Errorf("backend divergence for %s:\n kernel = %#v\n thrift = %#v", tc.name, kernel, thrift) + } + }) + } +} + +// Top-level (non-nested) scalars: assert the kernel ScanCell and the Thrift +// container render the same Go value. DECIMAL is deliberately excluded here — see +// TestTopLevelDecimalRendering for why the two *container-level* paths differ by +// arrow type while the actual driver results agree. +func TestArrowbasedKernelTopLevelScalarParity(t *testing.T) { + pool := memory.NewGoAllocator() + loc := time.UTC + + cases := []struct { + name string + build func() arrow.Array + }{ + {"int64", func() arrow.Array { + b := array.NewInt64Builder(pool) + b.Append(42) + return b.NewArray() + }}, + {"float32", func() arrow.Array { + b := array.NewFloat32Builder(pool) + b.Append(0.1) + return b.NewArray() + }}, + {"float64", func() arrow.Array { + b := array.NewFloat64Builder(pool) + b.Append(3.5) + return b.NewArray() + }}, + {"string", func() arrow.Array { + b := array.NewStringBuilder(pool) + b.Append("hi") + return b.NewArray() + }}, + {"timestamp", func() arrow.Array { + b := array.NewTimestampBuilder(pool, &arrow.TimestampType{Unit: arrow.Microsecond}) + b.Append(arrow.Timestamp(time.Date(2026, 7, 9, 12, 0, 0, 0, time.UTC).UnixMicro())) + return b.NewArray() + }}, + {"date32", func() arrow.Array { + b := array.NewDate32Builder(pool) + b.Append(arrow.Date32FromTime(time.Date(2026, 7, 9, 0, 0, 0, 0, time.UTC))) + return b.NewArray() + }}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + arr := tc.build() + defer arr.Release() + kernel, err := arrowscan.ScanCell(arr, 0, loc) + if err != nil { + t.Fatalf("arrowscan.ScanCell: %v", err) + } + thrift := renderViaArrowbased(t, arr, 0, loc) + if kernel != thrift { + t.Errorf("top-level %s divergence:\n kernel = %#v (%T)\n thrift = %#v (%T)", tc.name, kernel, kernel, thrift, thrift) + } + }) + } +} + +// TestTopLevelDecimalRendering documents the top-level DECIMAL story (review r5 +// M1-r5), which is subtler than "kernel string vs Thrift float64": +// +// - The kernel path always delivers DECIMAL as a native arrow decimal128 and +// renders it as an exact scale-applied string (arrowscan.ScanCell). +// - The Thrift *container* (decimal128Container.Value) converts a decimal128 to +// a lossy float64 — but that path is only reached when UseArrowNativeDecimal +// is on AND the value is read as a nested leaf. For a top-level column the +// scan uses DecimalStringValue (exact string) when the flag is on. +// - In the DEFAULT config the flag is off, so the server sends DECIMAL as a +// string column (DecimalAsArrow=false) — no decimal128 arrives at all, and the +// user gets the exact string. +// +// So across every real driver configuration a top-level DECIMAL comes back as the +// exact string on both backends (verified live). This test pins the kernel side +// (exact string) and the two Thrift container behaviors so the distinction can't +// silently drift into an actual result divergence. +func TestTopLevelDecimalRendering(t *testing.T) { + pool := memory.NewGoAllocator() + dt := &arrow.Decimal128Type{Precision: 38, Scale: 4} + b := array.NewDecimal128Builder(pool, dt) + defer b.Release() + n, _ := decimal128.FromString("1234567890123456.7890", dt.Precision, dt.Scale) + b.Append(n) + arr := b.NewArray() + defer arr.Release() + + // Kernel: exact string. + got, err := arrowscan.ScanCell(arr, 0, time.UTC) + if err != nil { + t.Fatal(err) + } + if got != "1234567890123456.7890" { + t.Errorf("kernel top-level decimal = %#v, want exact string", got) + } + + // Thrift container Value() is the lossy-float64 path (reached only with native + // decimal + nested read); DecimalStringValue is the exact top-level scan path. + // Assert both so a change to either is caught. + holder := &decimal128Container{scale: dt.Scale} + if err := holder.SetValueArray(arr.Data()); err != nil { + t.Fatal(err) + } + if s := holder.ValueString(0); s != "1234567890123456.7890" { + t.Errorf("Thrift DecimalStringValue = %q, want exact string", s) + } + if v, _ := holder.Value(0); v == "1234567890123456.7890" { + t.Errorf("Thrift container Value() unexpectedly exact; it is the documented lossy-float64 path") + } +} diff --git a/internal/rows/arrowbased/columnValues.go b/internal/rows/arrowbased/columnValues.go index 72e7f5f5..0e47d105 100644 --- a/internal/rows/arrowbased/columnValues.go +++ b/internal/rows/arrowbased/columnValues.go @@ -2,13 +2,12 @@ package arrowbased import ( "encoding/json" - "math/big" "strings" "time" "github.com/apache/arrow/go/v12/arrow" "github.com/apache/arrow/go/v12/arrow/array" - "github.com/apache/arrow/go/v12/arrow/decimal128" + "github.com/databricks/databricks-sql-go/internal/decimalfmt" "github.com/databricks/databricks-sql-go/internal/rows/rowscanner" dbsqllog "github.com/databricks/databricks-sql-go/logger" "github.com/pkg/errors" @@ -364,6 +363,7 @@ func (mvc *mapValueContainer) SetValueArray(colData arrow.ArrayData) error { type structValueContainer struct { structArray *array.Struct fieldNames []string + fieldKeys []string // precomputed JSON-escaped `"name":` prefixes, one per field complexValue []bool fieldValues []columnValues structArrayType *arrow.StructType @@ -371,11 +371,19 @@ type structValueContainer struct { var _ columnValues = (*structValueContainer)(nil) +// Value renders the struct at row i as a JSON object string. This grammar is +// mirrored by internal/arrowscan (the kernel backend); the two must render +// byte-identically — internal/rows/arrowbased/arrowscan_parity_test.go guards it. func (svc *structValueContainer) Value(i int) (any, error) { if i < svc.structArray.Len() { r := "{" for j := range svc.fieldValues { - r = r + "\"" + svc.fieldNames[j] + "\":" + // fieldKeys[j] is the precomputed JSON-escaped `"name":` prefix (built + // once at construction). Escaping matters — a raw `"` + name + `"` concat + // produces invalid JSON for a name with a quote/backslash/control char + // (`a"b` → `{"a"b":...}`) — but doing it per row would re-marshal the same + // constant strings N_rows × N_fields times on this hot path. + r = r + svc.fieldKeys[j] if svc.fieldValues[j].IsNull(int(i)) { r = r + "null" @@ -550,48 +558,9 @@ func (tvc *decimal128Container) Value(i int) (any, error) { // for backwards-compatible JSON rendering inside complex types. // See databricks/databricks-sql-go#274. func (tvc *decimal128Container) ValueString(i int) string { - return decimal128ToString(tvc.decimalArray.Value(i), tvc.scale) -} - -// decimal128ToString renders a decimal128 value exactly as a fixed-point -// string with `scale` fractional digits. It formats from the value's exact -// 128-bit unscaled integer (BigInt) and inserts the decimal point by string -// manipulation, so — unlike arrow's decimal128.Num.ToString, which divides via -// big.Float and rounds beyond ~17 significant digits — it never loses -// precision. See databricks/databricks-sql-go#274. -func decimal128ToString(n decimal128.Num, scale int32) string { - unscaled := n.BigInt() // exact signed unscaled integer - - neg := unscaled.Sign() < 0 - digits := new(big.Int).Abs(unscaled).String() - - var b strings.Builder - if neg { - b.WriteByte('-') - } - - if scale <= 0 { - // No fractional part (negative scale is not produced by the server, - // but treat it as scale 0 rather than panicking). - b.WriteString(digits) - return b.String() - } - - s := int(scale) - if len(digits) <= s { - // Pad with leading zeros so there are exactly `scale` fractional digits - // and a single leading integer zero, e.g. 5 with scale 3 -> "0.005". - b.WriteString("0.") - b.WriteString(strings.Repeat("0", s-len(digits))) - b.WriteString(digits) - } else { - intPart := digits[:len(digits)-s] - fracPart := digits[len(digits)-s:] - b.WriteString(intPart) - b.WriteByte('.') - b.WriteString(fracPart) - } - return b.String() + // Exact fixed-point rendering lives in internal/decimalfmt, shared with the + // kernel backend so both paths render DECIMAL identically. See #274. + return decimalfmt.ExactString(tvc.decimalArray.Value(i), tvc.scale) } func (tvc *decimal128Container) IsNull(i int) bool { diff --git a/kernel_backend.go b/kernel_backend.go new file mode 100644 index 00000000..412ffe78 --- /dev/null +++ b/kernel_backend.go @@ -0,0 +1,52 @@ +//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 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) { + // Reject options the kernel path can't honor yet + resolve the PAT. The + // validation is pure Go and lives in kernel_config.go (untagged) so its tests — + // including the exhaustiveness guard against a dropped Config field — run in the + // default CGO_ENABLED=0 build. + token, err := validateKernelConfig(cfg) + if err != nil { + return nil, err + } + + kc := kernel.Config{ + Host: cfg.Host, + HTTPPath: cfg.HTTPPath, + WarehouseID: cfg.WarehouseID, + Token: token, + Location: cfg.Location, + // Session confs (STATEMENT_TIMEOUT, QUERY_TAGS, TIMEZONE, …) — the same + // SessionParams map the Thrift backend forwards, so they flow to the + // 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 + // 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 (pure Go, no kernel dependency) lives in kernel_proxy.go so +// its test runs in the default CGO_ENABLED=0 build. diff --git a/kernel_backend_test.go b/kernel_backend_test.go new file mode 100644 index 00000000..eb3244c4 --- /dev/null +++ b/kernel_backend_test.go @@ -0,0 +1,40 @@ +//go:build cgo && databricks_kernel + +package dbsql + +import ( + "context" + "testing" + + "github.com/databricks/databricks-sql-go/internal/config" +) + +// newKernelBackend builds a backend for a supported config and propagates a +// validation error otherwise. The exhaustive per-option reject/auth assertions +// live in the untagged TestValidateKernelConfig (kernel_config_test.go), which +// runs in the default CGO_ENABLED=0 build; this tagged smoke test just confirms +// the tagged wrapper wires validateKernelConfig + the kernel.Config assembly. +func TestNewKernelBackend(t *testing.T) { + base := func() *config.Config { + c := config.WithDefaults() + c.Host = "h.databricks.com" + c.Port = 443 + c.HTTPPath = "/sql/1.0/warehouses/abc" + c.AccessToken = "dapi-x" + return c + } + + t.Run("supported config builds", func(t *testing.T) { + if _, err := newKernelBackend(context.Background(), base()); err != nil { + t.Errorf("a supported config should build cleanly, got %v", err) + } + }) + + t.Run("validation error propagates", func(t *testing.T) { + c := base() + c.Catalog = "main" // rejected by validateKernelConfig + if _, err := newKernelBackend(context.Background(), c); err == nil { + t.Error("newKernelBackend should propagate the validation error") + } + }) +} diff --git a/kernel_config.go b/kernel_config.go new file mode 100644 index 00000000..28cb35ad --- /dev/null +++ b/kernel_config.go @@ -0,0 +1,81 @@ +package dbsql + +import ( + "errors" + + "github.com/databricks/databricks-sql-go/auth/noop" + "github.com/databricks/databricks-sql-go/auth/pat" + "github.com/databricks/databricks-sql-go/internal/config" +) + +// This file is intentionally NOT behind the `cgo && databricks_kernel` build tag. +// The kernel backend's option-validation is pure Go (it reads config.Config and +// returns an error or a resolved PAT), so keeping it untagged lets its tests — +// including the reflective exhaustiveness check that guards against a future +// Config field being silently dropped — run under CGO_ENABLED=0. The tagged +// newKernelBackend calls validateKernelConfig, then assembles the cgo kernel.Config. + +// validateKernelConfig enforces the kernel backend's "nothing silently ignored" +// contract: it rejects every option the kernel path can't yet honor with a clear +// error (rather than dropping it, which would behave differently than Thrift) and +// resolves the PAT the kernel authenticates with. On success it returns the token +// to use. Options it does NOT reject are either forwarded by newKernelBackend or +// intentionally accepted-but-inert (documented in doc.go and asserted by +// TestKernelConfigFieldsClassified). +func validateKernelConfig(cfg *config.Config) (token string, err error) { + // Initial namespace (WithInitialNamespace): no kernel C-ABI setter yet, so the + // session would run in the default namespace and unqualified names would + // resolve differently than Thrift. + if cfg.Catalog != "" || cfg.Schema != "" { + return "", errors.New("databricks: WithInitialNamespace (catalog/schema) is not yet supported by the kernel backend; " + + "omit it or use the default (Thrift) backend") + } + // EnableMetricViewMetadata: maps to a server session conf we want to route + // backend-neutrally rather than duplicate here. + if cfg.EnableMetricViewMetadata { + return "", errors.New("databricks: WithEnableMetricViewMetadata is not yet supported by the kernel backend; " + + "omit it or use the default (Thrift) backend") + } + // Auth: the kernel backend authenticates with a PAT only. Any other + // authenticator sets cfg.Authenticator but leaves cfg.AccessToken empty, so an + // empty PAT would reach the kernel and fail with an opaque Unauthenticated + // error. Reject it here so the failure names the cause. + token = cfg.AccessToken + switch a := cfg.Authenticator.(type) { + case nil, *noop.NoopAuth: + // No explicit authenticator — token comes from cfg.AccessToken (may be + // empty; caught below). + case *pat.PATAuth: + // WithAccessToken sets both cfg.AccessToken and this authenticator, but + // WithAuthenticator(&pat.PATAuth{...}) sets only the authenticator and leaves + // cfg.AccessToken empty. Take the token from the authenticator when + // cfg.AccessToken didn't carry it, so both PAT paths work. + if token == "" { + token = a.AccessToken + } + default: + return "", errors.New("databricks: only personal access token (WithAccessToken) auth is supported by the kernel backend; " + + "OAuth (M2M/U2M), token-provider, external/static, and federated authenticators are not yet supported — " + + "use PAT or the default (Thrift) backend") + } + if token == "" { + return "", errors.New("databricks: the kernel backend requires a personal access token; " + + "set one with WithAccessToken (or a *pat.PATAuth via WithAuthenticator)") + } + // WithTimeout maps to a per-statement server timeout on Thrift + // (TExecuteStatementReq.QueryTimeout); the kernel C ABI exposes no equivalent, + // so reject it rather than run the query with no server-side timeout. + if cfg.QueryTimeout > 0 { + return "", errors.New("databricks: WithTimeout (server query timeout) is not yet supported by the kernel backend; " + + "omit it or use the default (Thrift) backend") + } + // WithRetries(-1) explicitly disables retries, but the kernel retries + // internally below the C ABI with no user-facing toggle — so a disable request + // would be silently violated. Reject it. Positive/default RetryMax is fine: the + // kernel provides retries (just not user-tunable), documented in doc.go. + if cfg.RetryMax < 0 { + return "", errors.New("databricks: disabling retries via WithRetries is not supported by the kernel backend " + + "(the kernel retries internally); omit it or use the default (Thrift) backend") + } + return token, nil +} diff --git a/kernel_config_test.go b/kernel_config_test.go new file mode 100644 index 00000000..9e1ff642 --- /dev/null +++ b/kernel_config_test.go @@ -0,0 +1,212 @@ +package dbsql + +import ( + "net/http" + "reflect" + "testing" + "time" + + "github.com/databricks/databricks-sql-go/auth/pat" + "github.com/databricks/databricks-sql-go/internal/config" +) + +// nonPATAuth stands in for any non-PAT authenticator (OAuth / token-provider / +// external / federated) — the kernel backend must reject it. +type nonPATAuth struct{} + +func (nonPATAuth) Authenticate(*http.Request) error { return nil } + +func baseKernelConfig() *config.Config { + c := config.WithDefaults() + c.Host = "h.databricks.com" + c.Port = 443 + c.HTTPPath = "/sql/1.0/warehouses/abc" + c.AccessToken = "dapi-x" + return c +} + +// validateKernelConfig enforces the kernel backend's "nothing silently ignored" +// contract: unsupported options are rejected loudly and the PAT is resolved. This +// is pure Go (no cgo), so these run in the default CGO_ENABLED=0 build. +func TestValidateKernelConfig(t *testing.T) { + t.Run("supported config ok", func(t *testing.T) { + c := baseKernelConfig() + c.SessionParams = map[string]string{"QUERY_TAGS": "a:1"} + if _, err := validateKernelConfig(c); err != nil { + t.Errorf("a supported config should validate, got %v", err) + } + }) + + t.Run("catalog rejected", func(t *testing.T) { + c := baseKernelConfig() + c.Catalog = "main" + if _, err := validateKernelConfig(c); err == nil { + t.Error("expected an error when a catalog is set") + } + }) + + t.Run("schema rejected", func(t *testing.T) { + c := baseKernelConfig() + c.Schema = "sys" + if _, err := validateKernelConfig(c); err == nil { + t.Error("expected an error when a schema is set") + } + }) + + t.Run("metric view rejected", func(t *testing.T) { + c := baseKernelConfig() + c.EnableMetricViewMetadata = true + if _, err := validateKernelConfig(c); err == nil { + t.Error("expected an error when metric-view metadata is enabled") + } + }) + + t.Run("PAT via WithAuthenticator resolves the token", func(t *testing.T) { + c := baseKernelConfig() + c.AccessToken = "" + c.Authenticator = &pat.PATAuth{AccessToken: "dapi-y"} + tok, err := validateKernelConfig(c) + if err != nil { + t.Fatalf("PAT via WithAuthenticator should validate, got %v", err) + } + if tok != "dapi-y" { + t.Errorf("token = %q, want dapi-y (sourced from the authenticator)", tok) + } + }) + + t.Run("empty token rejected", func(t *testing.T) { + c := baseKernelConfig() + c.AccessToken = "" + c.Authenticator = &pat.PATAuth{AccessToken: ""} + if _, err := validateKernelConfig(c); err == nil { + t.Error("expected an error when the resolved PAT is empty") + } + }) + + t.Run("non-PAT authenticator rejected", func(t *testing.T) { + c := baseKernelConfig() + c.Authenticator = nonPATAuth{} + if _, err := validateKernelConfig(c); err == nil { + t.Error("expected an error for a non-PAT authenticator") + } + }) + + t.Run("query timeout rejected", func(t *testing.T) { + c := baseKernelConfig() + c.QueryTimeout = 30 * time.Second + if _, err := validateKernelConfig(c); err == nil { + t.Error("expected an error when WithTimeout (query timeout) is set") + } + }) + + t.Run("disabling retries rejected", func(t *testing.T) { + c := baseKernelConfig() + c.RetryMax = -1 + if _, err := validateKernelConfig(c); err == nil { + t.Error("expected an error when retries are disabled (WithRetries(-1))") + } + }) + + t.Run("positive retry tuning + maxrows accepted", func(t *testing.T) { + c := baseKernelConfig() + c.RetryMax = 8 + c.MaxRows = 5000 + if _, err := validateKernelConfig(c); err != nil { + t.Errorf("positive retry/maxrows tuning should validate, got %v", err) + } + }) +} + +// kernelConfigFieldDisposition records, for every UserConfig field, how the kernel +// backend treats it. UserConfig is the user-facing option surface ("Only +// UserConfig are currently exposed to users"). Adding a field to config.UserConfig +// without classifying it here fails TestKernelConfigFieldsClassified — forcing a +// deliberate decision (forward it, reject it loudly, or accept it as inert) rather +// than silently dropping it on the kernel path. (TLSConfig and ArrowConfig live on +// the outer config.Config, not +// UserConfig; newKernelBackend reads TLSConfig.InsecureSkipVerify explicitly, and +// the kernel renders decimals exactly regardless of ArrowConfig.) +var kernelConfigFieldDisposition = map[string]string{ + // Forwarded to kernel.Config (see newKernelBackend). + "Host": "forwarded", + "HTTPPath": "forwarded", + "WarehouseID": "forwarded", + "AccessToken": "forwarded", // as the resolved PAT (kc.Token) + "Authenticator": "forwarded", // PAT authenticator resolved to the token + "Location": "forwarded", + "SessionParams": "forwarded", + "Port": "forwarded", // part of the endpoint URL (host:port) + "Protocol": "forwarded", // part of the endpoint URL (scheme) + "UseKernel": "forwarded", // the routing flag itself + + // Rejected loudly by validateKernelConfig. + "Catalog": "rejected", + "Schema": "rejected", + "EnableMetricViewMetadata": "rejected", + "QueryTimeout": "rejected", // when > 0 (WithTimeout) + "RetryMax": "rejected", // when < 0 (disable retries) + + // Accepted but intentionally inert on the kernel path (documented in doc.go): + // the kernel manages these internally, below the C ABI, with no user knob. + "MaxRows": "inert", + "RetryWaitMin": "inert", + "RetryWaitMax": "inert", + "UseLz4Compression": "inert", // kernel negotiates compression internally + + // Not applicable to the kernel path (Thrift/HTTP-transport or telemetry knobs + // that don't reach the kernel binding). + "UserAgentEntry": "inert", // TODO: forward once the kernel exposes a UA setter + "Transport": "inert", // custom RoundTripper; kernel uses its own HTTP stack + "EnableTelemetry": "inert", + "TelemetryBatchSize": "inert", + "TelemetryFlushInterval": "inert", + "UseArrowNativeDecimalDSN": "inert", // DSN carrier; kernel renders decimals exactly regardless + + // Fields promoted from the embedded CloudFetchConfig. The kernel does + // CloudFetch internally (below the C ABI), so none is forwarded — but each is + // classified individually so a new CloudFetch option can't slip the guard. + "UseCloudFetch": "inert", + "MaxDownloadThreads": "inert", + "MaxFilesInMemory": "inert", + "MinTimeToExpiry": "inert", + "CloudFetchSpeedThresholdMbps": "inert", + "HTTPClient": "inert", +} + +// kernelConfigClassifiedNames returns every UserConfig field name the kernel gate +// must account for, flattening anonymous embedded structs (e.g. CloudFetchConfig) +// into their promoted fields — reflect's NumField reports an embed as one field, +// which would hide new sub-fields from the drop guard below. +func kernelConfigClassifiedNames(t reflect.Type) []string { + var names []string + for i := 0; i < t.NumField(); i++ { + f := t.Field(i) + if f.Anonymous && f.Type.Kind() == reflect.Struct { + names = append(names, kernelConfigClassifiedNames(f.Type)...) + continue + } + names = append(names, f.Name) + } + return names +} + +func TestKernelConfigFieldsClassified(t *testing.T) { + names := kernelConfigClassifiedNames(reflect.TypeOf(config.UserConfig{})) + classified := make(map[string]bool, len(names)) + for _, name := range names { + classified[name] = true + if _, ok := kernelConfigFieldDisposition[name]; !ok { + t.Errorf("config.UserConfig field %q (incl. promoted embed fields) is not classified "+ + "for the kernel backend. Add it to kernelConfigFieldDisposition as "+ + "forwarded/rejected/inert (and wire it in validateKernelConfig / newKernelBackend if "+ + "it must be honored) so it isn't silently dropped on the kernel path.", name) + } + } + // Guard the reverse too: a disposition entry for a field that no longer exists + // is stale. + for name := range kernelConfigFieldDisposition { + if !classified[name] { + t.Errorf("kernelConfigFieldDisposition has %q but config.UserConfig (incl. embeds) no longer does; remove it", name) + } + } +} diff --git a/kernel_e2e_test.go b/kernel_e2e_test.go new file mode 100644 index 00000000..4b150b1d --- /dev/null +++ b/kernel_e2e_test.go @@ -0,0 +1,264 @@ +//go:build cgo && databricks_kernel + +package dbsql + +import ( + "bytes" + "context" + "database/sql" + "database/sql/driver" + "errors" + "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() + return kernelTestDBWith(t) +} + +// 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) + } +} + +// 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) + } +} + +// 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. +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 +// 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)}, + // 0.1 is NOT exactly representable, so a float32-vs-float64 mismatch would + // surface here (float64(float32(0.1)) != float32(0.1)); a bare FLOAT must + // scan to a native float32, matching Thrift. + {"float", "CAST(0.1 AS FLOAT)", float32(0.1)}, + {"boolean", "true", true}, + {"string", "'hi'", "hi"}, + {"binary", "CAST('abc' AS BINARY)", []byte("abc")}, + {"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}, + // 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 { + 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") + } + // It must be the deadline firing, not an unrelated failure (syntax error, + // transient network, server-side timeout) — otherwise the test would pass + // without proving cancellation works. + if !errors.Is(err, context.DeadlineExceeded) { + t.Errorf("expected context.DeadlineExceeded, got %v", err) + } + // Well under the query's natural runtime and not far past the 3s deadline; + // a no-op cancel that only returned when the query finished would blow past + // this. + if elapsed > 10*time.Second { + t.Errorf("cancellation took %v; expected it to abandon close to the 3s deadline", elapsed) + } + t.Logf("cancelled after %v with err=%v", elapsed, err) +} diff --git a/kernel_parity_test.go b/kernel_parity_test.go new file mode 100644 index 00000000..8f570302 --- /dev/null +++ b/kernel_parity_test.go @@ -0,0 +1,102 @@ +//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, " + + // A bare FLOAT at a non-exactly-representable value: catches the + // float32-vs-widened-float64 divergence (0.1 vs 0.10000000149011612). + "CAST(0.1 AS FLOAT), " + + "array(1, 2, 3), map('k', 9), named_struct('a', 1, 'b', 'x'), " + + `parse_json('{"a":1,"b":[2,3]}'), st_point(1, 2), ` + + // A decimal inside a struct: exercises exact-string nested-decimal + // rendering (19.99, not a lossy 19.990000000000002). + "named_struct('d', CAST(19.99 AS DECIMAL(5,2)))" + + 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_proxy.go b/kernel_proxy.go new file mode 100644 index 00000000..0e89377b --- /dev/null +++ b/kernel_proxy.go @@ -0,0 +1,45 @@ +package dbsql + +import ( + "net/http" + "net/url" + + "github.com/databricks/databricks-sql-go/internal/config" +) + +// This file is intentionally NOT behind the `cgo && databricks_kernel` build tag: +// proxy resolution is pure Go (no kernel C symbol), so keeping it in the default +// build lets its test run under CGO_ENABLED=0 rather than being dead behind the +// tag. newKernelBackend (tagged) calls proxyForEndpoint. + +// proxyForEndpoint resolves the proxy the Thrift path would use for this +// connection, via the same http.ProxyFromEnvironment (HTTP(S)_PROXY / NO_PROXY) +// the Thrift transport applies at request time. Building the endpoint request +// lets ProxyFromEnvironment apply the NO_PROXY rules for this exact host, so the +// kernel sees the same effective proxy decision — returning "" (direct) when +// NO_PROXY excludes the host or no proxy is set. No extra dependency: this is the +// stdlib function the driver already relies on. +func proxyForEndpoint(cfg *config.Config) string { + return proxyForEndpointFunc(cfg, http.ProxyFromEnvironment) +} + +// proxyForEndpointFunc is the testable core: it builds the endpoint request and +// asks resolve for the proxy, returning "" (direct) on any error, no proxy, or an +// unbuildable endpoint. resolve is http.ProxyFromEnvironment in production; tests +// inject a deterministic resolver to exercise proxy-set / NO_PROXY / direct +// without depending on http.ProxyFromEnvironment's process-wide env caching. +func proxyForEndpointFunc(cfg *config.Config, resolve func(*http.Request) (*url.URL, error)) string { + endpoint, err := cfg.ToEndpointURL() + if err != nil { + return "" + } + req, err := http.NewRequest(http.MethodPost, endpoint, nil) + if err != nil { + return "" + } + proxyURL, err := resolve(req) + if err != nil || proxyURL == nil { + return "" + } + return proxyURL.String() +} diff --git a/kernel_proxy_test.go b/kernel_proxy_test.go new file mode 100644 index 00000000..c7f73b13 --- /dev/null +++ b/kernel_proxy_test.go @@ -0,0 +1,78 @@ +package dbsql + +import ( + "errors" + "net/http" + "net/url" + "testing" + + "github.com/databricks/databricks-sql-go/internal/config" +) + +// proxyForEndpointFunc maps the injected resolver's decision to a proxy URL +// string (or "" for direct), and returns "" for an unbuildable endpoint. The +// production path uses http.ProxyFromEnvironment, whose env is snapshotted once +// per process (sync.Once) and so can't be re-driven mid-test; the resolver seam +// lets us assert every branch deterministically. Each case mirrors an +// httpproxy-style outcome: proxy set for this host, host excluded by NO_PROXY, +// no proxy configured, and an unbuildable endpoint. Pure Go, so it runs in the +// default CGO_ENABLED=0 build (no kernel lib required). +func TestProxyForEndpoint(t *testing.T) { + validCfg := func() *config.Config { + c := config.WithDefaults() + c.Host = "my-workspace.databricks.com" + c.Port = 443 + c.HTTPPath = "/sql/1.0/warehouses/abc" + return c + } + proxyURL, _ := url.Parse("http://corp-proxy:3128") + + cases := []struct { + name string + cfg *config.Config + resolve func(*http.Request) (*url.URL, error) + want string + }{ + { + name: "proxy set for host", + cfg: validCfg(), + resolve: func(*http.Request) (*url.URL, error) { return proxyURL, nil }, + want: "http://corp-proxy:3128", + }, + { + name: "host excluded by NO_PROXY -> direct", + cfg: validCfg(), + resolve: func(*http.Request) (*url.URL, error) { return nil, nil }, + want: "", + }, + { + name: "resolver error -> direct", + cfg: validCfg(), + resolve: func(*http.Request) (*url.URL, error) { return nil, errors.New("bad proxy url") }, + want: "", + }, + { + name: "unbuildable endpoint -> direct (resolver never consulted)", + cfg: func() *config.Config { + c := config.WithDefaults() + c.Host = "" + return c + }(), + resolve: func(*http.Request) (*url.URL, error) { + t.Error("resolver must not be called for an unbuildable endpoint") + return proxyURL, nil + }, + want: "", + }, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + if got := proxyForEndpointFunc(tc.cfg, tc.resolve); got != tc.want { + t.Errorf("proxyForEndpointFunc = %q, want %q", got, tc.want) + } + }) + } + + // The production wrapper wires http.ProxyFromEnvironment and must not panic. + _ = proxyForEndpoint(validCfg()) +} diff --git a/kernel_stub.go b/kernel_stub.go new file mode 100644 index 00000000..4774278b --- /dev/null +++ b/kernel_stub.go @@ -0,0 +1,20 @@ +//go:build !cgo || !databricks_kernel + +package dbsql + +import ( + "context" + "errors" + + "github.com/databricks/databricks-sql-go/internal/backend" + "github.com/databricks/databricks-sql-go/internal/config" +) + +// newKernelBackend is the stub compiled when the kernel backend is not built in. +// It fails loudly rather than silently falling back to Thrift, so a mismatch +// between WithUseKernel and the build tags surfaces at connect time. The real +// implementation is in kernel_backend.go. +func newKernelBackend(_ context.Context, _ *config.Config) (backend.Backend, error) { + return nil, errors.New("databricks: the SEA-via-kernel backend is not compiled into this binary; " + + "rebuild with -tags databricks_kernel and CGO_ENABLED=1, or unset WithUseKernel") +} diff --git a/kernel_stub_test.go b/kernel_stub_test.go new file mode 100644 index 00000000..9f8e8cde --- /dev/null +++ b/kernel_stub_test.go @@ -0,0 +1,33 @@ +//go:build !cgo || !databricks_kernel + +package dbsql + +import ( + "context" + "strings" + "testing" +) + +// In the default pure-Go build (no databricks_kernel tag) the kernel backend is +// not compiled in. Connecting with WithUseKernel(true) must fail loudly with a +// clear "not compiled in" error rather than silently falling back to Thrift or +// panicking on a nil backend. This guards the stub that fails closed. +func TestKernelBackendNotCompiledIn(t *testing.T) { + connector, err := NewConnector( + WithServerHostname("example.cloud.databricks.com"), + WithPort(443), + WithHTTPPath("/sql/1.0/endpoints/12346a5b5b0e123a"), + WithAccessToken("supersecret"), + WithUseKernel(true), + ) + if err != nil { + t.Fatalf("NewConnector: %v", err) + } + _, err = connector.Connect(context.Background()) + if err == nil { + t.Fatal("Connect with WithUseKernel(true) in a non-kernel build should error, got nil") + } + if !strings.Contains(err.Error(), "databricks_kernel") && !strings.Contains(err.Error(), "not compiled") { + t.Errorf("error should mention the missing build tag / not-compiled-in; got: %v", err) + } +}