Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 27 additions & 1 deletion internal/arrowscan/arrowscan.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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 {
Expand Down
34 changes: 34 additions & 0 deletions internal/arrowscan/arrowscan_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down
77 changes: 22 additions & 55 deletions internal/arrowscan/coltype.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
Expand All @@ -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:
Expand All @@ -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}
Expand All @@ -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)
Expand All @@ -112,34 +87,26 @@ 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.
return ColumnTypeInfo{DatabaseTypeName: "", ScanType: scanTypeUnknown}
}
}

// 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}
}
22 changes: 8 additions & 14 deletions internal/arrowscan/coltype_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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{})
Expand Down Expand Up @@ -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) {
Expand All @@ -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{
Expand Down
21 changes: 21 additions & 0 deletions internal/backend/kernel/affectedrows.go
Original file line number Diff line number Diff line change
@@ -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<i64> 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
}
24 changes: 24 additions & 0 deletions internal/backend/kernel/affectedrows_test.go
Original file line number Diff line number Diff line change
@@ -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)
}
})
}
}
14 changes: 14 additions & 0 deletions internal/backend/kernel/backend.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
4 changes: 4 additions & 0 deletions internal/backend/kernel/config.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
3 changes: 2 additions & 1 deletion internal/backend/kernel/operation.go
Original file line number Diff line number Diff line change
Expand Up @@ -194,7 +194,8 @@ 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))
// 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
// completed statement into a cancellation: for a non-idempotent DML that would
Expand Down
Loading
Loading