diff --git a/.github/workflows/trigger-integration-tests.yml b/.github/workflows/trigger-integration-tests.yml index 4da5deb1..db1bcf9f 100644 --- a/.github/workflows/trigger-integration-tests.yml +++ b/.github/workflows/trigger-integration-tests.yml @@ -9,9 +9,18 @@ name: Trigger Integration Tests # preview with the `integration-test` label (replay) or `integration-test-full` # (full passthrough), and merge queue runs the real required gate. # +# The suite runs against either driver backend, selected by the label: +# - `integration-test` (replay) / `integration-test-full` (passthrough) → Thrift +# - `integration-test-kernel` (passthrough) → SEA-via-kernel +# The kernel label adds `go_mode: sea` to the dispatch payload; driver-test then +# builds the kernel static lib and runs the tagged (databricks_kernel) leg. It uses +# passthrough (real warehouse) — the sea leg has no committed recordings to replay +# yet, so there is no sea replay label until those are captured. +# # Required external setup: # -# 1. `integration-test` and `integration-test-full` labels exist in this repo. +# 1. The three trigger labels exist in this repo: `integration-test`, +# `integration-test-full`, `integration-test-kernel`. # 2. `INTEGRATION_TEST_APP_ID` / `INTEGRATION_TEST_PRIVATE_KEY` repo secrets # are installed in this repo for the dispatcher GitHub App. # 3. The app is installed/granted on `databricks-driver-test` so this workflow @@ -46,7 +55,7 @@ jobs: with: script: | const present = context.payload.pull_request.labels.map((l) => l.name); - const triggerLabels = ['integration-test', 'integration-test-full']; + const triggerLabels = ['integration-test', 'integration-test-full', 'integration-test-kernel']; const removed = []; for (const name of triggerLabels) { if (!present.includes(name)) continue; @@ -111,7 +120,7 @@ jobs: completed_at: new Date().toISOString(), output: { title: 'Skipped on PR - runs in merge queue', - summary: 'Go integration tests are skipped on ordinary PR events and run as a required gate in the merge queue. Add the `integration-test` label to preview replay (or `integration-test-full` for the full passthrough suite) on this PR.', + summary: 'Go integration tests are skipped on ordinary PR events and run as a required gate in the merge queue. Preview on this PR by adding a label: `integration-test` (Thrift, replay) or `integration-test-full` (Thrift, passthrough); `integration-test-kernel` runs the SEA-via-kernel backend against the real warehouse (passthrough).', }, }); @@ -122,7 +131,8 @@ jobs: github.event_name == 'pull_request' && github.event.action == 'labeled' && (github.event.label.name == 'integration-test' || - github.event.label.name == 'integration-test-full') + github.event.label.name == 'integration-test-full' || + github.event.label.name == 'integration-test-kernel') runs-on: group: databricks-protected-runner-group labels: linux-ubuntu-latest @@ -134,14 +144,24 @@ jobs: pull-requests: write checks: write steps: - - name: Resolve proxy mode from label + - name: Resolve proxy mode + backend from label id: mode + # go_mode: integration-test-kernel runs the SEA-via-kernel backend; the others + # run Thrift. proxy_mode: integration-test-full and integration-test-kernel hit + # the real warehouse (passthrough); plain integration-test serves recordings + # (replay). The kernel label is passthrough because the sea leg has no committed + # recordings to replay yet. driver-test reads go_mode to decide whether to build + # the kernel lib and run the tagged leg. run: | - if [ "${{ github.event.label.name }}" = "integration-test-full" ]; then - echo "proxy_mode=passthrough" >> "$GITHUB_OUTPUT" - else - echo "proxy_mode=replay" >> "$GITHUB_OUTPUT" - fi + LABEL="${{ github.event.label.name }}" + case "$LABEL" in + integration-test-full|integration-test-kernel) echo "proxy_mode=passthrough" >> "$GITHUB_OUTPUT" ;; + *) echo "proxy_mode=replay" >> "$GITHUB_OUTPUT" ;; + esac + case "$LABEL" in + integration-test-kernel) echo "go_mode=sea" >> "$GITHUB_OUTPUT" ;; + *) echo "go_mode=thrift" >> "$GITHUB_OUTPUT" ;; + esac - name: Generate GitHub App token (driver-test repo) id: app-token @@ -156,6 +176,7 @@ jobs: uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0 env: PROXY_MODE: ${{ steps.mode.outputs.proxy_mode }} + GO_MODE: ${{ steps.mode.outputs.go_mode }} with: github-token: ${{ steps.app-token.outputs.token }} script: | @@ -170,9 +191,10 @@ jobs: pr_repo: context.repo.owner + '/' + context.repo.repo, pr_url: pr.html_url, proxy_mode: process.env.PROXY_MODE, + go_mode: process.env.GO_MODE, }, }); - core.info(`Dispatched go-pr-test (${process.env.PROXY_MODE}) for PR #${pr.number} @ ${pr.head.sha}`); + core.info(`Dispatched go-pr-test (${process.env.GO_MODE}/${process.env.PROXY_MODE}) for PR #${pr.number} @ ${pr.head.sha}`); - name: Fail check on dispatch error if: failure() @@ -198,13 +220,14 @@ jobs: uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 # v9.0.0 env: PROXY_MODE: ${{ steps.mode.outputs.proxy_mode }} + GO_MODE: ${{ steps.mode.outputs.go_mode }} with: script: | await github.rest.issues.createComment({ owner: context.repo.owner, repo: context.repo.repo, issue_number: context.issue.number, - body: `Go integration tests triggered (\`${process.env.PROXY_MODE}\`). [View workflow runs](https://github.com/databricks/databricks-driver-test/actions).`, + body: `Go integration tests triggered (\`${process.env.GO_MODE}\` / \`${process.env.PROXY_MODE}\`). [View workflow runs](https://github.com/databricks/databricks-driver-test/actions).`, }); # Merge queue: the required gate. Runs replay once before merge. @@ -261,9 +284,14 @@ jobs: pr_repo: context.repo.owner + '/' + context.repo.repo, pr_url: `${context.serverUrl}/${context.repo.owner}/${context.repo.repo}/pull/${prNumber}`, proxy_mode: 'replay', + // The required merge-queue gate runs the Thrift backend only. The + // kernel (sea) leg is previewed on demand via the kernel labels but + // is not a required gate — it becomes one once the SEA backend ships + // in a released driver and its recordings are captured. + go_mode: 'thrift', }, }); - core.info(`Merge-queue dispatch go-pr-test (replay) for PR #${prNumber} @ ${process.env.HEAD_SHA}`); + core.info(`Merge-queue dispatch go-pr-test (thrift/replay) for PR #${prNumber} @ ${process.env.HEAD_SHA}`); - name: Fail check on dispatch error if: failure() diff --git a/internal/backend/kernel/kernel_test.go b/internal/backend/kernel/kernel_test.go index a293f915..f535e013 100644 --- a/internal/backend/kernel/kernel_test.go +++ b/internal/backend/kernel/kernel_test.go @@ -14,6 +14,7 @@ import ( "github.com/databricks/databricks-sql-go/driverctx" dbsqlerr "github.com/databricks/databricks-sql-go/errors" "github.com/databricks/databricks-sql-go/internal/backend" + dbsqlrows "github.com/databricks/databricks-sql-go/internal/rows" "github.com/databricks/databricks-sql-go/logger" ) @@ -345,3 +346,92 @@ func strPtr(s string) *string { return &s } // 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. + +// kernelRows.Close() must fire the OnClose telemetry callback so the kernel path +// records CLOSE_STATEMENT / latency / statement success-or-failure — conn gates +// that recording on OnClose being called, and the Thrift path fires it. Before this +// wiring, kernel queries emitted no close telemetry (a production blind spot). A bare +// kernelRows is safe: Close() nil-guards cur/stream/op. +func TestKernelRowsCloseFiresOnClose(t *testing.T) { + t.Run("success path reports nil iterErr", func(t *testing.T) { + var got struct { + called bool + chunkCount int + iterErr, closeErr error + } + r := &kernelRows{ + chunkCount: 3, + callbacks: &dbsqlrows.TelemetryCallbacks{ + OnClose: func(latencyMs int64, chunkCount int, iterErr, closeErr error) { + got.called, got.chunkCount, got.iterErr, got.closeErr = true, chunkCount, iterErr, closeErr + }, + }, + } + if err := r.Close(); err != nil { + t.Fatalf("Close() = %v, want nil", err) + } + if !got.called { + t.Fatal("OnClose was not fired") + } + if got.chunkCount != 3 { + t.Errorf("OnClose chunkCount = %d, want 3", got.chunkCount) + } + if got.iterErr != nil || got.closeErr != nil { + t.Errorf("OnClose errs = (%v, %v), want (nil, nil)", got.iterErr, got.closeErr) + } + }) + + t.Run("iterationErr is reported", func(t *testing.T) { + sentinel := errors.New("boom") + var gotIter error + fired := 0 + r := &kernelRows{ + iterationErr: sentinel, + callbacks: &dbsqlrows.TelemetryCallbacks{ + OnClose: func(_ int64, _ int, iterErr, _ error) { fired++; gotIter = iterErr }, + }, + } + _ = r.Close() + if !errors.Is(gotIter, sentinel) { + t.Errorf("OnClose iterErr = %v, want %v", gotIter, sentinel) + } + // Idempotent: a second Close must not re-fire (conn/database-sql may double-close). + _ = r.Close() + if fired != 1 { + t.Errorf("OnClose fired %d times across two Close() calls, want 1", fired) + } + }) + + t.Run("nil callbacks is safe", func(t *testing.T) { + r := &kernelRows{} // no callbacks + if err := r.Close(); err != nil { + t.Errorf("Close() with nil callbacks = %v, want nil", err) + } + }) + + t.Run("construction-failure Close must not fire a success OnClose", func(t *testing.T) { + // Drive the real newKernelRows construction-failure path: a nil result + // stream makes kernel_result_stream_get_schema return a defined + // InvalidArgument error (the kernel null-checks the handle — never UB), so + // newKernelRows takes its cleanup r.Close() branch and returns an error. The + // callback must NOT have been armed yet, so the supplied OnClose must not + // fire a (falsely successful) CLOSE_STATEMENT for a statement that produced + // no rows. This is the invariant that keeping r.callbacks unset until after + // a successful build guarantees. + fired := false + cb := &dbsqlrows.TelemetryCallbacks{ + OnClose: func(int64, int, error, error) { fired = true }, + } + op := &kernelOp{backend: &KernelBackend{}} // for evictIfSessionFatal on the error path + rows, err := newKernelRows(context.Background(), op, nil /* stream */, cb) + if err == nil { + t.Fatal("newKernelRows(nil stream) = nil error, want a construction failure") + } + if rows != nil { + t.Errorf("newKernelRows on failure = %v rows, want nil", rows) + } + if fired { + t.Error("OnClose fired during construction-failure cleanup — callback armed too early") + } + }) +} diff --git a/internal/backend/kernel/rows.go b/internal/backend/kernel/rows.go index b03290f4..44557bf1 100644 --- a/internal/backend/kernel/rows.go +++ b/internal/backend/kernel/rows.go @@ -17,6 +17,7 @@ import ( "database/sql/driver" "fmt" "io" + "time" "unsafe" "github.com/apache/arrow/go/v12/arrow" @@ -47,6 +48,10 @@ type kernelRows struct { chunkCount int // cumulative batches fetched, for OnChunkFetched closed bool eof bool + // iterationErr is the first non-EOF error seen during Next(), reported to the + // OnClose telemetry callback so a failed statement is recorded (matching the + // Thrift path's rows.iterationErr). io.EOF is normal termination, not an error. + iterationErr error // 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. @@ -56,7 +61,14 @@ type kernelRows struct { // 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()} + // The telemetry callback is deliberately NOT set on r yet: the two cleanup + // r.Close() calls below run when construction FAILS (schema fetch/import), and a + // Close() with the callback set would fire OnClose as a *successful* close for a + // statement that never produced rows — masking the failure in CLOSE_STATEMENT + // telemetry. Assign it only on the success path so cleanup Close() on a + // schema/import failure does not record a falsely successful CLOSE_STATEMENT; the + // construction error itself is surfaced to and recorded by the conn execute path. + r := &kernelRows{ctx: ctx, op: op, stream: stream, keyCache: arrowscan.NewStructKeyCache()} var csch C.struct_ArrowSchema if err := call(func() C.KernelStatusCode { @@ -76,6 +88,9 @@ func newKernelRows(ctx context.Context, op *kernelOp, stream *C.kernel_result_st for i, f := range fields { r.cols[i] = f.Name } + // Construction succeeded — now arm the close telemetry callback so a normal + // Close() (after row iteration) records CLOSE_STATEMENT. + r.callbacks = cb klogCtx(ctx, "newKernelRows: %d columns", len(r.cols)) return r, nil } @@ -90,6 +105,7 @@ func (r *kernelRows) Close() error { return nil } r.closed = true + closeStart := time.Now() if r.cur != nil { r.cur.Release() r.cur = nil @@ -101,6 +117,14 @@ func (r *kernelRows) Close() error { if r.op != nil { r.op.close() } + // Fire the close telemetry callback so the kernel path records CLOSE_STATEMENT / + // execution latency / statement success-or-failure like the Thrift path does + // (conn gates this on OnClose being called). The kernel teardown has no fallible + // close RPC — the C stream/statement closes don't surface an error — so closeErr + // is nil; iterationErr carries any failure seen during Next(). + if r.callbacks != nil && r.callbacks.OnClose != nil { + r.callbacks.OnClose(time.Since(closeStart).Milliseconds(), r.chunkCount, r.iterationErr, nil) + } klogCtx(r.ctx, "kernelRows closed") return nil } @@ -108,6 +132,16 @@ func (r *kernelRows) Close() error { // 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 { + err := r.next(dest) + // Record the first non-EOF error for the OnClose telemetry callback (io.EOF is + // normal drain, not a failure). Mirrors the Thrift path's iterationErr capture. + if err != nil && err != io.EOF && r.iterationErr == nil { + r.iterationErr = err + } + return err +} + +func (r *kernelRows) next(dest []driver.Value) error { if r.closed { return io.EOF } diff --git a/kernel_config.go b/kernel_config.go index 168c3ab8..cf5fa0d0 100644 --- a/kernel_config.go +++ b/kernel_config.go @@ -129,6 +129,9 @@ func resolveKernelAuth(cfg *config.Config) (kernel.Auth, error) { } } if token == "" { + // Missing required config (not an unsupported-feature rejection), so this is + // intentionally NOT wrapped with ErrNotSupportedByKernel — a caller shouldn't + // fall back to Thrift for a forgotten token, it should supply one. return kernel.Auth{}, errors.New("databricks: the kernel backend requires a personal access token; " + "set one with WithAccessToken (or a *pat.PATAuth via WithAuthenticator)") } diff --git a/kernel_config_test.go b/kernel_config_test.go index 11c492c2..81e0e848 100644 --- a/kernel_config_test.go +++ b/kernel_config_test.go @@ -96,8 +96,8 @@ func TestValidateKernelConfig(t *testing.T) { // 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. (Catalog/schema/ - // metric-view moved to forwarded above; a non-PAT authenticator is rejected too but - // not sentinel-wrapped, so it's asserted separately below.) + // metric-view moved to forwarded above; a non-PAT authenticator is also + // sentinel-wrapped but needs its own AccessToken="" setup, so it's asserted separately below.) rejections := []struct { name string mut func(*config.Config)