diff --git a/.github/workflows/go.yml b/.github/workflows/go.yml index beff4157..98658dcf 100644 --- a/.github/workflows/go.yml +++ b/.github/workflows/go.yml @@ -90,6 +90,10 @@ jobs: fi go get -v -t -d ./... + # This matrix builds the default pure-Go driver (CGO_ENABLED=0), which does + # NOT compile the SEA-via-kernel backend (//go:build cgo && databricks_kernel). + # That opt-in path is exercised by the separate build-and-test-kernel job + # below, which builds the kernel static library and links it with CGO. - name: Test run: make test env: @@ -100,3 +104,149 @@ jobs: - name: Build run: make linux + + build-and-test-kernel: + name: Test (kernel backend) + # Exercises the opt-in SEA-via-kernel backend: builds the Rust kernel static + # lib from the pinned KERNEL_REV and runs the `databricks_kernel`-tagged unit + # tests with CGO. Separate from build-and-test so that job's CGO_ENABLED=0 + # pure-Go invariant is untouched. No warehouse creds here, so the live e2e / + # parity tests self-skip; only the tagged unit tests run. + # + # Bound the blast radius of the source build: it clones an external repo and + # cold-compiles ~200 Rust crates, so a network stall or hung cargo would + # otherwise hold a protected-runner slot up to the 360-min default. A tight + # per-ref concurrency group also collapses redundant heavy builds when the + # branch is pushed repeatedly (each push cancels the prior in-flight run). + timeout-minutes: 30 + concurrency: + group: kernel-build-${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + runs-on: + group: databricks-protected-runner-group + labels: linux-ubuntu-latest + + steps: + - name: Check out code into the Go module directory + uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 + + # Go module proxy (GOPROXY + ~/.netrc) via JFrog OIDC. Also exports + # JFROG_ACCESS_TOKEN, which the cargo step below reuses. + - name: Setup JFrog + uses: ./.github/actions/setup-jfrog + + # The protected runner blocks direct crates.io access (go/hardened-gha), + # so point cargo at the JFrog crates proxy. This repo's setup-jfrog is + # Go-only, so configure cargo inline here reusing its JFROG_ACCESS_TOKEN. + # The token stays in ~/.cargo/credentials.toml (not a CARGO*-prefixed env + # var) so rust-cache keys stay stable across runs. + - name: Configure cargo to use JFrog + shell: bash + run: | + set -euo pipefail + mkdir -p ~/.cargo + cat > ~/.cargo/config.toml << 'EOF' + [source.crates-io] + replace-with = "jfrog" + + [source.jfrog] + registry = "sparse+https://databricks.jfrog.io/artifactory/api/cargo/db-cargo-remote/index/" + + [registries.jfrog] + index = "sparse+https://databricks.jfrog.io/artifactory/api/cargo/db-cargo-remote/index/" + credential-provider = ["cargo:token"] + EOF + cat > ~/.cargo/credentials.toml << EOF + [registries.jfrog] + token = "Bearer ${JFROG_ACCESS_TOKEN}" + EOF + chmod 600 ~/.cargo/credentials.toml + + - name: Set up Go Toolchain + uses: actions/setup-go@924ae3a1cded613372ab5595356fb5720e22ba16 # v6.5.0 + with: + go-version: '1.25.x' + cache: false + + # Install the exact toolchain the kernel .a is built with. rust-toolchain.toml + # at the repo root is what actually governs the cargo build (it's a parent of + # build/kernel-src/, and the kernel repo pins no toolchain of its own), so a + # floating `stable` here would drift the archive under a fixed KERNEL_REV. + # Keep this in lockstep with rust-toolchain.toml's channel. + - name: Set up Rust Toolchain + uses: actions-rust-lang/setup-rust-toolchain@46268bd060767258de96ed93c1251119784f2ab6 # v1.16.1 + with: + toolchain: 1.96.1 + + # The kernel repo (databricks/databricks-sql-kernel) is private, and the + # hardened runner has no ambient git credentials, so kernel-lib.sh's fetch + # fails with "could not read Username". Mint a repo-scoped token from the + # INTEGRATION_TEST_APP GitHub App (installed on the org with kernel read + # access — the same mechanism the ODBC driver uses) and rewrite the kernel + # HTTPS URL to carry it. This is transparent to kernel-lib.sh. Retire once + # the kernel publishes a release artifact (KERNEL_LOCAL_A / download path). + - name: Generate GitHub App token for databricks-sql-kernel + id: kernel-token + uses: actions/create-github-app-token@f8d387b68d61c58ab83c6c016672934102569859 # v3.0.0 + with: + app-id: ${{ secrets.INTEGRATION_TEST_APP_ID }} + private-key: ${{ secrets.INTEGRATION_TEST_PRIVATE_KEY }} + owner: databricks + repositories: databricks-sql-kernel + + - name: Rewrite kernel repo URL to authenticated HTTPS + env: + TOKEN: ${{ steps.kernel-token.outputs.token }} + run: | + git config --global \ + url."https://x-access-token:${TOKEN}@github.com/databricks/".insteadOf \ + "https://github.com/databricks/" + + # Cache the built kernel .a keyed on KERNEL_REV: rebuild only when the pin + # moves. The ~85MB archive dwarfs a rebuild trigger, so keep the key tight. + - name: Cache kernel static lib + uses: actions/cache@55cc8345863c7cc4c66a329aec7e433d2d1c52a9 # v6.1.0 + with: + path: | + internal/backend/kernel/lib + internal/backend/kernel/include + key: ${{ runner.os }}-kernellib-${{ hashFiles('KERNEL_REV', 'rust-toolchain.toml') }} + + # Cache the cargo registry + kernel build tree so a cache miss on the .a + # is still an incremental Rust build, not a cold ~200-crate compile. + - name: Cache cargo + kernel build tree + uses: actions/cache@55cc8345863c7cc4c66a329aec7e433d2d1c52a9 # v6.1.0 + with: + path: | + ~/.cargo/registry + ~/.cargo/git + build/kernel-src/target + key: ${{ runner.os }}-kernel-cargo-${{ hashFiles('KERNEL_REV', 'rust-toolchain.toml') }} + restore-keys: | + ${{ runner.os }}-kernel-cargo- + + # make (for the recipe) and a C compiler (cgo compiles the kernel binding + # stubs and links libdatabricks_sql_kernel.a). cc is preinstalled on the + # protected runner — the kernel repo's own c-abi job relies on the same — + # so these conditional installs are a no-op guard that keeps the job robust + # to a future image change rather than a known gap. + - name: Install build prerequisites (make, C compiler) + run: | + if ! command -v make &> /dev/null ; then + apt-get update && apt-get install -y make + fi + if ! command -v cc &> /dev/null ; then + apt-get update && apt-get install -y build-essential + fi + + # make test-kernel builds the kernel lib (make kernel-lib) if the cache + # missed, then runs `CGO_ENABLED=1 go test -tags databricks_kernel ./...`. + # + # No warehouse creds here — this job runs only the tagged UNIT tests, so the + # live e2e / Thrift-parity tests self-skip (they need DATABRICKS_HOST / + # DATABRICKS_HTTP_PATH / DATABRICKS_TOKEN). This matches how build-and-test + # treats the Thrift path (pure-Go unit tests only); live integration tests for + # both backends run through the separate trigger-integration-tests.yml dispatch + # to databricks-driver-test, not in this workflow. + - name: Build kernel lib + run tagged tests + run: make test-kernel diff --git a/.gitignore b/.gitignore index 063361a8..970720b6 100644 --- a/.gitignore +++ b/.gitignore @@ -66,3 +66,8 @@ _tmp* .vscode/ __debug_bin .DS_Store + +# Kernel (SEA) backend build artifacts — produced by `make kernel-lib`, never committed. +/build/kernel-src/ +/internal/backend/kernel/lib/ +/internal/backend/kernel/include/ diff --git a/KERNEL_REV b/KERNEL_REV new file mode 100644 index 00000000..cfbf3024 --- /dev/null +++ b/KERNEL_REV @@ -0,0 +1 @@ +9b90406 diff --git a/Makefile b/Makefile index ebe22c05..6744c9f9 100644 --- a/Makefile +++ b/Makefile @@ -81,3 +81,46 @@ build: linux darwin ## Build the multi-arch binaries $(PLATFORMS): mkdir -p bin GOOS=$(os) GOARCH=amd64 $(GO) build $(GOBUILD_ARGS) -ldflags '$(LDFLAGS)' -o bin/$(BINARY)-$(os)-amd64 . + +# ── Kernel (SEA) backend ────────────────────────────────────────────────────── +# Opt-in cgo build that links the Rust SQL kernel's C ABI, gated behind the +# `databricks_kernel` build tag. These targets are wholly separate from the +# pure-Go defaults above (which stay CGO_ENABLED=0 and never touch the kernel). +KERNEL_REV = $(shell cat KERNEL_REV) +KERNEL_REPO ?= https://github.com/databricks/databricks-sql-kernel.git +KERNEL_SRC ?= build/kernel-src +# Target OS/arch (honors GOOS/GOARCH overrides) vs the actual host, so the build +# step can reject a source cross-build that would drop a host .a into a foreign +# target dir. Multi-OS support is via native per-OS runners (host == target) or +# staging a prebuilt .a with KERNEL_LOCAL_A — not cross-compiling from source. +KERNEL_GOOS = $(shell go env GOOS) +KERNEL_GOARCH = $(shell go env GOARCH) +KERNEL_GOHOSTOS = $(shell go env GOHOSTOS) +KERNEL_GOHOSTARCH = $(shell go env GOHOSTARCH) +KERNEL_LIB_DIR = internal/backend/kernel/lib/$(KERNEL_GOOS)_$(KERNEL_GOARCH) +KERNEL_INC_DIR = internal/backend/kernel/include +KERNEL_GO = CGO_ENABLED=1 go +KERNEL_TAGS = -tags databricks_kernel + +.PHONY: kernel-lib +kernel-lib: ## Build the pinned kernel static lib + header into the cgo link dir (source build). + KERNEL_REPO="$(KERNEL_REPO)" KERNEL_REV="$(KERNEL_REV)" KERNEL_SRC="$(KERNEL_SRC)" \ + KERNEL_LIB_DIR="$(KERNEL_LIB_DIR)" KERNEL_INC_DIR="$(KERNEL_INC_DIR)" \ + KERNEL_GOOS="$(KERNEL_GOOS)" KERNEL_GOARCH="$(KERNEL_GOARCH)" \ + KERNEL_GOHOSTOS="$(KERNEL_GOHOSTOS)" KERNEL_GOHOSTARCH="$(KERNEL_GOHOSTARCH)" \ + KERNEL_LOCAL_A="$(KERNEL_LOCAL_A)" KERNEL_LOCAL_HEADER="$(KERNEL_LOCAL_HEADER)" \ + ./build/kernel-lib.sh + +.PHONY: build-kernel +build-kernel: kernel-lib ## Build the driver with the kernel backend linked. + $(KERNEL_GO) build $(KERNEL_TAGS) ./... + +.PHONY: test-kernel +test-kernel: kernel-lib ## Run the kernel-tagged unit tests (no warehouse needed). + $(KERNEL_GO) test $(KERNEL_TAGS) ./... + +.PHONY: kernel-lib-download +kernel-lib-download: ## Prod mode (download prebuilt .a): blocked until the kernel publishes release artifacts. + @echo "kernel-lib-download: blocked — the kernel does not yet publish per-platform .a release artifacts." + @echo "Use 'make kernel-lib' (source build) meanwhile. See the distribution design doc." + @false diff --git a/build/kernel-lib.sh b/build/kernel-lib.sh new file mode 100755 index 00000000..ffc17e6f --- /dev/null +++ b/build/kernel-lib.sh @@ -0,0 +1,181 @@ +#!/usr/bin/env bash +# +# kernel-lib.sh — build (or copy in) the Databricks SQL kernel static library +# for the cgo `databricks_kernel` backend. +# +# Mode 1 (source build, the default and the only thing that works today): clone +# the kernel at the pinned KERNEL_REV, `cargo build` a self-contained static +# archive with pure-Rust TLS, and drop the .a + C header where the per-platform +# cgo_.go files link them (${SRCDIR}/lib/_ and ${SRCDIR}/include, +# both .gitignore'd). Kernel releases publish no binary assets yet, so building +# from a pinned commit is the current distribution model (mirrors the ODBC +# driver's Corrosion-from-source build). +# +# Local override: set KERNEL_LOCAL_A (and optionally KERNEL_LOCAL_HEADER) to +# copy an already-built archive instead of cloning + building — the analogue of +# the ODBC driver's KERNEL_LOCAL_PATH. +# +# Invoked by `make kernel-lib`, which supplies the environment below. It can +# also be run directly with the same variables set. +# +# Required env (Makefile provides these): +# KERNEL_REV kernel commit SHA to build (from the repo-root KERNEL_REV) +# KERNEL_REPO git URL of the kernel repo +# KERNEL_SRC working dir for the kernel checkout (gitignored) +# KERNEL_LIB_DIR dest dir for the .a (…/lib/_) +# KERNEL_INC_DIR dest dir for the header (…/include) +# Optional env: +# KERNEL_LOCAL_A path to a prebuilt libdatabricks_sql_kernel.a to copy +# KERNEL_LOCAL_HEADER path to a databricks_kernel.h to copy (defaults to a +# header next to KERNEL_LOCAL_A, else the checkout's) + +set -euo pipefail + +: "${KERNEL_LIB_DIR:?KERNEL_LIB_DIR must be set (run via 'make kernel-lib')}" +: "${KERNEL_INC_DIR:?KERNEL_INC_DIR must be set (run via 'make kernel-lib')}" + +LIB_NAME="libdatabricks_sql_kernel.a" +HEADER_NAME="databricks_kernel.h" + +log() { printf '[kernel-lib] %s\n' "$*" >&2; } + +emit_checksum() { + # A reproducibility breadcrumb: the archive path + its content hash. + local a="$KERNEL_LIB_DIR/$LIB_NAME" + if command -v sha256sum >/dev/null 2>&1; then + log "artifact: $a" + log "sha256: $(sha256sum "$a" | cut -d' ' -f1)" + elif command -v shasum >/dev/null 2>&1; then + log "artifact: $a" + log "sha256: $(shasum -a 256 "$a" | cut -d' ' -f1)" + fi +} + +mkdir -p "$KERNEL_LIB_DIR" "$KERNEL_INC_DIR" + +# ── Local override: copy a prebuilt archive, skip clone + cargo entirely. ────── +if [ -n "${KERNEL_LOCAL_A:-}" ]; then + [ -f "$KERNEL_LOCAL_A" ] || { log "KERNEL_LOCAL_A not found: $KERNEL_LOCAL_A"; exit 1; } + local_header="${KERNEL_LOCAL_HEADER:-}" + if [ -z "$local_header" ]; then + # Prefer a header next to the archive; fall back to the checkout's include/. + if [ -f "$(dirname "$KERNEL_LOCAL_A")/$HEADER_NAME" ]; then + local_header="$(dirname "$KERNEL_LOCAL_A")/$HEADER_NAME" + elif [ -n "${KERNEL_SRC:-}" ] && [ -f "$KERNEL_SRC/include/$HEADER_NAME" ]; then + local_header="$KERNEL_SRC/include/$HEADER_NAME" + fi + fi + [ -n "$local_header" ] && [ -f "$local_header" ] || { + log "header not found; set KERNEL_LOCAL_HEADER to a $HEADER_NAME"; exit 1; } + log "copying prebuilt archive from $KERNEL_LOCAL_A" + cp "$KERNEL_LOCAL_A" "$KERNEL_LIB_DIR/$LIB_NAME" + cp "$local_header" "$KERNEL_INC_DIR/$HEADER_NAME" + emit_checksum + exit 0 +fi + +# ── Mode 1: build from source at the pinned rev. ────────────────────────────── +: "${KERNEL_REV:?KERNEL_REV must be set (run via 'make kernel-lib')}" +: "${KERNEL_REPO:?KERNEL_REPO must be set (run via 'make kernel-lib')}" +: "${KERNEL_SRC:?KERNEL_SRC must be set (run via 'make kernel-lib')}" + +# Cache short-circuit: if a previously built archive + header are present AND the +# rev-stamp beside the archive matches KERNEL_REV, the artifact is already what +# this pin would produce — skip the git fetch + cargo build entirely (and don't +# even require a Rust toolchain). `kernel-lib` is .PHONY and CI re-invokes it on +# every run after restoring the .a cache (keyed on KERNEL_REV), so without this +# guard the build re-ran and overwrote the just-restored archive every time. +# The stamp (not just file existence) is what makes this safe against a local +# KERNEL_REV bump with a stale on-disk .a: a changed pin fails the match and +# forces a rebuild. In CI the .a-cache key already includes KERNEL_REV, so a pin +# change misses the cache and the artifact is simply absent — either way a bump +# rebuilds. +REV_STAMP="$KERNEL_LIB_DIR/.kernel-rev" +if [ -f "$KERNEL_LIB_DIR/$LIB_NAME" ] && [ -f "$KERNEL_INC_DIR/$HEADER_NAME" ] && + [ -f "$REV_STAMP" ] && [ "$(cat "$REV_STAMP")" = "$KERNEL_REV" ]; then + log "cache hit: $LIB_NAME already built for $KERNEL_REV — skipping source build" + emit_checksum + exit 0 +fi + +# Reject a source cross-build. `cargo build` below has no --target, so it emits +# the HOST triple's archive — but KERNEL_LIB_DIR is named for the TARGET +# (GOOS/GOARCH). If they differ, copying the host .a into the target dir would +# silently produce a wrong-arch artifact. Multi-OS is served by native per-OS +# runners (host == target, so this passes) or by staging a prebuilt cross-target +# .a via KERNEL_LOCAL_A (handled above, before this check) — not by cross- +# building the kernel from source, which the distribution design defers to the +# download path. Fail loud rather than mislink. Only enforced when the Makefile +# passes the host vars; a direct script call without them skips the check. +if [ -n "${KERNEL_GOOS:-}" ] && [ -n "${KERNEL_GOHOSTOS:-}" ] && + { [ "$KERNEL_GOOS" != "$KERNEL_GOHOSTOS" ] || [ "${KERNEL_GOARCH:-}" != "${KERNEL_GOHOSTARCH:-}" ]; }; then + log "refusing source cross-build: target ${KERNEL_GOOS}/${KERNEL_GOARCH} != host ${KERNEL_GOHOSTOS}/${KERNEL_GOHOSTARCH}." + log "build on a native ${KERNEL_GOOS} runner, or stage a prebuilt archive with KERNEL_LOCAL_A=." + exit 1 +fi + +command -v cargo >/dev/null 2>&1 || { + log "cargo not found — the source build needs a Rust toolchain." + log "install rustup, or use 'make kernel-lib KERNEL_LOCAL_A=' with a prebuilt .a." + exit 1 +} + +# Make $KERNEL_SRC a git repo pointed at the kernel remote, WITHOUT assuming an +# empty destination. CI caches build/kernel-src/target/ (for incremental Rust +# builds) but not .git, so on a cache hit the dir exists with a target/ subtree +# and no repo — a plain `git clone` would abort ("destination path already +# exists and is not an empty directory"). `git init` is idempotent and works +# whether the dir is absent, empty, or holds a restored target/; the kernel's +# own .gitignore excludes /target, so the later checkout leaves it untouched. +mkdir -p "$KERNEL_SRC" +if [ ! -d "$KERNEL_SRC/.git" ]; then + log "initializing git repo in $KERNEL_SRC (-> $KERNEL_REPO)" + git -C "$KERNEL_SRC" init --quiet +fi +if git -C "$KERNEL_SRC" remote get-url origin >/dev/null 2>&1; then + git -C "$KERNEL_SRC" remote set-url origin "$KERNEL_REPO" +else + git -C "$KERNEL_SRC" remote add origin "$KERNEL_REPO" +fi + +log "fetching + checking out $KERNEL_REV" +git -C "$KERNEL_SRC" fetch --all --tags --quiet || true +# checkout -f: this is a tool-managed, gitignored scratch checkout of a pinned +# third-party repo, so force past any leftover files (e.g. a source tree left +# behind if .git was lost). -f overwrites tracked-path collisions but leaves the +# gitignored target/ alone, so a cache-restored build tree survives. Without -f, +# an untracked-file conflict would masquerade as "commit not reachable" and send +# us down the PR-head fallback for the wrong reason. +# +# KERNEL_REV may be a bare commit that only lives on a PR ref (e.g. #163's head +# before it merges), so try the commit directly, then the PR head ref. +if ! git -C "$KERNEL_SRC" checkout -f --quiet "$KERNEL_REV" 2>/dev/null; then + log "commit not directly reachable; trying PR head refs" + git -C "$KERNEL_SRC" fetch --quiet origin '+refs/pull/*/head:refs/remotes/origin/pr/*' || true + git -C "$KERNEL_SRC" checkout -f --quiet "$KERNEL_REV" || { + log "could not check out $KERNEL_REV — is it pushed/fetchable?"; exit 1; } +fi +log "kernel at $(git -C "$KERNEL_SRC" rev-parse --short HEAD)" + +# --no-default-features --features tls-rustls: pure-Rust TLS (no system OpenSSL) +# keeps the archive self-contained and cross-compile-tractable. The kernel's +# default is tls-native, so the override is required. +# --locked: build against the kernel's committed Cargo.lock rather than +# re-resolving, so a fixed KERNEL_REV yields a fixed dependency graph (paired +# with the pinned rustc in rust-toolchain.toml, the .a is reproducible). Fails +# loud if the lock is stale instead of silently pulling newer deps. +log "cargo build --release --locked --no-default-features --features tls-rustls" +( cd "$KERNEL_SRC" && cargo build --release --locked --no-default-features --features tls-rustls ) + +src_a="$KERNEL_SRC/target/release/$LIB_NAME" +src_h="$KERNEL_SRC/include/$HEADER_NAME" +[ -f "$src_a" ] || { log "expected archive not produced: $src_a"; exit 1; } +[ -f "$src_h" ] || { log "expected header not found: $src_h"; exit 1; } + +cp "$src_a" "$KERNEL_LIB_DIR/$LIB_NAME" +cp "$src_h" "$KERNEL_INC_DIR/$HEADER_NAME" +# Record the rev this archive was built for so a later run can short-circuit +# (see the cache short-circuit above). Written last, only after a successful +# build, so a stamp never claims an artifact that isn't there. +printf '%s\n' "$KERNEL_REV" > "$REV_STAMP" +emit_checksum diff --git a/connector.go b/connector.go index e1701fcb..1dd89264 100644 --- a/connector.go +++ b/connector.go @@ -15,6 +15,7 @@ import ( "github.com/databricks/databricks-sql-go/auth/pat" "github.com/databricks/databricks-sql-go/auth/tokenprovider" "github.com/databricks/databricks-sql-go/driverctx" + "github.com/databricks/databricks-sql-go/internal/backend" "github.com/databricks/databricks-sql-go/internal/backend/thrift" "github.com/databricks/databricks-sql-go/internal/client" "github.com/databricks/databricks-sql-go/internal/config" @@ -32,9 +33,18 @@ type connector struct { func (c *connector) Connect(ctx context.Context) (driver.Conn, error) { defer debuglog.Track(ctx, "connector.Connect", "host=%s", c.cfg.Host)() - // Build the execution backend. Currently always the Thrift backend; a second - // SEA-via-kernel backend will be selected here by a connect-time flag. - be, err := thrift.New(ctx, c.cfg, c.client) + // Build the execution backend. Thrift is the default; the SEA-via-kernel + // backend is selected when UseKernel is set. newKernelBackend is build-tag + // gated: in the default pure-Go build it returns a clear "not linked in" + // error, so the kernel path compiles and links only under -tags + // databricks_kernel + CGO_ENABLED=1. + var be backend.Backend + var err error + if c.cfg.UseKernel { + be, err = newKernelBackend(ctx, c.cfg) + } else { + be, err = thrift.New(ctx, c.cfg, c.client) + } if err != nil { return nil, err } @@ -77,7 +87,14 @@ func (c *connector) Connect(ctx context.Context) (driver.Conn, error) { conn.telemetry.RecordOperation(ctx, conn.id, "", telemetry.OperationTypeCreateSession, sessionLatencyMs, nil) } - log.Info().Msgf("connect: host=%s port=%d httpPath=%s serverProtocolVersion=0x%X", c.cfg.Host, c.cfg.Port, c.cfg.HTTPPath, be.ServerProtocolVersion()) + // ServerProtocolVersion is Thrift-specific (not on the neutral backend + // interface); the kernel backend has no negotiated Thrift protocol, so log it + // only when present. + if tb, ok := be.(*thrift.Backend); ok { + log.Info().Msgf("connect: host=%s port=%d httpPath=%s serverProtocolVersion=0x%X", c.cfg.Host, c.cfg.Port, c.cfg.HTTPPath, tb.ServerProtocolVersion()) + } else { + log.Info().Msgf("connect: host=%s port=%d httpPath=%s backend=kernel", c.cfg.Host, c.cfg.Port, c.cfg.HTTPPath) + } return conn, nil } @@ -301,6 +318,27 @@ func WithHTTPPath(path string) ConnOption { } } +// WithUseKernel selects the SEA-via-kernel backend instead of the default +// Thrift backend. It has effect only in a build compiled with +// `-tags databricks_kernel` and CGO_ENABLED=1; in the default pure-Go build a +// connection made with this option set returns a clear error at connect time +// (the kernel backend is not linked in). +func WithUseKernel(useKernel bool) ConnOption { + return func(c *config.Config) { + c.UseKernel = useKernel + } +} + +// WithWarehouseID sets the bare SQL warehouse id. It has no effect unless +// WithUseKernel(true) is also set: the kernel backend addresses a warehouse by id +// (preferred over the http path when set), while the default Thrift backend +// ignores it entirely and continues to route by http path. +func WithWarehouseID(id string) ConnOption { + return func(c *config.Config) { + c.WarehouseID = id + } +} + // WithMaxRows sets up the max rows fetched per request. Default is 10000 func WithMaxRows(n int) ConnOption { return func(c *config.Config) { diff --git a/doc.go b/doc.go index 6ae9616b..c0d33113 100644 --- a/doc.go +++ b/doc.go @@ -44,6 +44,8 @@ Supported optional connection parameters can be specified in param=value and inc - accessToken: Personal access token. Required if authType set to Pat - clientID: Specifies the client ID to use with OauthM2M - clientSecret: Specifies the client secret to use with OauthM2M + - useKernel: Routes execution through the SEA-via-kernel backend instead of Thrift. Requires a build with -tags databricks_kernel (CGO_ENABLED=1); the default pure-Go build returns a clear error. Default is false. See the kernel-backend section below + - warehouseId: The bare SQL warehouse id, used by the kernel backend (which addresses a warehouse by id) in preference to the http path. Ignored by the Thrift backend Supported optional session parameters can be specified in param=value and include: @@ -90,6 +92,8 @@ Supported functional options include: - WithMaxDownloadThreads ( int). Sets up the max number of concurrent workers for cloud fetch. Default is 10. Optional - WithAuthenticator ( auth.Authenticator). Sets up authentication. Required if neither access token or client credentials are provided. - WithClientCredentials( string, string). Sets up Oauth M2M authentication. + - WithUseKernel( bool). Routes execution through the SEA-via-kernel backend instead of Thrift. Requires a build with -tags databricks_kernel (CGO_ENABLED=1); the default build returns a clear error. Default is false. See the kernel-backend section below. Optional + - WithWarehouseID( string). The bare SQL warehouse id used by the kernel backend in preference to the http path; ignored by the Thrift backend. Optional # Query cancellation and timeout @@ -155,6 +159,79 @@ The result log may look like this: {"level":"debug","connId":"01ed6545-5669-1ec7-8c7e-6d8a1ea0ab16","corrId":"workflow-example","queryId":"01ed6545-57cc-188a-bfc5-d9c0eaf8e189","time":1668558402,"message":"Run Main elapsed time: 1.298712292s"} +# SEA-via-kernel backend (experimental) + +By default the driver uses the Thrift/HiveServer2 backend and is pure Go +(CGO_ENABLED=0, go-gettable, cross-compilable). An experimental second backend +runs statements over the Statement Execution API via the Rust +databricks-sql-kernel, reached through a cgo C ABI. It is opt-in and compiled in +only under a build tag, so the default build is unchanged. + +To use it, build with the databricks_kernel tag and CGO enabled, and select it +per connection with WithUseKernel. The tagged build links the kernel static +library, which is not committed — run `make kernel-lib` first to produce it (and +the C header) under the cgo link path; `make build-kernel` does both steps: + + make kernel-lib # builds the kernel .a + header at the pinned KERNEL_REV + CGO_ENABLED=1 go build -tags databricks_kernel ./... + + connector, _ := dbsql.NewConnector( + dbsql.WithServerHostname(host), + dbsql.WithHTTPPath(httpPath), + dbsql.WithAccessToken(token), + dbsql.WithUseKernel(true), + ) + db := sql.OpenDB(connector) + +In a build without the tag, WithUseKernel(true) returns a clear error at connect +time rather than silently using Thrift. + +To see the kernel's own logs interleaved with the driver's, set DBSQL_KERNEL_DEBUG +to any non-empty value. That single flag turns on the driver's binding step tracer +AND installs the kernel's Rust log subscriber, so both write to the same stderr in +execution order. It is off by default (and must stay off during benchmarks — debug +logging perturbs latency). The kernel's verbosity is then controlled by RUST_LOG, +which the kernel honors directly; filter on the target databricks::sql::kernel +(note the colons — it is the kernel's explicit log target, NOT the crate module +path databricks_sql_kernel, which would match nothing): + + # kernel logs only: + DBSQL_KERNEL_DEBUG=1 RUST_LOG=databricks::sql::kernel=debug ./your_app 2>&1 + # kernel logs plus its HTTP stack (hyper/reqwest): + DBSQL_KERNEL_DEBUG=1 RUST_LOG=debug ./your_app 2>&1 + +The kernel backend currently supports PAT authentication; reading scalar, nested, +and complex-typed results (CloudFetch is handled transparently); context +cancellation during execute (a cancelled ctx fires a real server-side cancel; on +the read path cancellation is honored at result-batch boundaries, not mid-fetch); +and the TLS, proxy, and session-conf (query tags, statement timeout, time zone) +connection options. OAuth (M2M/U2M), initial catalog/schema, +WithEnableMetricViewMetadata, a non-default WithPort, and a non-https protocol +are not yet supported and return a clear error at connect time rather than being +silently ignored (the kernel backend connects over https:443 and has no port or +scheme setter); likewise WithTimeout (a server query timeout the kernel C ABI +can't set) and WithRetries used to disable retries (the kernel retries +internally). Bound query parameters and staging operations +(PUT/GET/REMOVE on a Unity Catalog volume, which need a local file transfer this +backend cannot perform) are likewise not yet supported and return a clear error +at execute time (they are per-statement, not connect-time). None of these is +silently ignored. (Metadata is issued as ordinary SQL — +SHOW/DESCRIBE/information_schema — and runs on this backend like any other +query.) + +WithMaxRows and retry tuning (WithRetries with a positive limit) are accepted but +not applied on the kernel path: the kernel manages result fetching and retries +internally, below the C ABI, with no user-facing knob. + +Two further kernel-backend caveats. The INTERVAL types (year-month / day-time) +listed in the type table below are not yet handled by the kernel scanner and +return a scan error; use the default (Thrift) backend for interval columns. And +the kernel backend does not yet surface a per-statement server query id on the +success path, so a QueryIdCallback (see below) fires with "" and no +EXECUTE_STATEMENT telemetry is emitted for kernel queries. (On a query failure the +server query id IS available: the returned error is a DBExecutionError whose +QueryId() carries it — see the Errors section.) + # Programmatically Retrieving Connection and Query Id Use the driverctx package under driverctx/ctx.go to add callbacks to the query context to receive the connection id and query id. @@ -226,6 +303,17 @@ Each type has a corresponding sentinel value which can be used with errors.Is() RequestError ExecutionError +The kernel backend (WithUseKernel, see above) additionally wraps every rejection +of an option or feature it cannot yet honor with the sentinel ErrNotSupportedByKernel. +A caller can detect this case with errors.Is(err, dbsqlerr.ErrNotSupportedByKernel) +— for example to fall back to the default (Thrift) backend — rather than matching +on the error message text: + + if errors.Is(err, dbsqlerr.ErrNotSupportedByKernel) { + // this option/statement isn't supported on the kernel backend yet; + // retry with the default backend (omit WithUseKernel). + } + Example usage: import ( diff --git a/errors/errors.go b/errors/errors.go index d6a19158..88bce5e2 100644 --- a/errors/errors.go +++ b/errors/errors.go @@ -59,6 +59,20 @@ var ExecutionError error = errors.New("Execution Error") // value to be used with errors.Is() to determine if an error chain contains any databricks error var DatabricksError error = errors.New("Databricks Error") +// value to be used with errors.Is() to determine if an error is a kernel-backend +// "option/feature not supported" rejection. The kernel backend (WithUseKernel) +// rejects options it cannot yet honor rather than silently ignoring them; wrapping +// them all with this sentinel lets a caller detect the case programmatically (e.g. +// to fall back to the default Thrift backend) instead of matching on message text. +var ErrNotSupportedByKernel error = errors.New("not supported by the kernel backend") + +// value to be used with errors.Is() to determine that WithUseKernel(true) was set +// but the binary was built without the kernel backend compiled in (missing the +// `databricks_kernel` build tag / CGO_ENABLED=1). Lets a caller detect the +// build-mismatch case programmatically — the same detection mechanism as +// ErrNotSupportedByKernel — instead of matching on message text. +var ErrKernelNotCompiled error = errors.New("the SEA-via-kernel backend is not compiled into this binary") + // Base interface for driver errors type DBError interface { // Descriptive message describing the error diff --git a/internal/arrowscan/arrowscan.go b/internal/arrowscan/arrowscan.go new file mode 100644 index 00000000..964db594 --- /dev/null +++ b/internal/arrowscan/arrowscan.go @@ -0,0 +1,385 @@ +// Package arrowscan converts Arrow array cells to database/sql driver.Values, +// with nested types (List/Map/Struct, and VARIANT which arrives nested) rendered +// to a JSON string byte-identical to the Thrift arrow path +// (internal/rows/arrowbased). It is pure Go (no cgo), so it is shared by the +// kernel backend and testable in the default CGO_ENABLED=0 build — the tests here +// are the regression guard for the exact rendering rules (native float32, exact +// decimals, time.Time formatting, JSON grammar) both backends must agree on. +// +// Rendering to JSON (not a Go map/slice) is deliberate: it is what the Thrift +// path returns, so a query's result is identical across backends. +// - list → [v0,v1,...] +// - map → {"k0":v0,"k1":v1,...} (keys stringified) +// - struct → {"field0":v0,...} +// - nested NULL → null +// - time.Time → quoted .String() (matches the Thrift marshal() special-case) +// - nested decimal → exact scale-applied JSON number literal (never a lossy +// float64), matching Thrift's marshalScalar → ValueString (#253/#274) +// - float32 → native float32 (not widened to float64), so JSON renders +// 3.14, not 3.140000104904175 +package arrowscan + +import ( + "database/sql/driver" + "encoding/json" + "fmt" + "strings" + "time" + + "github.com/apache/arrow/go/v12/arrow" + "github.com/apache/arrow/go/v12/arrow/array" + "github.com/databricks/databricks-sql-go/internal/decimalfmt" +) + +// ScanCell extracts one cell as a driver.Value. Scalars map to their Go value: +// bool, all int/uint widths, float (native float32/float64), string, binary, +// date, timestamp, and top-level decimal (as an exact fixed-point string, +// matching the Thrift path — a float64 would lose precision beyond ~17 digits; +// see databricks-sql-go#274). Nested types (List/Map/Struct, and VARIANT which +// arrives nested) render to a JSON string byte-identical to the Thrift path; +// GEOMETRY arrives as a WKB/WKT string and is handled by the string arm. NULLs +// map to nil. A genuinely unhandled type (e.g. interval/duration) returns an +// error rather than a silently wrong value. loc renders DATE / TIMESTAMP in the +// session time zone (nil = UTC, arrow's ToTime default). +func ScanCell(col arrow.Array, row int, loc *time.Location) (driver.Value, error) { + return ScanCellCached(col, row, loc, nil) +} + +// StructKeyCache memoizes the JSON-escaped `"name":` prefixes for a struct type, +// so writeStructJSON doesn't re-marshal constant field names on every row. It is +// caller-owned and must be scoped to a single result set (e.g. one driver.Rows) +// and discarded with it — NOT a process-global, which would leak. +// +// The Arrow C Data import allocates a fresh *StructType per batch, so a key is +// only ever hit within the batch that created it: across a multi-batch result the +// map would otherwise accrue one never-evicted entry per batch for the whole Rows +// lifetime. Callers should therefore Reset() the cache at each batch boundary — +// all rows of a batch share one imported Record, so the intra-batch win (escape +// each field name once per batch, not once per row) is fully preserved while the +// map stays bounded to a single batch's struct types. A nil cache is valid: +// rendering just recomputes the keys inline. +type StructKeyCache struct { + m map[*arrow.StructType][]string +} + +// NewStructKeyCache returns a cache ready to pass to ScanCellCached. +func NewStructKeyCache() *StructKeyCache { + return &StructKeyCache{m: make(map[*arrow.StructType][]string)} +} + +// Reset drops all memoized prefixes. Callers scope the cache to one batch by +// calling this when a new batch is imported (see StructKeyCache): the prior +// batch's *StructType keys can never be hit again, so keeping them only grows the +// map. Safe on a nil receiver. +func (c *StructKeyCache) Reset() { + if c == nil { + return + } + for k := range c.m { + delete(c.m, k) + } +} + +// keyPrefixes returns the escaped `"name":` prefix for each field of st, +// memoized. Safe on a nil receiver (computes without caching). +func (c *StructKeyCache) keyPrefixes(st *arrow.StructType) []string { + if c != nil { + if p, ok := c.m[st]; ok { + return p + } + } + fields := st.Fields() + prefixes := make([]string, len(fields)) + for i := range fields { + name, _ := json.Marshal(fields[i].Name) // JSON-escapes the field name + prefixes[i] = string(name) + ":" + } + if c != nil { + c.m[st] = prefixes + } + return prefixes +} + +// ScanCellCached is ScanCell with a caller-owned StructKeyCache (see +// StructKeyCache) so struct field-name keys are escaped once per result set +// rather than once per row. Pass nil for the un-memoized one-shot behavior. +func ScanCellCached(col arrow.Array, row int, loc *time.Location, keys *StructKeyCache) (driver.Value, error) { + if col.IsNull(row) { + return nil, nil + } + switch c := col.(type) { + case *array.Null: + return nil, nil + case *array.Boolean: + return c.Value(row), nil + case *array.Int8: + // Return the native width (int8/16/32), NOT a widened int64, to match the + // Thrift path — its columnValuesTyped[*array.Int8, int8] returns a raw int8, + // so a top-level TINYINT scanned into `any` is int8 on both backends. + // (database/sql's convertAssign still coerces these into a typed *int64.) + // NOTE: the driver.Value spec names only int64 for integers, so both + // backends are technically off-spec here; unifying on int64 across Thrift + + // kernel is a deliberate driver-wide follow-up (needs maintainer sign-off, + // as it changes the prod Thrift path's observable type) — matching Thrift is + // the correct choice until then, so the two backends stay identical. + return c.Value(row), nil + case *array.Int16: + return c.Value(row), nil + case *array.Int32: + return c.Value(row), nil + case *array.Int64: + return c.Value(row), nil + case *array.Uint8: + return int64(c.Value(row)), nil + case *array.Uint16: + return int64(c.Value(row)), nil + case *array.Uint32: + return int64(c.Value(row)), nil + case *array.Uint64: + // Databricks SQL has no unsigned types, so a Uint64 column does not occur + // in practice; this arm is defensive. driver.Value has no uint64 and the + // driver convention is int64 for integers, so a value above MaxInt64 would + // wrap — acceptable for an unreachable path. + return int64(c.Value(row)), nil // #nosec G115 -- see above; unreachable for Databricks types + case *array.Float32: + // Return the native float32, NOT a widened float64: the Thrift path returns + // a float32 driver.Value for a bare FLOAT column, and database/sql's + // asString formats it at bit-size 32 — so widening here would render + // CAST(0.1 AS FLOAT) as "0.10000000149011612" vs Thrift's "0.1". + return c.Value(row), nil + case *array.Float64: + return c.Value(row), nil + case *array.String: + return c.Value(row), nil + case *array.LargeString: + return c.Value(row), nil + case *array.Binary: + return c.Value(row), nil + case *array.Date32: + return inLocation(c.Value(row).ToTime(), loc), nil + case *array.Date64: + return inLocation(c.Value(row).ToTime(), loc), nil + case *array.Timestamp: + dt, ok := col.DataType().(*arrow.TimestampType) + if !ok { + return nil, fmt.Errorf("timestamp column has unexpected datatype %s", col.DataType()) + } + // Honor the schema's declared unit. This is intentionally stricter than the + // Thrift path (arrowRows hardcodes Timestamp_us): Databricks TIMESTAMP is + // always microseconds, so both agree in practice — but if a non-µs TIMESTAMP + // ever arrived, this renders the correct instant while Thrift would misread + // 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 + case *array.Decimal128: + dt := col.DataType().(*arrow.Decimal128Type) + return decimalfmt.ExactString(c.Value(row), dt.Scale), nil + case *array.List, *array.LargeList, *array.FixedSizeList, *array.Map, *array.Struct: + // Nested types (and VARIANT, which arrives as a nested value) render to a + // JSON string matching the Thrift path. + return renderJSONString(col, row, loc, keys) + default: + return nil, fmt.Errorf("scanning arrow type %s is not supported "+ + "(intervals are not yet handled)", col.DataType()) + } +} + +// inLocation renders t in loc, matching the Thrift path's .In(location); a nil +// loc leaves the value in UTC (arrow's ToTime default). +func inLocation(t time.Time, loc *time.Location) time.Time { + if loc == nil { + return t + } + return t.In(loc) +} + +// renderJSONString renders one nested cell as a JSON string (nil for a NULL +// cell). loc is applied to any timestamp/date leaves, matching the top-level +// scan and the Thrift path. +func renderJSONString(col arrow.Array, row int, loc *time.Location, keys *StructKeyCache) (driver.Value, error) { + if col.IsNull(row) { + return nil, nil + } + var b strings.Builder + if err := writeJSON(&b, col, row, loc, keys); err != nil { + return nil, err + } + return b.String(), nil +} + +// writeJSON writes the JSON form of col[row] into b, recursing for nested types. +func writeJSON(b *strings.Builder, col arrow.Array, row int, loc *time.Location, keys *StructKeyCache) error { + if col.IsNull(row) { + b.WriteString("null") + return nil + } + switch c := col.(type) { + case *array.List: + // Use the offset-aware ValueOffsets, NOT Offsets()[row]: arrow-go's + // Offsets() returns the full un-sliced buffer, so a List with a non-zero + // logical offset (a sliced array, or a List field of a struct — Struct.Field + // re-slices preserving data.offset) must index offsets[row+data.offset]. + // ValueOffsets does that; this mirrors the Thrift path's ValueOffsets use. + s, e := c.ValueOffsets(row) + return writeListJSON(b, c.ListValues(), int(s), int(e), loc, keys) + case *array.LargeList: + // arrow-go's LargeList.ValueOffsets (unlike List's) does NOT add data.offset, + // so add it by hand to stay offset-correct. + off := c.Data().Offset() + return writeListJSON(b, c.ListValues(), int(c.Offsets()[row+off]), int(c.Offsets()[row+off+1]), loc, keys) + case *array.FixedSizeList: + n := int(c.DataType().(*arrow.FixedSizeListType).Len()) + base := (row + c.Data().Offset()) * n + return writeListJSON(b, c.ListValues(), base, base+n, loc, keys) + case *array.Map: + return writeMapJSON(b, c, row, loc, keys) + case *array.Struct: + return writeStructJSON(b, c, row, loc, keys) + case *array.Decimal128: + // Emit the exact scale-applied decimal as a raw JSON number literal, not a + // float64 — a float64 would render DECIMAL(5,2) 19.99 as 19.990000000000002 + // and corrupt high-precision values. Matches the Thrift path's marshalScalar + // → ValueString (databricks-sql-go#253/#274). + b.WriteString(decimalfmt.ExactString(c.Value(row), col.DataType().(*arrow.Decimal128Type).Scale)) + return nil + case *array.Float32: + // Marshal the native float32, NOT a widened float64: json.Marshal(float32 + // (3.14)) is "3.14" but json.Marshal(float64(float32(3.14))) is + // "3.140000104904175". The Thrift nested path marshals the native float32, + // so widening here would break byte-parity for ARRAY/MAP/STRUCT<…FLOAT…>. + return writeScalarJSON(b, c.Value(row)) + default: + v, err := scalarForJSON(col, row, loc) + if err != nil { + return err + } + return writeScalarJSON(b, v) + } +} + +func writeListJSON(b *strings.Builder, values arrow.Array, start, end int, loc *time.Location, keys *StructKeyCache) error { + b.WriteByte('[') + for i := start; i < end; i++ { + if i > start { + b.WriteByte(',') + } + if err := writeJSON(b, values, i, loc, keys); err != nil { + return err + } + } + b.WriteByte(']') + return nil +} + +func writeMapJSON(b *strings.Builder, m *array.Map, row int, loc *time.Location, keys *StructKeyCache) error { + // Deferred (tracked): every entry key is written unconditionally in Arrow + // order, so an Arrow-legal MAP with duplicate keys renders non-unique JSON + // object keys (e.g. {"1":"a","1":"b"}). This mirrors the Thrift path exactly + // (parity holds), so it is intentionally NOT changed here — deciding dedup-last + // vs. error is a cross-backend contract change that belongs in its own PR so + // both renderers stay byte-identical. + // + // Map embeds *List, so ValueOffsets is offset-aware (adds data.offset); use it + // rather than Offsets()[row] for the same reason as writeJSON's List case. + s, e := m.ValueOffsets(row) + start, end := int(s), int(e) + mapKeys := m.Keys() + items := m.Items() + b.WriteByte('{') + for i := start; i < end; i++ { + if i > start { + b.WriteByte(',') + } + // JSON object keys must be strings; stringify the key value. + kv, err := scalarForJSON(mapKeys, i, loc) + if err != nil { + return err + } + writeJSONKey(b, kv) + b.WriteByte(':') + if err := writeJSON(b, items, i, loc, keys); err != nil { + return err + } + } + b.WriteByte('}') + return nil +} + +func writeStructJSON(b *strings.Builder, s *array.Struct, row int, loc *time.Location, keys *StructKeyCache) error { + // Field-name keys are constant across rows; keys.keyPrefixes memoizes them per + // result set so this per-row/per-cell path doesn't re-marshal them (a nil cache + // recomputes inline — correct, just not memoized). + prefixes := keys.keyPrefixes(s.DataType().(*arrow.StructType)) + b.WriteByte('{') + for f := 0; f < s.NumField(); f++ { + if f > 0 { + b.WriteByte(',') + } + b.WriteString(prefixes[f]) // pre-escaped `"name":` + if err := writeJSON(b, s.Field(f), row, loc, keys); err != nil { + return err + } + } + b.WriteByte('}') + return nil +} + +// writeJSONKey writes a value as a JSON object key (always a quoted string), +// matching the Thrift path's mapValueContainer: marshal the value the same way a +// leaf is marshaled (marshalScalar → json.Marshal, so a []byte key renders as +// base64 "YWJj", NOT fmt "%v" [97 98 99]), then quote-wrap the result if it is +// not already a JSON string (numbers/bools become "7"/"true" keys). +func writeJSONKey(b *strings.Builder, v any) { + // A time.Time key renders as its quoted .String() (same special-case as + // writeScalarJSON / the Thrift marshal()), not RFC3339. + if t, ok := v.(time.Time); ok { + kb, _ := json.Marshal(t.String()) + b.Write(kb) + return + } + kb, err := json.Marshal(v) + if err != nil { + // Unmarshalable key value — fall back to its stringified form, quoted. + kb, _ = json.Marshal(fmt.Sprintf("%v", v)) + } + if len(kb) > 0 && kb[0] == '"' { + b.Write(kb) // already a JSON string (string / []byte via marshal) + return + } + b.WriteByte('"') + b.Write(kb) + b.WriteByte('"') +} + +// writeScalarJSON writes a scalar leaf, mirroring the Thrift marshal(): a +// time.Time becomes a quoted .String(); everything else uses json.Marshal. +func writeScalarJSON(b *strings.Builder, v any) error { + if v == nil { + b.WriteString("null") + return nil + } + if t, ok := v.(time.Time); ok { + b.WriteByte('"') + b.WriteString(t.String()) + b.WriteByte('"') + return nil + } + vb, err := json.Marshal(v) + if err != nil { + return err + } + b.Write(vb) + return nil +} + +// scalarForJSON returns the Go value used for a nested leaf that is not itself a +// container — today only a map key (values are written directly by writeJSON). +// It reuses ScanCell's scalar arm, so a decimal key renders via the exact-string +// path (writeJSONKey then quotes it), never a lossy float64. +func scalarForJSON(col arrow.Array, row int, loc *time.Location) (any, error) { + if col.IsNull(row) { + return nil, nil + } + return ScanCell(col, row, loc) +} diff --git a/internal/arrowscan/arrowscan_test.go b/internal/arrowscan/arrowscan_test.go new file mode 100644 index 00000000..8ad4824b --- /dev/null +++ b/internal/arrowscan/arrowscan_test.go @@ -0,0 +1,509 @@ +package arrowscan + +import ( + "fmt" + "testing" + "time" + + "github.com/apache/arrow/go/v12/arrow" + "github.com/apache/arrow/go/v12/arrow/array" + "github.com/apache/arrow/go/v12/arrow/decimal128" + "github.com/apache/arrow/go/v12/arrow/memory" +) + +// ScanCell renders the supported scalar types and rejects an unsupported type +// (rather than returning a silently wrong value). +func TestScanCellScalars(t *testing.T) { + pool := memory.NewGoAllocator() + + t.Run("int64", func(t *testing.T) { + b := array.NewInt64Builder(pool) + defer b.Release() + b.Append(42) + arr := b.NewArray() + defer arr.Release() + v, err := ScanCell(arr, 0, nil) + if err != nil { + t.Fatal(err) + } + if v.(int64) != 42 { + t.Errorf("got %v", v) + } + }) + + t.Run("string", func(t *testing.T) { + b := array.NewStringBuilder(pool) + defer b.Release() + b.Append("hi") + arr := b.NewArray() + defer arr.Release() + v, err := ScanCell(arr, 0, nil) + if err != nil { + t.Fatal(err) + } + if v.(string) != "hi" { + t.Errorf("got %v", v) + } + }) + + t.Run("float32_native", func(t *testing.T) { + // A top-level FLOAT column must scan to a native float32, not a widened + // float64 — matching Thrift, so database/sql's asString renders + // CAST(0.1 AS FLOAT) as "0.1", not "0.10000000149011612". + b := array.NewFloat32Builder(pool) + defer b.Release() + b.Append(0.1) + arr := b.NewArray() + defer arr.Release() + v, err := ScanCell(arr, 0, nil) + if err != nil { + t.Fatal(err) + } + got, ok := v.(float32) + if !ok { + t.Fatalf("want float32, got %T", v) + } + if got != float32(0.1) { + t.Errorf("got %v, want 0.1", got) + } + }) + + t.Run("null", func(t *testing.T) { + b := array.NewInt64Builder(pool) + defer b.Release() + b.AppendNull() + arr := b.NewArray() + defer arr.Release() + v, err := ScanCell(arr, 0, nil) + if err != nil { + t.Fatal(err) + } + if v != nil { + t.Errorf("null should scan to nil, got %v", v) + } + }) + + t.Run("decimal_exact_string", func(t *testing.T) { + // 12345 at scale 2 = "123.45", exact (not a float64). + dt := &arrow.Decimal128Type{Precision: 10, Scale: 2} + b := array.NewDecimal128Builder(pool, dt) + defer b.Release() + b.Append(decimal128.FromU64(12345)) + arr := b.NewArray() + defer arr.Release() + v, err := ScanCell(arr, 0, nil) + if err != nil { + t.Fatal(err) + } + if v.(string) != "123.45" { + t.Errorf("got %v, want 123.45", v) + } + }) + + t.Run("unsupported_type_errors", func(t *testing.T) { + // A duration (INTERVAL) is not yet handled: must error, not return a + // wrong value. + b := array.NewDurationBuilder(pool, &arrow.DurationType{Unit: arrow.Microsecond}) + defer b.Release() + b.Append(1000) + arr := b.NewArray() + defer arr.Release() + if _, err := ScanCell(arr, 0, nil); err == nil { + t.Error("scanning a Duration should return an unsupported-type error") + } + }) +} + +// ScanCell renders DATE / TIMESTAMP in the requested location, matching the +// Thrift path's .In(location); a nil location leaves the value in UTC. +func TestScanCellTimestampLocation(t *testing.T) { + pool := memory.NewGoAllocator() + loc, err := time.LoadLocation("America/New_York") + if err != nil { + t.Skipf("tz database unavailable: %v", err) + } + + // 2026-07-09T12:00:00Z as microseconds since epoch. + utcTS := time.Date(2026, time.July, 9, 12, 0, 0, 0, time.UTC) + b := array.NewTimestampBuilder(pool, &arrow.TimestampType{Unit: arrow.Microsecond}) + defer b.Release() + b.Append(arrow.Timestamp(utcTS.UnixMicro())) + arr := b.NewArray() + defer arr.Release() + + t.Run("location applied", func(t *testing.T) { + v, err := ScanCell(arr, 0, loc) + if err != nil { + t.Fatal(err) + } + got := v.(time.Time) + if got.Location() != loc { + t.Errorf("location = %v, want %v", got.Location(), loc) + } + if !got.Equal(utcTS) { + t.Errorf("instant changed: got %v, want %v", got, utcTS) + } + }) + + t.Run("nil location is UTC", func(t *testing.T) { + v, err := ScanCell(arr, 0, nil) + if err != nil { + t.Fatal(err) + } + if v.(time.Time).Location() != time.UTC { + t.Errorf("nil location should render UTC, got %v", v.(time.Time).Location()) + } + }) +} + +// ScanCell honors the TIMESTAMP column's declared arrow unit (not a hardcoded +// microsecond assumption), so the same wall-clock instant is rendered whether the +// unit is s / ms / µs / ns. The cross-backend parity test can't cover this (it +// pins both sides to µs, matching the Databricks wire reality), so the unit +// arithmetic is verified here directly. +func TestScanCellTimestampUnits(t *testing.T) { + pool := memory.NewGoAllocator() + // One fixed instant, expressed in each unit's ticks-since-epoch. + inst := time.Date(2026, time.July, 9, 12, 0, 0, 0, time.UTC) + cases := []struct { + unit arrow.TimeUnit + ticks int64 + }{ + {arrow.Second, inst.Unix()}, + {arrow.Millisecond, inst.UnixMilli()}, + {arrow.Microsecond, inst.UnixMicro()}, + {arrow.Nanosecond, inst.UnixNano()}, + } + for _, tc := range cases { + t.Run(tc.unit.String(), func(t *testing.T) { + b := array.NewTimestampBuilder(pool, &arrow.TimestampType{Unit: tc.unit}) + defer b.Release() + b.Append(arrow.Timestamp(tc.ticks)) + 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(inst) { + t.Errorf("unit %s: got %v, want %v", tc.unit, got.UTC(), inst) + } + }) + } +} + +// ScanCell renders nested types (list/struct/map) to a JSON string matching the +// Thrift path. +func TestScanCellNested(t *testing.T) { + pool := memory.NewGoAllocator() + + t.Run("list", func(t *testing.T) { + b := array.NewListBuilder(pool, arrow.PrimitiveTypes.Int64) + defer b.Release() + vb := b.ValueBuilder().(*array.Int64Builder) + b.Append(true) + vb.Append(1) + vb.Append(2) + vb.Append(3) + arr := b.NewArray() + defer arr.Release() + v, err := ScanCell(arr, 0, nil) + if err != nil { + t.Fatal(err) + } + if v.(string) != "[1,2,3]" { + t.Errorf("got %q, want [1,2,3]", v) + } + }) + + t.Run("struct", func(t *testing.T) { + dt := arrow.StructOf( + arrow.Field{Name: "a", Type: arrow.PrimitiveTypes.Int64}, + arrow.Field{Name: "b", Type: arrow.BinaryTypes.String}, + ) + b := array.NewStructBuilder(pool, dt) + defer b.Release() + b.Append(true) + b.FieldBuilder(0).(*array.Int64Builder).Append(1) + b.FieldBuilder(1).(*array.StringBuilder).Append("x") + arr := b.NewArray() + defer arr.Release() + v, err := ScanCell(arr, 0, nil) + if err != nil { + t.Fatal(err) + } + if v.(string) != `{"a":1,"b":"x"}` { + t.Errorf("got %q, want {\"a\":1,\"b\":\"x\"}", v) + } + }) + + t.Run("map", func(t *testing.T) { + b := array.NewMapBuilder(pool, arrow.BinaryTypes.String, arrow.PrimitiveTypes.Int64, false) + defer b.Release() + kb := b.KeyBuilder().(*array.StringBuilder) + ib := b.ItemBuilder().(*array.Int64Builder) + b.Append(true) + kb.Append("k") + ib.Append(9) + arr := b.NewArray() + defer arr.Release() + v, err := ScanCell(arr, 0, nil) + if err != nil { + t.Fatal(err) + } + if v.(string) != `{"k":9}` { + t.Errorf("got %q, want {\"k\":9}", v) + } + }) + + t.Run("nested_decimal_exact", func(t *testing.T) { + // A decimal inside a struct must render as an exact JSON number, not a + // lossy float64 (19.99, not 19.990000000000002) — matching Thrift's + // marshalScalar. + dt := arrow.StructOf(arrow.Field{Name: "d", Type: &arrow.Decimal128Type{Precision: 5, Scale: 2}}) + b := array.NewStructBuilder(pool, dt) + defer b.Release() + b.Append(true) + b.FieldBuilder(0).(*array.Decimal128Builder).Append(decimal128.FromU64(1999)) + arr := b.NewArray() + defer arr.Release() + v, err := ScanCell(arr, 0, nil) + if err != nil { + t.Fatal(err) + } + if v.(string) != `{"d":19.99}` { + t.Errorf("got %q, want {\"d\":19.99}", v) + } + }) + + t.Run("nested_float32_exact", func(t *testing.T) { + // A float32 inside a struct must marshal as the native float32 (3.14), not + // a widened float64 (3.140000104904175) — matching Thrift's nested path, + // which marshals the native float32. + dt := arrow.StructOf(arrow.Field{Name: "f", Type: arrow.PrimitiveTypes.Float32}) + b := array.NewStructBuilder(pool, dt) + defer b.Release() + b.Append(true) + b.FieldBuilder(0).(*array.Float32Builder).Append(3.14) + arr := b.NewArray() + defer arr.Release() + v, err := ScanCell(arr, 0, nil) + if err != nil { + t.Fatal(err) + } + if v.(string) != `{"f":3.14}` { + t.Errorf("got %q, want {\"f\":3.14}", v) + } + }) + + t.Run("nested_null", func(t *testing.T) { + b := array.NewListBuilder(pool, arrow.PrimitiveTypes.Int64) + defer b.Release() + b.AppendNull() + arr := b.NewArray() + defer arr.Release() + v, err := ScanCell(arr, 0, nil) + if err != nil { + t.Fatal(err) + } + if v != nil { + t.Errorf("null list should scan to nil, got %v", v) + } + }) +} + +// A StructKeyCache must produce byte-identical output to the nil (recompute-inline) +// path, and reused across rows of the same struct type — so the round-trip +// through the cache can never silently diverge from the one-shot rendering. +func TestScanCellCachedMatchesUncached(t *testing.T) { + pool := memory.NewGoAllocator() + dt := arrow.StructOf( + arrow.Field{Name: "a", Type: arrow.PrimitiveTypes.Int64}, + arrow.Field{Name: `q"x`, Type: arrow.BinaryTypes.String}, // needs escaping + ) + b := array.NewStructBuilder(pool, dt) + defer b.Release() + for i := 0; i < 3; i++ { // multiple rows: exercises the memoized second+ hit + b.Append(true) + b.FieldBuilder(0).(*array.Int64Builder).Append(int64(i)) + b.FieldBuilder(1).(*array.StringBuilder).Append("v") + } + arr := b.NewArray() + defer arr.Release() + + cache := NewStructKeyCache() + for row := 0; row < 3; row++ { + uncached, err := ScanCell(arr, row, nil) + if err != nil { + t.Fatal(err) + } + cached, err := ScanCellCached(arr, row, nil, cache) + if err != nil { + t.Fatal(err) + } + if cached != uncached { + t.Errorf("row %d: cached %q != uncached %q", row, cached, uncached) + } + // Pin the escaped output to an independent expected value, not just + // cached==uncached: a wrong-but-consistent escaping rule would satisfy the + // equality check but fail here. The `q"x` field name must be JSON-escaped + // to q\"x in the key. + want := fmt.Sprintf(`{"a":%d,"q\"x":"v"}`, row) + if cached != want { + t.Errorf("row %d: got %q, want %q", row, cached, want) + } + } +} + +// StructKeyCache.Reset drops memoized entries (callers scope the cache to one +// batch this way) and stays correct afterward; Reset on a nil cache is a no-op. +func TestStructKeyCacheReset(t *testing.T) { + pool := memory.NewGoAllocator() + dt := arrow.StructOf(arrow.Field{Name: "a", Type: arrow.PrimitiveTypes.Int64}) + b := array.NewStructBuilder(pool, dt) + defer b.Release() + b.Append(true) + b.FieldBuilder(0).(*array.Int64Builder).Append(7) + arr := b.NewArray() + defer arr.Release() + + cache := NewStructKeyCache() + first, err := ScanCellCached(arr, 0, nil, cache) + if err != nil { + t.Fatal(err) + } + if len(cache.m) == 0 { + t.Fatal("expected the cache to memoize the struct type") + } + cache.Reset() + if len(cache.m) != 0 { + t.Errorf("Reset should empty the cache, got %d entries", len(cache.m)) + } + // Rendering after Reset must still be correct (re-memoizes on demand). + second, err := ScanCellCached(arr, 0, nil, cache) + if err != nil { + t.Fatal(err) + } + if first != second { + t.Errorf("render after Reset diverged: %q != %q", first, second) + } + + var nilCache *StructKeyCache + nilCache.Reset() // must not panic +} + +// LargeList and FixedSizeList are handled by arrowscan but NOT by the Thrift +// arrowbased renderer, so the cross-backend parity test can't reach them; their +// offset arithmetic (hand-rolled, distinct from List's ValueOffsets — arrow-go v12 +// LargeList.ValueOffsets does not add the logical offset, and FixedSizeList has no +// ValueOffsets at all) is verified here directly, including on a sliced array where +// a non-zero logical offset would expose an off-by-offset bug. +func TestScanCellLargeAndFixedSizeList(t *testing.T) { + pool := memory.NewGoAllocator() + + t.Run("large list", func(t *testing.T) { + b := array.NewLargeListBuilder(pool, arrow.PrimitiveTypes.Int64) + defer b.Release() + vb := b.ValueBuilder().(*array.Int64Builder) + for r := 0; r < 3; r++ { // 3 rows so a slice can drop the head + b.Append(true) + vb.Append(int64(r * 10)) + vb.Append(int64(r*10 + 1)) + } + full := b.NewArray() + defer full.Release() + sliced := array.NewSlice(full, 1, 3) // logical offset != 0 + defer sliced.Release() + + // row 0 of the slice is row 1 of full → [10,11] + got, err := ScanCell(sliced, 0, nil) + if err != nil { + t.Fatal(err) + } + if got != "[10,11]" { + t.Errorf("sliced large-list row 0 = %q, want [10,11]", got) + } + }) + + t.Run("fixed-size list", func(t *testing.T) { + b := array.NewFixedSizeListBuilder(pool, 2, arrow.PrimitiveTypes.Int64) + defer b.Release() + vb := b.ValueBuilder().(*array.Int64Builder) + for r := 0; r < 3; r++ { + b.Append(true) + vb.Append(int64(r * 10)) + vb.Append(int64(r*10 + 1)) + } + full := b.NewArray() + defer full.Release() + sliced := array.NewSlice(full, 1, 3) // offset != 0 → base = (row+offset)*n + defer sliced.Release() + + got, err := ScanCell(sliced, 0, nil) + if err != nil { + t.Fatal(err) + } + if got != "[10,11]" { + t.Errorf("sliced fixed-size-list row 0 = %q, want [10,11]", got) + } + }) +} + +// Empty-but-non-null collections must render as [] / {} (not null): the row is +// present, the collection just has zero elements. A nil/absent parent is a +// separate case (covered by the null-leaf tests). +func TestScanCellEmptyCollections(t *testing.T) { + pool := memory.NewGoAllocator() + + t.Run("empty list", func(t *testing.T) { + b := array.NewListBuilder(pool, arrow.PrimitiveTypes.Int64) + defer b.Release() + b.Append(true) // present row, no ValueBuilder appends → empty + arr := b.NewArray() + defer arr.Release() + got, err := ScanCell(arr, 0, nil) + if err != nil { + t.Fatal(err) + } + if got != "[]" { + t.Errorf("empty list = %q, want []", got) + } + }) + + t.Run("empty map", func(t *testing.T) { + b := array.NewMapBuilder(pool, arrow.BinaryTypes.String, arrow.PrimitiveTypes.Int64, false) + defer b.Release() + b.Append(true) // present row, no entries → empty + arr := b.NewArray() + defer arr.Release() + got, err := ScanCell(arr, 0, nil) + if err != nil { + t.Fatal(err) + } + if got != "{}" { + t.Errorf("empty map = %q, want {}", got) + } + }) +} + +// DECIMAL rendering at the len(digits) == scale boundary: a value whose digit +// count equals the scale must render as "0." (leading zero + full +// fractional part), not drop the integer-part zero or misplace the point. +func TestScanCellDecimalScaleBoundary(t *testing.T) { + pool := memory.NewGoAllocator() + // scale 2, value 0.05 → unscaled 5 (1 digit, scale 2): integer part is 0. + dt := &arrow.Decimal128Type{Precision: 5, Scale: 2} + b := array.NewDecimal128Builder(pool, dt) + defer b.Release() + b.Append(decimal128.FromU64(5)) + arr := b.NewArray() + defer arr.Release() + got, err := ScanCell(arr, 0, nil) + if err != nil { + t.Fatal(err) + } + if got != "0.05" { + t.Errorf("decimal 5 @ scale 2 = %q, want 0.05", got) + } +} diff --git a/internal/backend/kernel/backend.go b/internal/backend/kernel/backend.go new file mode 100644 index 00000000..3dad1ad6 --- /dev/null +++ b/internal/backend/kernel/backend.go @@ -0,0 +1,288 @@ +//go:build cgo && databricks_kernel + +package kernel + +/* +#include +#include "databricks_kernel.h" +*/ +import "C" + +import ( + "context" + "errors" + "fmt" + "sync/atomic" + "time" + + dbsqlerr "github.com/databricks/databricks-sql-go/errors" + "github.com/databricks/databricks-sql-go/internal/backend" +) + +// kernelSessionSeq is a process-wide monotonic counter for kernel session ids. +// The C ABI exposes no server session-id accessor, so we mint our own stable, +// collision-free id per OpenSession rather than deriving one from the session +// handle pointer (a freed handle's address can be reused, colliding telemetry / +// log correlation across sequential connections). +var kernelSessionSeq atomic.Uint64 + +// Config is the flat connection config for the kernel backend. The connector +// fills it from the driver's config so the user-facing options are unchanged +// (this mirrors how the kernel's pyo3/napi bindings take flat connection +// params). Zero-valued fields are simply not applied. +type Config struct { + Host string // workspace hostname, no scheme + HTTPPath string // e.g. /sql/1.0/warehouses/abc123 (carries ?o= org routing) + WarehouseID string // bare warehouse id; preferred over HTTPPath when set + Token string // PAT (dapi...) + + // SessionConf carries server-bound session confs verbatim — the same map the + // Thrift backend forwards (STATEMENT_TIMEOUT, QUERY_TAGS, TIMEZONE, …). + SessionConf map[string]string + + // TLSSkipVerify accepts any server cert (maps the driver's + // WithSkipTLSHostVerify / TLSConfig.InsecureSkipVerify). crypto/tls's + // InsecureSkipVerify disables both chain validation and the hostname check, + // so the kernel path relaxes both to match. + TLSSkipVerify bool + + // ProxyURL configures an HTTP proxy, already resolved for this endpoint from + // the same HTTP(S)_PROXY / NO_PROXY environment the Thrift path uses (NO_PROXY + // is applied during resolution). Empty leaves the kernel on a direct + // connection. + ProxyURL string + + // Location is the session time zone used to render DATE / TIMESTAMP values, + // matching the Thrift path which returns them in this location. nil means UTC. + Location *time.Location +} + +// KernelBackend implements backend.Backend over the kernel C ABI. One backend +// backs one conn, which database/sql serializes to a single goroutine at a time, +// so the kernel session inherits single-owner-ship and needs no locks; the only +// concurrency is the per-statement cancel watcher (see operation.go), which +// touches only the kernel's internal inflight-id slot. +type KernelBackend struct { + cfg Config + session *C.kernel_session_t + sessionID string + valid bool +} + +var _ backend.Backend = (*KernelBackend)(nil) + +// New builds a kernel backend without opening the session; the connector calls +// OpenSession immediately after, mirroring the Thrift backend's shape. +func New(cfg Config) *KernelBackend { + return &KernelBackend{cfg: cfg} +} + +// OpenSession builds a session config (warehouse/http-path + PAT), opens the +// session, and captures a per-conn id. Called once by the connector at connect +// time. The config handle is consumed by kernel_session_open on success and +// freed by us on any earlier failure. +func (k *KernelBackend) OpenSession(ctx context.Context) error { + // Fail fast on an already-cancelled context before the blocking kernel_session + // _open (which the C ABI does not let us interrupt mid-call). + // + // Deferred (tracked): this ctx is only checked here, at entry — once inside the + // blocking kernel_session_open there is no way to honor a deadline/cancel that + // fires mid-connect (a slow warehouse cold-start or a connect-time network + // partition blocks until the kernel returns on its own). Same class and same + // root cause as the CloseSession no-deadline note below (the kernel C ABI + // exposes no deadline/cancellation on the session-lifecycle calls); the fix is + // the same kernel-side change (a deadline arg or cancel handle), so it's grouped + // with that follow-up rather than fixed Go-side with a watchdog here. + if err := ctx.Err(); err != nil { + return err + } + initKernelLogging() + klog("OpenSession host=%s httpPath=%s warehouse=%s", k.cfg.Host, k.cfg.HTTPPath, k.cfg.WarehouseID) + + var cfg *C.KernelSessionConfig + if err := call(func() C.KernelStatusCode { return C.kernel_session_config_new(&cfg) }); err != nil { + return fmt.Errorf("kernel: config_new: %w", toConnError(err)) + } + // kernel_session_open consumes the config on EVERY path — success and + // failure alike (it reclaims the box up front). So we free the config + // ourselves only when we bail out BEFORE reaching kernel_session_open (a + // setter error below); once that call is made, ownership has transferred and + // a free here would double-free. + consumed := false + defer func() { + if !consumed { + C.kernel_session_config_free(cfg) + } + }() + + // Warehouse addressing: bare id when provided, else the http path (which also + // carries ?o= org routing for shared hosts). + host := newCStr(k.cfg.Host) + defer host.free() + if k.cfg.WarehouseID != "" { + wh := newCStr(k.cfg.WarehouseID) + defer wh.free() + if err := call(func() C.KernelStatusCode { + return C.kernel_session_config_set_warehouse(cfg, host.c, wh.c) + }); err != nil { + return fmt.Errorf("kernel: set_warehouse: %w", toConnError(err)) + } + } else { + path := newCStr(k.cfg.HTTPPath) + defer path.free() + if err := call(func() C.KernelStatusCode { + return C.kernel_session_config_set_http_path(cfg, host.c, path.c) + }); err != nil { + return fmt.Errorf("kernel: set_http_path: %w", toConnError(err)) + } + } + + tok := newCStr(k.cfg.Token) + defer tok.free() + if err := call(func() C.KernelStatusCode { + return C.kernel_session_config_set_auth_pat(cfg, tok.c) + }); err != nil { + return fmt.Errorf("kernel: set_auth_pat: %w", toConnError(err)) + } + + // TLS: crypto/tls's InsecureSkipVerify accepts any server cert, so relax both + // chain validation and the hostname check — mapping only one would leave the + // kernel path stricter than the Thrift path it mirrors (a self-signed cert + // would still be rejected). + if k.cfg.TLSSkipVerify { + if err := call(func() C.KernelStatusCode { + return C.kernel_session_config_set_tls_allow_self_signed(cfg, C.bool(true)) + }); err != nil { + return fmt.Errorf("kernel: set_tls_allow_self_signed: %w", toConnError(err)) + } + if err := call(func() C.KernelStatusCode { + return C.kernel_session_config_set_tls_skip_hostname_verification(cfg, C.bool(true)) + }); err != nil { + return fmt.Errorf("kernel: set_tls_skip_hostname_verification: %w", toConnError(err)) + } + } + + // Proxy: only when the environment configured one for this endpoint. NO_PROXY + // was already applied during resolution, so no bypass list is needed here; + // any credentials are carried in the URL userinfo (Go's proxy-env convention), + // so username/password are NULL. + if k.cfg.ProxyURL != "" { + url := newCStr(k.cfg.ProxyURL) + defer url.free() + if err := call(func() C.KernelStatusCode { + return C.kernel_session_config_set_proxy(cfg, url.c, nil, nil, nil) + }); err != nil { + return fmt.Errorf("kernel: set_proxy: %w", toConnError(err)) + } + } + + // Session confs (STATEMENT_TIMEOUT, QUERY_TAGS, TIMEZONE, …) — the same map + // the Thrift backend forwards, applied one key at a time. + for key, val := range k.cfg.SessionConf { + ck := newCStr(key) + cv := newCStr(val) + errSet := call(func() C.KernelStatusCode { + return C.kernel_session_config_set_session_conf(cfg, ck.c, cv.c) + }) + ck.free() + cv.free() + if errSet != nil { + return fmt.Errorf("kernel: set_session_conf[%s]: %w", key, toConnError(errSet)) + } + } + + // kernel_session_open takes ownership of cfg here. Its documented C-ABI + // contract (databricks_kernel.h: "CONSUMES config on both success and failure + // — do not use or free config afterwards") is what makes the unconditional + // consumed=true correct: mark it before checking the error so the deferred + // free never double-frees the already-consumed config. This rests on that + // header guarantee, NOT an assumption — if a future kernel revision validated + // args and returned before consuming, this would leak, so the contract must be + // re-verified against the header when KERNEL_REV is bumped. + var sess *C.kernel_session_t + err := call(func() C.KernelStatusCode { return C.kernel_session_open(cfg, &sess) }) + consumed = true + if err != nil { + return fmt.Errorf("kernel: session_open: %w", toConnError(err)) + } + k.session = sess + k.valid = true + // The C ABI exposes no server session-id accessor; mint a process-unique id + // for logging / telemetry correlation. NOT derived from the handle pointer — + // a freed pointer's address can be reused and collide across connections. + k.sessionID = fmt.Sprintf("kernel-%d", kernelSessionSeq.Add(1)) + klog("OpenSession OK session=%s", k.sessionID) + return nil +} + +// CloseSession tears down the server-side session. Best-effort: the kernel's +// close is async (see the C header), so an error is logged, not hard-failed. +// +// Deferred (tracked): this ignores ctx and blocks in the synchronous call() until +// kernel_session_close returns, with no deadline — a stalled kernel-side close +// (e.g. a shutdown-time network partition) can block database/sql pool cleanup. +// A bounded close needs either a kernel_session_close_blocking with a deadline or +// a Go-side watchdog; grouped with the kernel C-ABI follow-ups. Not fixed here to +// keep this PR scoped to the opt-in backend rather than add a watchdog. +func (k *KernelBackend) CloseSession(ctx context.Context) error { + if k.session == nil { + return nil + } + klog("CloseSession session=%s", k.sessionID) + err := call(func() C.KernelStatusCode { return C.kernel_session_close(k.session) }) + k.session = nil + k.valid = false + return toConnError(err) +} + +// SessionValid backs conn.IsValid → pool eviction. No I/O; inspects state +// captured at OpenSession and updated by markSessionDead. +func (k *KernelBackend) SessionValid() bool { return k.valid && k.session != nil } + +// markSessionDead marks the session unusable so SessionValid() → conn.IsValid() +// returns false and database/sql discards this conn on return to the pool. Called +// from the statement/read path when a kernel error is session-fatal (isSessionFatal): +// it evicts the dead conn WITHOUT returning driver.ErrBadConn, so the statement is +// never transparently re-run (no duplicate write). The backend +// is single-owner per conn (only the cancel watcher shares state, and it touches +// only the canceller's inflight slot), so this write is race-free. +func (k *KernelBackend) markSessionDead() { k.valid = false } + +// evictIfSessionFatal marks the session dead when err is (or wraps) a +// session-fatal KernelError, so a conn whose server session died mid-life is +// evicted from the pool. A no-op for nil, non-KernelError, or non-fatal errors. +// Uses errors.As, not a bare type assertion, so it still fires if a caller wraps +// the KernelError before passing it here. +func (k *KernelBackend) evictIfSessionFatal(err error) { + var ke *KernelError + if errors.As(err, &ke) && isSessionFatal(ke.Code) { + k.markSessionDead() + } +} + +// SessionID is the per-conn id (conn.id). Valid after OpenSession. +func (k *KernelBackend) SessionID() string { return k.sessionID } + +// Execute runs a statement to a terminal state via the blocking execute path. +// Per the Backend contract it returns a non-nil Operation even on error so the +// caller can read StatementID / wrap the error / Close uniformly. +func (k *KernelBackend) Execute(ctx context.Context, req backend.ExecRequest) (backend.Operation, error) { + // Bound parameters are not yet wired for the kernel backend. Reject them with + // a clear error rather than silently shipping the query with unbound + // placeholders (which would behave differently than Thrift). Parameters arrive + // per-query, so this is an execute-time error, not a connect-time one. Return a + // non-nil Operation per the Backend contract. + if len(req.Params) > 0 { + return &kernelOp{}, fmt.Errorf("databricks: query parameters are %w; "+ + "inline the values or use the default (Thrift) backend", dbsqlerr.ErrNotSupportedByKernel) + } + // Staging (Unity Catalog volume PUT/GET/REMOVE) needs a local file transfer the + // kernel path can't perform, and the kernel C ABI surfaces no IsStagingOperation + // signal to drive conn.execStagingOperation. Reject it here rather than let + // IsStaging return false and report success with no file moved (silent data loss). + if isStagingStatement(req.Query) { + return &kernelOp{}, fmt.Errorf("databricks: staging operations (PUT/GET/REMOVE on a volume) are %w; "+ + "use the default (Thrift) backend", dbsqlerr.ErrNotSupportedByKernel) + } + return k.execute(ctx, req) +} diff --git a/internal/backend/kernel/cgo.go b/internal/backend/kernel/cgo.go new file mode 100644 index 00000000..b3f905c3 --- /dev/null +++ b/internal/backend/kernel/cgo.go @@ -0,0 +1,212 @@ +//go:build cgo && databricks_kernel + +// Package kernel implements backend.Backend over the Databricks SQL kernel's C +// ABI (databricks-sql-kernel, src/c_abi). It is build-tag-gated behind +// `databricks_kernel` so the default driver build stays pure-Go (CGO_ENABLED=0, +// go-gettable, cross-compilable); only a build that opts in with +// `-tags databricks_kernel` and CGO_ENABLED=1 links the kernel static lib. +// +// This file holds the cgo plumbing: the C ABI include/link directives, the +// fallible-call helper that makes the kernel's thread-local last error readable, +// the error mapping to the driver's error surface, and the gated step logger. +// The backend, operation, and rows layers live in sibling files. +// +// Link contract (${SRCDIR}-relative, machine-independent). The header is +// included from ${SRCDIR}/include and the static lib is linked from +// ${SRCDIR}/lib/_; the per-platform link flags live in the +// cgo_.go files beside this one. Both directories are produced by the +// build step (`make kernel-lib`), which checks out the kernel at the commit +// pinned in the repo-root KERNEL_REV file and `cargo build`s a static lib — +// so the kernel revision is a reviewable pin, never baked into a #cgo line +// (those expand only ${SRCDIR} and cannot run git or read env). The dirs are +// .gitignore'd; nothing kernel-built is committed. For local development +// against an existing checkout, `make kernel-lib KERNEL_LOCAL_A=.a +// KERNEL_LOCAL_HEADER=databricks_kernel.h` copies those in instead of +// building. The eventual release path downloads a published .a at the pinned +// rev rather than building it (see the driver's distribution design). +package kernel + +/* +#cgo CFLAGS: -I${SRCDIR}/include +#include +#include "databricks_kernel.h" +*/ +import "C" + +import ( + "fmt" + "os" + "runtime" + "sync" + "unsafe" + + "github.com/databricks/databricks-sql-go/logger" +) + +// ─── Debug logging ─────────────────────────────────────────────────────────── +// +// Gated on DBSQL_KERNEL_DEBUG so it is OFF by default and, in particular, OFF +// during benchmarks (debug logging perturbs latency). Every binding step logs +// through klog when enabled — entry, status codes, handle addresses, batch +// counts — which is what makes a failing e2e cheap to diagnose. The same flag +// also installs the kernel's own Rust (tracing) logs via initKernelLogging, so +// binding and kernel logs interleave on stderr as one stream. +var kdebug = os.Getenv("DBSQL_KERNEL_DEBUG") != "" + +// Deferred (tracked): klog writes raw to stderr, NOT through the driver's +// logger.Logger, so its lines carry no connId/corrId/queryId and can't be +// correlated in a multi-conn process; it's also gated on its own +// DBSQL_KERNEL_DEBUG rather than the driver's DATABRICKS_LOG_LEVEL knob. Unifying +// kernel logging onto logger.Logger (which no-ops below its level, so no +// benchmark cost) is a tracked logging-unification follow-up, kept separate from +// this PR because it changes the debug-logging surface. +func klog(format string, args ...any) { + if !kdebug { + return + } + fmt.Fprintf(os.Stderr, "[kernel] "+format+"\n", args...) +} + +// KernelDebugEnabled reports whether binding-level debug logging is on (i.e. +// DBSQL_KERNEL_DEBUG is set). Exposed so a benchmark can assert the flag is off +// before measuring — debug logging perturbs latency; there is no such benchmark +// in-tree yet, so this currently has no callers. +func KernelDebugEnabled() bool { return kdebug } + +// initLoggingOnce guards kernel_init_logging, which is process-wide and +// first-call-wins in the kernel. We install the kernel subscriber lazily on the +// first session open rather than in init(), so a process that never opens a +// kernel session installs nothing. +var initLoggingOnce sync.Once + +// initKernelLogging turns on the kernel's own Rust (tracing) logs, gated on the +// same DBSQL_KERNEL_DEBUG flag as klog so both are OFF by default and, in +// particular, OFF during benchmarks — the subscriber is never installed in a +// benchmark run. level=NULL lets the kernel honor RUST_LOG (default warn); +// file_path=NULL sends kernel logs to stderr so they interleave with klog on one +// stream. Best-effort: Internal (e.g. the host already installed a global +// subscriber) is a documented, benign outcome — logged, never fatal to connect. +// +// Scope caveat: the kernel subscriber is PROCESS-WIDE, first-call-wins, and never +// uninstalled. Because kdebug is read once from DBSQL_KERNEL_DEBUG at package +// init, the switch is process-global, not per-connection: in a long-lived +// multi-tenant process, the first kernel session opened with the flag set +// installs the subscriber (and its stderr output) for ALL subsequent kernel +// sessions in that process, and there is no way to scope it to one connection or +// turn it off afterward. The "off during benchmarks" guarantee therefore depends +// on DBSQL_KERNEL_DEBUG being unset before package init. Making this a +// per-process runtime knob (rather than an init-time env read) is tracked with +// the logging-unification follow-up. +func initKernelLogging() { + if !kdebug { + return + } + initLoggingOnce.Do(func() { + if err := call(func() C.KernelStatusCode { + return C.kernel_init_logging(nil, nil) + }); err != nil { + klog("kernel_init_logging: %v (kernel logs unavailable; proceeding)", err) + } + }) +} + +// call runs a fallible kernel entry point and, on a non-Success status, reads +// the kernel's thread-local last error into a Go error. +// +// The kernel reports rich errors via a thread-local buffer read by a *second* +// call (kernel_get_last_error). Go's M:N scheduler can move a goroutine to a +// different OS thread between two cgo calls, so a naive call-then-read pair can +// observe the wrong thread's buffer. LockOSThread pins the goroutine to its OS +// thread across the call and its error read, closing that window. +func call(fn func() C.KernelStatusCode) error { + runtime.LockOSThread() + defer runtime.UnlockOSThread() + st := fn() + if st == C.KernelStatusCode_Success { + return nil + } + return lastError(st) +} + +// fireCancel dispatches a server-side cancel and reports whether the RPC was +// actually sent. The kernel returns dispatched=false while no server statement +// id has been observed yet (the window before the execute POST returns), and +// true once the cancel RPC goes out — which is the watcher's cue to stop +// re-firing. +func fireCancel(canceller *C.kernel_statement_canceller_t) bool { + var dispatched C.bool + C.kernel_statement_canceller_cancel(canceller, &dispatched) + return bool(dispatched) +} + +// lastError reads the kernel's thread-local last error and copies its string +// fields out immediately — the C `char*` fields are valid only until the next +// FFI call on this thread. Must run on the same OS thread as the failing call; +// call guarantees that via LockOSThread. +func lastError(code C.KernelStatusCode) *KernelError { + var e C.KernelError + if !bool(C.kernel_get_last_error(&e)) { + return &KernelError{Code: int(code), Message: fmt.Sprintf("kernel status %d (no detail)", int(code))} + } + ke := &KernelError{ + Code: int(e.code), + Message: C.GoString(e.message), + VendorCode: int32(e.vendor_code), + HTTPStatus: uint16(e.http_status), + Retryable: bool(e.retryable), + } + if e.sql_state != nil { + ke.SQLState = C.GoString(e.sql_state) + } + if e.query_id != nil { + ke.QueryID = C.GoString(e.query_id) + } + klog("kernel error: code=%d sqlstate=%q vendor=%d http=%d retryable=%v msg=%q", + ke.Code, ke.SQLState, ke.VendorCode, ke.HTTPStatus, ke.Retryable, ke.Message) + // Also emit through the driver's logger — no SQL text or PII, just the + // status/sqlstate/http fields plus the server query id (a correlation token, + // not PII) so on-call can pivot to server-side query history — so a kernel-path + // failure is visible without DBSQL_KERNEL_DEBUG. This is the error path only + // (never the hot per-row/per-batch path), so it does not perturb benchmarks. + // User/query faults (bad SQL, bad argument) are routine and log at Debug so they + // don't inflate the WARN rate on-call alerts key on; infra codes stay at Warn. + msg := "databricks: kernel call failed: code=%d sqlstate=%q vendor=%d http=%d retryable=%v queryId=%q" + if isUserFault(ke.Code) { + logger.Logger.Debug().Msgf(msg, ke.Code, ke.SQLState, ke.VendorCode, ke.HTTPStatus, ke.Retryable, ke.QueryID) + } else { + logger.Logger.Warn().Msgf(msg, ke.Code, ke.SQLState, ke.VendorCode, ke.HTTPStatus, ke.Retryable, ke.QueryID) + } + return ke +} + +// The plain-int status constants used by the untagged classifier logic +// (errors_classify.go) must stay in lockstep with the C enum in +// databricks_kernel.h. These compile-time assertions make a drift a build error +// under -tags databricks_kernel. Each converts BOTH directions of the difference +// to uint: if the Go constant and the C value disagree, one of a-b / b-a is a +// negative constant, and `uint()` is a hard compile error +// ("constant overflows uint") — so the file won't compile whether the header +// renumbers a code up OR down. (A one-sided `[a-b]struct{}` array size only +// catches C > Go: a downward renumber makes a-b positive, a legal array, and the +// drift slips through.) +const ( + _ = uint(statusInvalidArgument-int(C.KernelStatusCode_InvalidArgument)) | uint(int(C.KernelStatusCode_InvalidArgument)-statusInvalidArgument) + _ = uint(statusUnauthenticated-int(C.KernelStatusCode_Unauthenticated)) | uint(int(C.KernelStatusCode_Unauthenticated)-statusUnauthenticated) + _ = uint(statusUnavailable-int(C.KernelStatusCode_Unavailable)) | uint(int(C.KernelStatusCode_Unavailable)-statusUnavailable) + _ = uint(statusTimeout-int(C.KernelStatusCode_Timeout)) | uint(int(C.KernelStatusCode_Timeout)-statusTimeout) + _ = uint(statusNetworkError-int(C.KernelStatusCode_NetworkError)) | uint(int(C.KernelStatusCode_NetworkError)-statusNetworkError) + _ = uint(statusSqlError-int(C.KernelStatusCode_SqlError)) | uint(int(C.KernelStatusCode_SqlError)-statusSqlError) +) + +// cStr wraps C.CString with a guaranteed free. The kernel copies strings into +// owned Rust memory on receipt, so freeing immediately after the call is safe. +// Use: cs := newCStr(s); defer cs.free(); ...C.fn(cs.c)... +type cStr struct{ c *C.char } + +func newCStr(s string) cStr { return cStr{c: C.CString(s)} } + +func (s cStr) free() { + if s.c != nil { + C.free(unsafe.Pointer(s.c)) + } +} diff --git a/internal/backend/kernel/cgo_darwin.go b/internal/backend/kernel/cgo_darwin.go new file mode 100644 index 00000000..bc59d19e --- /dev/null +++ b/internal/backend/kernel/cgo_darwin.go @@ -0,0 +1,20 @@ +//go:build cgo && databricks_kernel && darwin && arm64 + +package kernel + +// Link flags for darwin/arm64. NOTE: this platform is not yet exercised in CI +// (M0 is linux/amd64); the flags below are the intended shape but must be +// validated on a mac before darwin is enabled. +// +// Two darwin-specific differences from linux: +// - Apple's ld64 does NOT accept the GNU `-l:.a` extension, so the +// archive is passed as a positional input by absolute ${SRCDIR} path +// instead. Since only the .a is placed under lib/darwin_arm64 (see +// kernel-lib.sh), there is no .so to accidentally prefer. +// - -lc++ (not -lstdc++) is the macOS C++ runtime; @loader_path keeps any +// dynamic reference resolvable relative to the built binary. + +/* +#cgo LDFLAGS: ${SRCDIR}/lib/darwin_arm64/libdatabricks_sql_kernel.a -lc++ -lm -Wl,-rpath,@loader_path +*/ +import "C" diff --git a/internal/backend/kernel/cgo_linux.go b/internal/backend/kernel/cgo_linux.go new file mode 100644 index 00000000..62274a34 --- /dev/null +++ b/internal/backend/kernel/cgo_linux.go @@ -0,0 +1,17 @@ +//go:build cgo && databricks_kernel && linux && amd64 + +package kernel + +// Link flags for linux/amd64. The static archive is forced with the +// -l:.a form (a GNU-ld extension) so the linker never prefers a +// same-named .so — the kernel's cargo build emits both a .a and a .so into the +// same dir, and a bare -ldatabricks_sql_kernel would pick the .so and bake in +// an rpath. -lstdc++/-lm/-ldl are the kernel's transitive system deps. +// +// The path is ${SRCDIR}-relative; `make kernel-lib` drops the archive at +// ${SRCDIR}/lib/linux_amd64/libdatabricks_sql_kernel.a. + +/* +#cgo LDFLAGS: -L${SRCDIR}/lib/linux_amd64 -l:libdatabricks_sql_kernel.a -lstdc++ -lm -ldl +*/ +import "C" diff --git a/internal/backend/kernel/cgo_unsupported.go b/internal/backend/kernel/cgo_unsupported.go new file mode 100644 index 00000000..bcc96910 --- /dev/null +++ b/internal/backend/kernel/cgo_unsupported.go @@ -0,0 +1,19 @@ +//go:build cgo && databricks_kernel && !(linux && amd64) && !(darwin && arm64) && !(windows && amd64) + +package kernel + +// This file is compiled only on GOOS/GOARCH combinations the kernel backend does +// not support. Per-platform link flags (cgo_.go) exist only for linux/amd64, +// darwin/arm64, and windows/amd64; on any other target there is no static archive +// to link, so cgo.go's C ABI calls would otherwise fail at the LINK step with an +// opaque "undefined reference to kernel_*". The Makefile's host==target guard +// does not catch this — it happily source-builds a host .a on e.g. linux/arm64 +// (Graviton) or an Intel Mac, only for the link to fall over with no matching +// LDFLAGS file. Referencing an undefined identifier here fails earlier, at +// COMPILE time, with a message that names the supported targets — a legible build +// error instead of a linker dump. +// +// Broader OS/arch coverage is tracked in the distribution design (native per-OS +// runners or a staged prebuilt .a); until then this guard makes the supported +// boundary explicit rather than latent. +const _ = kernel_backend_supports_only_linux_amd64_darwin_arm64_and_windows_amd64 diff --git a/internal/backend/kernel/cgo_windows.go b/internal/backend/kernel/cgo_windows.go new file mode 100644 index 00000000..ae650b9b --- /dev/null +++ b/internal/backend/kernel/cgo_windows.go @@ -0,0 +1,20 @@ +//go:build cgo && databricks_kernel && windows && amd64 + +package kernel + +// Link flags for windows/amd64. NOTE: this platform is not yet exercised in CI +// (M0 is linux/amd64); the flags below are the intended shape but must be +// validated on windows before it is enabled. +// +// cgo on windows uses the mingw/gcc toolchain, which links GNU archives (.a) — +// NOT the MSVC .lib that `cargo build --target x86_64-pc-windows-msvc` emits. +// So the kernel must be built for the windows-gnu target +// (`--target x86_64-pc-windows-gnu`) to produce a mingw-compatible .a; that +// target selection is the build step's responsibility. -lws2_32/-lwsock32 are +// the kernel's Winsock deps and -lrstrtmgr is Restart Manager (pulled in by the +// Rust std/dep graph on windows). + +/* +#cgo LDFLAGS: -L${SRCDIR}/lib/windows_amd64 -l:libdatabricks_sql_kernel.a -lws2_32 -lwsock32 -lrstrtmgr +*/ +import "C" diff --git a/internal/backend/kernel/errors_classify.go b/internal/backend/kernel/errors_classify.go new file mode 100644 index 00000000..d9a7e0c5 --- /dev/null +++ b/internal/backend/kernel/errors_classify.go @@ -0,0 +1,149 @@ +package kernel + +import ( + "fmt" + + dbsqlerrint "github.com/databricks/databricks-sql-go/internal/errors" +) + +// This file is intentionally NOT behind the `cgo && databricks_kernel` build tag. +// It holds the kernel error type + the connection-classification logic that +// enforces two safety-critical guarantees — a statement-path error must never +// become driver.ErrBadConn, so database/sql cannot silently re-run a possibly- +// committed statement, and a permanent auth failure is not retried. Keeping +// it pure Go lets its tests run in the default CGO_ENABLED=0 build; a future edit +// that reintroduced the duplicate-write bug (or the auth-retry storm) would then +// fail CI instead of shipping green. The cgo file (cgo.go) populates KernelError +// from the C struct and asserts these status constants stay in lockstep with the +// C enum at compile time (see the compile-time assertion block in cgo.go). + +// Status codes mirrored from the kernel C enum (KernelStatusCode in +// databricks_kernel.h) as plain Go ints, so this non-cgo file can classify errors +// without the C import. cgo.go asserts each equals its C.KernelStatusCode_* value +// at compile time, so drift from the header is a build error, not a latent bug. +const ( + statusInvalidArgument = 1 + statusUnauthenticated = 2 + statusUnavailable = 6 + statusTimeout = 7 + statusNetworkError = 12 + statusSqlError = 13 +) + +// KernelError is the Go-side structured error mapped from the kernel's KernelError +// struct. It carries the sqlstate so the backend's ExecutionError can attach it, +// matching the Thrift error surface. +type KernelError struct { + Code int + Message string + SQLState string + VendorCode int32 + HTTPStatus uint16 + Retryable bool + QueryID string +} + +func (e *KernelError) Error() string { + // Append the server query id when present — it is the one correlation handle + // to server-side query history, and StatementID() is "" on this backend. + q := "" + if e.QueryID != "" { + q = fmt.Sprintf(", queryId=%s", e.QueryID) + } + if e.SQLState != "" { + return fmt.Sprintf("kernel: %s (sqlstate=%s, code=%d%s)", e.Message, e.SQLState, e.Code, q) + } + return fmt.Sprintf("kernel: %s (code=%d%s)", e.Message, e.Code, q) +} + +// isBadConnection reports whether a status code is a *transient* connection +// failure — one where retrying on a fresh connection could succeed. On the +// session-lifecycle path this is wrapped as driver.ErrBadConn so database/sql +// retries connect. Unauthenticated is deliberately excluded: a wrong/expired PAT +// is permanent, so retrying it just burns connect attempts (and can worsen +// server-side auth rate-limiting) and fails identically — matching Thrift, which +// only treats an invalid session handle (a liveness signal), not a 401, as +// bad-conn. (Auth failure still marks the session dead for pool eviction — see +// isSessionFatal.) +func isBadConnection(code int) bool { + switch code { + case statusUnavailable, statusNetworkError: + return true + default: + return false + } +} + +// isSessionFatal reports whether a status code means the server-side session is no +// longer usable, so the conn must be evicted from the pool rather than reused for +// the next query. Broader than isBadConnection: it also covers Unauthenticated (an +// expired/revoked token kills the session) — but unlike isBadConnection it is NOT +// used to produce driver.ErrBadConn on the statement path, so eviction happens +// without database/sql replaying the statement (see toStatementError + the +// KernelBackend.markSessionDead call sites). +func isSessionFatal(code int) bool { + switch code { + case statusUnauthenticated, statusUnavailable, statusNetworkError: + return true + default: + return false + } +} + +// isUserFault reports whether a status code is a user/query fault (a bad SQL +// statement or a bad argument) rather than an infrastructure problem. Used to +// pick the log level for a kernel-call failure: user faults are routine — a +// fat-fingered query should not raise the driver's WARN rate and page on-call — +// so they log at Debug, while infra codes (Unavailable / NetworkError / Timeout / +// Unauthenticated) stay at Warn. Mirrors the Thrift path, which keeps user SQL +// errors out of operational-noise logging. +func isUserFault(code int) bool { + switch code { + case statusSqlError, statusInvalidArgument: + return true + default: + return false + } +} + +// toConnError classifies a kernel error on a SESSION-lifecycle path (open/close/ +// config, where nothing has executed): a status that means the session is unusable +// is wrapped as a bad-connection error, which identifies as driver.ErrBadConn so +// database/sql evicts the conn from the pool. Safe here because no statement ran, +// so there is nothing for database/sql to unsafely re-run. Other kernel errors — +// and plain (non-KernelError) errors — are returned unchanged, carrying sqlstate. +func toConnError(err error) error { + if err == nil { + return nil + } + ke, ok := err.(*KernelError) + if !ok { + return err + } + if isBadConnection(ke.Code) { + return dbsqlerrint.NewBadConnectionError(ke) + } + return ke +} + +// toStatementError classifies a kernel error on the STATEMENT path (execute and +// result read). It NEVER returns driver.ErrBadConn: once a statement has been +// sent, a network/unavailable failure surfaced afterward may have committed +// server-side, and driver.ErrBadConn would make database/sql transparently +// re-run the statement — a silent duplicate write for a non-idempotent +// INSERT/UPDATE/MERGE. This mirrors the kernel's own retry contract +// (ExecuteStatement is NonIdempotent, retried only on connect-phase failures) and +// the Thrift backend (ExecuteStatement is non-retryable), and honors Go's +// driver.ErrBadConn rule ("never return ErrBadConn if the server might have +// performed the operation"). The kernel has already exhausted its safe internal +// retries by the time we see the error. Returns the KernelError (or plain error) +// unchanged, carrying sqlstate. +func toStatementError(err error) error { + if err == nil { + return nil + } + if ke, ok := err.(*KernelError); ok { + return ke + } + return err +} diff --git a/internal/backend/kernel/errors_classify_test.go b/internal/backend/kernel/errors_classify_test.go new file mode 100644 index 00000000..e0180a51 --- /dev/null +++ b/internal/backend/kernel/errors_classify_test.go @@ -0,0 +1,142 @@ +package kernel + +import ( + "database/sql/driver" + "errors" + "testing" +) + +// These guard the safety contract (which errors may trigger database/sql's +// transparent statement replay, and which permanent failures must not be retried). +// They live in an untagged file so they run under CGO_ENABLED=0 — a future edit +// reintroducing the duplicate-write bug or the auth-retry storm fails CI here, +// rather than shipping green because the tagged tests never ran. + +// isBadConnection is the TRANSIENT-failure set (retrying a fresh connect could +// succeed) — used to produce driver.ErrBadConn on the session-lifecycle path. +// Unauthenticated is excluded: a wrong/expired PAT is permanent, not retryable. +func TestIsBadConnection(t *testing.T) { + transient := []int{statusUnavailable, statusNetworkError} + for _, code := range transient { + if !isBadConnection(code) { + t.Errorf("code %d should be a (transient) bad connection", code) + } + } + // Unauthenticated is session-fatal but NOT retryable, so it is not bad-conn. + notBad := []int{statusUnauthenticated, statusInvalidArgument, statusSqlError, statusTimeout} + for _, code := range notBad { + if isBadConnection(code) { + t.Errorf("code %d should not be a bad connection", code) + } + } +} + +// isSessionFatal is the broader "session is dead, evict the conn" set — it adds +// Unauthenticated (an expired token kills the session) on top of the transient set. +func TestIsSessionFatal(t *testing.T) { + fatal := []int{statusUnauthenticated, statusUnavailable, statusNetworkError} + for _, code := range fatal { + if !isSessionFatal(code) { + t.Errorf("code %d should be session-fatal", code) + } + } + notFatal := []int{statusInvalidArgument, statusSqlError, statusTimeout} + for _, code := range notFatal { + if isSessionFatal(code) { + t.Errorf("code %d should not be session-fatal", code) + } + } +} + +// isUserFault picks the log level for a kernel-call failure: user/query faults +// (bad SQL, bad argument) log at Debug so they don't inflate the WARN rate; +// infra codes stay at Warn. +func TestIsUserFault(t *testing.T) { + faults := []int{statusSqlError, statusInvalidArgument} + for _, code := range faults { + if !isUserFault(code) { + t.Errorf("code %d should be a user fault (log at Debug)", code) + } + } + infra := []int{statusUnavailable, statusNetworkError, statusTimeout, statusUnauthenticated} + for _, code := range infra { + if isUserFault(code) { + t.Errorf("code %d is infra-side and should NOT be a user fault (stays Warn)", code) + } + } +} + +// toConnError (session-lifecycle path) wraps a session-unusable KernelError as +// driver.ErrBadConn so database/sql evicts the conn, and leaves other errors and +// their sqlstate intact. +func TestToConnError(t *testing.T) { + if toConnError(nil) != nil { + t.Fatal("nil should map to nil") + } + + badConn := &KernelError{Code: statusUnavailable, Message: "gone"} + if !errors.Is(toConnError(badConn), driver.ErrBadConn) { + t.Errorf("unavailable kernel error on the session path should identify as driver.ErrBadConn") + } + + sqlErr := &KernelError{Code: statusSqlError, Message: "boom", SQLState: "42703"} + ke, ok := toConnError(sqlErr).(*KernelError) + if !ok { + t.Fatalf("sql error should remain a *KernelError, got %T", toConnError(sqlErr)) + } + if ke.SQLState != "42703" { + t.Errorf("sqlstate lost: got %q", ke.SQLState) + } +} + +// toStatementError (execute/read path) must NEVER return driver.ErrBadConn — even +// for a network/unavailable status — so database/sql cannot transparently re-run a +// statement that may have already executed server-side (silent duplicate write). +func TestToStatementErrorNeverBadConn(t *testing.T) { + if toStatementError(nil) != nil { + t.Fatal("nil should map to nil") + } + + for _, code := range []int{statusUnavailable, statusNetworkError, statusUnauthenticated} { + err := toStatementError(&KernelError{Code: code, Message: "post-execute failure"}) + if errors.Is(err, driver.ErrBadConn) { + t.Errorf("statement-path error (code=%d) must NOT identify as driver.ErrBadConn "+ + "(would let database/sql re-run a possibly-committed statement)", code) + } + } + + // sqlstate still preserved. + sqlErr := &KernelError{Code: statusSqlError, Message: "boom", SQLState: "42703"} + ke, ok := toStatementError(sqlErr).(*KernelError) + if !ok { + t.Fatalf("sql error should remain a *KernelError, got %T", toStatementError(sqlErr)) + } + if ke.SQLState != "42703" { + t.Errorf("sqlstate lost: got %q", ke.SQLState) + } +} + +// KernelError.Error() includes sqlstate and the server queryId when present. +func TestKernelErrorString(t *testing.T) { + withState := (&KernelError{Code: statusSqlError, Message: "boom", SQLState: "42703", QueryID: "q-1"}).Error() + for _, want := range []string{"boom", "sqlstate=42703", "code=13", "queryId=q-1"} { + if !contains(withState, want) { + t.Errorf("Error() = %q, missing %q", withState, want) + } + } + noState := (&KernelError{Code: statusUnavailable, Message: "gone"}).Error() + if contains(noState, "sqlstate=") || contains(noState, "queryId=") { + t.Errorf("Error() = %q, should omit empty sqlstate/queryId", noState) + } +} + +func contains(s, sub string) bool { return len(s) >= len(sub) && (s == sub || indexOf(s, sub) >= 0) } + +func indexOf(s, sub string) int { + for i := 0; i+len(sub) <= len(s); i++ { + if s[i:i+len(sub)] == sub { + return i + } + } + return -1 +} diff --git a/internal/backend/kernel/kernel_test.go b/internal/backend/kernel/kernel_test.go new file mode 100644 index 00000000..0656b302 --- /dev/null +++ b/internal/backend/kernel/kernel_test.go @@ -0,0 +1,149 @@ +//go:build cgo && databricks_kernel + +package kernel + +import ( + "context" + "database/sql/driver" + "errors" + "testing" + + dbsqlerr "github.com/databricks/databricks-sql-go/errors" + "github.com/databricks/databricks-sql-go/internal/backend" +) + +// The pure error-classifier tests (TestIsBadConnection, TestIsSessionFatal, +// TestToConnError, TestToStatementErrorNeverBadConn) live in the untagged +// errors_classify_test.go so they run under CGO_ENABLED=0. The tests below need a +// *KernelBackend, so they stay tagged. + +// evictIfSessionFatal flips SessionValid()→false on a session-fatal error (so the +// pool discards the conn) WITHOUT the error being driver.ErrBadConn (so the +// statement is never transparently re-run). +func TestEvictIfSessionFatal(t *testing.T) { + // valid tracks the session-dead flag SessionValid() gates on; the opaque + // session pointer is orthogonal here (can't construct the incomplete C type), + // so assert on k.valid directly. + k := &KernelBackend{valid: true} + + // Non-fatal (e.g. a SQL error) leaves the session valid. + k.evictIfSessionFatal(&KernelError{Code: statusSqlError}) + if !k.valid { + t.Error("a SQL error must not evict the session") + } + + // A session-fatal error marks the session dead, and the surfaced statement-path + // error is NOT driver.ErrBadConn (so database/sql won't re-run the statement). + fatal := &KernelError{Code: statusUnavailable, Message: "session gone"} + k.evictIfSessionFatal(fatal) + if k.valid { + t.Error("a session-fatal error must evict the session (valid=false)") + } + if errors.Is(toStatementError(fatal), driver.ErrBadConn) { + t.Error("the statement-path error must not be driver.ErrBadConn (no replay)") + } +} + +// Bound parameters are rejected up front by Execute with a clear error, before +// any session/C work — so this runs on a zero-value backend. The returned +// Operation must be non-nil (Backend contract) and its Close must report +// closed=false, since no server statement was ever created (a phantom +// CLOSE_STATEMENT would otherwise be recorded for it). +func TestExecuteRejectsParams(t *testing.T) { + k := &KernelBackend{} + op, err := k.Execute(context.Background(), backend.ExecRequest{ + Query: "SELECT ?", + Params: []backend.Param{{Name: "x"}}, + }) + if err == nil { + t.Fatal("expected an error for bound parameters, got nil") + } + if !errors.Is(err, dbsqlerr.ErrNotSupportedByKernel) { + t.Errorf("params rejection should wrap ErrNotSupportedByKernel, got %v", err) + } + if op == nil { + t.Fatal("Execute must return a non-nil Operation per the Backend contract") + } + closed, closeErr := op.Close(context.Background()) + if closeErr != nil { + t.Errorf("Close error = %v, want nil", closeErr) + } + if closed { + t.Error("Close on a handle-less op must report closed=false (no CLOSE_STATEMENT)") + } + if got := op.AffectedRows(); got != 0 { + t.Errorf("AffectedRows on a handle-less op = %d, want 0", got) + } +} + +// TestExecuteRejectsStaging drives Execute with a staging statement (not just the +// isStagingStatement detector in isolation) to pin the detector→Execute wiring: a +// refactor that dropped or reordered the check would silently reopen the +// silent-no-op data-loss path. Mirrors TestExecuteRejectsParams. +func TestExecuteRejectsStaging(t *testing.T) { + k := &KernelBackend{} + op, err := k.Execute(context.Background(), backend.ExecRequest{ + Query: "PUT '/tmp/f' INTO '/Volumes/main/s/e/f.csv'", + }) + if err == nil { + t.Fatal("expected an error for a staging statement, got nil") + } + if !errors.Is(err, dbsqlerr.ErrNotSupportedByKernel) { + t.Errorf("staging rejection should wrap ErrNotSupportedByKernel, got %v", err) + } + if op == nil { + t.Fatal("Execute must return a non-nil Operation per the Backend contract") + } + closed, closeErr := op.Close(context.Background()) + if closeErr != nil { + t.Errorf("Close error = %v, want nil", closeErr) + } + if closed { + t.Error("Close on a handle-less op must report closed=false (no CLOSE_STATEMENT)") + } +} + +// ExecutionError must satisfy the same public contract as the Thrift path so the +// errors.Is → errors.As → SqlState()/QueryId() recipe documented in doc.go works +// on the kernel backend too (it previously returned a bare *KernelError that +// matched none of it). +func TestExecutionErrorContract(t *testing.T) { + o := &kernelOp{} + + if got := o.ExecutionError(context.Background(), nil); got != nil { + t.Errorf("ExecutionError(nil) = %v, want nil", got) + } + + cause := &KernelError{Code: statusSqlError, Message: "boom", SQLState: "42000", QueryID: "q-123"} + err := o.ExecutionError(context.Background(), cause) + if err == nil { + t.Fatal("ExecutionError(cause) should not be nil") + } + if !errors.Is(err, dbsqlerr.ExecutionError) { + t.Errorf("kernel execution error should match dbsqlerr.ExecutionError; got %v", err) + } + var dbExec dbsqlerr.DBExecutionError + if !errors.As(err, &dbExec) { + t.Fatalf("kernel execution error should be a DBExecutionError; got %T", err) + } + if dbExec.SqlState() != "42000" { + t.Errorf("SqlState() = %q, want 42000 (from the KernelError)", dbExec.SqlState()) + } + // QueryId must come from the KernelError, not the (empty) ctx query id — the + // kernel path's StatementID() is "", so relying on ctx would drop the one + // server-side correlation handle. + if dbExec.QueryId() != "q-123" { + t.Errorf("QueryId() = %q, want q-123 (from the KernelError)", dbExec.QueryId()) + } + // The *KernelError cause stays reachable via Unwrap. + var ke *KernelError + if !errors.As(err, &ke) { + t.Error("the *KernelError cause should remain reachable via errors.As") + } +} + +// The cell/nested rendering (ScanCell and the JSON grammar) now lives in the +// untagged internal/arrowscan package, where its tests run in the default +// CGO_ENABLED=0 build; see arrowscan_test.go. The decimal formatter lives in +// internal/decimalfmt. This file keeps the kernel-specific tests: error mapping, +// bad-connection classification, and the bound-params rejection. diff --git a/internal/backend/kernel/operation.go b/internal/backend/kernel/operation.go new file mode 100644 index 00000000..eec325a6 --- /dev/null +++ b/internal/backend/kernel/operation.go @@ -0,0 +1,289 @@ +//go:build cgo && databricks_kernel + +package kernel + +/* +#include +#include "databricks_kernel.h" +*/ +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" +) + +// execute runs one statement to a terminal state on the blocking-execute path, +// with out-of-band cancellation from a watcher goroutine: +// 1. new_statement + set_sql +// 2. canceller_new BEFORE execute, so it can observe the server statement id +// 3. a watcher goroutine that fires the canceller on ctx.Done() +// 4. the single blocking kernel_statement_execute (inline/CloudFetch and +// long-query polling all happen inside the kernel, invisibly) +// 5. drain the watcher before returning, so a late cancel cannot land on a +// statement that reuses this handle +// +// Executes SQL text only; bound parameters are rejected up front by Execute. +func (k *KernelBackend) execute(ctx context.Context, req backend.ExecRequest) (backend.Operation, error) { + // Log the SQL length, not the text: query bodies can carry PII/secrets in + // WHERE/INSERT/SET, and this goes to stderr. Matches the driver's own + // debuglog convention (conn.ExecContext logs sql.len=%d). + klog("Execute sql.len=%d", len(req.Query)) + + var stmt *C.kernel_statement_t + if err := call(func() C.KernelStatusCode { + return C.kernel_session_new_statement(k.session, &stmt) + }); err != nil { + k.evictIfSessionFatal(err) + return &kernelOp{}, fmt.Errorf("kernel: new_statement: %w", toStatementError(err)) + } + + sql := newCStr(req.Query) + if err := call(func() C.KernelStatusCode { + return C.kernel_statement_set_sql(stmt, sql.c) + }); err != nil { + sql.free() + C.kernel_statement_close(stmt) + k.evictIfSessionFatal(err) + return &kernelOp{}, fmt.Errorf("kernel: set_sql: %w", toStatementError(err)) + } + sql.free() + + // Detached canceller, obtained before execute so it observes the server + // statement id the moment execute publishes it. Non-fatal on failure: proceed + // without cancellation rather than failing the query. + var canceller *C.kernel_statement_canceller_t + if err := call(func() C.KernelStatusCode { + return C.kernel_statement_canceller_new(stmt, &canceller) + }); err != nil { + klog("canceller_new failed (proceeding without cancel): %v", err) + canceller = nil + } + + // Watcher goroutine (only when there is both a cancellable ctx and a + // canceller). The server publishes the statement id to the canceller's + // inflight slot only when the initial POST returns — held up to the server's + // inline wait even for a long query — so a cancel fired before that is + // a no-op. Re-fire every 250ms after ctx.Done until the kernel reports the + // cancel RPC was actually dispatched (dispatched=true, i.e. the id appeared + // and the RPC went out), then stop: one real cancel is enough, and further + // fires would only hammer the server. Falls back to firing until execute + // returns (done) if the RPC never dispatches. + done := make(chan struct{}) + var watcherWg sync.WaitGroup + if canceller != nil && ctx.Done() != nil { + watcherWg.Add(1) + go func() { + defer watcherWg.Done() + select { + case <-ctx.Done(): + case <-done: + return + } + // Both ctx.Done() and done can be ready at once — execute just returned + // and the ctx was cancelled in the same window — and select picks + // randomly. If done is (also) closed, execute already completed, so skip + // firing: a cancel here would be a spurious RPC against a terminal + // statement (server no-op, but a wasted call + a misleading "cancelled" + // entry in query history for a query that actually returned). + select { + case <-done: + return + default: + } + klog("ctx.Done (%v) → firing canceller (with retry until dispatched)", ctx.Err()) + ticker := time.NewTicker(250 * time.Millisecond) + defer ticker.Stop() + if fireCancel(canceller) { + klog("cancel dispatched on first fire") + return + } + for { + select { + case <-done: + return + case <-ticker.C: + if fireCancel(canceller) { + klog("cancel dispatched, watcher stopping") + return + } + } + } + }() + } + + // The one blocking call. inline vs CloudFetch and long-query polling are all + // resolved inside the kernel; Go just waits here. + var exec *C.kernel_executed_statement_t + execErr := call(func() C.KernelStatusCode { + return C.kernel_statement_execute(stmt, &exec) + }) + + // Drain the watcher before returning so a late canceller fire cannot land on + // a subsequent statement reusing this handle. + close(done) + watcherWg.Wait() + if canceller != nil { + C.kernel_statement_canceller_free(canceller) + } + + op := &kernelOp{backend: k, stmt: stmt, location: k.cfg.Location} + if execErr != nil { + // Prefer the caller's ctx error when the ctx was cancelled (database/sql + // convention), keeping the kernel error as the cause. + if ctx.Err() != nil { + klog("Execute failed under cancelled ctx: kernelErr=%v ctxErr=%v", execErr, ctx.Err()) + op.close() + return op, fmt.Errorf("kernel: execute cancelled: %w", ctx.Err()) + } + klog("Execute failed: %v", execErr) + // A session-fatal status (expired token, dropped/unavailable session) means + // this conn is unusable: evict it so the pool doesn't hand it out again. We + // still return the PLAIN error (toStatementError, never ErrBadConn), so the + // conn is discarded without database/sql re-running the statement. + k.evictIfSessionFatal(execErr) + op.close() + return op, fmt.Errorf("kernel: execute: %w", toStatementError(execErr)) + } + op.exec = exec + // Capture the modified-row count now, while exec is live — the operation is + // closed (nulling exec) before AffectedRows is read on the ExecContext path. + op.affectedRows = int64(C.kernel_executed_statement_num_modified_rows(exec)) + // 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 + // report a committed write as cancelled, and a caller treating cancellation as + // retryable would double-write. This matches the Thrift path, which returns + // success on completion with no post-completion ctx re-check. A cancel that + // arrives before completion is still honored via the watcher → execErr branch + // above; one that loses the race to a committed statement yields that statement's + // result, same as Thrift. + klog("Execute OK stmt=%p exec=%p affectedRows=%d", stmt, exec, op.affectedRows) + return op, nil +} + +// kernelOp implements backend.Operation over a sync executed statement. +type kernelOp struct { + // backend is the owning connection's backend, held so a session-fatal error on + // the result-read path can evict the conn (markSessionDead), mirroring how the + // Thrift thriftOperation holds its *Backend. + backend *KernelBackend + stmt *C.kernel_statement_t + exec *C.kernel_executed_statement_t + closed bool + // affectedRows is the modified-row count captured at execute time. It is + // cached (not read live from exec) because the caller closes the operation — + // which nulls exec — before reading AffectedRows (see conn.ExecContext). + affectedRows int64 + // location renders DATE / TIMESTAMP values in the session time zone, matching + // the Thrift path; nil means UTC. Carried onto the rows built by Results. + location *time.Location +} + +var _ backend.Operation = (*kernelOp)(nil) + +// StatementID returns "": the kernel C ABI exposes no success-path statement/query +// id accessor (only the error path carries a query_id, on KernelError). Because +// conn.ExecContext/QueryContext gate per-statement telemetry on +// StatementID() != "", kernel queries currently emit no EXECUTE_STATEMENT metric +// and their query-id log field is empty — there is no session-id fallback on that +// gate. Surfacing an id needs a kernel accessor (e.g. +// kernel_executed_statement_query_id); tracked as a follow-up. +func (o *kernelOp) StatementID() string { return "" } + +// AffectedRows is the modified-row count for ExecContext. It returns the value +// cached at execute time, so it is correct even after the operation is closed +// (the ExecContext path closes the op before reading this). +func (o *kernelOp) AffectedRows() int64 { + return o.affectedRows +} + +// Results builds the driver.Rows over the executed statement's result stream. +// On the query path the returned Rows owns closing the server-side operation. +func (o *kernelOp) Results(ctx context.Context, callbacks *dbsqlrows.TelemetryCallbacks) (driver.Rows, error) { + if o.exec == nil { + return nil, fmt.Errorf("kernel: no executed statement") + } + var stream *C.kernel_result_stream_t + if err := call(func() C.KernelStatusCode { + return C.kernel_executed_statement_get_result_stream(o.exec, &stream) + }); err != nil { + // No Rows is returned to own teardown, and the query path does not call + // Operation.Close on a Results error — so close the handles here to avoid + // leaking the statement / executed handle (and its server operation). + o.backend.evictIfSessionFatal(err) + o.close() + return nil, fmt.Errorf("kernel: get_result_stream: %w", toStatementError(err)) + } + return newKernelRows(ctx, o, stream, callbacks) +} + +// IsStaging reports whether this is a staging (PUT/GET/REMOVE) operation. Always +// false on the kernel path: staging statements are rejected up front in +// Execute (isStagingStatement), so no kernelOp is ever produced for one — there +// is nothing for conn.execStagingOperation to act on here. +func (o *kernelOp) IsStaging(ctx context.Context) (bool, error) { return false, nil } + +// Close best-effort closes the executed statement and its statement handle. It is +// idempotent (a second call, or a call after Rows.Close already tore the +// operation down, is a no-op) per the backend.Operation contract. closed reports +// whether a close was actually issued. +func (o *kernelOp) Close(ctx context.Context) (bool, error) { + return o.close(), nil +} + +// close is the shared idempotent teardown used by both Operation.Close and +// Rows.Close. Closing the executed handle first, then the statement, matches the +// C ABI teardown order (result stream is closed by Rows before this runs). +func (o *kernelOp) close() bool { + if o.closed { + return false + } + o.closed = true + // Report closed=true only when there was actually a handle to tear down. A + // handle-less op (e.g. the &kernelOp{} returned on the bound-params error + // path) reaches here via conn.ExecContext's unconditional Close; returning + // true would record a phantom CLOSE_STATEMENT for a statement that never hit + // the server. Matches the backend.Operation contract (closed=false when the + // operation had no handle) and the Thrift backend's !hasHandle() behavior. + didClose := o.exec != nil || o.stmt != nil + if o.exec != nil { + C.kernel_executed_statement_close(o.exec) + o.exec = nil + } + if o.stmt != nil { + C.kernel_statement_close(o.stmt) + o.stmt = nil + } + klog("kernelOp closed (didClose=%v)", didClose) + return didClose +} + +// ExecutionError wraps cause as the driver's execution error so it satisfies the +// same public contract as the Thrift path: errors.Is(err, dbsqlerr.ExecutionError) +// matches, and errors.As(err, &dbExecErr) exposes SqlState()/QueryId() — the recipe +// documented in doc.go. The kernel carries BOTH sqlState and the server query id in +// its own *KernelError (not a Thrift TGetOperationStatusResp, and unlike Thrift the +// kernel path has no ctx query id — StatementID() is ""), so pull both out and hand +// them to NewExecutionErrorWithState, keeping the *KernelError as the cause so its +// detail stays reachable via Unwrap. +func (o *kernelOp) ExecutionError(ctx context.Context, cause error) error { + if cause == nil { + return nil + } + sqlState, queryID := "", "" + var ke *KernelError + if errors.As(cause, &ke) { + sqlState, queryID = ke.SQLState, ke.QueryID + } + return dbsqlerrint.NewExecutionErrorWithState(ctx, dbsqlerr.ErrQueryExecution, cause, queryID, sqlState) +} diff --git a/internal/backend/kernel/rows.go b/internal/backend/kernel/rows.go new file mode 100644 index 00000000..82df1920 --- /dev/null +++ b/internal/backend/kernel/rows.go @@ -0,0 +1,194 @@ +//go:build cgo && databricks_kernel + +package kernel + +/* +#include +#include "databricks_kernel.h" +// Forward-declare the Arrow C Data Interface structs so cgo can take their +// addresses; arrow-go's cdata package reinterprets these via unsafe.Pointer. +struct ArrowSchema; +struct ArrowArray; +*/ +import "C" + +import ( + "context" + "database/sql/driver" + "fmt" + "io" + "unsafe" + + "github.com/apache/arrow/go/v12/arrow" + "github.com/apache/arrow/go/v12/arrow/cdata" + "github.com/databricks/databricks-sql-go/internal/arrowscan" + dbsqlrows "github.com/databricks/databricks-sql-go/internal/rows" +) + +var _ driver.Rows = (*kernelRows)(nil) + +// kernelRows implements driver.Rows over the kernel result stream. It pulls one +// Arrow RecordBatch at a time via kernel_result_stream_next_batch (inline and +// CloudFetch are transparent below the C ABI), imports it zero-copy through the +// Arrow C Data Interface, and scans rows out on demand. +// +// The batch/row split mirrors the kernel's own ResultStream: next_batch does the +// per-batch network/decode work once; row reads walk the already-imported +// arrow.Record with O(1) indexing. +type kernelRows struct { + ctx context.Context + op *kernelOp + stream *C.kernel_result_stream_t + callbacks *dbsqlrows.TelemetryCallbacks + + cols []string + cur arrow.Record // current batch (nil until first Next) + rowInCur int // next row index within cur + chunkCount int // cumulative batches fetched, for OnChunkFetched + closed bool + eof bool + // keyCache memoizes struct field-name JSON keys for this result set so + // per-row rendering doesn't re-marshal constant names. Scoped to this Rows + // (freed with it) — not a process-global, which would leak. + keyCache *arrowscan.StructKeyCache +} + +// newKernelRows fetches the schema up front (for Columns()) and returns the row +// iterator; batches are pulled lazily on Next. +func newKernelRows(ctx context.Context, op *kernelOp, stream *C.kernel_result_stream_t, cb *dbsqlrows.TelemetryCallbacks) (driver.Rows, error) { + r := &kernelRows{ctx: ctx, op: op, stream: stream, callbacks: cb, keyCache: arrowscan.NewStructKeyCache()} + + var csch C.struct_ArrowSchema + if err := call(func() C.KernelStatusCode { + return C.kernel_result_stream_get_schema(stream, &csch) + }); err != nil { + op.backend.evictIfSessionFatal(err) + r.Close() + return nil, fmt.Errorf("kernel: get_schema: %w", toStatementError(err)) + } + sch, err := cdata.ImportCArrowSchema((*cdata.CArrowSchema)(unsafe.Pointer(&csch))) + if err != nil { + r.Close() + return nil, fmt.Errorf("kernel: import schema: %w", err) + } + fields := sch.Fields() + r.cols = make([]string, len(fields)) + for i, f := range fields { + r.cols[i] = f.Name + } + klog("newKernelRows: %d columns", len(r.cols)) + return r, nil +} + +// Columns returns the result-set column names. +func (r *kernelRows) Columns() []string { return r.cols } + +// Close releases the current batch, the kernel result stream, and (query-path +// ownership) the server operation. Idempotent. +func (r *kernelRows) Close() error { + if r.closed { + return nil + } + r.closed = true + if r.cur != nil { + r.cur.Release() + r.cur = nil + } + if r.stream != nil { + C.kernel_result_stream_close(r.stream) + r.stream = nil + } + if r.op != nil { + r.op.close() + } + klog("kernelRows closed") + return nil +} + +// Next fills dest with the next row's values, advancing across batches. Returns +// io.EOF when the stream is drained. +func (r *kernelRows) Next(dest []driver.Value) error { + if r.closed { + return io.EOF + } + for r.cur == nil || r.rowInCur >= int(r.cur.NumRows()) { + if r.eof { + return io.EOF + } + if err := r.nextBatch(); err != nil { + return err + } + } + rec := r.cur + for c := 0; c < len(dest); c++ { + v, err := arrowscan.ScanCellCached(rec.Column(c), r.rowInCur, r.op.location, r.keyCache) + if err != nil { + return fmt.Errorf("kernel: scan col %d (%s): %w", c, r.cols[c], err) + } + dest[c] = v + } + r.rowInCur++ + return nil +} + +// nextBatch pulls the next Arrow batch. A released array (release==NULL) is the +// kernel's end-of-stream sentinel. +func (r *kernelRows) nextBatch() error { + // Honor cancellation at batch boundaries: check ctx before entering the + // blocking C fetch (which cannot itself observe ctx). This does NOT interrupt + // a fetch already in flight — and database/sql's own cancel watcher can't + // either: its Rows.Close takes rs.closemu.Lock(), which blocks until the + // in-progress Next (holding the RLock) returns, so the stream close waits for + // the C call to finish on its own. A single hung CloudFetch batch is therefore + // uninterruptible (the kernel exposes no per-download timeout); mid-fetch + // cancellation would need the execute path's watcher/canceller applied here. + if r.ctx != nil { + if err := r.ctx.Err(); err != nil { + return err + } + } + if r.cur != nil { + r.cur.Release() + r.cur = nil + } + var carr C.struct_ArrowArray + var csch C.struct_ArrowSchema + if err := call(func() C.KernelStatusCode { + return C.kernel_result_stream_next_batch(r.stream, &carr, &csch) + }); err != nil { + r.op.backend.evictIfSessionFatal(err) + return fmt.Errorf("kernel: next_batch: %w", toStatementError(err)) + } + if carr.release == nil { + r.eof = true + klog("nextBatch: EOF") + return io.EOF + } + // Zero-copy import. The kernel exports self-contained batches (Rust to_ffi + // moves the Arc-owned buffers in), so the arrow.Record safely outlives the + // stream; we still Release each batch explicitly as we advance. + rec, err := cdata.ImportCRecordBatch( + (*cdata.CArrowArray)(unsafe.Pointer(&carr)), + (*cdata.CArrowSchema)(unsafe.Pointer(&csch))) + if err != nil { + return fmt.Errorf("kernel: import batch: %w", err) + } + r.cur = rec + r.rowInCur = 0 + // Scope the struct-key cache to this batch: the C Data import mints a fresh + // *StructType per batch, so the prior batch's cached prefixes can never be hit + // again — resetting keeps the cache from growing one entry per batch over the + // whole Rows lifetime (the intra-batch memoization win is unaffected). + r.keyCache.Reset() + r.chunkCount++ + if r.callbacks != nil && r.callbacks.OnChunkFetched != nil { + // chunkCount is cumulative (per the callback contract). bytesDownloaded, + // chunkIndex, and latency are left 0: the kernel does CloudFetch/decompress + // internally and hands back ready Arrow batches, so the Go side never sees + // the compressed wire bytes or per-chunk fetch latency the Thrift path + // reports — a fabricated number would be worse than a truthful zero. + r.callbacks.OnChunkFetched(r.chunkCount, 0, 0, 0, 0) + } + klog("nextBatch: %d rows (chunk %d)", rec.NumRows(), r.chunkCount) + return nil +} diff --git a/internal/backend/kernel/staging.go b/internal/backend/kernel/staging.go new file mode 100644 index 00000000..c614169c --- /dev/null +++ b/internal/backend/kernel/staging.go @@ -0,0 +1,67 @@ +package kernel + +import "strings" + +// This file is intentionally NOT behind the `cgo && databricks_kernel` build tag. +// Staging detection is pure Go (string inspection of the SQL), so keeping it +// untagged lets its test run under CGO_ENABLED=0. + +// isStagingStatement reports whether sql is a Unity Catalog volume staging +// command (PUT / GET / REMOVE). The Thrift backend learns a statement is staging +// from server result metadata (IsStagingOperation); the kernel C ABI exposes no +// such signal, and the kernel path cannot perform the local file transfer these +// commands require. We therefore detect them from the SQL and reject at execute +// time (see Execute) rather than silently returning success with no file moved — +// which is what returning false from IsStaging would do. Case-insensitive; skips +// leading whitespace AND leading SQL comments (`-- ...` and `/* ... */`), so a +// comment-prefixed staging command is still caught rather than slipping through +// as a silent no-op. Errs toward matching only a clear leading keyword — a false +// negative just yields the kernel's normal "unsupported syntax" server error, +// whereas a false positive would reject a legitimate query. +func isStagingStatement(sql string) bool { + s := stripLeadingSQLComments(sql) + for _, kw := range []string{"PUT", "GET", "REMOVE"} { + if len(s) > len(kw) && + strings.EqualFold(s[:len(kw)], kw) && + isSQLTokenBreak(s[len(kw)]) { + return true + } + } + return false +} + +// stripLeadingSQLComments removes leading whitespace and leading SQL comments +// (both `-- ...` line comments and `/* ... */` block comments), repeatedly, so +// the first real token is exposed. An unterminated block comment consumes the +// rest of the string (yielding ""). Comment bodies are not scanned for nested +// markers — SQL block comments do not nest — and this only touches the LEADING +// run, so a comment later in the statement is untouched. +func stripLeadingSQLComments(sql string) string { + s := strings.TrimSpace(sql) + for { + switch { + case strings.HasPrefix(s, "--"): + if i := strings.IndexByte(s, '\n'); i >= 0 { + s = s[i+1:] + } else { + return "" // line comment to end of input + } + case strings.HasPrefix(s, "/*"): + if i := strings.Index(s[2:], "*/"); i >= 0 { + s = s[2+i+2:] + } else { + return "" // unterminated block comment + } + default: + return s + } + s = strings.TrimSpace(s) + } +} + +// isSQLTokenBreak reports whether b ends a leading keyword — a space or a quote +// (staging commands take a quoted path, e.g. PUT '' INTO ''), so +// "PUT '..." matches but an identifier like "PUTS" or "GETTER" does not. +func isSQLTokenBreak(b byte) bool { + return b == ' ' || b == '\t' || b == '\n' || b == '\r' || b == '\'' || b == '"' +} diff --git a/internal/backend/kernel/staging_test.go b/internal/backend/kernel/staging_test.go new file mode 100644 index 00000000..91d69794 --- /dev/null +++ b/internal/backend/kernel/staging_test.go @@ -0,0 +1,41 @@ +package kernel + +import "testing" + +func TestIsStagingStatement(t *testing.T) { + staging := []string{ + "PUT '/tmp/f.csv' INTO '/Volumes/main/s/e/f.csv' OVERWRITE", + "GET '/Volumes/main/s/e/f.csv' TO 'local.csv'", + "REMOVE '/Volumes/main/s/e/f.csv'", + " put '/tmp/f' into '/Volumes/x'", // leading space + lowercase + "Get\t'/Volumes/x' to 'y'", // tab after keyword, mixed case + "/* audit */ PUT '/local' INTO '/Volumes/main/s/v'", // leading block comment + "/* a */ /* b */ PUT '/l' INTO '/Volumes/x'", // multiple block comments + "-- upload the file\nPUT '/l' INTO '/Volumes/x'", // leading line comment, then staging + " /* c */\n -- d\n REMOVE '/Volumes/x'", // mixed comments + whitespace + } + for _, sql := range staging { + if !isStagingStatement(sql) { + t.Errorf("expected staging: %q", sql) + } + } + + notStaging := []string{ + "SELECT 1", + "PUTS ON A SHOW", // PUT is a prefix of an identifier, not a token + "GETTER()", // likewise for GET + "REMOVED_AT FROM t", // likewise for REMOVE + "INSERT INTO t VALUES (1)", // contains INTO but is not staging + "", // empty + "-- PUT '/x' INTO '/y'", // whole line commented out; nothing after + "/* PUT '/x' INTO '/y' */", // whole statement inside a block comment + "SELECT 'PUT ' AS col", // keyword only inside a literal + "/* unterminated PUT '/x'", // unterminated block comment consumes the rest + "/* c */ SELECT 1", // real leading comment, but not a staging command + } + for _, sql := range notStaging { + if isStagingStatement(sql) { + t.Errorf("did not expect staging: %q", sql) + } + } +} diff --git a/internal/config/config.go b/internal/config/config.go index fb7c6afb..12b59395 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -125,6 +125,13 @@ type UserConfig struct { // See databricks/databricks-sql-go#274. UseArrowNativeDecimalDSN bool CloudFetchConfig + // UseKernel selects the SEA-via-kernel backend instead of Thrift. See the + // WithUseKernel connector option for the build requirements. DSN: useKernel=true. + UseKernel bool + // WarehouseID is the bare SQL warehouse id, used by the kernel backend (which + // addresses a warehouse by id) in preference to HTTPPath. The Thrift backend + // ignores it and routes by HTTPPath. DSN: warehouseId=. + WarehouseID string } // DeepCopy returns a true deep copy of UserConfig @@ -171,6 +178,8 @@ func (ucfg UserConfig) DeepCopy() UserConfig { EnableTelemetry: ucfg.EnableTelemetry, TelemetryBatchSize: ucfg.TelemetryBatchSize, TelemetryFlushInterval: ucfg.TelemetryFlushInterval, + UseKernel: ucfg.UseKernel, + WarehouseID: ucfg.WarehouseID, } } @@ -320,6 +329,17 @@ func ParseDSN(dsn string) (UserConfig, error) { ucfg.UseArrowNativeDecimalDSN = useArrowNativeDecimal } + // Kernel backend parameters + if useKernel, ok, err := params.extractAsBool("useKernel"); ok { + if err != nil { + return UserConfig{}, err + } + ucfg.UseKernel = useKernel + } + if warehouseID, ok := params.extract("warehouseId"); ok { + ucfg.WarehouseID = warehouseID + } + // Telemetry parameters if enableTelemetry, ok, err := params.extractAsBool("enableTelemetry"); ok { if err != nil { diff --git a/internal/config/config_test.go b/internal/config/config_test.go index d3574fd1..4c70e993 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -586,6 +586,34 @@ func TestParseConfig(t *testing.T) { wantCfg: UserConfig{}, wantErr: true, }, + { + name: "kernel backend params useKernel and warehouseId", + args: args{dsn: "token:supersecret@example.cloud.databricks.com:443/sql/1.0/endpoints/12346a5b5b0e123a?useKernel=true&warehouseId=abc123"}, + wantCfg: UserConfig{ + Protocol: "https", + Host: "example.cloud.databricks.com", + Port: 443, + MaxRows: defaultMaxRows, + Authenticator: &pat.PATAuth{AccessToken: "supersecret"}, + AccessToken: "supersecret", + HTTPPath: "/sql/1.0/endpoints/12346a5b5b0e123a", + SessionParams: make(map[string]string), + RetryMax: 4, + RetryWaitMin: 1 * time.Second, + RetryWaitMax: 30 * time.Second, + CloudFetchConfig: defCloudConfig, + UseKernel: true, + WarehouseID: "abc123", + }, + wantURL: "https://example.cloud.databricks.com:443/sql/1.0/endpoints/12346a5b5b0e123a", + wantErr: false, + }, + { + name: "malformed useKernel is an error", + args: args{dsn: "token:supersecret@example.cloud.databricks.com:443/sql/1.0/endpoints/12346a5b5b0e123a?useKernel=notabool"}, + wantCfg: UserConfig{}, + wantErr: true, + }, } for i, tt := range tests { fmt.Println(i) @@ -649,6 +677,8 @@ func TestUserConfig_DeepCopy(t *testing.T) { UserAgentEntry: "test", Location: location, SessionParams: map[string]string{"a": "32", "b": "4"}, + UseKernel: true, + WarehouseID: "wh-abc123", } cfg_copy := cfg.DeepCopy() diff --git a/internal/decimalfmt/decimalfmt.go b/internal/decimalfmt/decimalfmt.go new file mode 100644 index 00000000..7d21ebf5 --- /dev/null +++ b/internal/decimalfmt/decimalfmt.go @@ -0,0 +1,47 @@ +// Package decimalfmt renders Arrow decimal128 values as exact fixed-point +// strings, shared by the Thrift (arrowbased) and kernel result paths so a +// DECIMAL renders identically regardless of backend. Keeping one implementation +// means a precision fix (e.g. databricks-sql-go#274) lands in both at once. +package decimalfmt + +import ( + "math/big" + "strings" + + "github.com/apache/arrow/go/v12/arrow/decimal128" +) + +// ExactString renders an Arrow decimal128 as an exact fixed-point string, +// applying scale by string placement rather than float conversion. This +// preserves precision a float64 would lose beyond ~17 significant digits +// (databricks-sql-go#274). A negative scale is not produced by the server and is +// treated as scale 0 rather than panicking. +func ExactString(n decimal128.Num, scale int32) string { + unscaled := n.BigInt() // exact signed unscaled integer + neg := unscaled.Sign() < 0 + digits := new(big.Int).Abs(unscaled).String() + + var b strings.Builder + if neg { + b.WriteByte('-') + } + + if scale <= 0 { + b.WriteString(digits) + return b.String() + } + + s := int(scale) + if len(digits) <= s { + // Pad with leading zeros so there are exactly `scale` fractional digits + // and a single leading integer zero, e.g. 5 with scale 3 -> "0.005". + b.WriteString("0.") + b.WriteString(strings.Repeat("0", s-len(digits))) + b.WriteString(digits) + } else { + b.WriteString(digits[:len(digits)-s]) + b.WriteByte('.') + b.WriteString(digits[len(digits)-s:]) + } + return b.String() +} diff --git a/internal/decimalfmt/decimalfmt_test.go b/internal/decimalfmt/decimalfmt_test.go new file mode 100644 index 00000000..7765c7c6 --- /dev/null +++ b/internal/decimalfmt/decimalfmt_test.go @@ -0,0 +1,43 @@ +package decimalfmt + +import ( + "math/big" + "testing" + + "github.com/apache/arrow/go/v12/arrow/decimal128" +) + +// ExactString applies scale by string placement, preserving digits a float64 +// would lose beyond ~17 significant figures. +func TestExactString(t *testing.T) { + cases := []struct { + unscaled int64 + scale int32 + want string + }{ + {12345, 2, "123.45"}, + {5, 3, "0.005"}, + {100, 0, "100"}, + {1999, 2, "19.99"}, + {-12345, 2, "-123.45"}, + {-5, 3, "-0.005"}, + {0, 2, "0.00"}, + } + for _, c := range cases { + got := ExactString(decimal128.FromI64(c.unscaled), c.scale) + if got != c.want { + t.Errorf("unscaled=%d scale=%d: got %q want %q", c.unscaled, c.scale, got, c.want) + } + } +} + +// A value beyond float64's exact integer range must survive intact — the whole +// point of formatting from the 128-bit unscaled integer rather than a float. +func TestExactStringHighPrecision(t *testing.T) { + // 12345678901234567890 (20 digits) — past float64's 2^53 exact-integer limit. + big20, _ := new(big.Int).SetString("12345678901234567890", 10) + got := ExactString(decimal128.FromBigInt(big20), 4) + if want := "1234567890123456.7890"; got != want { + t.Errorf("got %q want %q", got, want) + } +} diff --git a/internal/errors/err.go b/internal/errors/err.go index 7d6765d6..1cba5a93 100644 --- a/internal/errors/err.go +++ b/internal/errors/err.go @@ -185,6 +185,21 @@ 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 +// queryId and sqlState, for backends that don't have a Thrift +// TGetOperationStatusResp (e.g. the kernel backend, which carries both in its own +// KernelError). Same shape as NewExecutionError otherwise. An empty queryId falls +// back to the ctx value, so a caller that only has the ctx id (or none) still +// behaves like NewExecutionError. +func NewExecutionErrorWithState(ctx context.Context, msg string, err error, queryId, sqlState string) *executionError { + dbErr := newDatabricksError(ctx, msg, err) + dbErr.errType = "execution error" + if queryId == "" { + queryId = driverctx.QueryIdFromContext(ctx) + } + return &executionError{databricksError: dbErr, queryId: queryId, 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/rows/arrowbased/arrowRows.go b/internal/rows/arrowbased/arrowRows.go index b7a03fb5..4b99eb9d 100644 --- a/internal/rows/arrowbased/arrowRows.go +++ b/internal/rows/arrowbased/arrowRows.go @@ -4,6 +4,7 @@ import ( "bytes" "context" "database/sql/driver" + "encoding/json" "io" "time" @@ -686,11 +687,16 @@ func (vcm *arrowValueContainerMaker) makeColumnValueContainer(t arrow.DataType, case *arrow.StructType: svc := &structValueContainer{structArrayType: t} - svc.fieldNames = make([]string, len(t.Fields())) + svc.fieldKeys = make([]string, len(t.Fields())) svc.fieldValues = make([]columnValues, len(t.Fields())) svc.complexValue = make([]bool, len(t.Fields())) for i, f := range t.Fields() { - svc.fieldNames[i] = f.Name + // Precompute the JSON-escaped `"name":` key once — field names are + // invariant across rows, so escaping them per row (structValueContainer + // .Value runs per row) would re-marshal constant strings. json.Marshal a + // string never errors. + keyBytes, _ := json.Marshal(f.Name) + svc.fieldKeys[i] = string(keyBytes) + ":" c, err := vcm.makeColumnValueContainer(f.Type, location, toTimestampFn, nil) if err != nil { return nil, err diff --git a/internal/rows/arrowbased/arrowscan_parity_test.go b/internal/rows/arrowbased/arrowscan_parity_test.go new file mode 100644 index 00000000..762d11fe --- /dev/null +++ b/internal/rows/arrowbased/arrowscan_parity_test.go @@ -0,0 +1,445 @@ +package arrowbased + +import ( + "testing" + "time" + + "github.com/apache/arrow/go/v12/arrow" + "github.com/apache/arrow/go/v12/arrow/array" + "github.com/apache/arrow/go/v12/arrow/decimal128" + "github.com/apache/arrow/go/v12/arrow/memory" + + "github.com/databricks/databricks-sql-go/internal/arrowscan" +) + +// The kernel backend (internal/arrowscan) and the Thrift backend (this package) +// each render Arrow cells to driver.Values — nested types to a JSON string — and +// the "identical results across backends" contract requires byte-for-byte +// agreement. That contract used to be guarded only by TestKernelThriftParity, +// which is build-tagged AND needs a live warehouse, so it never runs in CI. This +// test feeds the same arrow.Array through both renderers with no cgo and no +// warehouse, so a divergence (e.g. struct-key JSON escaping) fails a default +// CGO_ENABLED=0 run. +// +// It lives in package arrowbased because the container factory +// (makeColumnValueContainer) is package-private; arrowscan is a pure leaf import +// (no cycle). +func renderViaArrowbased(t *testing.T, arr arrow.Array, row int, loc *time.Location) any { + t.Helper() + maker := &arrowValueContainerMaker{} + holder, err := maker.makeColumnValueContainer(arr.DataType(), loc, func(ts arrow.Timestamp) time.Time { + return ts.ToTime(arrow.Microsecond) + }, nil) + if err != nil { + t.Fatalf("makeColumnValueContainer(%s): %v", arr.DataType(), err) + } + if err := holder.SetValueArray(arr.Data()); err != nil { + t.Fatalf("SetValueArray(%s): %v", arr.DataType(), err) + } + if holder.IsNull(row) { + return nil + } + v, err := holder.Value(row) + if err != nil { + t.Fatalf("Value(%s): %v", arr.DataType(), err) + } + return v +} + +func TestArrowbasedKernelRenderParity(t *testing.T) { + pool := memory.NewGoAllocator() + + // structWithKey builds a single-row STRUCT whose field is `name`. + structWithKey := func(name string) arrow.Array { + dt := arrow.StructOf(arrow.Field{Name: name, Type: arrow.PrimitiveTypes.Int64}) + b := array.NewStructBuilder(pool, dt) + b.Append(true) + b.FieldBuilder(0).(*array.Int64Builder).Append(1) + return b.NewArray() + } + + cases := []struct { + name string + build func() arrow.Array + }{ + {"list_int", func() arrow.Array { + b := array.NewListBuilder(pool, arrow.PrimitiveTypes.Int64) + vb := b.ValueBuilder().(*array.Int64Builder) + b.Append(true) + vb.Append(1) + vb.Append(2) + return b.NewArray() + }}, + {"struct_simple", func() arrow.Array { return structWithKey("a") }}, + // The escaping divergence: a field name with a quote must render as valid, + // identically-escaped JSON on both backends. + {"struct_key_with_quote", func() arrow.Array { return structWithKey(`a"b`) }}, + {"struct_key_with_backslash", func() arrow.Array { return structWithKey(`a\b`) }}, + {"struct_key_with_newline", func() arrow.Array { return structWithKey("a\nb") }}, + {"map_string_key", func() arrow.Array { + b := array.NewMapBuilder(pool, arrow.BinaryTypes.String, arrow.PrimitiveTypes.Int64, false) + b.Append(true) // open row 0's map + b.KeyBuilder().(*array.StringBuilder).Append("k") + b.ItemBuilder().(*array.Int64Builder).Append(9) + return b.NewArray() + }}, + {"map_special_string_key", func() arrow.Array { + b := array.NewMapBuilder(pool, arrow.BinaryTypes.String, arrow.PrimitiveTypes.Int64, false) + b.Append(true) + b.KeyBuilder().(*array.StringBuilder).Append(`k"x`) + b.ItemBuilder().(*array.Int64Builder).Append(9) + return b.NewArray() + }}, + {"map_int_key", func() arrow.Array { + b := array.NewMapBuilder(pool, arrow.PrimitiveTypes.Int64, arrow.PrimitiveTypes.Int64, false) + b.Append(true) + b.KeyBuilder().(*array.Int64Builder).Append(7) + b.ItemBuilder().(*array.Int64Builder).Append(9) + return b.NewArray() + }}, + {"map_binary_key", func() arrow.Array { + // []byte key: json.Marshal → base64 "YWJj", NOT fmt "%v" [97 98 99]. + b := array.NewMapBuilder(pool, arrow.BinaryTypes.Binary, arrow.PrimitiveTypes.Int64, false) + b.Append(true) + b.KeyBuilder().(*array.BinaryBuilder).Append([]byte("abc")) + b.ItemBuilder().(*array.Int64Builder).Append(9) + return b.NewArray() + }}, + {"map_date_key", func() arrow.Array { + b := array.NewMapBuilder(pool, arrow.FixedWidthTypes.Date32, arrow.PrimitiveTypes.Int64, false) + b.Append(true) + b.KeyBuilder().(*array.Date32Builder).Append(arrow.Date32FromTime(time.Date(2026, time.July, 9, 0, 0, 0, 0, time.UTC))) + b.ItemBuilder().(*array.Int64Builder).Append(9) + return b.NewArray() + }}, + {"nested_float32", func() arrow.Array { + dt := arrow.StructOf(arrow.Field{Name: "f", Type: arrow.PrimitiveTypes.Float32}) + b := array.NewStructBuilder(pool, dt) + b.Append(true) + b.FieldBuilder(0).(*array.Float32Builder).Append(0.1) + return b.NewArray() + }}, + {"nested_decimal", func() arrow.Array { + dt := arrow.StructOf(arrow.Field{Name: "d", Type: &arrow.Decimal128Type{Precision: 5, Scale: 2}}) + b := array.NewStructBuilder(pool, dt) + b.Append(true) + b.FieldBuilder(0).(*array.Decimal128Builder).Append(decimal128.FromU64(1999)) + return b.NewArray() + }}, + // Highest-drift shapes: recursive nesting, a nested timestamp leaf (the + // time.Time → quoted .String() special-case), and a null leaf inside a + // container (vs a null container, which is already covered). + {"array_of_struct", func() arrow.Array { + elem := arrow.StructOf(arrow.Field{Name: "a", Type: arrow.PrimitiveTypes.Int64}) + b := array.NewListBuilder(pool, elem) + sb := b.ValueBuilder().(*array.StructBuilder) + b.Append(true) + sb.Append(true) + sb.FieldBuilder(0).(*array.Int64Builder).Append(1) + sb.Append(true) + sb.FieldBuilder(0).(*array.Int64Builder).Append(2) + return b.NewArray() + }}, + {"struct_of_list", func() arrow.Array { + dt := arrow.StructOf(arrow.Field{Name: "xs", Type: arrow.ListOf(arrow.PrimitiveTypes.Int64)}) + b := array.NewStructBuilder(pool, dt) + b.Append(true) + lb := b.FieldBuilder(0).(*array.ListBuilder) + lb.Append(true) + vb := lb.ValueBuilder().(*array.Int64Builder) + vb.Append(1) + vb.Append(2) + return b.NewArray() + }}, + {"map_of_struct_value", func() arrow.Array { + valT := arrow.StructOf(arrow.Field{Name: "a", Type: arrow.PrimitiveTypes.Int64}) + b := array.NewMapBuilder(pool, arrow.BinaryTypes.String, valT, false) + b.Append(true) + b.KeyBuilder().(*array.StringBuilder).Append("k") + sb := b.ItemBuilder().(*array.StructBuilder) + sb.Append(true) + sb.FieldBuilder(0).(*array.Int64Builder).Append(5) + return b.NewArray() + }}, + {"nested_timestamp", func() arrow.Array { + dt := arrow.StructOf(arrow.Field{Name: "ts", Type: &arrow.TimestampType{Unit: arrow.Microsecond}}) + b := array.NewStructBuilder(pool, dt) + b.Append(true) + ts := time.Date(2026, time.July, 9, 12, 34, 56, 0, time.UTC) + b.FieldBuilder(0).(*array.TimestampBuilder).Append(arrow.Timestamp(ts.UnixMicro())) + return b.NewArray() + }}, + {"null_leaf_in_struct", func() arrow.Array { + dt := arrow.StructOf( + arrow.Field{Name: "a", Type: arrow.PrimitiveTypes.Int64}, + arrow.Field{Name: "b", Type: arrow.PrimitiveTypes.Int64}, + ) + b := array.NewStructBuilder(pool, dt) + b.Append(true) + b.FieldBuilder(0).(*array.Int64Builder).Append(1) + b.FieldBuilder(1).(*array.Int64Builder).AppendNull() + return b.NewArray() + }}, + // Multi-entry map: with only one entry the leading-comma (kernel) and + // trailing-comma (Thrift) separators coincide by accident; 2+ entries is what + // actually exercises the separator, and the per-row loop covers rows 1+. + {"map_multi_entry", func() arrow.Array { + b := array.NewMapBuilder(pool, arrow.BinaryTypes.String, arrow.PrimitiveTypes.Int64, false) + b.Append(true) + kb := b.KeyBuilder().(*array.StringBuilder) + ib := b.ItemBuilder().(*array.Int64Builder) + kb.Append("a") + ib.Append(1) + kb.Append("b") + ib.Append(2) + kb.Append("c") + ib.Append(3) + return b.NewArray() + }}, + // Multi-row list: rows 1+ must be indexed via the per-row value offsets, not + // offsets[0]. Different lengths per row so a wrong index is visible. + {"list_multi_row", func() arrow.Array { + b := array.NewListBuilder(pool, arrow.PrimitiveTypes.Int64) + vb := b.ValueBuilder().(*array.Int64Builder) + b.Append(true) + vb.Append(1) + vb.Append(2) + b.Append(true) + vb.Append(3) + b.Append(true) + vb.Append(4) + vb.Append(5) + vb.Append(6) + return b.NewArray() + }}, + // Empty-but-non-null collections: the row is present with zero elements, so + // both backends must render [] / {} (not null). Distinct from a null parent. + {"empty_list", func() arrow.Array { + b := array.NewListBuilder(pool, arrow.PrimitiveTypes.Int64) + b.Append(true) // present, no element appends + return b.NewArray() + }}, + {"empty_map", func() arrow.Array { + b := array.NewMapBuilder(pool, arrow.BinaryTypes.String, arrow.PrimitiveTypes.Int64, false) + b.Append(true) // present, no entries + return b.NewArray() + }}, + } + + // Sliced-array cases: a slice sets a non-zero logical offset, so an + // offset-unaware Offsets()[row] reads the wrong element range or panics. Build a + // 3-row array, then slice off the head so row 0 of the slice sits at a non-zero + // underlying offset — the direct regression guard for the ValueOffsets fix. + slicedCases := []struct { + name string + build func() arrow.Array + }{ + {"sliced_list", func() arrow.Array { + b := array.NewListBuilder(pool, arrow.PrimitiveTypes.Int64) + vb := b.ValueBuilder().(*array.Int64Builder) + for r := 0; r < 3; r++ { + b.Append(true) + vb.Append(int64(r * 10)) + vb.Append(int64(r*10 + 1)) + } + full := b.NewArray() + defer full.Release() + return array.NewSlice(full, 1, 3) // drop row 0 → offset != 0 + }}, + {"sliced_map", func() arrow.Array { + b := array.NewMapBuilder(pool, arrow.BinaryTypes.String, arrow.PrimitiveTypes.Int64, false) + kb := b.KeyBuilder().(*array.StringBuilder) + ib := b.ItemBuilder().(*array.Int64Builder) + for r := 0; r < 3; r++ { + b.Append(true) + kb.Append("k") + ib.Append(int64(r)) + kb.Append("j") + ib.Append(int64(r + 100)) + } + full := b.NewArray() + defer full.Release() + return array.NewSlice(full, 1, 3) + }}, + {"sliced_struct_of_list", func() arrow.Array { + dt := arrow.StructOf(arrow.Field{Name: "xs", Type: arrow.ListOf(arrow.PrimitiveTypes.Int64)}) + b := array.NewStructBuilder(pool, dt) + for r := 0; r < 3; r++ { + b.Append(true) + lb := b.FieldBuilder(0).(*array.ListBuilder) + lb.Append(true) + vb := lb.ValueBuilder().(*array.Int64Builder) + vb.Append(int64(r)) + vb.Append(int64(r + 1)) + } + full := b.NewArray() + defer full.Release() + return array.NewSlice(full, 1, 3) // struct offset propagates to the list field + }}, + } + cases = append(cases, slicedCases...) + + // A non-UTC location, so the timestamp/date rendering path (both backends + // apply .In(loc)) is actually exercised rather than a UTC no-op. + loc, err := time.LoadLocation("America/New_York") + if err != nil { + t.Skipf("tz database unavailable: %v", err) + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + arr := tc.build() + defer arr.Release() + + // Render EVERY row, not just row 0: multi-row cases exercise per-row + // offset indexing (rows 1+), which a single-row check can't reach — the + // exact gap that let the un-offset List/Map indexing bug through. + for row := 0; row < arr.Len(); row++ { + kernel, err := arrowscan.ScanCell(arr, row, loc) + if err != nil { + t.Fatalf("arrowscan.ScanCell row %d: %v", row, err) + } + thrift := renderViaArrowbased(t, arr, row, loc) + + if kernel != thrift { + t.Errorf("backend divergence for %s row %d:\n kernel = %#v\n thrift = %#v", tc.name, row, kernel, thrift) + } + } + }) + } +} + +// Top-level (non-nested) scalars: assert the kernel ScanCell and the Thrift +// container render the same Go value. DECIMAL is deliberately excluded here — see +// TestTopLevelDecimalRendering for why the two *container-level* paths differ by +// arrow type while the actual driver results agree. +func TestArrowbasedKernelTopLevelScalarParity(t *testing.T) { + pool := memory.NewGoAllocator() + loc := time.UTC + + cases := []struct { + name string + build func() arrow.Array + }{ + {"tinyint", func() arrow.Array { + b := array.NewInt8Builder(pool) + b.Append(1) + return b.NewArray() + }}, + {"smallint", func() arrow.Array { + b := array.NewInt16Builder(pool) + b.Append(2) + return b.NewArray() + }}, + {"int", func() arrow.Array { + b := array.NewInt32Builder(pool) + b.Append(3) + return b.NewArray() + }}, + {"int64", func() arrow.Array { + b := array.NewInt64Builder(pool) + b.Append(42) + return b.NewArray() + }}, + {"float32", func() arrow.Array { + b := array.NewFloat32Builder(pool) + b.Append(0.1) + return b.NewArray() + }}, + {"float64", func() arrow.Array { + b := array.NewFloat64Builder(pool) + b.Append(3.5) + return b.NewArray() + }}, + {"string", func() arrow.Array { + b := array.NewStringBuilder(pool) + b.Append("hi") + return b.NewArray() + }}, + {"timestamp", func() arrow.Array { + b := array.NewTimestampBuilder(pool, &arrow.TimestampType{Unit: arrow.Microsecond}) + b.Append(arrow.Timestamp(time.Date(2026, 7, 9, 12, 0, 0, 0, time.UTC).UnixMicro())) + return b.NewArray() + }}, + {"date32", func() arrow.Array { + b := array.NewDate32Builder(pool) + b.Append(arrow.Date32FromTime(time.Date(2026, 7, 9, 0, 0, 0, 0, time.UTC))) + return b.NewArray() + }}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + arr := tc.build() + defer arr.Release() + kernel, err := arrowscan.ScanCell(arr, 0, loc) + if err != nil { + t.Fatalf("arrowscan.ScanCell: %v", err) + } + thrift := renderViaArrowbased(t, arr, 0, loc) + if kernel != thrift { + t.Errorf("top-level %s divergence:\n kernel = %#v (%T)\n thrift = %#v (%T)", tc.name, kernel, kernel, thrift, thrift) + } + }) + } +} + +// TestTopLevelDecimalRendering documents the top-level DECIMAL story, which is +// subtler than "kernel string vs Thrift float64": +// +// - The kernel path always delivers DECIMAL as a native arrow decimal128 and +// renders it as an exact scale-applied string (arrowscan.ScanCell). +// - The Thrift *container* (decimal128Container.Value) converts a decimal128 to +// a lossy float64 — but that path is only reached when UseArrowNativeDecimal +// is on AND the value is read as a nested leaf. For a top-level column the +// scan uses DecimalStringValue (exact string) when the flag is on. +// - In the DEFAULT config the flag is off, so the server sends DECIMAL as a +// string column (DecimalAsArrow=false) — no decimal128 arrives at all, and the +// user gets the exact string. +// +// So across every real driver configuration a top-level DECIMAL comes back as the +// exact string on both backends (verified live). This test pins the kernel side +// (exact string) and the two Thrift container behaviors so the distinction can't +// silently drift into an actual result divergence. +func TestTopLevelDecimalRendering(t *testing.T) { + pool := memory.NewGoAllocator() + dt := &arrow.Decimal128Type{Precision: 38, Scale: 4} + b := array.NewDecimal128Builder(pool, dt) + defer b.Release() + n, _ := decimal128.FromString("1234567890123456.7890", dt.Precision, dt.Scale) + b.Append(n) + arr := b.NewArray() + defer arr.Release() + + // Kernel: exact string. + got, err := arrowscan.ScanCell(arr, 0, time.UTC) + if err != nil { + t.Fatal(err) + } + if got != "1234567890123456.7890" { + t.Errorf("kernel top-level decimal = %#v, want exact string", got) + } + + // Thrift container Value() is the lossy-float64 path (reached only with native + // decimal + nested read); DecimalStringValue is the exact top-level scan path. + // Assert both so a change to either is caught. + holder := &decimal128Container{scale: dt.Scale} + if err := holder.SetValueArray(arr.Data()); err != nil { + t.Fatal(err) + } + if s := holder.ValueString(0); s != "1234567890123456.7890" { + t.Errorf("Thrift DecimalStringValue = %q, want exact string", s) + } + // Value() is the lossy path: it returns a float64, and a 20-digit decimal + // cannot survive it — assert the actual rounded value, not just the type. + v, err := holder.Value(0) + if err != nil { + t.Fatal(err) + } + f, ok := v.(float64) + if !ok { + t.Fatalf("Thrift container Value() = %T, want the documented lossy float64", v) + } + if f != 1234567890123456.7890 { + t.Errorf("Thrift container Value() = %v, want the float64 rounding of the decimal", f) + } +} diff --git a/internal/rows/arrowbased/columnValues.go b/internal/rows/arrowbased/columnValues.go index 72e7f5f5..73aeb01f 100644 --- a/internal/rows/arrowbased/columnValues.go +++ b/internal/rows/arrowbased/columnValues.go @@ -2,13 +2,12 @@ package arrowbased import ( "encoding/json" - "math/big" "strings" "time" "github.com/apache/arrow/go/v12/arrow" "github.com/apache/arrow/go/v12/arrow/array" - "github.com/apache/arrow/go/v12/arrow/decimal128" + "github.com/databricks/databricks-sql-go/internal/decimalfmt" "github.com/databricks/databricks-sql-go/internal/rows/rowscanner" dbsqllog "github.com/databricks/databricks-sql-go/logger" "github.com/pkg/errors" @@ -363,7 +362,7 @@ func (mvc *mapValueContainer) SetValueArray(colData arrow.ArrayData) error { type structValueContainer struct { structArray *array.Struct - fieldNames []string + fieldKeys []string // precomputed JSON-escaped `"name":` prefixes, one per field complexValue []bool fieldValues []columnValues structArrayType *arrow.StructType @@ -371,11 +370,19 @@ type structValueContainer struct { var _ columnValues = (*structValueContainer)(nil) +// Value renders the struct at row i as a JSON object string. This grammar is +// mirrored by internal/arrowscan (the kernel backend); the two must render +// byte-identically — internal/rows/arrowbased/arrowscan_parity_test.go guards it. func (svc *structValueContainer) Value(i int) (any, error) { if i < svc.structArray.Len() { r := "{" for j := range svc.fieldValues { - r = r + "\"" + svc.fieldNames[j] + "\":" + // fieldKeys[j] is the precomputed JSON-escaped `"name":` prefix (built + // once at construction). Escaping matters — a raw `"` + name + `"` concat + // produces invalid JSON for a name with a quote/backslash/control char + // (`a"b` → `{"a"b":...}`) — but doing it per row would re-marshal the same + // constant strings N_rows × N_fields times on this hot path. + r = r + svc.fieldKeys[j] if svc.fieldValues[j].IsNull(int(i)) { r = r + "null" @@ -550,48 +557,9 @@ func (tvc *decimal128Container) Value(i int) (any, error) { // for backwards-compatible JSON rendering inside complex types. // See databricks/databricks-sql-go#274. func (tvc *decimal128Container) ValueString(i int) string { - return decimal128ToString(tvc.decimalArray.Value(i), tvc.scale) -} - -// decimal128ToString renders a decimal128 value exactly as a fixed-point -// string with `scale` fractional digits. It formats from the value's exact -// 128-bit unscaled integer (BigInt) and inserts the decimal point by string -// manipulation, so — unlike arrow's decimal128.Num.ToString, which divides via -// big.Float and rounds beyond ~17 significant digits — it never loses -// precision. See databricks/databricks-sql-go#274. -func decimal128ToString(n decimal128.Num, scale int32) string { - unscaled := n.BigInt() // exact signed unscaled integer - - neg := unscaled.Sign() < 0 - digits := new(big.Int).Abs(unscaled).String() - - var b strings.Builder - if neg { - b.WriteByte('-') - } - - if scale <= 0 { - // No fractional part (negative scale is not produced by the server, - // but treat it as scale 0 rather than panicking). - b.WriteString(digits) - return b.String() - } - - s := int(scale) - if len(digits) <= s { - // Pad with leading zeros so there are exactly `scale` fractional digits - // and a single leading integer zero, e.g. 5 with scale 3 -> "0.005". - b.WriteString("0.") - b.WriteString(strings.Repeat("0", s-len(digits))) - b.WriteString(digits) - } else { - intPart := digits[:len(digits)-s] - fracPart := digits[len(digits)-s:] - b.WriteString(intPart) - b.WriteByte('.') - b.WriteString(fracPart) - } - return b.String() + // Exact fixed-point rendering lives in internal/decimalfmt, shared with the + // kernel backend so both paths render DECIMAL identically. See #274. + return decimalfmt.ExactString(tvc.decimalArray.Value(i), tvc.scale) } func (tvc *decimal128Container) IsNull(i int) bool { diff --git a/kernel_backend.go b/kernel_backend.go new file mode 100644 index 00000000..412ffe78 --- /dev/null +++ b/kernel_backend.go @@ -0,0 +1,52 @@ +//go:build cgo && databricks_kernel + +package dbsql + +import ( + "context" + + "github.com/databricks/databricks-sql-go/internal/backend" + "github.com/databricks/databricks-sql-go/internal/backend/kernel" + "github.com/databricks/databricks-sql-go/internal/config" +) + +// newKernelBackend builds the SEA-via-kernel backend from the driver config; the +// connector opens the session right after, matching the Thrift path. It reads the +// same config fields Thrift does and translates them to the kernel's flat +// connection config, so the user-facing options are unchanged — only the routing +// differs. The public API adds nothing beyond WithUseKernel. +func newKernelBackend(_ context.Context, cfg *config.Config) (backend.Backend, error) { + // Reject options the kernel path can't honor yet + resolve the PAT. The + // validation is pure Go and lives in kernel_config.go (untagged) so its tests — + // including the exhaustiveness guard against a dropped Config field — run in the + // default CGO_ENABLED=0 build. + token, err := validateKernelConfig(cfg) + if err != nil { + return nil, err + } + + kc := kernel.Config{ + Host: cfg.Host, + HTTPPath: cfg.HTTPPath, + WarehouseID: cfg.WarehouseID, + Token: token, + Location: cfg.Location, + // Session confs (STATEMENT_TIMEOUT, QUERY_TAGS, TIMEZONE, …) — the same + // SessionParams map the Thrift backend forwards, so they flow to the + // server identically with no per-backend translation. SPOG org routing + // rides in HTTPPath's ?o= and is parsed kernel-side. + SessionConf: cfg.SessionParams, + } + // TLS: the driver honors TLSConfig only for InsecureSkipVerify (see + // internal/client), so map exactly that knob to the kernel. + if cfg.TLSConfig != nil && cfg.TLSConfig.InsecureSkipVerify { + kc.TLSSkipVerify = true + } + // Proxy: the Thrift path uses http.ProxyFromEnvironment; mirror it by reading + // the same HTTP(S)_PROXY / NO_PROXY environment for the kernel. + kc.ProxyURL = proxyForEndpoint(cfg) + return kernel.New(kc), nil +} + +// proxyForEndpoint (pure Go, no kernel dependency) lives in kernel_proxy.go so +// its test runs in the default CGO_ENABLED=0 build. diff --git a/kernel_backend_test.go b/kernel_backend_test.go new file mode 100644 index 00000000..eb3244c4 --- /dev/null +++ b/kernel_backend_test.go @@ -0,0 +1,40 @@ +//go:build cgo && databricks_kernel + +package dbsql + +import ( + "context" + "testing" + + "github.com/databricks/databricks-sql-go/internal/config" +) + +// newKernelBackend builds a backend for a supported config and propagates a +// validation error otherwise. The exhaustive per-option reject/auth assertions +// live in the untagged TestValidateKernelConfig (kernel_config_test.go), which +// runs in the default CGO_ENABLED=0 build; this tagged smoke test just confirms +// the tagged wrapper wires validateKernelConfig + the kernel.Config assembly. +func TestNewKernelBackend(t *testing.T) { + base := func() *config.Config { + c := config.WithDefaults() + c.Host = "h.databricks.com" + c.Port = 443 + c.HTTPPath = "/sql/1.0/warehouses/abc" + c.AccessToken = "dapi-x" + return c + } + + t.Run("supported config builds", func(t *testing.T) { + if _, err := newKernelBackend(context.Background(), base()); err != nil { + t.Errorf("a supported config should build cleanly, got %v", err) + } + }) + + t.Run("validation error propagates", func(t *testing.T) { + c := base() + c.Catalog = "main" // rejected by validateKernelConfig + if _, err := newKernelBackend(context.Background(), c); err == nil { + t.Error("newKernelBackend should propagate the validation error") + } + }) +} diff --git a/kernel_config.go b/kernel_config.go new file mode 100644 index 00000000..ebc7c059 --- /dev/null +++ b/kernel_config.go @@ -0,0 +1,111 @@ +package dbsql + +import ( + "fmt" + + "github.com/databricks/databricks-sql-go/auth/noop" + "github.com/databricks/databricks-sql-go/auth/pat" + dbsqlerr "github.com/databricks/databricks-sql-go/errors" + "github.com/databricks/databricks-sql-go/internal/config" +) + +// This file is intentionally NOT behind the `cgo && databricks_kernel` build tag. +// The kernel backend's option-validation is pure Go (it reads config.Config and +// returns an error or a resolved PAT), so keeping it untagged lets its tests — +// including the reflective exhaustiveness check that guards against a future +// Config field being silently dropped — run under CGO_ENABLED=0. The tagged +// newKernelBackend calls validateKernelConfig, then assembles the cgo kernel.Config. + +// validateKernelConfig enforces the kernel backend's "nothing silently ignored" +// contract: it rejects every option the kernel path can't yet honor with a clear +// error (rather than dropping it, which would behave differently than Thrift) and +// resolves the PAT the kernel authenticates with. On success it returns the token +// to use. Options it does NOT reject are either forwarded by newKernelBackend or +// intentionally accepted-but-inert (documented in doc.go and asserted by +// TestKernelConfigFieldsClassified). +// +// Every rejection wraps errors.ErrNotSupportedByKernel so a caller can detect the +// "kernel can't honor this option" case with errors.Is (e.g. to fall back to the +// default backend) instead of matching on message text. +func validateKernelConfig(cfg *config.Config) (token string, err error) { + // Initial namespace (WithInitialNamespace): no kernel C-ABI setter yet, so the + // session would run in the default namespace and unqualified names would + // resolve differently than Thrift. + if cfg.Catalog != "" || cfg.Schema != "" { + return "", fmt.Errorf("databricks: WithInitialNamespace (catalog/schema) is %w; "+ + "omit it or use the default (Thrift) backend", dbsqlerr.ErrNotSupportedByKernel) + } + // EnableMetricViewMetadata: maps to a server session conf we want to route + // backend-neutrally rather than duplicate here. + if cfg.EnableMetricViewMetadata { + return "", fmt.Errorf("databricks: WithEnableMetricViewMetadata is %w; "+ + "omit it or use the default (Thrift) backend", dbsqlerr.ErrNotSupportedByKernel) + } + // Port / Protocol: the kernel C ABI takes only a bare host and connects over + // https:443; it has no port or scheme setter. The Thrift path honors a custom + // port/scheme via ToEndpointURL, so a non-default value here would be silently + // ignored on the kernel path (it would just hit 443) — reject it instead, per + // the "nothing silently ignored" contract. Defaults (https/443) are fine. + if cfg.Protocol != "" && cfg.Protocol != "https" { + return "", fmt.Errorf("databricks: a non-https protocol is %w "+ + "(it connects over https); use the default (Thrift) backend", dbsqlerr.ErrNotSupportedByKernel) + } + if cfg.Port != 0 && cfg.Port != 443 { + return "", fmt.Errorf("databricks: a non-default port (WithPort) is %w "+ + "(it connects on 443); omit it or use the default (Thrift) backend", dbsqlerr.ErrNotSupportedByKernel) + } + // Transport (WithTransport, a custom http.RoundTripper carrying a custom CA + // bundle / mTLS / proxy): the kernel uses its own Rust HTTP stack below the C + // ABI and never sees a Go RoundTripper, so a custom Transport would be silently + // ignored. Reject it per the "nothing silently ignored" contract. (The kernel + // does honor HTTPS_PROXY and InsecureSkipVerify through their own mappings; only + // a wholesale custom Transport is unsupported.) + if cfg.Transport != nil { + return "", fmt.Errorf("databricks: a custom WithTransport (RoundTripper) is %w "+ + "(the kernel uses its own HTTP stack); use the default (Thrift) backend", dbsqlerr.ErrNotSupportedByKernel) + } + // Auth: the kernel backend authenticates with a PAT only. Any other + // authenticator sets cfg.Authenticator but leaves cfg.AccessToken empty, so an + // empty PAT would reach the kernel and fail with an opaque Unauthenticated + // error. Reject it here so the failure names the cause. + token = cfg.AccessToken + switch a := cfg.Authenticator.(type) { + case nil, *noop.NoopAuth: + // No explicit authenticator — token comes from cfg.AccessToken (may be + // empty; caught below). + case *pat.PATAuth: + // WithAccessToken sets both cfg.AccessToken and this authenticator, but + // WithAuthenticator(&pat.PATAuth{...}) sets only the authenticator and leaves + // cfg.AccessToken empty. Take the token from the authenticator when + // cfg.AccessToken didn't carry it, so both PAT paths work. + if token == "" { + token = a.AccessToken + } + default: + return "", fmt.Errorf("databricks: only personal access token (WithAccessToken) auth is supported by the kernel backend; "+ + "OAuth (M2M/U2M), token-provider, external/static, and federated authenticators are %w — "+ + "use PAT or the default (Thrift) backend", dbsqlerr.ErrNotSupportedByKernel) + } + if token == "" { + // Missing required config (not an unsupported-feature rejection), so this is + // intentionally NOT wrapped with ErrNotSupportedByKernel. + return "", fmt.Errorf("databricks: the kernel backend requires a personal access token; " + + "set one with WithAccessToken (or a *pat.PATAuth via WithAuthenticator)") + } + // WithTimeout maps to a per-statement server timeout on Thrift + // (TExecuteStatementReq.QueryTimeout); the kernel C ABI exposes no equivalent, + // so reject it rather than run the query with no server-side timeout. + if cfg.QueryTimeout > 0 { + return "", fmt.Errorf("databricks: WithTimeout (server query timeout) is %w; "+ + "omit it or use the default (Thrift) backend", dbsqlerr.ErrNotSupportedByKernel) + } + // WithRetries(-1) explicitly disables retries, but the kernel retries + // internally below the C ABI with no user-facing toggle — so a disable request + // would be silently violated. Reject it. Positive/default RetryMax is fine: the + // kernel provides retries (just not user-tunable), documented in doc.go. + if cfg.RetryMax < 0 { + return "", fmt.Errorf("databricks: disabling retries via WithRetries is %w "+ + "(the kernel retries internally); omit it or use the default (Thrift) backend", dbsqlerr.ErrNotSupportedByKernel) + } + return token, nil +} diff --git a/kernel_config_test.go b/kernel_config_test.go new file mode 100644 index 00000000..1f3a0649 --- /dev/null +++ b/kernel_config_test.go @@ -0,0 +1,207 @@ +package dbsql + +import ( + "errors" + "net/http" + "reflect" + "testing" + "time" + + "github.com/databricks/databricks-sql-go/auth/pat" + dbsqlerr "github.com/databricks/databricks-sql-go/errors" + "github.com/databricks/databricks-sql-go/internal/config" +) + +// nonPATAuth stands in for any non-PAT authenticator (OAuth / token-provider / +// external / federated) — the kernel backend must reject it. +type nonPATAuth struct{} + +func (nonPATAuth) Authenticate(*http.Request) error { return nil } + +func baseKernelConfig() *config.Config { + c := config.WithDefaults() + c.Host = "h.databricks.com" + c.Port = 443 + c.HTTPPath = "/sql/1.0/warehouses/abc" + c.AccessToken = "dapi-x" + return c +} + +// validateKernelConfig enforces the kernel backend's "nothing silently ignored" +// contract: unsupported options are rejected loudly and the PAT is resolved. This +// is pure Go (no cgo), so these run in the default CGO_ENABLED=0 build. +func TestValidateKernelConfig(t *testing.T) { + t.Run("supported config ok", func(t *testing.T) { + c := baseKernelConfig() + c.SessionParams = map[string]string{"QUERY_TAGS": "a:1"} + if _, err := validateKernelConfig(c); err != nil { + t.Errorf("a supported config should validate, got %v", err) + } + }) + + // Every rejection must (a) error and (b) wrap ErrNotSupportedByKernel, since + // that sentinel is the documented programmatic fallback-detection contract — + // asserting only err != nil would let a dropped or malformed %w wrap ship green. + // Table-driven so a new rejection is covered by adding one row. + rejections := []struct { + name string + mut func(*config.Config) + }{ + {"catalog", func(c *config.Config) { c.Catalog = "main" }}, + {"schema", func(c *config.Config) { c.Schema = "sys" }}, + {"metric view", func(c *config.Config) { c.EnableMetricViewMetadata = true }}, + {"non-PAT authenticator", func(c *config.Config) { c.Authenticator = nonPATAuth{} }}, + {"query timeout", func(c *config.Config) { c.QueryTimeout = 30 * time.Second }}, + {"disable retries", func(c *config.Config) { c.RetryMax = -1 }}, + {"non-https protocol", func(c *config.Config) { c.Protocol = "http" }}, + {"non-default port", func(c *config.Config) { c.Port = 8443 }}, + {"custom transport", func(c *config.Config) { c.Transport = http.DefaultTransport }}, + } + for _, tc := range rejections { + t.Run(tc.name+" rejected", func(t *testing.T) { + c := baseKernelConfig() + tc.mut(c) + _, err := validateKernelConfig(c) + if err == nil { + t.Fatalf("expected an error when %s is set", tc.name) + } + if !errors.Is(err, dbsqlerr.ErrNotSupportedByKernel) { + t.Errorf("%s rejection should wrap ErrNotSupportedByKernel, got %v", tc.name, err) + } + }) + } + + t.Run("PAT via WithAuthenticator resolves the token", func(t *testing.T) { + c := baseKernelConfig() + c.AccessToken = "" + c.Authenticator = &pat.PATAuth{AccessToken: "dapi-y"} + tok, err := validateKernelConfig(c) + if err != nil { + t.Fatalf("PAT via WithAuthenticator should validate, got %v", err) + } + if tok != "dapi-y" { + t.Errorf("token = %q, want dapi-y (sourced from the authenticator)", tok) + } + }) + + t.Run("empty token rejected", func(t *testing.T) { + // Missing-required-config, NOT an unsupported-feature rejection, so this is + // intentionally NOT expected to wrap ErrNotSupportedByKernel. + c := baseKernelConfig() + c.AccessToken = "" + c.Authenticator = &pat.PATAuth{AccessToken: ""} + if _, err := validateKernelConfig(c); err == nil { + t.Error("expected an error when the resolved PAT is empty") + } + }) + + t.Run("positive retry tuning + maxrows accepted", func(t *testing.T) { + c := baseKernelConfig() + c.RetryMax = 8 + c.MaxRows = 5000 + if _, err := validateKernelConfig(c); err != nil { + t.Errorf("positive retry/maxrows tuning should validate, got %v", err) + } + }) + + t.Run("default https/443 accepted", func(t *testing.T) { + c := baseKernelConfig() // WithDefaults sets Protocol=https, Port=443 + if _, err := validateKernelConfig(c); err != nil { + t.Errorf("the default https/443 endpoint should validate, got %v", err) + } + }) +} + +// kernelConfigFieldDisposition records, for every UserConfig field, how the kernel +// backend treats it. UserConfig is the user-facing option surface ("Only +// UserConfig are currently exposed to users"). Adding a field to config.UserConfig +// without classifying it here fails TestKernelConfigFieldsClassified — forcing a +// deliberate decision (forward it, reject it loudly, or accept it as inert) rather +// than silently dropping it on the kernel path. (TLSConfig and ArrowConfig live on +// the outer config.Config, not +// UserConfig; newKernelBackend reads TLSConfig.InsecureSkipVerify explicitly, and +// the kernel renders decimals exactly regardless of ArrowConfig.) +var kernelConfigFieldDisposition = map[string]string{ + // Forwarded to kernel.Config (see newKernelBackend). + "Host": "forwarded", + "HTTPPath": "forwarded", + "WarehouseID": "forwarded", + "AccessToken": "forwarded", // as the resolved PAT (kc.Token) + "Authenticator": "forwarded", // PAT authenticator resolved to the token + "Location": "forwarded", + "SessionParams": "forwarded", + "UseKernel": "forwarded", // the routing flag itself + + // Rejected loudly by validateKernelConfig. + "Catalog": "rejected", + "Schema": "rejected", + "EnableMetricViewMetadata": "rejected", + "QueryTimeout": "rejected", // when > 0 (WithTimeout) + "RetryMax": "rejected", // when < 0 (disable retries) + "Protocol": "rejected", // kernel is https-only; non-default rejected + "Port": "rejected", // kernel connects on 443; non-default rejected + "Transport": "rejected", // custom RoundTripper; kernel uses its own HTTP stack, so reject rather than drop + + // Accepted but intentionally inert on the kernel path (documented in doc.go): + // the kernel manages these internally, below the C ABI, with no user knob. + "MaxRows": "inert", + "RetryWaitMin": "inert", + "RetryWaitMax": "inert", + "UseLz4Compression": "inert", // kernel negotiates compression internally + + // Not applicable to the kernel path (Thrift/HTTP-transport or telemetry knobs + // that don't reach the kernel binding). + "UserAgentEntry": "inert", // TODO: forward once the kernel exposes a UA setter + "EnableTelemetry": "inert", + "TelemetryBatchSize": "inert", + "TelemetryFlushInterval": "inert", + "UseArrowNativeDecimalDSN": "inert", // DSN carrier; kernel renders decimals exactly regardless + + // Fields promoted from the embedded CloudFetchConfig. The kernel does + // CloudFetch internally (below the C ABI), so none is forwarded — but each is + // classified individually so a new CloudFetch option can't slip the guard. + "UseCloudFetch": "inert", + "MaxDownloadThreads": "inert", + "MaxFilesInMemory": "inert", + "MinTimeToExpiry": "inert", + "CloudFetchSpeedThresholdMbps": "inert", + "HTTPClient": "inert", +} + +// kernelConfigClassifiedNames returns every UserConfig field name the kernel gate +// must account for, flattening anonymous embedded structs (e.g. CloudFetchConfig) +// into their promoted fields — reflect's NumField reports an embed as one field, +// which would hide new sub-fields from the drop guard below. +func kernelConfigClassifiedNames(t reflect.Type) []string { + var names []string + for i := 0; i < t.NumField(); i++ { + f := t.Field(i) + if f.Anonymous && f.Type.Kind() == reflect.Struct { + names = append(names, kernelConfigClassifiedNames(f.Type)...) + continue + } + names = append(names, f.Name) + } + return names +} + +func TestKernelConfigFieldsClassified(t *testing.T) { + names := kernelConfigClassifiedNames(reflect.TypeOf(config.UserConfig{})) + classified := make(map[string]bool, len(names)) + for _, name := range names { + classified[name] = true + if _, ok := kernelConfigFieldDisposition[name]; !ok { + t.Errorf("config.UserConfig field %q (incl. promoted embed fields) is not classified "+ + "for the kernel backend. Add it to kernelConfigFieldDisposition as "+ + "forwarded/rejected/inert (and wire it in validateKernelConfig / newKernelBackend if "+ + "it must be honored) so it isn't silently dropped on the kernel path.", name) + } + } + // Guard the reverse too: a disposition entry for a field that no longer exists + // is stale. + for name := range kernelConfigFieldDisposition { + if !classified[name] { + t.Errorf("kernelConfigFieldDisposition has %q but config.UserConfig (incl. embeds) no longer does; remove it", name) + } + } +} diff --git a/kernel_e2e_test.go b/kernel_e2e_test.go new file mode 100644 index 00000000..7d847bb4 --- /dev/null +++ b/kernel_e2e_test.go @@ -0,0 +1,267 @@ +//go:build cgo && databricks_kernel + +package dbsql + +import ( + "bytes" + "context" + "database/sql" + "database/sql/driver" + "errors" + "os" + "testing" + "time" +) + +// kernelTestDB opens a kernel-backed *sql.DB from DATABRICKS_HOST / +// DATABRICKS_HTTP_PATH / DATABRICKS_TOKEN, or skips when they are unset. It goes +// through the standard connector with WithUseKernel(true) — the same path a real +// consumer uses — not a kernel-only connector. +func kernelTestDB(t *testing.T) *sql.DB { + t.Helper() + return kernelTestDBWith(t) +} + +// TestKernelE2ESelect1 is the smallest end-to-end proof: PAT session over the +// kernel, execute, scan one scalar row. +func TestKernelE2ESelect1(t *testing.T) { + db := kernelTestDB(t) + defer db.Close() + + var got int64 + if err := db.QueryRowContext(context.Background(), "SELECT 1").Scan(&got); err != nil { + t.Fatalf("query: %v", err) + } + if got != 1 { + t.Errorf("SELECT 1 = %d, want 1", got) + } +} + +// kernelTestDBWith opens a kernel-backed *sql.DB with extra connector options on +// top of the base host/path/PAT, or skips when creds are unset. It is the config +// counterpart to kernelTestDB. +func kernelTestDBWith(t *testing.T, extra ...ConnOption) *sql.DB { + t.Helper() + host := os.Getenv("DATABRICKS_HOST") + httpPath := os.Getenv("DATABRICKS_HTTP_PATH") + token := os.Getenv("DATABRICKS_TOKEN") + if host == "" || httpPath == "" || token == "" { + t.Skip("set DATABRICKS_HOST / DATABRICKS_HTTP_PATH / DATABRICKS_TOKEN for the kernel e2e") + } + opts := append([]ConnOption{ + WithServerHostname(host), + WithHTTPPath(httpPath), + WithAccessToken(token), + WithUseKernel(true), + }, extra...) + connector, err := NewConnector(opts...) + if err != nil { + t.Fatalf("NewConnector: %v", err) + } + return sql.OpenDB(connector) +} + +// TestKernelE2EQueryTags proves session confs reach the server: WithQueryTags +// (the same option the Thrift path uses) is routed to the kernel and read back +// via SET, which echoes each tag by key. +func TestKernelE2EQueryTags(t *testing.T) { + db := kernelTestDBWith(t, WithQueryTags(map[string]string{"team": "peco"})) + defer db.Close() + + var key, val string + if err := db.QueryRowContext(context.Background(), "SET query_tags").Scan(&key, &val); err != nil { + t.Fatalf("SET query_tags: %v", err) + } + if key != "team" || val != "peco" { + t.Errorf("query tag read back as %q=%q, want team=peco", key, val) + } +} + +// TestKernelE2EStatementTimeout proves a STATEMENT_TIMEOUT session param (via +// WithSessionParams) is applied on the kernel session and read back via SET. +func TestKernelE2EStatementTimeout(t *testing.T) { + db := kernelTestDBWith(t, WithSessionParams(map[string]string{"STATEMENT_TIMEOUT": "300"})) + defer db.Close() + + var key, val string + if err := db.QueryRowContext(context.Background(), "SET statement_timeout").Scan(&key, &val); err != nil { + t.Fatalf("SET statement_timeout: %v", err) + } + if val != "300" { + t.Errorf("statement_timeout read back as %q, want 300", val) + } +} + +// TestKernelE2ETimeZone proves the session time zone (WithSessionParams +// timezone) is applied to scanned TIMESTAMP values, matching the Thrift path — +// the returned time.Time carries the configured location, not UTC. +func TestKernelE2ETimeZone(t *testing.T) { + const tz = "America/New_York" + db := kernelTestDBWith(t, WithSessionParams(map[string]string{"timezone": tz})) + defer db.Close() + + var ts time.Time + if err := db.QueryRowContext(context.Background(), + "SELECT CAST('2026-07-09 12:00:00' AS TIMESTAMP)").Scan(&ts); err != nil { + t.Fatalf("query: %v", err) + } + if ts.Location().String() != tz { + t.Errorf("timestamp location = %q, want %q", ts.Location(), tz) + } +} + +// TestKernelE2ETLSSkipVerify checks that WithSkipTLSHostVerify (a relaxation +// knob) is accepted on the kernel path; the connection must still succeed +// against the warehouse's valid certificate. +func TestKernelE2ETLSSkipVerify(t *testing.T) { + db := kernelTestDBWith(t, WithSkipTLSHostVerify()) + defer db.Close() + + var got int64 + if err := db.QueryRowContext(context.Background(), "SELECT 1").Scan(&got); err != nil { + t.Fatalf("query with TLS skip-verify: %v", err) + } + if got != 1 { + t.Errorf("SELECT 1 = %d, want 1", got) + } +} + +// TestKernelE2EDataTypes scans each supported scalar type in its own subtest, so +// a failure names the exact type rather than being masked by others in a shared +// row. Each case selects a single value and compares the scanned result. NULL is +// covered as its own case. +func TestKernelE2EDataTypes(t *testing.T) { + db := kernelTestDB(t) + defer db.Close() + + cases := []struct { + name string + expr string // the single SELECT expression + want driver.Value // expected scanned value (nil for SQL NULL) + }{ + // Integer widths scan to their native Go type (int8/16/32/64), matching the + // Thrift backend. (Both are technically off the driver.Value spec, which + // names only int64; unifying both backends on int64 is a tracked follow-up.) + {"bigint", "CAST(42 AS BIGINT)", int64(42)}, + {"int", "CAST(7 AS INT)", int32(7)}, + {"smallint", "CAST(3 AS SMALLINT)", int16(3)}, + {"tinyint", "CAST(1 AS TINYINT)", int8(1)}, + {"double", "CAST(3.5 AS DOUBLE)", float64(3.5)}, + // 0.1 is NOT exactly representable, so a float32-vs-float64 mismatch would + // surface here (float64(float32(0.1)) != float32(0.1)); a bare FLOAT must + // scan to a native float32, matching Thrift. + {"float", "CAST(0.1 AS FLOAT)", float32(0.1)}, + {"boolean", "true", true}, + {"string", "'hi'", "hi"}, + {"binary", "CAST('abc' AS BINARY)", []byte("abc")}, + {"decimal_exact", "CAST(1.25 AS DECIMAL(5,2))", "1.25"}, + {"date", "CAST('2026-07-09' AS DATE)", time.Date(2026, time.July, 9, 0, 0, 0, 0, time.UTC)}, + {"timestamp", "CAST('2026-07-09 12:34:56' AS TIMESTAMP)", time.Date(2026, time.July, 9, 12, 34, 56, 0, time.UTC)}, + {"null", "CAST(NULL AS STRING)", nil}, + // Nested types render to a JSON string; VARIANT arrives nested, GEOMETRY + // as a WKT/WKB string. + {"array", "array(1, 2, 3)", "[1,2,3]"}, + {"map", "map('k', 9)", `{"k":9}`}, + {"struct", "named_struct('a', 1, 'b', 'x')", `{"a":1,"b":"x"}`}, + {"variant", `parse_json('{"a":1,"b":[2,3]}')`, `{"a":1,"b":[2,3]}`}, + {"geometry", "st_point(1, 2)", "POINT(1 2)"}, + } + + for _, c := range cases { + t.Run(c.name, func(t *testing.T) { + var got any + err := db.QueryRowContext(context.Background(), "SELECT "+c.expr).Scan(&got) + if err != nil { + t.Fatalf("scan %s: %v", c.expr, err) + } + if !dataTypeEqual(got, c.want) { + t.Errorf("%s = %#v (%T), want %#v (%T)", c.expr, got, got, c.want, c.want) + } + }) + } +} + +// dataTypeEqual compares scanned values, handling the two non-comparable cases: +// []byte (bytes.Equal) and time.Time (Equal, which is instant-based and ignores +// the location the value was materialized in). +func dataTypeEqual(got, want driver.Value) bool { + switch w := want.(type) { + case nil: + return got == nil + case []byte: + g, ok := got.([]byte) + return ok && bytes.Equal(g, w) + case time.Time: + g, ok := got.(time.Time) + return ok && g.Equal(w) + default: + return got == want + } +} + +// TestKernelE2ECloudFetch streams a CloudFetch-sized result end to end. CloudFetch +// is internal to the kernel, so "it works" means many batches stream and scan +// correctly — which also exercises the per-batch release/lifetime path. +func TestKernelE2ECloudFetch(t *testing.T) { + db := kernelTestDB(t) + defer db.Close() + + const want = 1_000_000 + rows, err := db.QueryContext(context.Background(), + "SELECT id FROM range(0, 1000000)") + if err != nil { + t.Fatalf("query: %v", err) + } + defer rows.Close() + + var count, last int64 + for rows.Next() { + var id int64 + if err := rows.Scan(&id); err != nil { + t.Fatalf("scan at row %d: %v", count, err) + } + count++ + last = id + } + if err := rows.Err(); err != nil { + t.Fatalf("iteration: %v", err) + } + if count != want { + t.Errorf("row count = %d, want %d", count, want) + } + if last != want-1 { + t.Errorf("last id = %d, want %d", last, want-1) + } +} + +// TestKernelE2ECancellation cancels a long-running query via ctx and asserts it +// returns well before its uncancelled runtime. +func TestKernelE2ECancellation(t *testing.T) { + db := kernelTestDB(t) + defer db.Close() + + ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + defer cancel() + + start := time.Now() + // A query that would run far longer than the 3s deadline. + _, err := db.QueryContext(ctx, "SELECT count(*) FROM range(0, 100000000000) WHERE id % 7 = 0") + elapsed := time.Since(start) + + if err == nil { + t.Fatal("expected a cancellation error, got nil") + } + // It must be the deadline firing, not an unrelated failure (syntax error, + // transient network, server-side timeout) — otherwise the test would pass + // without proving cancellation works. + if !errors.Is(err, context.DeadlineExceeded) { + t.Errorf("expected context.DeadlineExceeded, got %v", err) + } + // Well under the query's natural runtime and not far past the 3s deadline; + // a no-op cancel that only returned when the query finished would blow past + // this. + if elapsed > 10*time.Second { + t.Errorf("cancellation took %v; expected it to abandon close to the 3s deadline", elapsed) + } + t.Logf("cancelled after %v with err=%v", elapsed, err) +} diff --git a/kernel_parity_test.go b/kernel_parity_test.go new file mode 100644 index 00000000..8f570302 --- /dev/null +++ b/kernel_parity_test.go @@ -0,0 +1,102 @@ +//go:build cgo && databricks_kernel + +package dbsql + +import ( + "context" + "database/sql" + "os" + "testing" +) + +// thriftTestDB opens the same warehouse over the default Thrift backend, for +// parity comparison against the kernel backend. +func thriftTestDB(t *testing.T) *sql.DB { + t.Helper() + host := os.Getenv("DATABRICKS_HOST") + httpPath := os.Getenv("DATABRICKS_HTTP_PATH") + token := os.Getenv("DATABRICKS_TOKEN") + if host == "" || httpPath == "" || token == "" { + t.Skip("set DATABRICKS_HOST / DATABRICKS_HTTP_PATH / DATABRICKS_TOKEN for the parity test") + } + connector, err := NewConnector( + WithServerHostname(host), + WithHTTPPath(httpPath), + WithAccessToken(token), + ) + if err != nil { + t.Fatalf("NewConnector (thrift): %v", err) + } + return sql.OpenDB(connector) +} + +// TestKernelThriftParity runs the same scalar query through both backends and +// asserts identical row output. The two scanners are intentionally separate, so +// this golden comparison is the guarantee that they render the scalar types +// equivalently. +func TestKernelThriftParity(t *testing.T) { + const query = "SELECT CAST(7 AS BIGINT), CAST(2.5 AS DOUBLE), 'parity', " + + "CAST(NULL AS STRING), CAST(9.99 AS DECIMAL(5,2)), true, " + + // A bare FLOAT at a non-exactly-representable value: catches the + // float32-vs-widened-float64 divergence (0.1 vs 0.10000000149011612). + "CAST(0.1 AS FLOAT), " + + "array(1, 2, 3), map('k', 9), named_struct('a', 1, 'b', 'x'), " + + `parse_json('{"a":1,"b":[2,3]}'), st_point(1, 2), ` + + // A decimal inside a struct: exercises exact-string nested-decimal + // rendering (19.99, not a lossy 19.990000000000002). + "named_struct('d', CAST(19.99 AS DECIMAL(5,2)))" + + kernelDB := kernelTestDB(t) + defer kernelDB.Close() + thriftDB := thriftTestDB(t) + defer thriftDB.Close() + + kernelRow := scanOneRowAsStrings(t, kernelDB, query) + thriftRow := scanOneRowAsStrings(t, thriftDB, query) + + if len(kernelRow) != len(thriftRow) { + t.Fatalf("column count differs: kernel=%d thrift=%d", len(kernelRow), len(thriftRow)) + } + for i := range kernelRow { + if kernelRow[i] != thriftRow[i] { + t.Errorf("col %d differs: kernel=%q thrift=%q", i, kernelRow[i], thriftRow[i]) + } + } +} + +// scanOneRowAsStrings scans the first row into a []string via sql.RawBytes, so a +// NULL renders as "" and every value is compared in its wire form, +// independent of Go-type coercion differences between the backends. +func scanOneRowAsStrings(t *testing.T, db *sql.DB, query string) []string { + t.Helper() + rows, err := db.QueryContext(context.Background(), query) + if err != nil { + t.Fatalf("query: %v", err) + } + defer rows.Close() + + cols, err := rows.Columns() + if err != nil { + t.Fatalf("columns: %v", err) + } + if !rows.Next() { + t.Fatalf("no row: %v", rows.Err()) + } + raw := make([]sql.RawBytes, len(cols)) + dest := make([]any, len(cols)) + for i := range raw { + dest[i] = &raw[i] + } + if err := rows.Scan(dest...); err != nil { + t.Fatalf("scan: %v", err) + } + out := make([]string, len(cols)) + for i, b := range raw { + if b == nil { + out[i] = "" + } else { + out[i] = string(b) + } + } + return out +} diff --git a/kernel_proxy.go b/kernel_proxy.go new file mode 100644 index 00000000..9975e0ab --- /dev/null +++ b/kernel_proxy.go @@ -0,0 +1,59 @@ +package dbsql + +import ( + "fmt" + "net/http" + "net/url" + + "github.com/databricks/databricks-sql-go/internal/config" +) + +// This file is intentionally NOT behind the `cgo && databricks_kernel` build tag: +// proxy resolution is pure Go (no kernel C symbol), so keeping it in the default +// build lets its test run under CGO_ENABLED=0 rather than being dead behind the +// tag. newKernelBackend (tagged) calls proxyForEndpoint. + +// proxyForEndpoint resolves the proxy the Thrift path would use for this +// connection, via the same http.ProxyFromEnvironment (HTTP(S)_PROXY / NO_PROXY) +// the Thrift transport applies at request time. Building the endpoint request +// lets ProxyFromEnvironment apply the NO_PROXY rules for this exact host, so the +// kernel sees the same effective proxy decision — returning "" (direct) when +// NO_PROXY excludes the host or no proxy is set. No extra dependency: this is the +// stdlib function the driver already relies on. +func proxyForEndpoint(cfg *config.Config) string { + return proxyForEndpointFunc(cfg, http.ProxyFromEnvironment) +} + +// proxyForEndpointFunc is the testable core: it builds the endpoint request and +// asks resolve for the proxy, returning "" (direct) on any error, no proxy, or an +// unbuildable endpoint. resolve is http.ProxyFromEnvironment in production; tests +// inject a deterministic resolver to exercise proxy-set / NO_PROXY / direct +// without depending on http.ProxyFromEnvironment's process-wide env caching. +func proxyForEndpointFunc(cfg *config.Config, resolve func(*http.Request) (*url.URL, error)) string { + // Build the endpoint from scheme/host/port directly rather than via + // cfg.ToEndpointURL(): that requires a non-empty HTTPPath and errors in the + // warehouse-id addressing mode the kernel prefers (HTTPPath == ""), which would + // swallow the error and silently return "" (direct) — ignoring HTTPS_PROXY. + // Only scheme+host matter for the proxy/NO_PROXY decision; the path is + // irrelevant, so omitting it is correct here. + if cfg.Host == "" { + return "" + } + scheme := cfg.Protocol + if scheme == "" { + scheme = "https" + } + endpoint := fmt.Sprintf("%s://%s", scheme, cfg.Host) + if cfg.Port != 0 { + endpoint = fmt.Sprintf("%s://%s:%d", scheme, cfg.Host, cfg.Port) + } + req, err := http.NewRequest(http.MethodPost, endpoint, nil) + if err != nil { + return "" + } + proxyURL, err := resolve(req) + if err != nil || proxyURL == nil { + return "" + } + return proxyURL.String() +} diff --git a/kernel_proxy_test.go b/kernel_proxy_test.go new file mode 100644 index 00000000..0da7c72c --- /dev/null +++ b/kernel_proxy_test.go @@ -0,0 +1,93 @@ +package dbsql + +import ( + "errors" + "net/http" + "net/url" + "testing" + + "github.com/databricks/databricks-sql-go/internal/config" +) + +// proxyForEndpointFunc maps the injected resolver's decision to a proxy URL +// string (or "" for direct), and returns "" for an unbuildable endpoint. The +// production path uses http.ProxyFromEnvironment, whose env is snapshotted once +// per process (sync.Once) and so can't be re-driven mid-test; the resolver seam +// lets us assert every branch deterministically. Each case mirrors an +// httpproxy-style outcome: proxy set for this host, host excluded by NO_PROXY, +// no proxy configured, and an unbuildable endpoint. Pure Go, so it runs in the +// default CGO_ENABLED=0 build (no kernel lib required). +func TestProxyForEndpoint(t *testing.T) { + validCfg := func() *config.Config { + c := config.WithDefaults() + c.Host = "my-workspace.databricks.com" + c.Port = 443 + c.HTTPPath = "/sql/1.0/warehouses/abc" + return c + } + proxyURL, _ := url.Parse("http://corp-proxy:3128") + + cases := []struct { + name string + cfg *config.Config + resolve func(*http.Request) (*url.URL, error) + want string + }{ + { + name: "proxy set for host", + cfg: validCfg(), + resolve: func(*http.Request) (*url.URL, error) { return proxyURL, nil }, + want: "http://corp-proxy:3128", + }, + { + // Warehouse-id addressing mode (HTTPPath == "") is what the kernel + // prefers; the proxy must still resolve. Previously ToEndpointURL errored + // on the empty path and the proxy was silently dropped to "" (direct). + name: "warehouse-id mode (no HTTPPath) still resolves proxy", + cfg: func() *config.Config { + c := config.WithDefaults() + c.Host = "my-workspace.databricks.com" + c.Port = 443 + c.WarehouseID = "abc" // no HTTPPath + return c + }(), + resolve: func(*http.Request) (*url.URL, error) { return proxyURL, nil }, + want: "http://corp-proxy:3128", + }, + { + name: "host excluded by NO_PROXY -> direct", + cfg: validCfg(), + resolve: func(*http.Request) (*url.URL, error) { return nil, nil }, + want: "", + }, + { + name: "resolver error -> direct", + cfg: validCfg(), + resolve: func(*http.Request) (*url.URL, error) { return nil, errors.New("bad proxy url") }, + want: "", + }, + { + name: "unbuildable endpoint -> direct (resolver never consulted)", + cfg: func() *config.Config { + c := config.WithDefaults() + c.Host = "" + return c + }(), + resolve: func(*http.Request) (*url.URL, error) { + t.Error("resolver must not be called for an unbuildable endpoint") + return proxyURL, nil + }, + want: "", + }, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + if got := proxyForEndpointFunc(tc.cfg, tc.resolve); got != tc.want { + t.Errorf("proxyForEndpointFunc = %q, want %q", got, tc.want) + } + }) + } + + // The production wrapper wires http.ProxyFromEnvironment and must not panic. + _ = proxyForEndpoint(validCfg()) +} diff --git a/kernel_stub.go b/kernel_stub.go new file mode 100644 index 00000000..4ea1dde5 --- /dev/null +++ b/kernel_stub.go @@ -0,0 +1,23 @@ +//go:build !cgo || !databricks_kernel + +package dbsql + +import ( + "context" + "fmt" + + dbsqlerr "github.com/databricks/databricks-sql-go/errors" + "github.com/databricks/databricks-sql-go/internal/backend" + "github.com/databricks/databricks-sql-go/internal/config" +) + +// newKernelBackend is the stub compiled when the kernel backend is not built in. +// It fails loudly rather than silently falling back to Thrift, so a mismatch +// between WithUseKernel and the build tags surfaces at connect time. The error +// wraps ErrKernelNotCompiled so a caller can detect the build mismatch with +// errors.Is (the same mechanism as ErrNotSupportedByKernel) rather than matching +// message text. The real implementation is in kernel_backend.go. +func newKernelBackend(_ context.Context, _ *config.Config) (backend.Backend, error) { + return nil, fmt.Errorf("databricks: %w; rebuild with -tags databricks_kernel and "+ + "CGO_ENABLED=1, or unset WithUseKernel", dbsqlerr.ErrKernelNotCompiled) +} diff --git a/kernel_stub_test.go b/kernel_stub_test.go new file mode 100644 index 00000000..919b4362 --- /dev/null +++ b/kernel_stub_test.go @@ -0,0 +1,36 @@ +//go:build !cgo || !databricks_kernel + +package dbsql + +import ( + "context" + "errors" + "testing" + + dbsqlerr "github.com/databricks/databricks-sql-go/errors" +) + +// In the default pure-Go build (no databricks_kernel tag) the kernel backend is +// not compiled in. Connecting with WithUseKernel(true) must fail loudly with a +// clear "not compiled in" error rather than silently falling back to Thrift or +// panicking on a nil backend. This guards the stub that fails closed. +func TestKernelBackendNotCompiledIn(t *testing.T) { + connector, err := NewConnector( + WithServerHostname("example.cloud.databricks.com"), + WithPort(443), + WithHTTPPath("/sql/1.0/endpoints/12346a5b5b0e123a"), + WithAccessToken("supersecret"), + WithUseKernel(true), + ) + if err != nil { + t.Fatalf("NewConnector: %v", err) + } + _, err = connector.Connect(context.Background()) + if err == nil { + t.Fatal("Connect with WithUseKernel(true) in a non-kernel build should error, got nil") + } + // Detect the build mismatch via the exported sentinel, not message text. + if !errors.Is(err, dbsqlerr.ErrKernelNotCompiled) { + t.Errorf("error should wrap ErrKernelNotCompiled; got: %v", err) + } +} diff --git a/rust-toolchain.toml b/rust-toolchain.toml new file mode 100644 index 00000000..f3f40f0c --- /dev/null +++ b/rust-toolchain.toml @@ -0,0 +1,10 @@ +# Pins the Rust toolchain used to build the SEA kernel static lib, so the archive +# is reproducible: a floating `stable` drifts the .a (and its sha256) under a fixed +# KERNEL_REV as rustc/std advance. rustup honors this file for any cargo invocation +# in this dir or below — including the kernel checkout under build/kernel-src/ — so +# CI and local `make kernel-lib` build with the same compiler. Bump deliberately, +# in lockstep with the kernel's MSRV (kernel Cargo.toml edition 2021; no pin of its +# own). Only the opt-in `databricks_kernel` build reads this; the pure-Go default +# build has no Rust dependency. +[toolchain] +channel = "1.96.1"