diff --git a/KERNEL_REV b/KERNEL_REV index cfbf3024..73ce4cfa 100644 --- a/KERNEL_REV +++ b/KERNEL_REV @@ -1 +1 @@ -9b90406 +53171cd diff --git a/doc.go b/doc.go index a8b9626f..a056084f 100644 --- a/doc.go +++ b/doc.go @@ -211,21 +211,30 @@ time zone) connection options. WithTimeout (a server query timeout the kernel C can't set) and WithRetries used to disable retries (the kernel retries internally) return a clear error at connect time rather than being silently ignored; token- provider, external/static, and federated authenticators are likewise not supported -and rejected loudly. Bound query parameters are not yet supported and return a clear -error at execute time (they arrive per-query, not at connect). None of these is -silently ignored. (Metadata is issued as ordinary SQL — SHOW/DESCRIBE/ -information_schema — and runs on this backend like any other query.) +and rejected loudly. None of these is silently ignored. Bound query parameters +(positional and named) are supported — bound over the C ABI to match the Thrift +path. (Metadata is issued as ordinary SQL — SHOW/DESCRIBE/information_schema — and +runs on this backend like any other query.) WithMaxRows and retry tuning (WithRetries with a positive limit) are accepted but not applied on the kernel path: the kernel manages result fetching and retries internally, below the C ABI, with no user-facing knob. -Two further kernel-backend caveats. The INTERVAL types (year-month / day-time) -listed in the type table below are not yet handled by the kernel scanner and -return a scan error; use the default (Thrift) backend for interval columns. And -the kernel backend does not yet surface a per-statement server query id, so a -QueryIdCallback (see below) fires with "" and no EXECUTE_STATEMENT telemetry is -emitted for kernel queries. +Features that live above the backend seam are inherited unchanged: the +database/sql connection pool (each connection wraps one kernel session), +per-connection telemetry (CREATE_SESSION, EXECUTE_STATEMENT, and DELETE_SESSION +are recorded for kernel connections just like Thrift), and the telemetry +exporter's circuit breaker are all backend-agnostic. Result types are rendered to +match the Thrift backend byte-for-byte: scalars, DECIMAL (exact string), +TIMESTAMP / TIMESTAMP_NTZ (both shifted into the session time zone, as Thrift +does), INTERVAL year-month and day-time, nested Array/Map/Struct and VARIANT (as +JSON), and GEOMETRY (WKT). The per-statement server query id is surfaced on the +success path, so a QueryIdCallback (see below) fires with the real id and +EXECUTE_STATEMENT telemetry carries it. + +One kernel-backend limitation remains on the read path: context cancellation is +honored at result-batch boundaries, not mid-fetch (an in-flight CloudFetch batch +runs to completion before the cancel takes effect). # Programmatically Retrieving Connection and Query Id diff --git a/internal/arrowscan/arrowscan.go b/internal/arrowscan/arrowscan.go index b072d6b0..5cc983aa 100644 --- a/internal/arrowscan/arrowscan.go +++ b/internal/arrowscan/arrowscan.go @@ -37,10 +37,12 @@ import ( // matching the Thrift path — a float64 would lose precision beyond ~17 digits; // see databricks-sql-go#274). Nested types (List/Map/Struct, and VARIANT which // arrives nested) render to a JSON string byte-identical to the Thrift path; -// GEOMETRY arrives as a WKB/WKT string and is handled by the string arm. NULLs -// map to nil. A genuinely unhandled type (e.g. interval/duration) returns an -// error rather than a silently wrong value. loc renders DATE / TIMESTAMP in the -// session time zone (nil = UTC, arrow's ToTime default). +// GEOMETRY arrives as a WKB/WKT string and is handled by the string arm. INTERVAL +// day-time/year-month arrive as native arrow duration/month-interval and format to +// the same string the Thrift path receives pre-formatted from the server. NULLs +// map to nil. A genuinely unhandled type returns an error rather than a silently +// wrong value. loc renders DATE / TIMESTAMP in the session time zone (nil = UTC, +// arrow's ToTime default). func ScanCell(col arrow.Array, row int, loc *time.Location) (driver.Value, error) { return ScanCellCached(col, row, loc, nil) } @@ -148,16 +150,84 @@ func ScanCellCached(col arrow.Array, row int, loc *time.Location, keys *StructKe case *array.Decimal128: dt := col.DataType().(*arrow.Decimal128Type) return decimalfmt.ExactString(c.Value(row), dt.Scale), nil + case *array.Duration: + // INTERVAL DAY TO SECOND arrives as an arrow duration. The kernel returns the + // native arrow value, so we format it Go-side to the same "D HH:MM:SS.nnnnnnnnn" + // string the Thrift path gets pre-formatted from the server (its native-interval + // config is off in prod, so it never scans a duration array — hence there is no + // 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.MonthInterval: + // INTERVAL YEAR TO MONTH arrives as a month count; Thrift's server string is + // "years-months". + return formatYearMonthInterval(int32(c.Value(row))), nil case *array.List, *array.LargeList, *array.FixedSizeList, *array.Map, *array.Struct: // Nested types (and VARIANT, which arrives as a nested value) render to a // JSON string matching the Thrift path. return renderJSONString(col, row, loc, keys) default: - return nil, fmt.Errorf("scanning arrow type %s is not supported "+ - "(intervals are not yet handled)", col.DataType()) + return nil, fmt.Errorf("scanning arrow type %s is not supported", col.DataType()) } } +// formatDayTimeInterval renders an arrow duration (in the given time unit) as the +// Thrift path's "D HH:MM:SS.nnnnnnnnn" — days, then zero-padded hours:minutes:seconds +// with 9 fractional digits, negated with a leading '-'. +func formatDayTimeInterval(v int64, unit arrow.TimeUnit) string { + neg := v < 0 + if neg { + v = -v + } + // Split into whole seconds + a sub-second nanosecond remainder working in the + // native unit. We must NOT scale the full magnitude up to nanoseconds first: + // Spark day-time intervals run up to ~Long.MaxValue microseconds (~292 years), + // so v*1e3 (or *1e6/*1e9) would overflow int64 and silently produce a wrong + // (often negative) string. Deriving seconds by dividing keeps every + // intermediate in range; only the bounded sub-second remainder is scaled up. + var secs, frac int64 + switch unit { + case arrow.Second: + secs = v + case arrow.Millisecond: + secs = v / 1e3 + frac = (v % 1e3) * 1e6 + case arrow.Microsecond: + secs = v / 1e6 + frac = (v % 1e6) * 1e3 + default: // Nanosecond + secs = v / 1e9 + frac = v % 1e9 + } + days := secs / 86400 + rem := secs % 86400 + h := rem / 3600 + rem %= 3600 + m := rem / 60 + s := rem % 60 + sign := "" + if neg { + sign = "-" + } + return fmt.Sprintf("%s%d %02d:%02d:%02d.%09d", sign, days, h, m, s, frac) +} + +// formatYearMonthInterval renders a month count as the Thrift path's "years-months", +// negated with a leading '-'. +func formatYearMonthInterval(months int32) string { + neg := months < 0 + if neg { + months = -months + } + y := months / 12 + mo := months % 12 + sign := "" + if neg { + sign = "-" + } + return fmt.Sprintf("%s%d-%d", sign, y, mo) +} + // 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 089cea73..f8f5f7f3 100644 --- a/internal/arrowscan/arrowscan_test.go +++ b/internal/arrowscan/arrowscan_test.go @@ -100,19 +100,88 @@ func TestScanCellScalars(t *testing.T) { }) t.Run("unsupported_type_errors", func(t *testing.T) { - // A duration (INTERVAL) is not yet handled: must error, not return a - // wrong value. - b := array.NewDurationBuilder(pool, &arrow.DurationType{Unit: arrow.Microsecond}) + // An unhandled arrow type must error, not return a silently wrong value. + // Duration/MonthInterval are now handled (see TestScanCellInterval); use a + // type with no scan arm. + b := array.NewTime32Builder(pool, &arrow.Time32Type{Unit: arrow.Second}) defer b.Release() b.Append(1000) arr := b.NewArray() defer arr.Release() if _, err := ScanCell(arr, 0, nil); err == nil { - t.Error("scanning a Duration should return an unsupported-type error") + t.Error("scanning an unhandled arrow type should return an error") } }) } +// INTERVAL day-time (arrow duration) and year-month (arrow month-interval) arrive +// as native arrow values on the kernel path and must format to the exact string the +// Thrift path receives pre-formatted from the server: "D HH:MM:SS.nnnnnnnnn" and +// "years-months", with negatives signed. (These formatters were validated live +// kernel==Thrift in the PuPr POC; this is the regression guard.) +func TestScanCellInterval(t *testing.T) { + pool := memory.NewGoAllocator() + + dayTime := []struct { + name string + unit arrow.TimeUnit + v int64 + want string + }{ + {"one_day_us", arrow.Microsecond, 86400 * 1_000_000, "1 00:00:00.000000000"}, + {"day_to_sec_us", arrow.Microsecond, 90061_500000, "1 01:01:01.500000000"}, + {"seconds_unit", arrow.Second, 3661, "0 01:01:01.000000000"}, + {"negative_us", arrow.Microsecond, -90061_500000, "-1 01:01:01.500000000"}, + // A large microsecond magnitude (~106.75M days, near Long.MaxValue μs) must + // NOT overflow int64 while scaling to nanoseconds — regression guard for the + // prior multiply-first bug that produced a wrong/negative string here. + {"large_us_no_overflow", arrow.Microsecond, 9223372036854775807, "106751991 04:00:54.775807000"}, + } + for _, tc := range dayTime { + t.Run("daytime_"+tc.name, func(t *testing.T) { + b := array.NewDurationBuilder(pool, &arrow.DurationType{Unit: tc.unit}) + defer b.Release() + b.Append(arrow.Duration(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) + } + }) + } + + yearMonth := []struct { + name string + months int32 + want string + }{ + {"two_years", 24, "2-0"}, + {"year_and_month", 13, "1-1"}, + {"months_only", 5, "0-5"}, + {"negative", -13, "-1-1"}, + } + for _, tc := range yearMonth { + t.Run("yearmonth_"+tc.name, func(t *testing.T) { + b := array.NewMonthIntervalBuilder(pool) + defer b.Release() + b.Append(arrow.MonthInterval(tc.months)) + 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 // Thrift path's .In(location); a nil location leaves the value in UTC. func TestScanCellTimestampLocation(t *testing.T) { diff --git a/internal/backend/kernel/backend.go b/internal/backend/kernel/backend.go index c3aa7502..2b72d063 100644 --- a/internal/backend/kernel/backend.go +++ b/internal/backend/kernel/backend.go @@ -10,7 +10,6 @@ import "C" import ( "context" - "errors" "fmt" "strings" "time" @@ -343,16 +342,9 @@ func (k *KernelBackend) SessionID() string { return k.sessionID } // Execute runs a statement to a terminal state via the blocking execute path. // Per the Backend contract it returns a non-nil Operation even on error so the -// caller can read StatementID / wrap the error / Close uniformly. +// caller can read StatementID / wrap the error / Close uniformly. Bound +// parameters (req.Params) are bound onto the statement in execute (see +// bindParams). func (k *KernelBackend) Execute(ctx context.Context, req backend.ExecRequest) (backend.Operation, error) { - // Bound parameters are not yet wired for the kernel backend. Reject them with - // a clear error rather than silently shipping the query with unbound - // placeholders (which would behave differently than Thrift). Parameters arrive - // per-query, so this is an execute-time error, not a connect-time one. Return a - // non-nil Operation per the Backend contract. - if len(req.Params) > 0 { - return &kernelOp{}, errors.New("databricks: query parameters are not yet supported by the kernel backend; " + - "inline the values or use the default (Thrift) backend") - } return k.execute(ctx, req) } diff --git a/internal/backend/kernel/kernel_test.go b/internal/backend/kernel/kernel_test.go index 8746789a..b4fcbcf3 100644 --- a/internal/backend/kernel/kernel_test.go +++ b/internal/backend/kernel/kernel_test.go @@ -75,14 +75,19 @@ func TestEvictIfSessionFatal(t *testing.T) { // Operation must be non-nil (Backend contract) and its Close must report // closed=false, since no server statement was ever created (a phantom // CLOSE_STATEMENT would otherwise be recorded for it). -func TestExecuteRejectsParams(t *testing.T) { - k := &KernelBackend{} +// When Execute fails before it acquires a statement handle (here: a nil session +// makes new_statement fail), it must still honor the Backend contract — a non-nil, +// handle-less Operation that Closes as a no-op (closed=false, no CLOSE_STATEMENT) +// and reports zero AffectedRows. (Parameter binding is exercised live in +// TestKernelE2EParams; a nil-session unit test can't reach the bind path.) +func TestExecuteHandleLessOpContract(t *testing.T) { + k := &KernelBackend{} // nil session → new_statement fails op, err := k.Execute(context.Background(), backend.ExecRequest{ Query: "SELECT ?", - Params: []backend.Param{{Name: "x"}}, + Params: []backend.Param{{Name: "x", Type: "STRING", Value: strPtr("v")}}, }) if err == nil { - t.Fatal("expected an error for bound parameters, got nil") + t.Fatal("expected an error from Execute on a nil-session backend, got nil") } if op == nil { t.Fatal("Execute must return a non-nil Operation per the Backend contract") @@ -99,6 +104,8 @@ func TestExecuteRejectsParams(t *testing.T) { } } +func strPtr(s string) *string { return &s } + // The cell/nested rendering (ScanCell and the JSON grammar) now lives in the // untagged internal/arrowscan package, where its tests run in the default // CGO_ENABLED=0 build; see arrowscan_test.go. The decimal formatter lives in diff --git a/internal/backend/kernel/operation.go b/internal/backend/kernel/operation.go index 0255fb61..16ef7a1d 100644 --- a/internal/backend/kernel/operation.go +++ b/internal/backend/kernel/operation.go @@ -11,11 +11,14 @@ import "C" import ( "context" "database/sql/driver" + "errors" "fmt" "sync" "time" + dbsqlerr "github.com/databricks/databricks-sql-go/errors" "github.com/databricks/databricks-sql-go/internal/backend" + dbsqlerrint "github.com/databricks/databricks-sql-go/internal/errors" dbsqlrows "github.com/databricks/databricks-sql-go/internal/rows" ) @@ -55,6 +58,17 @@ func (k *KernelBackend) execute(ctx context.Context, req backend.ExecRequest) (b } sql.free() + // Bind parameters. The driver hands us backend.Param{Name, Type, Value}: the + // value is already stringified, Type is the Databricks SQL type name, and a nil + // Value is SQL NULL (Type "VOID"). Each maps 1:1 onto + // kernel_statement_bind_parameter, which builds the SEA wire parameter directly + // (name/ordinal + type + string), matching the Thrift path's toSparkParameters. + if err := bindParams(stmt, req.Params); err != nil { + C.kernel_statement_close(stmt) + k.evictIfSessionFatal(err) + return &kernelOp{}, fmt.Errorf("kernel: bind params: %w", toStatementError(err)) + } + // Detached canceller, obtained before execute so it observes the server // statement id the moment execute publishes it. Non-fatal on failure: proceed // without cancellation rather than failing the query. @@ -141,13 +155,44 @@ func (k *KernelBackend) execute(ctx context.Context, req backend.ExecRequest) (b return op, fmt.Errorf("kernel: execute: %w", toStatementError(execErr)) } op.exec = exec - // Capture the modified-row count now, while exec is live — the operation is - // closed (nulling exec) before AffectedRows is read on the ExecContext path. + // 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)) - klog("Execute OK stmt=%p exec=%p affectedRows=%d", stmt, exec, op.affectedRows) + if qid := C.kernel_executed_statement_query_id(exec); qid != nil { + op.statementID = C.GoString(qid) // deep-copies out of the borrowed C string + } + klog("Execute OK stmt=%p exec=%p affectedRows=%d statementID=%q", stmt, exec, op.affectedRows, op.statementID) return op, nil } +// bindParams binds the driver's backend.Param list onto the statement via the +// kernel's raw-param bind. Each Param is already stringified with its Databricks +// SQL type name; an empty Name is a positional param (ordinal assigned kernel-side +// in push order) and a nil Value is SQL NULL (Type "VOID"). Runs before execute, +// so the params are set on the fresh statement (set_sql clears any prior binds). +func bindParams(stmt *C.kernel_statement_t, params []backend.Param) error { + for i, p := range params { + name := newCStrOrNull(p.Name) // empty Name → NULL → positional + typ := newCStr(p.Type) + val := newCStrOrNull("") + if p.Value != nil { + val = newCStr(*p.Value) // non-nil, possibly empty string, is a real value + } + err := call(func() C.KernelStatusCode { + return C.kernel_statement_bind_parameter(stmt, name.c, typ.c, val.c) + }) + name.free() + typ.free() + val.free() + if err != nil { + return fmt.Errorf("param %d (name=%q type=%q): %w", i, p.Name, p.Type, err) + } + klog("bound param %d name=%q type=%q null=%v", i, p.Name, p.Type, p.Value == nil) + } + return nil +} + // kernelOp implements backend.Operation over a sync executed statement. type kernelOp struct { // backend is the owning connection's backend, held so a session-fatal error on @@ -161,6 +206,11 @@ type kernelOp struct { // cached (not read live from exec) because the caller closes the operation — // which nulls exec — before reading AffectedRows (see conn.ExecContext). affectedRows int64 + // statementID is the server query id, captured at execute time. Cached (not + // read live) because kernel_executed_statement_query_id returns a pointer + // borrowed from the exec handle, valid only while exec is alive — the same + // lifetime discipline as affectedRows. + statementID string // location renders DATE / TIMESTAMP values in the session time zone, matching // the Thrift path; nil means UTC. Carried onto the rows built by Results. location *time.Location @@ -168,14 +218,10 @@ type kernelOp struct { var _ backend.Operation = (*kernelOp)(nil) -// StatementID returns "": the kernel C ABI exposes no success-path statement/query -// id accessor (only the error path carries a query_id, on KernelError). Because -// conn.ExecContext/QueryContext gate per-statement telemetry on -// StatementID() != "", kernel queries currently emit no EXECUTE_STATEMENT metric -// and their query-id log field is empty — there is no session-id fallback on that -// gate. Surfacing an id needs a kernel accessor (e.g. -// kernel_executed_statement_query_id); tracked as a follow-up. -func (o *kernelOp) StatementID() string { return "" } +// StatementID returns the server query id captured at execute time (empty on a +// handle-less op that never executed). A non-empty id ungates EXECUTE_STATEMENT +// telemetry and drives QueryIdCallback, matching the Thrift path. +func (o *kernelOp) StatementID() string { return o.statementID } // AffectedRows is the modified-row count for ExecContext. It returns the value // cached at execute time, so it is correct even after the operation is closed @@ -243,9 +289,21 @@ func (o *kernelOp) close() bool { return didClose } -// ExecutionError wraps cause as the driver's execution error. The kernel error -// already carries the sqlstate (see KernelError), so this returns cause as-is -// (nil when cause is nil), matching the neutral contract. +// ExecutionError wraps cause as the driver's execution error, attaching the +// kernel's sqlstate so kernel query failures surface with the same +// DBExecutionError shape (SqlState(), QueryId(), IsRetryable()) as the Thrift +// path — the parity property consumers rely on via errors.As. The sqlstate is +// dug out of the underlying *KernelError when present; the query id still comes +// from ctx (StatementID() is "" on this backend until the kernel exposes a +// success-path accessor). Returns nil when cause is nil. func (o *kernelOp) ExecutionError(ctx context.Context, cause error) error { - return toStatementError(cause) + if cause == nil { + return nil + } + sqlState := "" + var ke *KernelError + if errors.As(cause, &ke) { + sqlState = ke.SQLState + } + return dbsqlerrint.NewExecutionErrorWithState(ctx, dbsqlerr.ErrQueryExecution, cause, sqlState) } diff --git a/internal/errors/err.go b/internal/errors/err.go index 7d6765d6..7d391845 100644 --- a/internal/errors/err.go +++ b/internal/errors/err.go @@ -185,6 +185,16 @@ func NewExecutionError(ctx context.Context, msg string, err error, opStatusResp return &executionError{databricksError: dbErr, queryId: driverctx.QueryIdFromContext(ctx), sqlState: sqlState} } +// NewExecutionErrorWithState builds an execution error from an already-extracted +// sqlState string, for backends that don't have a Thrift TGetOperationStatusResp +// (e.g. the kernel backend, which carries sqlState in its own KernelError). Same +// shape as NewExecutionError otherwise; queryId still comes from ctx. +func NewExecutionErrorWithState(ctx context.Context, msg string, err error, sqlState string) *executionError { + dbErr := newDatabricksError(ctx, msg, err) + dbErr.errType = "execution error" + return &executionError{databricksError: dbErr, queryId: driverctx.QueryIdFromContext(ctx), sqlState: sqlState} +} + // wraps an error and adds trace if not already present func WrapErr(err error, msg string) error { var st stackTracer diff --git a/internal/errors/err_test.go b/internal/errors/err_test.go index e62e0ff2..ff05445c 100644 --- a/internal/errors/err_test.go +++ b/internal/errors/err_test.go @@ -36,6 +36,23 @@ func TestDbSqlErrors(t *testing.T) { assert.Equal(t, ee, execError) }) + t.Run("NewExecutionErrorWithState carries the sqlstate", func(t *testing.T) { + // The kernel backend has no Thrift TGetOperationStatusResp; it passes an + // already-extracted sqlstate. The value must satisfy DBExecutionError and + // surface that sqlstate — the parity property with NewExecutionError. + cause := errors.New("cause") + execError := NewExecutionErrorWithState(context.TODO(), "exec error", cause, "42P01") + e := errors.Wrap(execError, "is wrapped") + + assert.Equal(t, "is wrapped: databricks: execution error: exec error: cause", e.Error()) + assert.True(t, errors.Is(e, dbsqlerr.ExecutionError)) + assert.True(t, errors.Is(e, cause)) + + var ee dbsqlerr.DBExecutionError + assert.True(t, errors.As(e, &ee)) + assert.Equal(t, "42P01", ee.SqlState()) + }) + t.Run("errors.Is/As works with driver error values", func(t *testing.T) { // Create a driver error and wrap it in a regular error cause := errors.New("cause") diff --git a/internal/rows/arrowbased/arrowscan_parity_test.go b/internal/rows/arrowbased/arrowscan_parity_test.go index 2fbf0260..053d081d 100644 --- a/internal/rows/arrowbased/arrowscan_parity_test.go +++ b/internal/rows/arrowbased/arrowscan_parity_test.go @@ -169,6 +169,25 @@ func TestArrowbasedKernelRenderParity(t *testing.T) { b.FieldBuilder(0).(*array.TimestampBuilder).Append(arrow.Timestamp(ts.UnixMicro())) return b.NewArray() }}, + // VARIANT and GEOMETRY arrive over Arrow as plain STRING columns (verified + // live: a top-level VARIANT is the JSON *string* "{\"a\":1}", GEOMETRY is the + // WKT string "POINT(1 2)"). Nested inside a container they are string leaves, + // so they must render as a quoted, JSON-escaped string on both backends — the + // variant's own JSON is NOT re-parsed, it is escaped as text. + {"variant_string_leaf_in_array", func() arrow.Array { + b := array.NewListBuilder(pool, arrow.BinaryTypes.String) + vb := b.ValueBuilder().(*array.StringBuilder) + b.Append(true) + vb.Append(`{"a":1,"b":[2,3]}`) // a VARIANT element, delivered as a string + return b.NewArray() + }}, + {"geometry_wkt_leaf_in_array", func() arrow.Array { + b := array.NewListBuilder(pool, arrow.BinaryTypes.String) + vb := b.ValueBuilder().(*array.StringBuilder) + b.Append(true) + vb.Append("POINT(1 2)") // GEOMETRY delivered as a WKT string + return b.NewArray() + }}, {"null_leaf_in_struct", func() arrow.Array { dt := arrow.StructOf( arrow.Field{Name: "a", Type: arrow.PrimitiveTypes.Int64}, @@ -281,6 +300,87 @@ func TestArrowbasedKernelTopLevelScalarParity(t *testing.T) { } } +// TestArrowbasedKernelTimestampTZParity pins the TIMESTAMP vs TIMESTAMP_NTZ +// rendering contract: over Arrow, a TIMESTAMP arrives with TimeZone "UTC" and a +// TIMESTAMP_NTZ with an empty TimeZone (kernel json.rs:171 — "TIMESTAMP carries a +// tz on the Arrow side; TIMESTAMP_NTZ does not"). Both backends deliberately IGNORE +// that tz field: each renders via ToTime + .In(loc), so the LTZ-vs-NTZ difference is +// carried entirely by the instant value the server sends, NOT by the client +// inspecting the tz. Verified live on both backends (NY + Kolkata, incl. a DST-gap +// literal): kernel == Thrift byte-for-byte for both types. +// +// This test locks that in so a future "semantically-correct" change that skips +// .In(loc) for TimeZone=="" (which would look right in isolation) fails CI — it +// would make the kernel diverge from the Thrift path, which shifts NTZ too. +func TestArrowbasedKernelTimestampTZParity(t *testing.T) { + pool := memory.NewGoAllocator() + loc, err := time.LoadLocation("America/New_York") + if err != nil { + t.Skipf("tz database unavailable: %v", err) + } + // Same instant for both; only the arrow TimeZone field differs (UTC vs ""). + instant := time.Date(2026, 7, 9, 12, 0, 0, 0, time.UTC) + build := func(tz string) arrow.Array { + b := array.NewTimestampBuilder(pool, &arrow.TimestampType{Unit: arrow.Microsecond, TimeZone: tz}) + b.Append(arrow.Timestamp(instant.UnixMicro())) + return b.NewArray() + } + for _, tc := range []struct { + name string + tz string + }{ + {"timestamp_ltz_utc_zone", "UTC"}, + {"timestamp_ntz_empty_zone", ""}, + } { + t.Run(tc.name, func(t *testing.T) { + arr := build(tc.tz) + defer arr.Release() + kernel, err := arrowscan.ScanCell(arr, 0, loc) + if err != nil { + t.Fatalf("arrowscan.ScanCell: %v", err) + } + thrift := renderViaArrowbased(t, arr, 0, loc) + if kernel != thrift { + t.Errorf("TimeZone=%q divergence:\n kernel = %#v\n thrift = %#v", tc.tz, kernel, thrift) + } + }) + } +} + +// TestArrowbasedKernelVariantGeometryParity documents that VARIANT and GEOMETRY +// need no special rendering: over Arrow both arrive as plain STRING columns (a +// top-level VARIANT is its JSON text "{\"a\":1}", GEOMETRY is its WKT "POINT(1 2)"), +// so the existing string arm on both backends already produces identical output. +// Verified live on both backends. (GEOGRAPHY is intentionally not covered — it is +// not enabled on the benchmark warehouse and no consumer has asked for it.) +func TestArrowbasedKernelVariantGeometryParity(t *testing.T) { + pool := memory.NewGoAllocator() + loc := time.UTC + cases := []struct { + name, val string + }{ + {"variant_object", `{"a":1,"b":[2,3]}`}, + {"variant_scalar", "42"}, + {"geometry_wkt", "POINT(1 2)"}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + b := array.NewStringBuilder(pool) + b.Append(tc.val) + arr := b.NewArray() + defer arr.Release() + kernel, err := arrowscan.ScanCell(arr, 0, loc) + if err != nil { + t.Fatalf("arrowscan.ScanCell: %v", err) + } + thrift := renderViaArrowbased(t, arr, 0, loc) + if kernel != thrift || kernel != tc.val { + t.Errorf("%s divergence:\n kernel = %#v\n thrift = %#v\n want = %#v", tc.name, kernel, thrift, tc.val) + } + }) + } +} + // TestTopLevelDecimalRendering documents the top-level DECIMAL story, which is // subtler than "kernel string vs Thrift float64": // diff --git a/kernel_e2e_test.go b/kernel_e2e_test.go index d5744bc3..e99a089b 100644 --- a/kernel_e2e_test.go +++ b/kernel_e2e_test.go @@ -9,8 +9,13 @@ import ( "database/sql/driver" "errors" "os" + "sync" + "sync/atomic" "testing" "time" + + "github.com/databricks/databricks-sql-go/driverctx" + dbsqlerr "github.com/databricks/databricks-sql-go/errors" ) // kernelTestDB opens a kernel-backed *sql.DB from DATABRICKS_HOST / @@ -110,6 +115,50 @@ func TestKernelE2ETimeZone(t *testing.T) { } } +// TestKernelE2ETimestampNTZ proves TIMESTAMP and TIMESTAMP_NTZ round-trip +// byte-identically to the Thrift path at a non-UTC session tz. The kernel delivers +// TIMESTAMP with an Arrow tz and TIMESTAMP_NTZ without one, but — like the Thrift +// path — the driver ignores that field and renders via .In(loc). So the naive +// wall-clock of an NTZ literal is treated as a UTC instant and shifted into the +// session tz (12:00 NTZ in America/New_York → the -04:00 local of 12:00Z = 08:00), +// exactly as Thrift does. Verified live on both backends; this pins the round-trip +// so a one-sided "don't shift NTZ" change is caught. +func TestKernelE2ETimestampNTZ(t *testing.T) { + const tz = "America/New_York" + db := kernelTestDBWith(t, WithSessionParams(map[string]string{"timezone": tz})) + defer db.Close() + + loc, err := time.LoadLocation(tz) + if err != nil { + t.Skipf("tz database unavailable: %v", err) + } + ctx := context.Background() + + var ltz time.Time + if err := db.QueryRowContext(ctx, + "SELECT CAST('2026-07-09 12:00:00' AS TIMESTAMP)").Scan(<z); err != nil { + t.Fatalf("TIMESTAMP query: %v", err) + } + // A zoned TIMESTAMP literal is the wall-clock in the session tz. + if want := time.Date(2026, 7, 9, 12, 0, 0, 0, loc); !ltz.Equal(want) { + t.Errorf("TIMESTAMP = %s, want %s", ltz, want) + } + + var ntz time.Time + if err := db.QueryRowContext(ctx, + "SELECT CAST('2026-07-09 12:00:00' AS TIMESTAMP_NTZ)").Scan(&ntz); err != nil { + t.Fatalf("TIMESTAMP_NTZ query: %v", err) + } + // The NTZ wall-clock is treated as a UTC instant then shifted into the session + // tz: 12:00Z == 08:00 in America/New_York (-04:00 in July). + if want := time.Date(2026, 7, 9, 8, 0, 0, 0, loc); !ntz.Equal(want) { + t.Errorf("TIMESTAMP_NTZ = %s, want %s (UTC instant %s)", ntz, want, ntz.UTC()) + } + if ntz.Location().String() != tz { + t.Errorf("TIMESTAMP_NTZ location = %q, want %q", ntz.Location(), tz) + } +} + // TestKernelE2ETLSSkipVerify checks that WithSkipTLSHostVerify (a relaxation // knob) is accepted on the kernel path; the connection must still succeed // against the warehouse's valid certificate. @@ -199,6 +248,38 @@ func dataTypeEqual(got, want driver.Value) bool { } } +// TestKernelE2EErrorSurface proves a failed query surfaces as a DBExecutionError +// carrying a SqlState, the same error shape as the Thrift path (consumers rely on +// errors.As(err, &DBExecutionError) + SqlState()). Verified live: unknown table → +// 42P01, unknown column → 42703 — byte-identical sqlstate to Thrift, though the +// kernel's message is richer (it includes the SQL error class + suggestions). +func TestKernelE2EErrorSurface(t *testing.T) { + db := kernelTestDB(t) + defer db.Close() + + cases := []struct { + name, query, wantState string + }{ + {"unknown_table", "SELECT * FROM this_table_does_not_exist_xyz", "42P01"}, + {"unknown_column", "SELECT no_such_col FROM range(1)", "42703"}, + } + for _, c := range cases { + t.Run(c.name, func(t *testing.T) { + _, err := db.QueryContext(context.Background(), c.query) + if err == nil { + t.Fatalf("expected an error for %q, got none", c.query) + } + var de dbsqlerr.DBExecutionError + if !errors.As(err, &de) { + t.Fatalf("error is not a DBExecutionError: %v", err) + } + if de.SqlState() != c.wantState { + t.Errorf("sqlState = %q, want %q (err: %v)", de.SqlState(), c.wantState, err) + } + }) + } +} + // TestKernelE2ECloudFetch streams a CloudFetch-sized result end to end. CloudFetch // is internal to the kernel, so "it works" means many batches stream and scan // correctly — which also exercises the per-batch release/lifetime path. @@ -234,6 +315,68 @@ func TestKernelE2ECloudFetch(t *testing.T) { } } +// TestKernelE2EStatementID proves the kernel backend surfaces the server query id +// on the success path: a registered QueryIdCallback fires with a non-empty id +// (driven by kernelOp.StatementID() → kernel_executed_statement_query_id). This is +// the observable end of the EXECUTE_STATEMENT telemetry / query-history path that +// was previously dark on the kernel backend (StatementID() returned ""). +func TestKernelE2EStatementID(t *testing.T) { + db := kernelTestDB(t) + defer db.Close() + + var gotID string + var fired bool + ctx := driverctx.NewContextWithQueryIdCallback(context.Background(), func(id string) { + fired = true + gotID = id + }) + + var v int64 + if err := db.QueryRowContext(ctx, "SELECT 1").Scan(&v); err != nil { + t.Fatalf("query: %v", err) + } + if !fired { + t.Fatal("QueryIdCallback did not fire") + } + if gotID == "" { + t.Error("kernel backend surfaced an empty query id; want the server statement id") + } +} + +// TestKernelE2EConnectionPool proves the database/sql connection pool works with +// the kernel backend: the pool lives above the backend seam (each connection wraps +// one kernel session), so many concurrent queries over a capped pool must all +// succeed with checkout/return and per-conn single-session isolation intact. +func TestKernelE2EConnectionPool(t *testing.T) { + db := kernelTestDB(t) + defer db.Close() + db.SetMaxOpenConns(8) + db.SetMaxIdleConns(8) + + const n = 40 + var errs, ok int64 + var wg sync.WaitGroup + for i := 0; i < n; i++ { + wg.Add(1) + go func() { + defer wg.Done() + ctx, cancel := context.WithTimeout(context.Background(), 90*time.Second) + defer cancel() + var v int + if err := db.QueryRowContext(ctx, "SELECT 1").Scan(&v); err != nil || v != 1 { + atomic.AddInt64(&errs, 1) + return + } + atomic.AddInt64(&ok, 1) + }() + } + wg.Wait() + if errs != 0 { + t.Fatalf("connection pool: %d/%d queries failed", errs, n) + } + t.Logf("connection pool OK: %d/%d queries succeeded over pool cap 8", ok, n) +} + // TestKernelE2ECancellation cancels a long-running query via ctx and asserts it // returns well before its uncancelled runtime. func TestKernelE2ECancellation(t *testing.T) { diff --git a/kernel_parity_test.go b/kernel_parity_test.go index 8f570302..62d2576a 100644 --- a/kernel_parity_test.go +++ b/kernel_parity_test.go @@ -64,12 +64,67 @@ func TestKernelThriftParity(t *testing.T) { } } +// paramCase is one parameterized-query parity case: the same SQL + args run on +// both backends must yield byte-identical output. +type paramCase struct { + name string + sql string + args []any +} + +// The kernel binds these via kernel_statement_bind_parameter (the driver's +// backend.Param{Name, Type, Value}); Thrift binds via toSparkParameters. Covers +// positional (?) and named (:n) markers, each scalar type, SQL NULL, multi-param, +// and a predicate — mirroring the C2 POC gate. +var paramCases = []paramCase{ + {"pos_int", "SELECT ? AS v", []any{int64(42)}}, + {"pos_string", "SELECT ? AS v", []any{"hello"}}, + {"pos_double", "SELECT ? AS v", []any{3.5}}, + {"pos_bool", "SELECT ? AS v", []any{true}}, + {"pos_null", "SELECT ? AS v", []any{nil}}, + {"pos_two", "SELECT ? + ? AS v", []any{int64(2), int64(40)}}, + {"pos_in_predicate", "SELECT count(*) AS v FROM range(100) WHERE id < ?", []any{int64(10)}}, + {"named_int", "SELECT :n AS v", []any{sql.Named("n", int64(7))}}, + {"named_string", "SELECT :s AS v", []any{sql.Named("s", "world")}}, + {"named_two", "SELECT :a AS a, :b AS b", []any{sql.Named("a", int64(1)), sql.Named("b", "x")}}, +} + +// TestKernelParamsVsThrift asserts parameterized queries produce byte-identical +// output on the kernel and Thrift backends — the bound-parameter acceptance gate. +func TestKernelParamsVsThrift(t *testing.T) { + kernelDB := kernelTestDB(t) + defer kernelDB.Close() + thriftDB := thriftTestDB(t) + defer thriftDB.Close() + + for _, c := range paramCases { + t.Run(c.name, func(t *testing.T) { + kernelRow := scanOneRowAsStringsArgs(t, kernelDB, c.sql, c.args...) + thriftRow := scanOneRowAsStringsArgs(t, thriftDB, c.sql, c.args...) + if len(kernelRow) != len(thriftRow) { + t.Fatalf("column count differs: kernel=%d thrift=%d", len(kernelRow), len(thriftRow)) + } + for i := range kernelRow { + if kernelRow[i] != thriftRow[i] { + t.Errorf("col %d differs: kernel=%q thrift=%q (sql=%q args=%v)", i, kernelRow[i], thriftRow[i], c.sql, c.args) + } + } + }) + } +} + // scanOneRowAsStrings scans the first row into a []string via sql.RawBytes, so a // NULL renders as "" and every value is compared in its wire form, // independent of Go-type coercion differences between the backends. func scanOneRowAsStrings(t *testing.T, db *sql.DB, query string) []string { + return scanOneRowAsStringsArgs(t, db, query) +} + +// scanOneRowAsStringsArgs is scanOneRowAsStrings with query arguments (bound +// parameters). Same wire-form comparison contract. +func scanOneRowAsStringsArgs(t *testing.T, db *sql.DB, query string, args ...any) []string { t.Helper() - rows, err := db.QueryContext(context.Background(), query) + rows, err := db.QueryContext(context.Background(), query, args...) if err != nil { t.Fatalf("query: %v", err) }