From 465c68c1c59d473ddb8d2f4916d783d85796218e Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Mon, 20 Jul 2026 17:59:31 +0000 Subject: [PATCH 01/11] fix(kernel): render out-of-nanosecond-range TIMESTAMPs without wrapping MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The kernel Rows path converted arrow timestamps with arrow's ToTime, which forms an int64-nanosecond intermediate (value * unit-multiplier). For a microsecond TIMESTAMP — Databricks' wire unit — that overflows int64 outside ~1678-2262 and silently wraps a valid instant to a wrong one (TIMESTAMP '0001-01-01' scanned back as 1754-08-30, '9999-12-31' as 1816-03-30). Convert via the unit-specific time.Unix/UnixMilli/UnixMicro constructors, which split into (seconds, nanoseconds) internally and represent the full 0001-9999 range Databricks TIMESTAMP allows. The default Thrift path is unaffected — it receives TIMESTAMP as a preformatted server string and never runs this multiply — so this closes a kernel-only divergence from Thrift. TestScanCellTimestampOutOfNanoRange pins both repro endpoints plus the int64-ns boundary years; the existing unit and location tests still pass. Co-authored-by: Isaac Signed-off-by: Mani Kaustubh Mathur --- internal/arrowscan/arrowscan.go | 28 ++++++++++++++++++++++- internal/arrowscan/arrowscan_test.go | 34 ++++++++++++++++++++++++++++ 2 files changed, 61 insertions(+), 1 deletion(-) diff --git a/internal/arrowscan/arrowscan.go b/internal/arrowscan/arrowscan.go index cee634e2..2d70360a 100644 --- a/internal/arrowscan/arrowscan.go +++ b/internal/arrowscan/arrowscan.go @@ -173,7 +173,15 @@ func ScanCellCached(col arrow.Array, row int, loc *time.Location, keys *StructKe // it. The cross-backend parity test pins both sides to µs (matching the wire // reality), so it can't exercise a non-µs value; TestScanCellTimestampUnits // covers this arm's unit-correctness directly instead. - return inLocation(c.Value(row).ToTime(dt.Unit), loc), nil + // + // Convert via the unit-specific time.Unix* constructors rather than arrow's + // ToTime: ToTime forms an int64-nanosecond intermediate (value*multiplier), + // which overflows for TIMESTAMPs outside ~1678–2262 and silently wraps a valid + // instant to a wrong one (e.g. TIMESTAMP '0001-01-01' → 1754). time.UnixMicro/ + // UnixMilli/Unix split into (sec, nsec) internally, so the full 0001–9999 range + // round-trips. The Thrift default path never hits this — it receives TIMESTAMP + // as a preformatted server string — so this divergence is kernel-only. + return inLocation(timestampToTime(int64(c.Value(row)), dt.Unit), loc), nil case *array.Decimal128: dt := col.DataType().(*arrow.Decimal128Type) return decimalfmt.ExactString(c.Value(row), dt.Scale), nil @@ -276,6 +284,24 @@ func abs64(x int64) int64 { return x } +// timestampToTime converts an arrow timestamp value in the given unit to a UTC +// time.Time without arrow's ToTime int64-nanosecond overflow (see the *array.Timestamp +// scan arm). Each constructor takes its argument in native units and splits into +// (seconds, nanoseconds) internally, so the whole 0001–9999 range is representable; +// only the nanosecond arm passes ns directly (already the finest unit, cannot overflow). +func timestampToTime(v int64, unit arrow.TimeUnit) time.Time { + switch unit { + case arrow.Second: + return time.Unix(v, 0).UTC() + case arrow.Millisecond: + return time.UnixMilli(v).UTC() + case arrow.Microsecond: + return time.UnixMicro(v).UTC() + default: // arrow.Nanosecond + return time.Unix(0, v).UTC() + } +} + // inLocation renders t in loc, matching the Thrift path's .In(location); a nil // loc leaves the value in UTC (arrow's ToTime default). func inLocation(t time.Time, loc *time.Location) time.Time { diff --git a/internal/arrowscan/arrowscan_test.go b/internal/arrowscan/arrowscan_test.go index 957a8f47..f58740ef 100644 --- a/internal/arrowscan/arrowscan_test.go +++ b/internal/arrowscan/arrowscan_test.go @@ -274,6 +274,40 @@ func TestScanCellTimestampUnits(t *testing.T) { } } +// TestScanCellTimestampOutOfNanoRange pins the fix for the microsecond→nanosecond +// overflow: arrow's ToTime forms value*multiplier as an int64 ns count, which +// overflows for instants outside ~1678–2262 and silently wraps to a wrong time +// (TIMESTAMP '0001-01-01' scanned back as 1754, '9999-12-31' as 1816). The +// unit-split constructors used by ScanCell must round-trip the full 0001–9999 +// range that Databricks TIMESTAMP allows. +func TestScanCellTimestampOutOfNanoRange(t *testing.T) { + pool := memory.NewGoAllocator() + cases := []time.Time{ + time.Date(1, time.January, 1, 0, 0, 0, 0, time.UTC), // TIMESTAMP '0001-01-01 00:00:00' + time.Date(9999, time.December, 31, 23, 59, 59, 0, time.UTC), // TIMESTAMP '9999-12-31 23:59:59' + time.Date(1677, time.January, 1, 0, 0, 0, 0, time.UTC), // just below the int64-ns floor + time.Date(2263, time.January, 1, 0, 0, 0, 0, time.UTC), // just above the int64-ns ceiling + } + for _, want := range cases { + t.Run(want.Format("2006-01-02"), func(t *testing.T) { + // Databricks TIMESTAMP is always microseconds on the wire. + b := array.NewTimestampBuilder(pool, &arrow.TimestampType{Unit: arrow.Microsecond}) + defer b.Release() + b.Append(arrow.Timestamp(want.UnixMicro())) + arr := b.NewArray() + defer arr.Release() + + v, err := ScanCell(arr, 0, nil) + if err != nil { + t.Fatal(err) + } + if got := v.(time.Time); !got.Equal(want) { + t.Errorf("got %v, want %v (arrow ToTime would have wrapped this)", got.UTC(), want) + } + }) + } +} + // ScanCell renders nested types (list/struct/map) to a JSON string matching the // Thrift path. func TestScanCellNested(t *testing.T) { From 21ad6fad8419527cabb0619d9089d0ea7685760a Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Mon, 20 Jul 2026 17:59:40 +0000 Subject: [PATCH 02/11] fix(kernel): report 0 affected rows for DDL to match Thrift MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The C ABI returns -1 for "not applicable / unknown" affected-row counts — DDL, SELECT, or a warehouse that doesn't surface the counter (the kernel's num_modified_rows Option None). The Thrift path reports 0 in those cases, so Result.RowsAffected() diverged: a CREATE/DROP SCHEMA returned -1 on the kernel backend and 0 on Thrift. Fold the -1 sentinel to 0 so RowsAffected() is identical across backends; a real DML count (>= 0) passes through unchanged. The rule lives in a pure, untagged normalizeAffectedRows so it is unit-testable in the default CGO_ENABLED=0 build alongside the other kernel decision helpers. Co-authored-by: Isaac Signed-off-by: Mani Kaustubh Mathur --- internal/backend/kernel/affectedrows.go | 21 +++++++++++++++++ internal/backend/kernel/affectedrows_test.go | 24 ++++++++++++++++++++ internal/backend/kernel/operation.go | 9 +++++++- 3 files changed, 53 insertions(+), 1 deletion(-) create mode 100644 internal/backend/kernel/affectedrows.go create mode 100644 internal/backend/kernel/affectedrows_test.go diff --git a/internal/backend/kernel/affectedrows.go b/internal/backend/kernel/affectedrows.go new file mode 100644 index 00000000..ca2bd564 --- /dev/null +++ b/internal/backend/kernel/affectedrows.go @@ -0,0 +1,21 @@ +package kernel + +// This file is intentionally NOT behind the `cgo && databricks_kernel` build tag. +// It holds the pure normalization of the C ABI's affected-row count so the rule is +// unit-testable in the default CGO_ENABLED=0 build (like paramBindArg in bindparams.go), +// separate from the cgo call site in operation.go. + +// normalizeAffectedRows maps the kernel C ABI's modified-row count onto the value +// database/sql's Result.RowsAffected() reports, matching the Thrift path. +// +// The C ABI returns -1 for "not applicable / unknown": DDL, SELECT, or a warehouse +// that doesn't surface the counter (the kernel's num_modified_rows Option is +// None). The Thrift path reports 0 in exactly those cases (TGetOperationStatusResp +// defaults NumModifiedRows to 0), so the -1 sentinel is folded to 0 for parity. +// A real DML count (>= 0) passes through unchanged. +func normalizeAffectedRows(n int64) int64 { + if n < 0 { + return 0 + } + return n +} diff --git a/internal/backend/kernel/affectedrows_test.go b/internal/backend/kernel/affectedrows_test.go new file mode 100644 index 00000000..9df1b5f0 --- /dev/null +++ b/internal/backend/kernel/affectedrows_test.go @@ -0,0 +1,24 @@ +package kernel + +import "testing" + +func TestNormalizeAffectedRows(t *testing.T) { + cases := []struct { + name string + in int64 + want int64 + }{ + {"ddl or select sentinel (-1) folds to Thrift's 0", -1, 0}, + {"any negative folds to 0", -42, 0}, + {"zero real count is preserved", 0, 0}, + {"positive DML count passes through", 7, 7}, + {"large DML count passes through", 1_000_000, 1_000_000}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + if got := normalizeAffectedRows(tc.in); got != tc.want { + t.Errorf("normalizeAffectedRows(%d) = %d, want %d", tc.in, got, tc.want) + } + }) + } +} diff --git a/internal/backend/kernel/operation.go b/internal/backend/kernel/operation.go index 4e0228b9..0bd82c58 100644 --- a/internal/backend/kernel/operation.go +++ b/internal/backend/kernel/operation.go @@ -194,7 +194,14 @@ func (k *KernelBackend) execute(ctx context.Context, req backend.ExecRequest) (b // Capture the modified-row count and server query id now, while exec is live — // the operation is closed (nulling exec) before these are read on the // ExecContext path, and the query-id pointer is only valid while exec lives. - op.affectedRows = int64(C.kernel_executed_statement_num_modified_rows(exec)) + // + // The C ABI returns -1 for "not applicable / unknown" (DDL, SELECT, or a + // warehouse that doesn't surface the counter — the kernel's Option None). + // The Thrift path reports 0 in that case (TGetOperationStatusResp defaults + // NumModifiedRows to 0), so normalize the sentinel to 0 to keep RowsAffected() + // identical across backends; real DML counts (>= 0) pass through unchanged. + // normalizeAffectedRows is the pure (CGO_ENABLED=0-testable) form of this rule. + op.affectedRows = normalizeAffectedRows(int64(C.kernel_executed_statement_num_modified_rows(exec))) // A nil execErr means the statement reached a terminal state server-side (it // committed). We deliberately do NOT re-check ctx.Err() here to convert a // completed statement into a cancellation: for a non-idempotent DML that would From c0772716302e4db232b4a0d016aec74bb0de962e Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Wed, 22 Jul 2026 04:07:59 +0000 Subject: [PATCH 03/11] fix(kernel): forward User-Agent so query history attributes the Go driver The kernel backend never forwarded the driver's User-Agent, so the kernel's built-in UA leaked through and query history misattributed SEA-path queries. Forward the same composed UA the Thrift path sends via set_custom_header, so both backends are attributed alike. Default-on; WithUserAgentEntry still customizes it. Co-authored-by: Isaac Signed-off-by: Mani Kaustubh Mathur --- internal/backend/kernel/backend.go | 14 ++++++++++++++ internal/backend/kernel/config.go | 4 ++++ kernel_config.go | 3 +++ kernel_config_test.go | 4 +++- 4 files changed, 24 insertions(+), 1 deletion(-) diff --git a/internal/backend/kernel/backend.go b/internal/backend/kernel/backend.go index febf37e4..5dedc4b3 100644 --- a/internal/backend/kernel/backend.go +++ b/internal/backend/kernel/backend.go @@ -132,6 +132,20 @@ func (k *KernelBackend) OpenSession(ctx context.Context) error { return err } + // User-Agent so query history attributes the kernel path to this driver. + if k.cfg.UserAgent != "" { + name := newCStr("User-Agent") + val := newCStr(k.cfg.UserAgent) + errSet := call(func() C.KernelStatusCode { + return C.kernel_session_config_set_custom_header(cfg, name.c, val.c) + }) + name.free() + val.free() + if errSet != nil { + return fmt.Errorf("kernel: set_custom_header[User-Agent]: %w", toConnError(errSet)) + } + } + // 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 diff --git a/internal/backend/kernel/config.go b/internal/backend/kernel/config.go index 5727407c..e0577891 100644 --- a/internal/backend/kernel/config.go +++ b/internal/backend/kernel/config.go @@ -17,6 +17,10 @@ type Config struct { WarehouseID string // bare warehouse id; preferred over HTTPPath when set Auth Auth // PAT / OAuth M2M / OAuth U2M + // UserAgent is forwarded as the User-Agent header so the kernel path is + // attributed to this driver (not the kernel's built-in UA). Empty leaves it unset. + UserAgent string + // SessionConf carries server-bound session confs verbatim — the same map the // Thrift backend forwards (STATEMENT_TIMEOUT, QUERY_TAGS, TIMEZONE, …). SessionConf map[string]string diff --git a/kernel_config.go b/kernel_config.go index e4bf996f..72443729 100644 --- a/kernel_config.go +++ b/kernel_config.go @@ -11,6 +11,7 @@ import ( "github.com/databricks/databricks-sql-go/auth/pat" dbsqlerr "github.com/databricks/databricks-sql-go/errors" "github.com/databricks/databricks-sql-go/internal/backend/kernel" + "github.com/databricks/databricks-sql-go/internal/client" "github.com/databricks/databricks-sql-go/internal/config" ) @@ -116,6 +117,8 @@ func buildKernelConfig(cfg *config.Config, kauth kernel.Auth) kernel.Config { WarehouseID: cfg.WarehouseID, Auth: kauth, Location: cfg.Location, + // Same UA the Thrift path sends, so query history attributes both alike. + UserAgent: client.BuildUserAgent(cfg), // Initial namespace: no kernel config setter, so the kernel backend applies // these post-connect via USE CATALOG / USE SCHEMA. Catalog: cfg.Catalog, diff --git a/kernel_config_test.go b/kernel_config_test.go index c152eeab..3a5a2de2 100644 --- a/kernel_config_test.go +++ b/kernel_config_test.go @@ -329,9 +329,11 @@ var kernelConfigFieldDisposition = map[string]string{ "MaxRows": "inert", "UseLz4Compression": "inert", // kernel negotiates compression internally + // Rides in the forwarded User-Agent header (set_custom_header). + "UserAgentEntry": "forwarded", + // 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 "EnableTelemetry": "inert", "TelemetryBatchSize": "inert", "TelemetryFlushInterval": "inert", From 8815db9c35477a5db9bb546a43eb3a0054e000a9 Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Wed, 22 Jul 2026 18:16:27 +0000 Subject: [PATCH 04/11] fix(kernel): report VOID/NULL columns as STRING to match Thrift The server stringifies VOID over the Thrift wire (verified live: even bare SELECT NULL reports STRING_TYPE), same as intervals, but hands the kernel a native arrow.NULL schema. Map arrow.NULL to STRING so both backends report identical column metadata; correct the two parity tests that wrongly pinned NULL_TYPE (derived from the enum name, never captured live). Co-authored-by: Isaac Signed-off-by: Mani Kaustubh Mathur --- internal/arrowscan/coltype.go | 77 ++++++--------------- internal/arrowscan/coltype_test.go | 22 +++--- internal/rows/coltype_thrift_parity_test.go | 37 ++++------ 3 files changed, 42 insertions(+), 94 deletions(-) diff --git a/internal/arrowscan/coltype.go b/internal/arrowscan/coltype.go index e795fcfc..7cdba4d7 100644 --- a/internal/arrowscan/coltype.go +++ b/internal/arrowscan/coltype.go @@ -10,32 +10,18 @@ import ( ) // ColumnTypeInfo is the per-column metadata database/sql surfaces through -// sql.ColumnType — the DatabaseTypeName, ScanType, and Length the optional -// driver.RowsColumnType* interfaces return. The kernel backend derives it from -// the result's Arrow schema; ColumnTypeInfoFor is the single mapping both the -// value scanner (ScanCellCached) and the type-metadata reporter agree on, so a -// column's reported type can never drift from what a row actually scans into. -// -// The mapping mirrors the Thrift backend (internal/rows/rows.go getScanType / -// ColumnTypeDatabaseTypeName / ColumnTypeLength) so a query reports byte-identical -// column metadata on either backend — the same "identical results across backends" -// contract the value renderers hold. The pure-Go guards are TestColumnTypeInfoFor, -// TestColumnTypeInfoScanTypeCoversScanner, and TestColumnTypeInfoMatchesThriftMapping -// (which cross-checks this mapping against the Thrift functions directly); the live -// TestKernelThriftColumnTypeParity confirms it against a real warehouse. +// sql.ColumnType. The kernel derives it from the result's Arrow schema via +// ColumnTypeInfoFor — the mapping the value scanner and type reporter share, kept +// byte-identical to the Thrift backend (guarded by the coltype parity tests). type ColumnTypeInfo struct { - // DatabaseTypeName is the Databricks type name (e.g. "BIGINT", "DECIMAL", - // "ARRAY"), matching what the Thrift path reports; "" for a type with no - // Databricks name. + // DatabaseTypeName is the Databricks type name (e.g. "BIGINT", "DECIMAL"), + // matching the Thrift path; "" for a type with no Databricks name. DatabaseTypeName string // ScanType is the Go type database/sql recommends scanning the column into, - // matching the Thrift path. nil (only for the NULL type) makes database/sql - // fall back to interface{}, exactly as the Thrift path's nil scan type does. + // matching the Thrift path. ScanType reflect.Type - // Length / HasLength report a variable-length column's length. As on the - // Thrift path, only variable-length types (string/binary/nested, and the - // server-stringified interval/geo types) report a length, and the reported - // value is math.MaxInt64 (unbounded); fixed-width types report (0, false). + // Length / HasLength report a variable-length column's unbounded length + // (math.MaxInt64), matching Thrift; fixed-width types report (0, false). Length int64 HasLength bool } @@ -58,14 +44,8 @@ var ( // ColumnTypeInfoFor maps an Arrow column type to the metadata database/sql // exposes, matching the Thrift backend for every Databricks type. The Arrow types -// listed here are exactly those ScanCellCached scans, so the type reported for a -// column and the value produced for its cells stay in lockstep. -// -// Notably: DECIMAL reports sql.RawBytes (the Thrift scan type for DECIMAL) even -// though the value is rendered as an exact string — matching Thrift, which also -// scans DECIMAL into RawBytes; and the interval / geo types report STRING because -// the Thrift server pre-formats them to strings and the kernel formats them -// Go-side to the same string, so both are indistinguishable to a caller. +// here are exactly those ScanCellCached scans, so a column's reported type and its +// scanned value stay in lockstep. func ColumnTypeInfoFor(dt arrow.DataType) ColumnTypeInfo { switch dt.ID() { case arrow.BOOL: @@ -79,11 +59,8 @@ func ColumnTypeInfoFor(dt arrow.DataType) ColumnTypeInfo { case arrow.INT64: return ColumnTypeInfo{DatabaseTypeName: "BIGINT", ScanType: scanTypeInt64} case arrow.UINT8, arrow.UINT16, arrow.UINT32, arrow.UINT64: - // Databricks SQL has no unsigned types, so these do not occur in practice; - // this arm is defensive and stays in lockstep with ScanCellCached, which - // widens every unsigned integer to int64 (driver.Value has no uint64). Report - // BIGINT/int64 to match that scanned value rather than falling through to the - // generic *interface{} default. + // Databricks SQL has no unsigned types; defensive arm matching ScanCellCached, + // which widens unsigned ints to int64 (driver.Value has no uint64). return ColumnTypeInfo{DatabaseTypeName: "BIGINT", ScanType: scanTypeInt64} case arrow.FLOAT32: return ColumnTypeInfo{DatabaseTypeName: "FLOAT", ScanType: scanTypeFloat32} @@ -100,10 +77,8 @@ func ColumnTypeInfoFor(dt arrow.DataType) ColumnTypeInfo { case arrow.TIMESTAMP: return ColumnTypeInfo{DatabaseTypeName: "TIMESTAMP", ScanType: scanTypeDateTime} case arrow.DECIMAL128: - // Thrift scans DECIMAL into sql.RawBytes; match it even though the kernel - // renders the value as an exact fixed-point string (both convert cleanly to - // a caller's *string/*[]byte, and reporting the same scan type keeps the - // backends indistinguishable). + // Match Thrift's sql.RawBytes scan type even though the kernel renders the + // value as an exact string — both convert cleanly to a caller's *string/*[]byte. return ColumnTypeInfo{DatabaseTypeName: "DECIMAL", ScanType: scanTypeRawBytes} case arrow.LIST, arrow.LARGE_LIST, arrow.FIXED_SIZE_LIST: return varLen("ARRAY", scanTypeRawBytes) @@ -112,24 +87,17 @@ func ColumnTypeInfoFor(dt arrow.DataType) ColumnTypeInfo { case arrow.STRUCT: return varLen("STRUCT", scanTypeRawBytes) case arrow.DURATION: - // INTERVAL DAY TO SECOND. Parity target is what the Thrift backend REPORTS, - // which is config-dependent: in the prod default (native-interval Arrow off) - // the server pre-formats intervals to text and declares the Thrift column - // STRING_TYPE, so Thrift's GetDBTypeName yields "STRING" and MaxInt64 length - // — verified live against both backends. We therefore report STRING here even - // though the kernel receives a native arrow.DURATION (which it formats Go-side - // to the identical string). If a warehouse ever enables native-interval Thrift - // the server would instead declare INTERVAL_DAY_TIME and Thrift would report - // that; matching STRING is correct for the default path the parity test pins, - // and the scanned VALUE is identical either way — only this label would differ. + // INTERVAL DAY TO SECOND: Thrift's prod default pre-formats intervals to text + // and declares STRING_TYPE (verified live), so match STRING; the kernel formats + // the native arrow.DURATION Go-side to the identical string. return varLen("STRING", scanTypeString) case arrow.INTERVAL_MONTHS: // INTERVAL YEAR TO MONTH — same server-config reasoning as arrow.DURATION. return varLen("STRING", scanTypeString) case arrow.NULL: - // The NULL type has no scan type on the Thrift path (nil → database/sql - // falls back to interface{}); mirror that. - return ColumnTypeInfo{DatabaseTypeName: "NULL", ScanType: nil} + // VOID/NULL columns: the server stringifies them over Thrift (verified live: + // even bare SELECT NULL reports STRING), same as intervals — so match STRING. + return varLen("STRING", scanTypeString) default: // A type ScanCellCached does not handle: report the Thrift default scan type // (*interface{}) and no database name, rather than inventing one. @@ -137,9 +105,8 @@ func ColumnTypeInfoFor(dt arrow.DataType) ColumnTypeInfo { } } -// varLen builds a ColumnTypeInfo for a variable-length type, which reports an -// unbounded length (math.MaxInt64) just as the Thrift path does for -// string/binary/nested columns. +// varLen builds a ColumnTypeInfo for a variable-length type, reporting the +// unbounded length (math.MaxInt64) the Thrift path uses for such columns. func varLen(name string, scan reflect.Type) ColumnTypeInfo { return ColumnTypeInfo{DatabaseTypeName: name, ScanType: scan, Length: math.MaxInt64, HasLength: true} } diff --git a/internal/arrowscan/coltype_test.go b/internal/arrowscan/coltype_test.go index ac71ad0c..3878f79f 100644 --- a/internal/arrowscan/coltype_test.go +++ b/internal/arrowscan/coltype_test.go @@ -10,13 +10,10 @@ import ( "github.com/apache/arrow/go/v12/arrow" ) -// TestColumnTypeInfoFor pins the Arrow → column-metadata mapping the kernel -// backend reports through sql.ColumnType, so it stays byte-identical to the -// Thrift path (internal/rows/rows.go). This is a pure-Go, no-warehouse guard: the -// gap it regresses (PECOBLR-3692) was invisible to the value-parity suites -// because they compare scanned VALUES, not Rows.ColumnType* metadata. The -// expected DatabaseTypeName / ScanType / Length values are the ground truth -// captured from the Thrift backend on a live warehouse. +// TestColumnTypeInfoFor pins the Arrow → column-metadata mapping so it stays +// byte-identical to the Thrift path — a pure-Go guard the value-parity suites miss +// (they compare scanned values, not ColumnType metadata). Expected values are +// ground truth captured from Thrift on a live warehouse. func TestColumnTypeInfoFor(t *testing.T) { str := reflect.TypeOf("") raw := reflect.TypeOf(sql.RawBytes{}) @@ -53,7 +50,8 @@ func TestColumnTypeInfoFor(t *testing.T) { {"struct", arrow.StructOf(arrow.Field{Name: "x", Type: arrow.PrimitiveTypes.Int64}), "STRUCT", raw, math.MaxInt64, true}, {"duration", &arrow.DurationType{Unit: arrow.Microsecond}, "STRING", str, math.MaxInt64, true}, {"month_interval", arrow.FixedWidthTypes.MonthInterval, "STRING", str, math.MaxInt64, true}, - {"null", arrow.Null, "NULL", nil, 0, false}, + // VOID/NULL: Thrift stringifies to STRING on the wire (verified live), like intervals. + {"null", arrow.Null, "STRING", str, math.MaxInt64, true}, } for _, c := range cases { t.Run(c.name, func(t *testing.T) { @@ -72,12 +70,8 @@ func TestColumnTypeInfoFor(t *testing.T) { } // TestColumnTypeInfoScanTypeCoversScanner is the lockstep guard: every Arrow type -// the value scanner (ScanCellCached) handles must also have a non-fallback entry -// in ColumnTypeInfoFor, so a future scalar type added to the scanner without a -// matching type-metadata entry is caught here rather than silently reporting the -// generic *interface{} scan type at runtime. It checks the representative arrays -// the scanner switches on; the NULL type is intentionally excluded (its nil scan -// type is the correct, Thrift-matching value). +// ScanCellCached handles must have a non-fallback ColumnTypeInfoFor entry, so a new +// scanner type without matching metadata is caught here, not at runtime. func TestColumnTypeInfoScanTypeCoversScanner(t *testing.T) { unknown := reflect.TypeOf(new(any)) scannerTypes := []arrow.DataType{ diff --git a/internal/rows/coltype_thrift_parity_test.go b/internal/rows/coltype_thrift_parity_test.go index be178a64..7635ec00 100644 --- a/internal/rows/coltype_thrift_parity_test.go +++ b/internal/rows/coltype_thrift_parity_test.go @@ -10,24 +10,12 @@ import ( "github.com/databricks/databricks-sql-go/internal/rows/rowscanner" ) -// TestColumnTypeInfoMatchesThriftMapping is the pure-Go drift guard the kernel -// backend's column-metadata mapping (internal/arrowscan.ColumnTypeInfoFor) was -// missing. That mapping restates the Thrift path's per-type decisions — getScanType -// (this package), rowscanner.GetDBTypeName, and ColumnTypeLength — but pinned its own -// hardcoded expectations, so a change to the Thrift mapping (e.g. DECIMAL's scan -// type) would leave arrowscan silently divergent while CI stayed green; parity was -// then enforced only by the warehouse-gated, build-tagged TestKernelThriftColumnTypeParity -// (nightly, not PR CI). -// -// This test cross-checks the two mappings DIRECTLY, with no cgo and no warehouse, so -// drift fails a default CGO_ENABLED=0 PR run. The two switch on different type systems -// (Arrow IDs vs Thrift TTypeIds), so it can't share a map — instead it enumerates the -// shared Databricks types as {Arrow type, equivalent Thrift TColumnDesc} pairs and -// asserts ColumnTypeInfoFor(arrow) reports exactly what the Thrift functions produce -// for the paired column. -// -// It lives in package rows (white-box) to reach the unexported getScanType; -// arrowscan is a pure leaf import (no cycle). +// TestColumnTypeInfoMatchesThriftMapping cross-checks arrowscan.ColumnTypeInfoFor +// against the Thrift mapping (getScanType / GetDBTypeName / ColumnTypeLength) +// directly, so a change to one that isn't mirrored in the other fails a pure-Go PR +// run — not just the nightly warehouse-gated parity test. It enumerates each shared +// type as an {Arrow type, equivalent Thrift TColumnDesc} pair and asserts they agree. +// White-box (package rows) to reach the unexported getScanType. func TestColumnTypeInfoMatchesThriftMapping(t *testing.T) { // thriftCol builds the minimal TColumnDesc the Thrift mapping functions read. thriftCol := func(id cli_service.TTypeId) *cli_service.TColumnDesc { @@ -63,15 +51,14 @@ func TestColumnTypeInfoMatchesThriftMapping(t *testing.T) { {"array", arrow.ListOf(arrow.PrimitiveTypes.Int64), cli_service.TTypeId_ARRAY_TYPE}, {"map", arrow.MapOf(arrow.BinaryTypes.String, arrow.PrimitiveTypes.Int64), cli_service.TTypeId_MAP_TYPE}, {"struct", arrow.StructOf(arrow.Field{Name: "x", Type: arrow.PrimitiveTypes.Int64}), cli_service.TTypeId_STRUCT_TYPE}, - // Interval types: in the prod default (native-interval Arrow off) the server - // pre-formats intervals to text and declares the column STRING_TYPE, so the - // kernel — which receives native arrow.DURATION / arrow.INTERVAL_MONTHS and - // formats them Go-side to the same string — must report the STRING metadata the - // Thrift path reports for STRING_TYPE. Pairing them with STRING_TYPE here pins - // exactly that (see the arrowscan DURATION / INTERVAL_MONTHS arms). + // Intervals: the server pre-formats them to text and declares STRING_TYPE over + // Thrift, so pair with STRING_TYPE (the kernel formats the native arrow value to + // the same string). See the arrowscan DURATION / INTERVAL_MONTHS arms. {"interval_day_time", arrow.FixedWidthTypes.Duration_us, cli_service.TTypeId_STRING_TYPE}, {"interval_year_month", arrow.FixedWidthTypes.MonthInterval, cli_service.TTypeId_STRING_TYPE}, - {"null", arrow.Null, cli_service.TTypeId_NULL_TYPE}, + // VOID/NULL: the server stringifies it over Thrift as STRING_TYPE (verified + // live — even bare SELECT NULL), same as the interval types above. + {"null", arrow.Null, cli_service.TTypeId_STRING_TYPE}, } for _, c := range cases { From 0163bcdcd14c01b9145e053fad1d16ed1a746ad3 Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Wed, 22 Jul 2026 20:38:48 +0000 Subject: [PATCH 05/11] test(kernel): cover UA forwarding + VOID column-type parity MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Address review of the parity fixes: assert kc.UserAgent == BuildUserAgent (non-empty) in TestBuildKernelConfig so deleting the forwarding or an empty UA fails a pure-Go run; add a CAST(NULL AS VOID) column to the live column-type parity query so the arrow.NULL→STRING mapping is guarded against the real server, not only the pure-Go tests. Trim the duplicated -1→0 rationale at the affected-rows call site to a one-liner. Co-authored-by: Isaac Signed-off-by: Mani Kaustubh Mathur --- internal/backend/kernel/operation.go | 8 +------- kernel_config_test.go | 6 ++++++ kernel_parity_test.go | 4 +++- 3 files changed, 10 insertions(+), 8 deletions(-) diff --git a/internal/backend/kernel/operation.go b/internal/backend/kernel/operation.go index 0bd82c58..cc60a5bc 100644 --- a/internal/backend/kernel/operation.go +++ b/internal/backend/kernel/operation.go @@ -194,13 +194,7 @@ func (k *KernelBackend) execute(ctx context.Context, req backend.ExecRequest) (b // Capture the modified-row count and server query id now, while exec is live — // the operation is closed (nulling exec) before these are read on the // ExecContext path, and the query-id pointer is only valid while exec lives. - // - // The C ABI returns -1 for "not applicable / unknown" (DDL, SELECT, or a - // warehouse that doesn't surface the counter — the kernel's Option None). - // The Thrift path reports 0 in that case (TGetOperationStatusResp defaults - // NumModifiedRows to 0), so normalize the sentinel to 0 to keep RowsAffected() - // identical across backends; real DML counts (>= 0) pass through unchanged. - // normalizeAffectedRows is the pure (CGO_ENABLED=0-testable) form of this rule. + // normalizeAffectedRows folds the C ABI's -1 sentinel to 0 for Thrift parity. op.affectedRows = normalizeAffectedRows(int64(C.kernel_executed_statement_num_modified_rows(exec))) // A nil execErr means the statement reached a terminal state server-side (it // committed). We deliberately do NOT re-check ctx.Err() here to convert a diff --git a/kernel_config_test.go b/kernel_config_test.go index 3a5a2de2..c2795bf4 100644 --- a/kernel_config_test.go +++ b/kernel_config_test.go @@ -11,6 +11,7 @@ import ( "github.com/databricks/databricks-sql-go/auth/pat" dbsqlerr "github.com/databricks/databricks-sql-go/errors" "github.com/databricks/databricks-sql-go/internal/backend/kernel" + "github.com/databricks/databricks-sql-go/internal/client" "github.com/databricks/databricks-sql-go/internal/config" ) @@ -443,6 +444,11 @@ func TestBuildKernelConfig(t *testing.T) { if kc.Auth.Mode != kauth.Mode || kc.Auth.Token != kauth.Token { t.Errorf("auth not forwarded: got %+v, want %+v", kc.Auth, kauth) } + // UserAgent must be the driver's composed UA, non-empty — else query + // history mis-attributes SEA-path queries to the kernel's built-in UA. + if want := client.BuildUserAgent(c); kc.UserAgent == "" || kc.UserAgent != want { + t.Errorf("UserAgent not forwarded: got %q, want %q", kc.UserAgent, want) + } }) t.Run("MaxChunksInMemory injected into kernel SessionConf", func(t *testing.T) { diff --git a/kernel_parity_test.go b/kernel_parity_test.go index 4987e152..ed9e52ed 100644 --- a/kernel_parity_test.go +++ b/kernel_parity_test.go @@ -116,7 +116,9 @@ func TestKernelThriftColumnTypeParity(t *testing.T) { "CAST('2020-01-01 00:00:00' AS TIMESTAMP) a_ts, CAST('abc' AS BINARY) a_binary, " + "array(1,2,3) a_array, map('k',1) a_map, named_struct('x',1) a_struct, " + `parse_json('{"a":1}') a_variant, INTERVAL '1' DAY a_iv_dt, ` + - "INTERVAL '1' MONTH a_iv_ym, st_point(1,2) a_geom" + // VOID: the server declares it STRING_TYPE over Thrift; asserts the kernel + // arrow.NULL→STRING mapping matches live, not just in the pure-Go tests. + "INTERVAL '1' MONTH a_iv_ym, st_point(1,2) a_geom, CAST(NULL AS VOID) a_void" kernelDB := kernelTestDB(t) defer kernelDB.Close() From 531f6c53ccf769149a48b58d2980117c6217fbe5 Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Fri, 24 Jul 2026 18:41:40 +0000 Subject: [PATCH 06/11] fix(kernel): expose GetArrowBatches on the public Rows interface kernelRows implemented only driver.Rows + the ColumnType* interfaces, so a Conn.Raw + rows.(dbsqlrows.Rows) assertion failed on the kernel backend (it succeeds on Thrift). Implement GetArrowBatches by wrapping the existing zero-copy next_batch pull, and reject GetArrowIPCStreams (the kernel C ABI exports Arrow C Data, not IPC bytes). A compile-time dbrows.Rows assertion guards the gap. Co-authored-by: Isaac Signed-off-by: Mani Kaustubh Mathur --- internal/backend/kernel/rows.go | 91 +++++++++++++++++++++++++++++++++ 1 file changed, 91 insertions(+) diff --git a/internal/backend/kernel/rows.go b/internal/backend/kernel/rows.go index 2020dbda..c556ea3c 100644 --- a/internal/backend/kernel/rows.go +++ b/internal/backend/kernel/rows.go @@ -23,8 +23,10 @@ import ( "github.com/apache/arrow/go/v12/arrow" "github.com/apache/arrow/go/v12/arrow/cdata" + dbsqlerr "github.com/databricks/databricks-sql-go/errors" "github.com/databricks/databricks-sql-go/internal/arrowscan" dbsqlrows "github.com/databricks/databricks-sql-go/internal/rows" + dbrows "github.com/databricks/databricks-sql-go/rows" ) // kernelRows implements the same optional column-type interfaces the Thrift path @@ -37,6 +39,9 @@ var ( _ driver.RowsColumnTypeDatabaseTypeName = (*kernelRows)(nil) _ driver.RowsColumnTypeNullable = (*kernelRows)(nil) _ driver.RowsColumnTypeLength = (*kernelRows)(nil) + // Public batch API, matching the Thrift path (rows.rows). GetArrowBatches is + // real; GetArrowIPCStreams rejects (kernel C ABI exports C Data, not IPC bytes). + _ dbrows.Rows = (*kernelRows)(nil) ) // kernelRows implements driver.Rows over the kernel result stream. It pulls one @@ -55,6 +60,7 @@ type kernelRows struct { cols []string colTypes []arrowscan.ColumnTypeInfo // per-column type metadata (PECOBLR-3692) + schema *arrow.Schema // result-set schema, for GetArrowBatches().Schema() cur arrow.Record // current batch (nil until first Next) rowInCur int // next row index within cur chunkCount int // cumulative batches fetched, for OnChunkFetched @@ -95,6 +101,7 @@ func newKernelRows(ctx context.Context, op *kernelOp, stream *C.kernel_result_st r.Close() return nil, fmt.Errorf("kernel: import schema: %w", err) } + r.schema = sch fields := sch.Fields() r.cols = make([]string, len(fields)) // Derive per-column type metadata from the Arrow schema up front (the same @@ -284,3 +291,87 @@ func (r *kernelRows) nextBatch() error { klogCtx(r.ctx, "nextBatch: %d rows (chunk %d)", rec.NumRows(), r.chunkCount) return nil } + +// GetArrowBatches exposes the kernel stream as an arrow.Record iterator (the +// public zero-copy batch API), reusing the row scanner's next_batch pull. Callers +// MUST Release() each record. +func (r *kernelRows) GetArrowBatches(context.Context) (dbrows.ArrowBatchIterator, error) { + return &kernelBatchIterator{r: r}, nil +} + +// GetArrowIPCStreams is rejected on the kernel path: the C ABI exports Arrow via +// the C Data Interface (next_batch), not IPC bytes. Use GetArrowBatches, or the +// Thrift backend for IPC streams. +func (r *kernelRows) GetArrowIPCStreams(context.Context) (dbrows.ArrowIPCStreamIterator, error) { + return nil, fmt.Errorf("databricks: GetArrowIPCStreams is %w (kernel exports Arrow C Data, not IPC bytes); use GetArrowBatches", + dbsqlerr.ErrNotSupportedByKernel) +} + +// kernelBatchIterator adapts the kernelRows pull loop to the public +// ArrowBatchIterator: it prefetches one batch (so HasNext is exact) and transfers +// each record's ownership to the caller, who must Release it. +type kernelBatchIterator struct { + r *kernelRows + pending arrow.Record // prefetched, not yet handed out + done bool // stream drained (io.EOF seen) + err error // sticky fetch error, surfaced by Next +} + +var _ dbrows.ArrowBatchIterator = (*kernelBatchIterator)(nil) + +// fill buffers one batch into pending, or sets done/err. It re-uses nextBatch and +// takes ownership of r.cur so the row scanner and Close can't double-release it. +func (it *kernelBatchIterator) fill() { + if it.pending != nil || it.done || it.err != nil { + return + } + if it.r.closed || it.r.eof { // nothing left to pull + it.done = true + return + } + switch err := it.r.nextBatch(); err { + case nil: + it.pending, it.r.cur = it.r.cur, nil + case io.EOF: + it.done = true + default: + it.err = err + } +} + +// Next returns the next record (caller owns it) or io.EOF when the stream drains. +func (it *kernelBatchIterator) Next() (arrow.Record, error) { + it.fill() + if it.err != nil { + return nil, it.err + } + if it.pending == nil { + return nil, io.EOF + } + rec := it.pending + it.pending = nil + return rec, nil +} + +// HasNext reports whether a following Next would yield a record or an error. +func (it *kernelBatchIterator) HasNext() bool { + it.fill() + return it.pending != nil || it.err != nil +} + +// Close releases any buffered record and tears down the underlying stream. +func (it *kernelBatchIterator) Close() { + if it.pending != nil { + it.pending.Release() + it.pending = nil + } + it.r.Close() //nolint:errcheck +} + +// Schema returns the result-set schema captured at construction. +func (it *kernelBatchIterator) Schema() (*arrow.Schema, error) { + if it.r.schema == nil { + return nil, fmt.Errorf("kernel: no schema available") + } + return it.r.schema, nil +} From 5c25584b4e5ba27409a8f3786c90eae88b138387 Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Fri, 24 Jul 2026 19:09:20 +0000 Subject: [PATCH 07/11] fix(auth): route M2M token logs through the driver logger m2m.Authenticate logged via the global zerolog logger, whose default level is Debug, so "OAuth token fetched successfully" printed regardless of the driver's configured log level (default Warn) on both backends. Route it through the driver's logger package so it honors WithLogLevel like every other line. Co-authored-by: Isaac Signed-off-by: Mani Kaustubh Mathur --- auth/oauth/m2m/m2m.go | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/auth/oauth/m2m/m2m.go b/auth/oauth/m2m/m2m.go index 9e7c9c20..f9c8509b 100644 --- a/auth/oauth/m2m/m2m.go +++ b/auth/oauth/m2m/m2m.go @@ -8,7 +8,7 @@ import ( "github.com/databricks/databricks-sql-go/auth" "github.com/databricks/databricks-sql-go/auth/oauth" - "github.com/rs/zerolog/log" + "github.com/databricks/databricks-sql-go/logger" "golang.org/x/oauth2" "golang.org/x/oauth2/clientcredentials" ) @@ -72,9 +72,11 @@ func (c *authClient) Authenticate(r *http.Request) error { c.tokenSource = GetTokenSource(config) token, err := c.tokenSource.Token() - log.Debug().Msgf("databricks OAuth token fetched successfully") + // Log via the driver's configurable logger (defaults to Warn) rather than the + // global zerolog logger, so this line honors WithLogLevel like every other. + logger.Debug().Msgf("databricks OAuth token fetched successfully") if err != nil { - log.Err(err).Msg("failed to get token") + logger.Err(err).Msg("failed to get token") return err } From ba072afd9e2aea2083ac4b59a54687b3a2d1bff0 Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Fri, 24 Jul 2026 19:52:13 +0000 Subject: [PATCH 08/11] fix(auth): route remaining auth logs through the driver logger MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The U2M, OAuth-config, and token-provider paths still logged via the global zerolog logger (default level Debug), so their lines printed as raw JSON regardless of the driver's configured level — the same leak the m2m fix addressed. Route all of them through the driver's logger package so every auth log honors WithLogLevel and uses the driver's format; auth logic is untouched. This closes the last global-zerolog call sites in the driver. Co-authored-by: Isaac Signed-off-by: Mani Kaustubh Mathur --- auth/oauth/oauth.go | 8 ++++---- auth/oauth/u2m/authenticator.go | 14 +++++++------- auth/tokenprovider/authenticator.go | 4 ++-- auth/tokenprovider/cached.go | 6 +++--- auth/tokenprovider/exchange.go | 18 +++++++++--------- 5 files changed, 25 insertions(+), 25 deletions(-) diff --git a/auth/oauth/oauth.go b/auth/oauth/oauth.go index 0525e8cd..d0fd3f66 100644 --- a/auth/oauth/oauth.go +++ b/auth/oauth/oauth.go @@ -11,7 +11,7 @@ import ( "time" "github.com/coreos/go-oidc/v3/oidc" - "github.com/rs/zerolog/log" + "github.com/databricks/databricks-sql-go/logger" "golang.org/x/oauth2" ) @@ -99,7 +99,7 @@ func resolveOIDCIssuer(ctx context.Context, client *http.Client, hostName string cfgURL := fmt.Sprintf("https://%s/.well-known/databricks-config", hostName) meta, ok := fetchHostMetadata(ctx, client, cfgURL) if !ok || meta.OIDCEndpoint == "" { - log.Debug().Msgf("oauth: no usable databricks-config for %q; using bare-host OIDC issuer", hostName) + logger.Debug().Msgf("oauth: no usable databricks-config for %q; using bare-host OIDC issuer", hostName) return fallback } @@ -107,13 +107,13 @@ func resolveOIDCIssuer(ctx context.Context, client *http.Client, hostName string // placeholder would resolve to a malformed ".../accounts/" issuer. Fall back // rather than emit it (the function's documented contract). if strings.Contains(meta.OIDCEndpoint, accountIDPlaceholder) && meta.AccountID == "" { - log.Warn().Msgf("oauth: databricks-config for %q has an %s placeholder but empty account_id; using bare-host OIDC issuer", hostName, accountIDPlaceholder) + logger.Warn().Msgf("oauth: databricks-config for %q has an %s placeholder but empty account_id; using bare-host OIDC issuer", hostName, accountIDPlaceholder) return fallback } issuer := substituteAccountID(meta) if !isValidDatabricksIssuer(issuer) { - log.Warn().Msgf("oauth: databricks-config for %q advertised an unusable oidc_endpoint %q; using bare-host OIDC issuer", hostName, issuer) + logger.Warn().Msgf("oauth: databricks-config for %q advertised an unusable oidc_endpoint %q; using bare-host OIDC issuer", hostName, issuer) return fallback } return issuer diff --git a/auth/oauth/u2m/authenticator.go b/auth/oauth/u2m/authenticator.go index 1d5b46bf..c73e5ace 100644 --- a/auth/oauth/u2m/authenticator.go +++ b/auth/oauth/u2m/authenticator.go @@ -20,7 +20,7 @@ import ( "github.com/databricks/databricks-sql-go/auth" "github.com/databricks/databricks-sql-go/auth/oauth" - "github.com/rs/zerolog/log" + "github.com/databricks/databricks-sql-go/logger" "golang.org/x/oauth2" ) @@ -160,7 +160,7 @@ func (tsp *tokenSourceProvider) GetTokenSource() (oauth2.TokenSource, error) { loginURL := tsp.config.AuthCodeURL(state, challenge, challengeMethod) tsp.state = state - log.Info().Msgf("listening on %s://%s/", tsp.redirectURL.Scheme, tsp.redirectURL.Host) + logger.Info().Msgf("listening on %s://%s/", tsp.redirectURL.Scheme, tsp.redirectURL.Host) listener, err := net.Listen("tcp", tsp.redirectURL.Host) if err != nil { return nil, err @@ -225,28 +225,28 @@ func (tsp *tokenSourceProvider) ServeHTTP(w http.ResponseWriter, r *http.Request // Do some checking of the response here to show more relevant content if resp.err != "" { - log.Error().Msg(resp.err) + logger.Error().Msg(resp.err) w.WriteHeader(http.StatusBadRequest) _, err := w.Write([]byte(errorHTML("Identity Provider returned an error: " + resp.err))) //nolint:gosec // XSS not a concern for local OAuth callback if err != nil { - log.Error().Err(err).Msg("unable to write error response") + logger.Error().Err(err).Msg("unable to write error response") } return } if resp.state != tsp.state && r.URL.String() != "/favicon.ico" { msg := "Authentication state received did not match original request. Please try to login again." - log.Error().Msg(msg) + logger.Error().Msg(msg) w.WriteHeader(http.StatusBadRequest) _, err := w.Write([]byte(errorHTML(msg))) if err != nil { - log.Error().Err(err).Msg("unable to write error response") + logger.Error().Err(err).Msg("unable to write error response") } return } _, err := w.Write([]byte(infoHTML("CLI Login Success", "You may close this window anytime now and go back to terminal"))) if err != nil { - log.Error().Err(err).Msg("unable to write success response") + logger.Error().Err(err).Msg("unable to write success response") } } diff --git a/auth/tokenprovider/authenticator.go b/auth/tokenprovider/authenticator.go index b4ef57ff..9452811c 100644 --- a/auth/tokenprovider/authenticator.go +++ b/auth/tokenprovider/authenticator.go @@ -6,7 +6,7 @@ import ( "net/http" "github.com/databricks/databricks-sql-go/auth" - "github.com/rs/zerolog/log" + "github.com/databricks/databricks-sql-go/logger" ) // TokenProviderAuthenticator implements auth.Authenticator using a TokenProvider. @@ -47,7 +47,7 @@ func (a *TokenProviderAuthenticator) Authenticate(r *http.Request) error { } token.SetAuthHeader(r) - log.Debug().Msgf("token provider authenticator: authenticated using provider %s", a.provider.Name()) + logger.Debug().Msgf("token provider authenticator: authenticated using provider %s", a.provider.Name()) return nil } diff --git a/auth/tokenprovider/cached.go b/auth/tokenprovider/cached.go index fdb86c89..5c4de8b7 100644 --- a/auth/tokenprovider/cached.go +++ b/auth/tokenprovider/cached.go @@ -6,7 +6,7 @@ import ( "sync" "time" - "github.com/rs/zerolog/log" + "github.com/databricks/databricks-sql-go/logger" ) // CachedTokenProvider wraps another provider and caches tokens @@ -43,7 +43,7 @@ func (p *CachedTokenProvider) GetToken(ctx context.Context) (*Token, error) { // If cache is valid and not being refreshed, return a copy if cached != nil && !needsRefresh { - log.Debug().Msgf("cached token provider: using cached token for provider %s", p.provider.Name()) + logger.Debug().Msgf("cached token provider: using cached token for provider %s", p.provider.Name()) // Return a copy to avoid concurrent modification issues return copyToken(cached), nil } @@ -75,7 +75,7 @@ func (p *CachedTokenProvider) GetToken(ctx context.Context) (*Token, error) { p.mutex.Unlock() // Fetch new token WITHOUT holding the lock - log.Debug().Msgf("cached token provider: fetching new token from provider %s", p.provider.Name()) + logger.Debug().Msgf("cached token provider: fetching new token from provider %s", p.provider.Name()) token, err := p.provider.GetToken(ctx) // Update cache with result diff --git a/auth/tokenprovider/exchange.go b/auth/tokenprovider/exchange.go index 09c9268d..c0c5161e 100644 --- a/auth/tokenprovider/exchange.go +++ b/auth/tokenprovider/exchange.go @@ -10,8 +10,8 @@ import ( "strings" "time" + "github.com/databricks/databricks-sql-go/logger" "github.com/golang-jwt/jwt/v5" - "github.com/rs/zerolog/log" ) // FederationProvider wraps another token provider and automatically handles token exchange @@ -61,16 +61,16 @@ func (p *FederationProvider) GetToken(ctx context.Context) (*Token, error) { // Check if token is a JWT and needs exchange if p.needsTokenExchange(baseToken.AccessToken) { - log.Debug().Msgf("federation provider: attempting token exchange for %s", p.baseProvider.Name()) + logger.Debug().Msgf("federation provider: attempting token exchange for %s", p.baseProvider.Name()) // Try token exchange exchangedToken, err := p.tryTokenExchange(ctx, baseToken.AccessToken) if err != nil { - log.Warn().Err(err).Msg("federation provider: token exchange failed, using original token") + logger.Warn().Err(err).Msg("federation provider: token exchange failed, using original token") return baseToken, nil // Fall back to original token } - log.Debug().Msg("federation provider: token exchange successful") + logger.Debug().Msg("federation provider: token exchange successful") return exchangedToken, nil } @@ -87,7 +87,7 @@ func (p *FederationProvider) needsTokenExchange(tokenString string) bool { // 3. Token validation will be done by Databricks during exchange token, _, err := new(jwt.Parser).ParseUnverified(tokenString, jwt.MapClaims{}) if err != nil { - log.Debug().Err(err).Msg("federation provider: not a JWT token, skipping exchange") + logger.Debug().Err(err).Msg("federation provider: not a JWT token, skipping exchange") return false } @@ -114,7 +114,7 @@ func (p *FederationProvider) tryTokenExchange(ctx context.Context, subjectToken exchangeURL = "https://" + exchangeURL } else if strings.HasPrefix(exchangeURL, "http://") { // Warn if using insecure HTTP for token exchange - log.Warn().Msgf("federation provider: using insecure HTTP for token exchange: %s", exchangeURL) + logger.Warn().Msgf("federation provider: using insecure HTTP for token exchange: %s", exchangeURL) } if !strings.HasSuffix(exchangeURL, "/") { exchangeURL += "/" @@ -179,7 +179,7 @@ func (p *FederationProvider) tryTokenExchange(ctx context.Context, subjectToken return nil, fmt.Errorf("token exchange returned empty access token") } if tokenResp.TokenType == "" { - log.Debug().Msg("token exchange: token_type not specified, defaulting to Bearer") + logger.Debug().Msg("token exchange: token_type not specified, defaulting to Bearer") tokenResp.TokenType = "Bearer" } if tokenResp.ExpiresIn < 0 { @@ -211,7 +211,7 @@ func (p *FederationProvider) isSameHost(url1, url2 string) bool { u2, err2 := url.Parse(parsedURL2) if err1 != nil || err2 != nil { - log.Debug().Msgf("federation provider: failed to parse URLs for comparison: url1=%s err1=%v, url2=%s err2=%v", + logger.Debug().Msgf("federation provider: failed to parse URLs for comparison: url1=%s err1=%v, url2=%s err2=%v", url1, err1, parsedURL2, err2) return false } @@ -219,7 +219,7 @@ func (p *FederationProvider) isSameHost(url1, url2 string) bool { // Use Hostname() instead of Host to ignore port differences // This handles cases like "host.com:443" == "host.com" for HTTPS isSame := u1.Hostname() == u2.Hostname() - log.Debug().Msgf("federation provider: host comparison: %s vs %s = %v", u1.Hostname(), u2.Hostname(), isSame) + logger.Debug().Msgf("federation provider: host comparison: %s vs %s = %v", u1.Hostname(), u2.Hostname(), isSame) return isSame } From e4dbfb59832a9b333f932b9e9ba6ad4c2b8fe963 Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Fri, 24 Jul 2026 20:00:18 +0000 Subject: [PATCH 09/11] fix(kernel): forward Thrift-parity U2M scopes instead of kernel default MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The kernel U2M path forwarded only the client id, so the kernel applied its own default scope set (all-apis + offline_access). A workspace whose built-in databricks-sql-connector client isn't granted all-apis rejects that with access_denied. resolveKernelAuth now populates Auth.Scopes from oauth.GetScopes — the same function the Thrift path uses (offline_access + sql on AWS/GCP, offline_access + /user_impersonation on Azure) — so both backends authorize identically. The setAuth -> set_auth_u2m wiring already forwarded Scopes; only the population was missing. No kernel/C-ABI change. Co-authored-by: Isaac Signed-off-by: Mani Kaustubh Mathur --- internal/backend/kernel/auth.go | 18 ++++++++---------- kernel_auth_real_test.go | 6 ++++++ kernel_config.go | 17 ++++++++++------- kernel_config_test.go | 7 +++++++ 4 files changed, 31 insertions(+), 17 deletions(-) diff --git a/internal/backend/kernel/auth.go b/internal/backend/kernel/auth.go index 29563b28..ddef6510 100644 --- a/internal/backend/kernel/auth.go +++ b/internal/backend/kernel/auth.go @@ -20,21 +20,19 @@ const ( // validateKernelConfig); OpenSession maps it to exactly one // kernel_session_config_set_auth_* call. // Scopes and RedirectPort map to the optional args of set_auth_u2m and are wired -// through to it by setAuth, but no Go path populates them today: the driver exposes -// no user option for U2M scopes or redirect port on either backend (the native -// Thrift path hardcodes both), so resolveKernelAuth leaves them zero and the kernel -// applies its defaults. They are kept — rather than dropped and the setter hardcoded -// to NULL/0 — so kernel.Auth models the full set_auth_u2m surface: adding a future -// WithOAuthRedirectPort / scopes option becomes populating these, not re-plumbing -// the setter. TestSetAuthByMode's "U2M full" case pins that marshalling so the -// dormant path stays correct. +// through to it by setAuth. resolveKernelAuth populates Scopes with the same +// cloud-specific set the Thrift path requests (via oauth.GetScopes) so both +// backends authorize identically; RedirectPort stays zero (no user option, kernel +// default 8020) but is kept so kernel.Auth models the full set_auth_u2m surface — +// a future WithOAuthRedirectPort becomes populating it, not re-plumbing the setter. +// TestSetAuthByMode's "U2M full" case pins the marshalling of both. type Auth struct { Mode AuthMode Token string // PAT ClientID string // M2M + U2M (U2M: the cloud-inferred Go client id) ClientSecret string // M2M - Scopes []string // U2M — dormant (see note above); nil → kernel default scopes - RedirectPort uint16 // U2M — dormant (see note above); 0 → kernel default port (8020) + Scopes []string // U2M — Thrift-parity scopes from oauth.GetScopes; nil → kernel default + RedirectPort uint16 // U2M — no user option today; 0 → kernel default port (8020) } // M2MCredentialsProvider is implemented by the OAuth M2M authenticator to expose diff --git a/kernel_auth_real_test.go b/kernel_auth_real_test.go index d89ab30a..1e677725 100644 --- a/kernel_auth_real_test.go +++ b/kernel_auth_real_test.go @@ -1,8 +1,10 @@ package dbsql import ( + "reflect" "testing" + "github.com/databricks/databricks-sql-go/auth/oauth" "github.com/databricks/databricks-sql-go/auth/oauth/m2m" "github.com/databricks/databricks-sql-go/auth/oauth/u2m" "github.com/databricks/databricks-sql-go/internal/backend/kernel" @@ -64,5 +66,9 @@ func TestResolveKernelAuthRealAuthenticators(t *testing.T) { if got.Mode != kernel.AuthU2M || got.ClientID == "" { t.Errorf("auth = %+v, want mode=U2M with a non-empty (cloud-inferred) clientID", got) } + // Scopes are forwarded from oauth.GetScopes(cfg.Host) for Thrift parity. + if want := oauth.GetScopes(c.Host, nil); !reflect.DeepEqual(got.Scopes, want) { + t.Errorf("U2M scopes = %v, want %v (Thrift parity)", got.Scopes, want) + } }) } diff --git a/kernel_config.go b/kernel_config.go index 72443729..0a8053a6 100644 --- a/kernel_config.go +++ b/kernel_config.go @@ -8,6 +8,7 @@ import ( "time" "github.com/databricks/databricks-sql-go/auth/noop" + "github.com/databricks/databricks-sql-go/auth/oauth" "github.com/databricks/databricks-sql-go/auth/pat" dbsqlerr "github.com/databricks/databricks-sql-go/errors" "github.com/databricks/databricks-sql-go/internal/backend/kernel" @@ -307,13 +308,15 @@ func resolveKernelAuth(cfg *config.Config) (kernel.Auth, error) { clientID, clientSecret := a.M2MCredentials() return kernel.Auth{Mode: kernel.AuthM2M, ClientID: clientID, ClientSecret: clientSecret}, nil case kernel.U2MCredentialsProvider: - // Go sources only the client id for U2M; kernel.Auth.Scopes / RedirectPort - // are left zero so setAuth passes the kernel's defaults. Go exposes no - // user-facing option for U2M scopes or redirect port on either backend today - // (the native Thrift path hardcodes both), so there is nothing to forward — - // but the kernel.Auth fields + setAuth wiring model the full set_auth_u2m - // surface for if/when such an option is added (see kernel.Auth docs). - return kernel.Auth{Mode: kernel.AuthU2M, ClientID: a.U2MClientID()}, nil + // Forward the SAME cloud-specific scopes the Thrift path requests via + // oauth.GetScopes (offline_access + sql on AWS/GCP, offline_access + + // /user_impersonation on Azure), so both backends authorize against + // the built-in databricks-sql-connector client identically. Without this the + // kernel applied its own default set (all-apis + offline_access), which a + // workspace whose public client isn't granted all-apis rejects with + // access_denied. RedirectPort is still left zero (no user option; kernel + // default 8020). Passing nil to GetScopes yields the pure cloud-default set. + return kernel.Auth{Mode: kernel.AuthU2M, ClientID: a.U2MClientID(), Scopes: oauth.GetScopes(cfg.Host, nil)}, nil case nil, *noop.NoopAuth, *pat.PATAuth: // PAT (or no explicit authenticator). WithAccessToken sets both // cfg.AccessToken and a *pat.PATAuth, but WithAuthenticator(&pat.PATAuth{...}) diff --git a/kernel_config_test.go b/kernel_config_test.go index c2795bf4..bf0c97c1 100644 --- a/kernel_config_test.go +++ b/kernel_config_test.go @@ -8,6 +8,7 @@ import ( "testing" "time" + "github.com/databricks/databricks-sql-go/auth/oauth" "github.com/databricks/databricks-sql-go/auth/pat" dbsqlerr "github.com/databricks/databricks-sql-go/errors" "github.com/databricks/databricks-sql-go/internal/backend/kernel" @@ -182,6 +183,12 @@ func TestValidateKernelConfig(t *testing.T) { if a.Mode != kernel.AuthU2M || a.ClientID != "databricks-sql-connector" { t.Errorf("auth = %+v, want mode=U2M clientID=databricks-sql-connector", a) } + // Scopes must match what the Thrift path requests for this host, so both + // backends authorize against the same client identically (not the kernel's + // all-apis default). baseKernelConfig's host is AWS → [offline_access, sql]. + if want := oauth.GetScopes(c.Host, nil); !reflect.DeepEqual(a.Scopes, want) { + t.Errorf("U2M scopes = %v, want %v (Thrift parity)", a.Scopes, want) + } }) t.Run("last-applied auth wins: M2M then PAT resolves to PAT", func(t *testing.T) { From c9e758137e8057eb9658bfc816e2577236ce287d Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Fri, 24 Jul 2026 21:03:48 +0000 Subject: [PATCH 10/11] fix(kernel): map INTERVAL_DAY_TIME to STRING for Thrift parity MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit DAY-TIME intervals can arrive from the kernel as arrow.INTERVAL_DAY_TIME (*array.DayTimeInterval), not just arrow.DURATION. ColumnTypeInfoFor lacked the case, so dt_interval columns reported ""/*interface{} vs Thrift's STRING/string (comparator diff across 5 configs); the value scanner also lacked it, so a real DT-interval row would have errored. Add both — metadata → STRING and the scanner folding {Days,Milliseconds} to ms through the same renderer — mirroring the existing DURATION/INTERVAL_MONTHS handling. Kernel-path only; no other driver and no kernel/C-ABI change. Co-authored-by: Isaac Signed-off-by: Mani Kaustubh Mathur --- internal/arrowscan/arrowscan.go | 7 +++++++ internal/arrowscan/arrowscan_test.go | 28 ++++++++++++++++++++++++++++ internal/arrowscan/coltype.go | 8 +++++--- internal/arrowscan/coltype_test.go | 2 ++ 4 files changed, 42 insertions(+), 3 deletions(-) diff --git a/internal/arrowscan/arrowscan.go b/internal/arrowscan/arrowscan.go index 2d70360a..cb24b0a8 100644 --- a/internal/arrowscan/arrowscan.go +++ b/internal/arrowscan/arrowscan.go @@ -193,6 +193,13 @@ func ScanCellCached(col arrow.Array, row int, loc *time.Location, keys *StructKe // shared renderer to reuse, and this stays kernel-side). dt := col.DataType().(*arrow.DurationType) return formatDayTimeInterval(int64(c.Value(row)), dt.Unit), nil + case *array.DayTimeInterval: + // INTERVAL DAY TO SECOND can also arrive as arrow's DayTimeInterval ({Days, + // Milliseconds}) rather than a Duration. Fold it to total milliseconds and reuse + // the same renderer. int32 days × 86_400_000 stays well within int64, so the + // day→ms scaling cannot overflow. + v := c.Value(row) + return formatDayTimeInterval(int64(v.Days)*86_400_000+int64(v.Milliseconds), arrow.Millisecond), nil case *array.MonthInterval: // INTERVAL YEAR TO MONTH arrives as a month count; Thrift's server string is // "years-months". diff --git a/internal/arrowscan/arrowscan_test.go b/internal/arrowscan/arrowscan_test.go index f58740ef..b8d31bc3 100644 --- a/internal/arrowscan/arrowscan_test.go +++ b/internal/arrowscan/arrowscan_test.go @@ -193,6 +193,34 @@ func TestScanCellInterval(t *testing.T) { } }) } + + // DAY-TIME intervals can arrive as arrow's DayTimeInterval ({Days, Milliseconds}) + // instead of a Duration; it must render to the identical "D HH:MM:SS.nnnnnnnnn". + dayTimeInterval := []struct { + name string + v arrow.DayTimeInterval + want string + }{ + {"one_day", arrow.DayTimeInterval{Days: 1, Milliseconds: 0}, "1 00:00:00.000000000"}, + {"day_hms_millis", arrow.DayTimeInterval{Days: 1, Milliseconds: 3661_500}, "1 01:01:01.500000000"}, + {"negative", arrow.DayTimeInterval{Days: -1, Milliseconds: -3661_500}, "-1 01:01:01.500000000"}, + } + for _, tc := range dayTimeInterval { + t.Run("daytimeinterval_"+tc.name, func(t *testing.T) { + b := array.NewDayTimeIntervalBuilder(pool) + defer b.Release() + b.Append(tc.v) + arr := b.NewArray() + defer arr.Release() + v, err := ScanCell(arr, 0, nil) + if err != nil { + t.Fatal(err) + } + if v.(string) != tc.want { + t.Errorf("got %q, want %q", v, tc.want) + } + }) + } } // ScanCell renders DATE / TIMESTAMP in the requested location, matching the diff --git a/internal/arrowscan/coltype.go b/internal/arrowscan/coltype.go index 7cdba4d7..a0a5548e 100644 --- a/internal/arrowscan/coltype.go +++ b/internal/arrowscan/coltype.go @@ -86,10 +86,12 @@ func ColumnTypeInfoFor(dt arrow.DataType) ColumnTypeInfo { return varLen("MAP", scanTypeRawBytes) case arrow.STRUCT: return varLen("STRUCT", scanTypeRawBytes) - case arrow.DURATION: + case arrow.DURATION, arrow.INTERVAL_DAY_TIME: // INTERVAL DAY TO SECOND: Thrift's prod default pre-formats intervals to text - // and declares STRING_TYPE (verified live), so match STRING; the kernel formats - // the native arrow.DURATION Go-side to the identical string. + // and declares STRING_TYPE (verified live), so match STRING; the kernel returns + // it as either arrow.DURATION or arrow.INTERVAL_DAY_TIME, both formatted Go-side + // to the identical string (see ScanCellCached). Without INTERVAL_DAY_TIME here + // its metadata fell through to the default ""/interface{} (comparator diff). return varLen("STRING", scanTypeString) case arrow.INTERVAL_MONTHS: // INTERVAL YEAR TO MONTH — same server-config reasoning as arrow.DURATION. diff --git a/internal/arrowscan/coltype_test.go b/internal/arrowscan/coltype_test.go index 3878f79f..0dbd2f8a 100644 --- a/internal/arrowscan/coltype_test.go +++ b/internal/arrowscan/coltype_test.go @@ -49,6 +49,7 @@ func TestColumnTypeInfoFor(t *testing.T) { {"map", arrow.MapOf(arrow.BinaryTypes.String, arrow.PrimitiveTypes.Int64), "MAP", raw, math.MaxInt64, true}, {"struct", arrow.StructOf(arrow.Field{Name: "x", Type: arrow.PrimitiveTypes.Int64}), "STRUCT", raw, math.MaxInt64, true}, {"duration", &arrow.DurationType{Unit: arrow.Microsecond}, "STRING", str, math.MaxInt64, true}, + {"day_time_interval", arrow.FixedWidthTypes.DayTimeInterval, "STRING", str, math.MaxInt64, true}, {"month_interval", arrow.FixedWidthTypes.MonthInterval, "STRING", str, math.MaxInt64, true}, // VOID/NULL: Thrift stringifies to STRING on the wire (verified live), like intervals. {"null", arrow.Null, "STRING", str, math.MaxInt64, true}, @@ -88,6 +89,7 @@ func TestColumnTypeInfoScanTypeCoversScanner(t *testing.T) { arrow.MapOf(arrow.BinaryTypes.String, arrow.PrimitiveTypes.Int64), arrow.StructOf(arrow.Field{Name: "x", Type: arrow.PrimitiveTypes.Int64}), &arrow.DurationType{Unit: arrow.Microsecond}, + arrow.FixedWidthTypes.DayTimeInterval, arrow.FixedWidthTypes.MonthInterval, } for _, dt := range scannerTypes { From 4fde591295bb98e1c59407a9ca463fd321d3aa14 Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Sat, 25 Jul 2026 15:17:46 +0000 Subject: [PATCH 11/11] fix(kernel): address review findings on the Arrow batch / U2M-scope path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Detach kernelRows.ctx from the caller's QueryContext (context2.WithoutCancel + Close-driven cancel), matching the Thrift path: a deadline that only gated statement submission no longer truncates a large CloudFetch stream. Seed iterationErr on batch-path fetch errors so a failed stream is recorded as a failed statement in telemetry, not a success. Add kernel-tagged coverage for the previously-untested batch API: white-box unit tests (IPC rejection, schema, drained-stream, sticky fetch-error → iterationErr) and a live GetArrowBatches draining e2e with an exact row count. Add a literal-assertion TestGetScopes (AWS/GCP/Azure) so a scope regression — the cause of the kernel U2M access_denied incident — is caught rather than asserted tautologically. Fix stale comments: the U2M doc now reflects that scopes are forwarded, m2m references SetLogLevel (not the nonexistent WithLogLevel), and the m2m success log moved after the error check. Co-authored-by: Isaac Signed-off-by: Mani Kaustubh Mathur --- auth/oauth/m2m/m2m.go | 6 +- auth/oauth/oauth_test.go | 23 +++++ auth/oauth/u2m/authenticator.go | 12 +-- internal/backend/kernel/batchiterator_test.go | 91 +++++++++++++++++++ internal/backend/kernel/rows.go | 30 ++++-- kernel_e2e_test.go | 49 ++++++++++ 6 files changed, 190 insertions(+), 21 deletions(-) create mode 100644 internal/backend/kernel/batchiterator_test.go diff --git a/auth/oauth/m2m/m2m.go b/auth/oauth/m2m/m2m.go index f9c8509b..d08b6b89 100644 --- a/auth/oauth/m2m/m2m.go +++ b/auth/oauth/m2m/m2m.go @@ -72,14 +72,14 @@ func (c *authClient) Authenticate(r *http.Request) error { c.tokenSource = GetTokenSource(config) token, err := c.tokenSource.Token() - // Log via the driver's configurable logger (defaults to Warn) rather than the - // global zerolog logger, so this line honors WithLogLevel like every other. - logger.Debug().Msgf("databricks OAuth token fetched successfully") if err != nil { logger.Err(err).Msg("failed to get token") return err } + // Log via the driver's configurable logger (defaults to Warn) rather than the + // global zerolog logger, so this line honors SetLogLevel like every other. + logger.Debug().Msgf("databricks OAuth token fetched successfully") token.SetAuthHeader(r) return nil diff --git a/auth/oauth/oauth_test.go b/auth/oauth/oauth_test.go index e9e64291..96d71a2f 100644 --- a/auth/oauth/oauth_test.go +++ b/auth/oauth/oauth_test.go @@ -4,6 +4,7 @@ import ( "context" "net/http" "net/http/httptest" + "reflect" "strings" "testing" ) @@ -62,6 +63,28 @@ func TestGetAzureDnsZone(t *testing.T) { } } +// TestGetScopes pins the literal cloud-specific scope sets (not GetScopes's own +// output) so a regression in the AWS/GCP-vs-Azure branching — the source of the +// kernel U2M access_denied incident — is caught rather than asserted tautologically. +func TestGetScopes(t *testing.T) { + cases := []struct { + name string + host string + want []string + }{ + {"AWS", "dbc-1234.cloud.databricks.com", []string{"offline_access", "sql"}}, + {"GCP", "x.gcp.databricks.com", []string{"offline_access", "sql"}}, + {"Azure", "adb-123.azuredatabricks.net", []string{"offline_access", "2ff814a6-3304-4ab8-85cb-cd0e6f879c1d/user_impersonation"}}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + if got := GetScopes(tc.host, nil); !reflect.DeepEqual(got, tc.want) { + t.Fatalf("GetScopes(%q, nil) = %v, want %v", tc.host, got, tc.want) + } + }) + } +} + // TestResolveOIDCIssuer drives resolveOIDCIssuer end-to-end against an httptest // server: the server stands in for the connection host, so a 200 with a given // databricks-config body exercises the real fetch + substitution + validation + diff --git a/auth/oauth/u2m/authenticator.go b/auth/oauth/u2m/authenticator.go index c73e5ace..834265c0 100644 --- a/auth/oauth/u2m/authenticator.go +++ b/auth/oauth/u2m/authenticator.go @@ -84,15 +84,9 @@ type u2mAuthenticator struct { // mode. It structurally satisfies the U2MCredentialsProvider interface the kernel // backend asserts (defined in internal/backend/kernel, off the public API). // -// Only the client id crosses to the kernel — NOT the scopes. The Thrift path -// requests offline_access + sql (AWS/GCP) or offline_access + / -// user_impersonation (Azure) via oauth.GetScopes, while the kernel's U2M flow -// applies its own default set (all-apis + offline_access). Neither backend exposes -// a user-facing U2M-scopes option, so this is a fixed-vs-fixed difference, not a -// dropped caller choice; against the built-in databricks-sql-connector public -// client (which all-apis targets) both authorize successfully. If a U2M-scopes -// option is ever added, forward it via kernel.Auth.Scopes (already wired through -// set_auth_u2m) so the two paths request the same set. +// Only the client id is read here; resolveKernelAuth pairs it with the same +// oauth.GetScopes set the Thrift path requests, so both backends authorize against +// the built-in databricks-sql-connector client identically (not the kernel default). func (c *u2mAuthenticator) U2MClientID() string { return c.clientID } // Auth will start the OAuth Authorization Flow to authenticate the cli client diff --git a/internal/backend/kernel/batchiterator_test.go b/internal/backend/kernel/batchiterator_test.go new file mode 100644 index 00000000..833b60b9 --- /dev/null +++ b/internal/backend/kernel/batchiterator_test.go @@ -0,0 +1,91 @@ +//go:build cgo && databricks_kernel + +package kernel + +import ( + "context" + "errors" + "io" + "testing" + + "github.com/apache/arrow/go/v12/arrow" + dbsqlerr "github.com/databricks/databricks-sql-go/errors" + dbsqlrows "github.com/databricks/databricks-sql-go/internal/rows" + dbrows "github.com/databricks/databricks-sql-go/rows" +) + +// TestGetArrowIPCStreamsRejected pins the discoverable sentinel so callers can +// branch on it: the kernel C ABI exports Arrow C Data, not IPC bytes. +func TestGetArrowIPCStreamsRejected(t *testing.T) { + r := &kernelRows{} + it, err := r.GetArrowIPCStreams(context.Background()) + if it != nil { + t.Errorf("GetArrowIPCStreams = %v iterator, want nil", it) + } + if !errors.Is(err, dbsqlerr.ErrNotSupportedByKernel) { + t.Errorf("GetArrowIPCStreams err = %v, want wrapping ErrNotSupportedByKernel", err) + } +} + +// TestBatchIteratorSchema returns the schema captured at construction, and errors +// (rather than panics) when none was captured. +func TestBatchIteratorSchema(t *testing.T) { + sch := arrow.NewSchema([]arrow.Field{{Name: "c", Type: arrow.PrimitiveTypes.Int64}}, nil) + it := &kernelBatchIterator{r: &kernelRows{schema: sch}} + if got, err := it.Schema(); err != nil || got != sch { + t.Errorf("Schema() = (%v, %v), want (%v, nil)", got, err, sch) + } + itNil := &kernelBatchIterator{r: &kernelRows{}} + if _, err := itNil.Schema(); err == nil { + t.Error("Schema() with no schema = nil error, want a failure") + } +} + +// TestBatchIteratorDrainedStream: once the stream is drained (eof), fill reports +// done, so HasNext is false and Next returns io.EOF without touching the C stream. +func TestBatchIteratorDrainedStream(t *testing.T) { + it := &kernelBatchIterator{r: &kernelRows{eof: true}} + if it.HasNext() { + t.Error("HasNext on a drained stream = true, want false") + } + if rec, err := it.Next(); rec != nil || err != io.EOF { + t.Errorf("Next on a drained stream = (%v, %v), want (nil, io.EOF)", rec, err) + } +} + +// TestBatchIteratorFetchErrorSeedsIterationErr is the regression for the telemetry +// gap: the batch path bypasses Next(), so a fetch error must still seed iterationErr +// or OnClose records the failed stream as a successful statement. A cancelled r.ctx +// makes nextBatch return a non-EOF error at the boundary without a live C stream. +func TestBatchIteratorFetchErrorSeedsIterationErr(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + var gotIter error + r := &kernelRows{ + ctx: ctx, + cancel: cancel, + callbacks: &dbsqlrows.TelemetryCallbacks{ + OnClose: func(_ int64, _ int, iterErr, _ error) { gotIter = iterErr }, + }, + } + it := &kernelBatchIterator{r: r} + + _, err := it.Next() + if err == nil { + t.Fatal("Next with a cancelled ctx = nil error, want the fetch error") + } + // The error is sticky: a second Next re-returns it (it never clears). + if _, err2 := it.Next(); err2 != err { + t.Errorf("second Next = %v, want the same sticky error %v", err2, err) + } + if r.iterationErr == nil { + t.Error("iterationErr not seeded on batch-path fetch error — OnClose would report success") + } + it.Close() + if !errors.Is(gotIter, err) { + t.Errorf("OnClose iterErr = %v, want the fetch error %v", gotIter, err) + } +} + +// compile-time guard: kernelBatchIterator must satisfy the public iterator surface. +var _ dbrows.ArrowBatchIterator = (*kernelBatchIterator)(nil) diff --git a/internal/backend/kernel/rows.go b/internal/backend/kernel/rows.go index c556ea3c..d04a884e 100644 --- a/internal/backend/kernel/rows.go +++ b/internal/backend/kernel/rows.go @@ -25,6 +25,7 @@ import ( "github.com/apache/arrow/go/v12/arrow/cdata" dbsqlerr "github.com/databricks/databricks-sql-go/errors" "github.com/databricks/databricks-sql-go/internal/arrowscan" + context2 "github.com/databricks/databricks-sql-go/internal/compat/context" dbsqlrows "github.com/databricks/databricks-sql-go/internal/rows" dbrows "github.com/databricks/databricks-sql-go/rows" ) @@ -53,7 +54,11 @@ var ( // per-batch network/decode work once; row reads walk the already-imported // arrow.Record with O(1) indexing. type kernelRows struct { + // ctx is detached from the caller's QueryContext (via context2.WithoutCancel) + // so a submit-gating deadline can't truncate a large CloudFetch stream, matching + // the Thrift path (see ES-1934053); cancel makes it abortable from Close. ctx context.Context + cancel context.CancelFunc op *kernelOp stream *C.kernel_result_stream_t callbacks *dbsqlrows.TelemetryCallbacks @@ -86,7 +91,10 @@ func newKernelRows(ctx context.Context, op *kernelOp, stream *C.kernel_result_st // telemetry. Assign it only on the success path so cleanup Close() on a // schema/import failure does not record a falsely successful CLOSE_STATEMENT; the // construction error itself is surfaced to and recorded by the conn execute path. - r := &kernelRows{ctx: ctx, op: op, stream: stream, keyCache: arrowscan.NewStructKeyCache()} + // Detach from the caller's cancellation but keep its values (auth/logging), so a + // deadline that only gated statement submission can't truncate the result stream. + resultsCtx, resultsCancel := context.WithCancel(context2.WithoutCancel(ctx)) + r := &kernelRows{ctx: resultsCtx, cancel: resultsCancel, op: op, stream: stream, keyCache: arrowscan.NewStructKeyCache()} var csch C.struct_ArrowSchema if err := call(func() C.KernelStatusCode { @@ -171,6 +179,10 @@ func (r *kernelRows) Close() error { } r.closed = true closeStart := time.Now() + // Release the detached results context so an aborted Close doesn't leak it. + if r.cancel != nil { + r.cancel() + } if r.cur != nil { r.cur.Release() r.cur = nil @@ -233,14 +245,9 @@ func (r *kernelRows) next(dest []driver.Value) error { // nextBatch pulls the next Arrow batch. A released array (release==NULL) is the // kernel's end-of-stream sentinel. func (r *kernelRows) nextBatch() error { - // 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. + // Stop pulling once Close cancels r.ctx; r.ctx is detached from the caller's + // deadline, so this fires on abort, not on a submit-gating timeout. It does NOT + // interrupt an in-flight C fetch (the kernel exposes no per-download timeout). if r.ctx != nil { if err := r.ctx.Err(); err != nil { return err @@ -336,6 +343,11 @@ func (it *kernelBatchIterator) fill() { it.done = true default: it.err = err + // The batch path bypasses Next(), so seed iterationErr here too; else OnClose + // records the failed stream as a successful statement in telemetry. + if it.r.iterationErr == nil { + it.r.iterationErr = err + } } } diff --git a/kernel_e2e_test.go b/kernel_e2e_test.go index 1abaea10..ee2e53a7 100644 --- a/kernel_e2e_test.go +++ b/kernel_e2e_test.go @@ -8,6 +8,7 @@ import ( "database/sql" "database/sql/driver" "errors" + "io" "os" "sync" "sync/atomic" @@ -16,6 +17,7 @@ import ( "github.com/databricks/databricks-sql-go/driverctx" dbsqlerr "github.com/databricks/databricks-sql-go/errors" + dbsqlrows "github.com/databricks/databricks-sql-go/rows" ) // kernelTestDB opens a kernel-backed *sql.DB from the DATABRICKS_PECOTESTING_* @@ -549,3 +551,50 @@ func TestKernelE2EM2M(t *testing.T) { t.Errorf("SELECT 1 = %d, want 1", got) } } + +// TestKernelE2EGetArrowBatches drains the public Arrow batch API over the kernel +// backend end to end — the path driver_e2e_test only exercises through Thrift — +// asserting an exact row count and that each record is released by the caller. +func TestKernelE2EGetArrowBatches(t *testing.T) { + db := kernelTestDB(t) + defer db.Close() + + conn, err := db.Conn(context.Background()) + if err != nil { + t.Fatalf("conn: %v", err) + } + defer conn.Close() + + const want = 100_000 + var driverRows driver.Rows + if err := conn.Raw(func(d any) error { + var qErr error + driverRows, qErr = d.(driver.QueryerContext).QueryContext(context.Background(), "SELECT id FROM range(0, 100000)", nil) + return qErr + }); err != nil { + t.Fatalf("query: %v", err) + } + defer driverRows.Close() + + batches, err := driverRows.(dbsqlrows.Rows).GetArrowBatches(context.Background()) + if err != nil { + t.Fatalf("GetArrowBatches: %v", err) + } + defer batches.Close() + + var rowCount int64 + for batches.HasNext() { + rec, err := batches.Next() + if err == io.EOF { + break + } + if err != nil { + t.Fatalf("batch Next at row %d: %v", rowCount, err) + } + rowCount += rec.NumRows() + rec.Release() // caller owns each record + } + if rowCount != want { + t.Errorf("GetArrowBatches row count = %d, want %d", rowCount, want) + } +}