diff --git a/KERNEL_REV b/KERNEL_REV index 0fc1de14..d82b4bdc 100644 --- a/KERNEL_REV +++ b/KERNEL_REV @@ -1 +1 @@ -4d301455f7e70de9cb747e0f1143b8a3df8c5b48 +c19e5d1ea1a587d16e4ac54c8afd520eca56d00a diff --git a/connector.go b/connector.go index 3838cb9d..24149a89 100644 --- a/connector.go +++ b/connector.go @@ -45,14 +45,16 @@ func (c *connector) Connect(ctx context.Context) (driver.Conn, error) { if c.cfg.UseKernel { be, err = newKernelBackend(ctx, c.cfg) } else { - // The experimental WithKernel* TLS options have no Thrift-path equivalent — - // reject them loudly rather than silently ignore, so a caller who sets a - // trusted-CA bundle / an independent hostname skip and forgets WithUseKernel - // learns the option had no effect instead of connecting with a - // weaker-than-intended (or unconfigured) TLS trust store. + // The experimental WithKernel* options have no Thrift-path equivalent — + // reject them loudly rather than silently ignore, so a caller who sets one + // (a trusted-CA bundle, an independent hostname skip, an mTLS client + // certificate, or a CloudFetch toggle) and forgets WithUseKernel learns the + // option had no effect instead of connecting with a weaker-than-intended + // (or unconfigured) setup. if c.cfg.KernelExperimental != nil { return nil, fmt.Errorf("databricks: a WithKernel* option "+ - "(WithKernelTrustedCerts / WithKernelSkipHostnameVerify) %w; "+ + "(WithKernelTrustedCerts / WithKernelSkipHostnameVerify / "+ + "WithKernelClientCertificate / WithKernelCloudFetch) %w; "+ "add WithUseKernel(true) or remove it", dbsqlerr.ErrRequiresKernelBackend) } be, err = thrift.New(ctx, c.cfg, c.client) @@ -603,3 +605,51 @@ func WithKernelSkipHostnameVerify() ConnOption { kernelExperimental(c).TLSSkipHostnameVerify = true } } + +// WithKernelClientCertificate configures a client certificate + private key for +// mutual TLS (mTLS) on the kernel backend. cert is a PEM leaf certificate, +// optionally followed by its intermediate chain; key is the matching PEM private +// key (use PKCS#8 for portability). Both halves are required together — they map +// to the single paired kernel_session_config_set_tls_client_certificate, so +// cert-without-key is unrepresentable. The private key is never logged. +// +// EXPERIMENTAL, kernel-only: the default (Thrift) backend rejects this at connect. +func WithKernelClientCertificate(cert, key []byte) ConnOption { + return func(c *config.Config) { + // Copy defensively (matching KernelExperimentalConfig.DeepCopy) so a + // caller mutating the buffers between NewConnector and Connect can't change + // the identity out from under us. + k := kernelExperimental(c) + // Record that mTLS was explicitly requested, independent of whether the + // caller passed non-empty bytes. Without this marker an empty cert+key + // pair (e.g. from a failed PEM load) is indistinguishable from the option + // never being called, and validateKernelConfig would accept it while + // applyKernelTLS silently skipped the setter — connecting with no client + // identity. The marker lets validation reject the incomplete request loudly. + k.TLSClientCertConfigured = true + if len(cert) > 0 { + k.TLSClientCertPEM = append([]byte(nil), cert...) + } else { + k.TLSClientCertPEM = cert + } + if len(key) > 0 { + k.TLSClientKeyPEM = append([]byte(nil), key...) + } else { + k.TLSClientKeyPEM = key + } + } +} + +// WithKernelCloudFetch toggles CloudFetch (cloud-storage-backed download of large +// result sets) on the kernel backend. Passing false forces the inline-Arrow path +// for all result sizes. Unlike the backend-neutral WithCloudFetch (a plain bool +// whose unset state is indistinguishable from false), this is tri-state: not +// calling it leaves the kernel default (CloudFetch on), and it maps to +// kernel_session_config_set_cloudfetch_enabled only when set. +// +// EXPERIMENTAL, kernel-only: the default (Thrift) backend rejects this at connect. +func WithKernelCloudFetch(enabled bool) ConnOption { + return func(c *config.Config) { + kernelExperimental(c).CloudFetchEnabled = &enabled + } +} diff --git a/doc.go b/doc.go index 2febbf55..ed0d543f 100644 --- a/doc.go +++ b/doc.go @@ -195,30 +195,44 @@ level the lines are suppressed with no cost. dbsql.SetLogLevel("debug") // or DATABRICKS_LOG_LEVEL=debug -The same level is mapped into the kernel's internal (Rust) log subscriber, with two -caveats specific to the Rust lines: they go to stderr directly (not affected by -logger.SetLogOutput), and their verbosity is fixed when the first kernel session in -the process is opened — set the level before that first connect to control them. +The same level is mapped into the kernel's internal (Rust) log subscriber, which is +routed into the driver's logger via a reverse-call callback (kernel_set_log_callback): +kernel-internal tracing events are forwarded into the SAME unified sink as the Go +binding lines — including a custom logger.SetLogOutput — rather than going to stderr. +One caveat specific to the Rust lines: BOTH their verbosity AND their output +destination are captured when the first kernel session in the process is opened (the +forwarding sink snapshots the logger then). A later dbsql.SetLogLevel or +logger.SetLogOutput re-targets the Go binding lines but NOT the already-forwarded +kernel lines, so set the level and any custom output before that first connect to +govern them. For finer control of the Rust verbosity independent of the driver level, set -DBSQL_KERNEL_DEBUG to any non-empty value: it forces the kernel subscriber on and -defers to RUST_LOG. Filter on the target databricks::sql::kernel (note the colons): +DBSQL_KERNEL_DEBUG to any non-empty value: the callback then defers its level to +RUST_LOG. Filter on the target databricks::sql::kernel (note the colons): - # kernel logs only, at the kernel's own verbosity: - DBSQL_KERNEL_DEBUG=1 RUST_LOG=databricks::sql::kernel=debug ./your_app 2>&1 + # kernel logs at the kernel's own verbosity, forwarded into the driver logger: + DBSQL_KERNEL_DEBUG=1 RUST_LOG=databricks::sql::kernel=debug ./your_app # kernel logs plus its HTTP stack: - DBSQL_KERNEL_DEBUG=1 RUST_LOG=debug ./your_app 2>&1 + DBSQL_KERNEL_DEBUG=1 RUST_LOG=debug ./your_app + +The kernel exposes one process-global tracing subscriber reachable two mutually +exclusive ways over the C ABI (kernel_init_logging → stderr, kernel_set_log_callback +→ the driver logger); the driver installs the callback. A host that has already +installed its own global tracing subscriber will see the registration report a +benign, logged failure and kernel events simply won't forward. Supported on the kernel backend: PAT and OAuth (M2M via WithClientCredentials, U2M via the authType=oauthU2M DSN param); reading scalar, nested, and complex-typed results (CloudFetch is transparent); bound query parameters (positional and named); -context cancellation during execute; the initial namespace (WithInitialNamespace, -applied post-connect via USE CATALOG / USE SCHEMA); metric-view metadata -(WithEnableMetricViewMetadata); and the TLS, proxy, and session-conf (query tags, -statement timeout, time zone) options. Nothing is silently ignored: WithTimeout, a -retries-disabling WithRetries, token-provider / external / federated -authenticators, and custom M2M OAuth scopes (the kernel applies its own) are -rejected at connect; staging (PUT/GET/REMOVE on a Unity Catalog volume) is +context cancellation during execute and mid-fetch; the initial namespace +(WithInitialNamespace, applied post-connect via USE CATALOG / USE SCHEMA); +metric-view metadata (WithEnableMetricViewMetadata); and the TLS, proxy, and +session-conf (query tags, statement timeout, time zone) options. Nothing is silently +ignored: WithTimeout, a retries-disabling WithRetries, a CloudFetch-disabling +WithCloudFetch(false) (use WithKernelCloudFetch(false) instead), a non-default +WithPort, a non-https protocol, a custom WithTransport, token-provider / external / +federated authenticators, and custom M2M OAuth scopes (the kernel applies its own) +are rejected at connect; staging (PUT/GET/REMOVE on a Unity Catalog volume) is rejected at execute. WithMaxRows and positive-limit WithRetries are accepted but inert (the kernel manages fetching and retries below the C ABI). @@ -233,15 +247,22 @@ backend exposes a U2M-scopes option, so this is a fixed difference between the t default sets, not a dropped setting; both authorize against the built-in public client. -Experimental kernel-only TLS options (rejected by the default backend; the -WithKernel* prefix marks them experimental): +Experimental kernel-only options (rejected by the default backend; the WithKernel* +prefix marks them experimental): - WithKernelTrustedCerts(pem) adds a PEM CA bundle on top of the system roots (for a re-signing proxy or on-prem CA). Required because the kernel's TLS stack does not read SSL_CERT_FILE. - WithKernelSkipHostnameVerify() skips only the hostname check while keeping chain validation (finer-grained than WithSkipTLSHostVerify). - -Setting either without WithUseKernel fails Connect with an error wrapping the + - WithKernelClientCertificate(cert, key) configures a client certificate + private + key for mutual TLS (mTLS). Both PEM halves are required together; the key is + never logged. + - WithKernelCloudFetch(enabled) toggles CloudFetch. Tri-state: leaving it unset + keeps the kernel default (on); false forces the inline-Arrow path for all result + sizes. Distinct from the backend-neutral WithCloudFetch, whose plain-bool unset + state can't be told apart from false. + +Setting any of these without WithUseKernel fails Connect with an error wrapping the sentinel ErrRequiresKernelBackend, detectable with errors.Is. Features above the backend seam are inherited unchanged: the database/sql connection @@ -253,9 +274,11 @@ and GEOMETRY (WKT). The server query id is surfaced on the success path, so a QueryIdCallback (see below) fires with the real id and EXECUTE_STATEMENT telemetry carries it. -On the read path, context cancellation is honored at result-batch boundaries, not -mid-fetch: an in-flight CloudFetch batch runs to completion before the cancel takes -effect. +On the read path, context cancellation is honored both at result-batch boundaries +and mid-fetch: a cancel or deadline firing while a CloudFetch batch is in flight +aborts that fetch (a hung S3 / pre-signed-URL GET does not block the caller past its +deadline), rather than running the batch to completion first. A non-cancellable +context (e.g. context.Background) takes the plain fetch path with no added overhead. # Programmatically Retrieving Connection and Query Id diff --git a/internal/backend/kernel/abi.go b/internal/backend/kernel/abi.go new file mode 100644 index 00000000..bbf8c3c8 --- /dev/null +++ b/internal/backend/kernel/abi.go @@ -0,0 +1,23 @@ +package kernel + +import "fmt" + +// compareABI reports whether the C-ABI version compiled into the linked kernel +// library (got) matches the version the driver's header declares (want): it +// returns a descriptive, operator-facing error on mismatch and nil when they +// agree. +// +// It is split out from checkABIVersion (in the cgo-tagged cgo.go) as pure, +// cgo-free logic so the mismatch verdict and its message are unit-testable under +// the default CGO_ENABLED=0 build. The negative path can't otherwise be reached +// from a test: a real build always links a .a and header produced together, and +// abiVersions() reads fixed cgo constants with no injection seam. checkABIVersion +// runs abiVersions() through this once per process (sync.Once-cached), on the +// first connect. +func compareABI(got, want uint32) error { + if got != want { + return fmt.Errorf("databricks: kernel ABI version mismatch: linked library reports %d, "+ + "driver header expects %d; rebuild the kernel static lib and header together (make kernel-lib)", got, want) + } + return nil +} diff --git a/internal/backend/kernel/abi_test.go b/internal/backend/kernel/abi_test.go new file mode 100644 index 00000000..0059a4c3 --- /dev/null +++ b/internal/backend/kernel/abi_test.go @@ -0,0 +1,37 @@ +package kernel + +import ( + "strings" + "testing" +) + +// TestCompareABI covers both verdicts of the ABI-version handshake without the +// cgo symbols, so it runs under the default CGO_ENABLED=0 build. The matching +// (happy) path is additionally proven against the real linked library by +// TestABIVersionMatches in the cgo build; the mismatch path — the runtime hazard +// the check exists for (a driver header linked against a differently-built +// prebuilt .a) — is otherwise unreachable from a test, since a real build always +// links a matching .a + header. +func TestCompareABI(t *testing.T) { + // Matching versions must produce no error: checkABIVersion opens the session. + for _, v := range []uint32{0, 1, 42} { + if err := compareABI(v, v); err != nil { + t.Errorf("compareABI(%d, %d) = %v, want nil for matching versions", v, v, err) + } + } + + // Mismatched versions must produce an error so checkABIVersion refuses to open + // (rather than silently misread status codes / error-struct fields). The + // message must name both versions and the remediation so an operator hitting a + // stale prebuilt .a knows what to rebuild. + err := compareABI(2, 1) + if err == nil { + t.Fatal("compareABI(2, 1) = nil, want an error for mismatched versions") + } + msg := err.Error() + for _, want := range []string{"ABI version mismatch", "reports 2", "expects 1", "make kernel-lib"} { + if !strings.Contains(msg, want) { + t.Errorf("mismatch error %q does not contain %q", msg, want) + } + } +} diff --git a/internal/backend/kernel/backend.go b/internal/backend/kernel/backend.go index 361d8a2b..ecd772db 100644 --- a/internal/backend/kernel/backend.go +++ b/internal/backend/kernel/backend.go @@ -56,20 +56,21 @@ func New(cfg Config) *KernelBackend { // 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. + // Fail fast on an already-cancelled context before doing any work, then honor + // a deadline that fires mid-connect via the ctxWatcher below: kernel_session + // _open blocks the calling thread inside the C ABI, so a slow warehouse + // cold-start or a connect-time network partition can't be interrupted by the + // caller's ctx unless a cancel token is fired from another thread. if err := ctx.Err(); err != nil { return err } + // Verify the linked kernel library's C-ABI version matches the header the + // driver compiled against before making any other kernel call — a mismatch + // means every status code / error struct we read afterward could be + // misinterpreted. Runs once per process; cheap and cached thereafter. + if err := checkABIVersion(); err != nil { + return err + } initKernelLogging() klogCtx(ctx, "OpenSession host=%s httpPath=%s warehouse=%s", k.cfg.Host, k.cfg.HTTPPath, k.cfg.WarehouseID) @@ -154,6 +155,19 @@ func (k *KernelBackend) OpenSession(ctx context.Context) error { } } + // Experimental kernel-only CloudFetch toggle (WithKernelCloudFetch). Tri-state: + // only set it when the caller did, so an unset value leaves the kernel default + // (CloudFetch on). Do NOT route this through set_session_conf — the kernel owns + // the can_cloud_download conf and the server rejects it as not user-settable. + if k.cfg.CloudFetchEnabled != nil { + enabled := C.bool(*k.cfg.CloudFetchEnabled) + if err := call(func() C.KernelStatusCode { + return C.kernel_session_config_set_cloudfetch_enabled(cfg, enabled) + }); err != nil { + return fmt.Errorf("kernel: set_cloudfetch_enabled: %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 { @@ -178,9 +192,25 @@ func (k *KernelBackend) OpenSession(ctx context.Context) error { // 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) }) + // Bridge ctx onto a cancel token so a deadline firing mid-connect drops the + // in-flight connect request rather than blocking until the kernel returns. A + // non-cancellable ctx yields a nil watcher → NULL token → the plain open path, + // so there is no watcher overhead on a background context. + watcher := newCtxWatcher(ctx) + defer watcher.stop() + err := call(func() C.KernelStatusCode { + return C.kernel_session_open_cancellable(cfg, &sess, watcher.tokenPtr()) + }) consumed = true if err != nil { + // Prefer the caller's ctx error when the connect was interrupted by the + // deadline; cancelledErr holds the shared dual-%w wrap (see its doc) so + // errors.Is still matches the ctx error AND the connect failure's server + // diagnostics stay reachable via errors.As. + if ctxErr := ctx.Err(); ctxErr != nil { + klogCtx(ctx, "OpenSession interrupted by ctx: kernelErr=%v ctxErr=%v", err, ctxErr) + return cancelledErr("session_open", ctxErr, toConnError(err)) + } return fmt.Errorf("kernel: session_open: %w", toConnError(err)) } k.session = sess @@ -213,10 +243,9 @@ func (k *KernelBackend) OpenSession(ctx context.Context) error { } // applyKernelTLS forwards the experimental kernel-only TLS knobs to the session -// config: a custom CA bundle and an independent hostname-skip. Each is a no-op -// when its field is unset, so this is safe to call unconditionally. (mTLS client -// cert/key is a separate follow-up — it needs a kernel C-ABI setter that is not -// yet on kernel main.) +// config: a custom CA bundle, an independent hostname-skip, and an mTLS client +// certificate + key. Each is a no-op when its field is unset, so this is safe to +// call unconditionally. func (k *KernelBackend) applyKernelTLS(cfg *C.KernelSessionConfig) error { if len(k.cfg.TLSTrustedCertsPEM) > 0 { ca := newCBytes(k.cfg.TLSTrustedCertsPEM) @@ -234,6 +263,20 @@ func (k *KernelBackend) applyKernelTLS(cfg *C.KernelSessionConfig) error { return fmt.Errorf("kernel: set_tls_skip_hostname_verification: %w", toConnError(err)) } } + // mTLS client identity. The cert and key travel as a pair (WithKernelClientCertificate + // sets both or neither), forwarded via the single paired setter; checking the + // cert is enough. The key bytes go to owned Rust memory and are never logged. + if len(k.cfg.TLSClientCertPEM) > 0 { + cert := newCBytes(k.cfg.TLSClientCertPEM) + defer cert.free() + key := newCBytes(k.cfg.TLSClientKeyPEM) + defer key.free() + if err := call(func() C.KernelStatusCode { + return C.kernel_session_config_set_tls_client_certificate(cfg, cert.ptr, cert.len, key.ptr, key.len) + }); err != nil { + return fmt.Errorf("kernel: set_tls_client_certificate: %w", toConnError(err)) + } + } return nil } @@ -345,14 +388,17 @@ func (k *KernelBackend) runNamespaceStmt(ctx context.Context, sql string) error return closeErr } -// 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. +// CloseSession tears down the server-side session. Best-effort: kernel_session +// _close initiates the delete without waiting (see the C header), so it does not +// block and 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. +// Stays fire-and-forget deliberately. The C ABI also offers +// kernel_session_close_blocking (awaits DeleteSession, for Python/Node parity), +// but adopting it here would make close a blocking network round-trip with no +// deadline honored — a stalled close (shutdown-time network partition) would then +// block database/sql pool cleanup. Swapping to it is grouped with the +// cancellable-close follow-up so the blocking close ships with a ctx bridge; until +// then fire-and-forget is the safer default. func (k *KernelBackend) CloseSession(ctx context.Context) error { if k.session == nil { return nil diff --git a/internal/backend/kernel/cancel.go b/internal/backend/kernel/cancel.go new file mode 100644 index 00000000..7554c915 --- /dev/null +++ b/internal/backend/kernel/cancel.go @@ -0,0 +1,148 @@ +//go:build cgo && databricks_kernel + +package kernel + +/* +#include +#include "databricks_kernel.h" +*/ +import "C" + +import ( + "context" + "fmt" + "sync" +) + +// ctxWatcher bridges a Go context deadline/cancellation onto a kernel cancel +// token so a blocking cgo call (connect or a hung result-stream fetch) can be +// interrupted when the caller's ctx fires. +// +// The kernel's C ABI cannot observe a Go ctx mid-call: kernel_session_open and +// kernel_result_stream_next_batch block the calling OS thread inside Rust until +// the operation completes. The *_cancellable variants take a +// kernel_cancel_token_t that a *different* thread fires; firing drops the raced +// future in the kernel and unblocks the calling thread promptly. On connect that +// is a real abort — kernel_session_open awaits its request inline, so the +// in-flight reqwest future cancels on drop. On mid-fetch it is a prompt +// stop-waiting: the CloudFetch download runs on a detached kernel task, so the +// call returns at the deadline but that chunk's download drains in the +// background (bounded by its ~60s read-timeout). Either way the caller's OS +// thread / goroutine is freed at the deadline. This helper owns that token plus +// a watcher goroutine that fires it on ctx.Done(). +// +// It mirrors the execute-path canceller watcher in operation.go (same +// done-channel + WaitGroup drain shape), but drives the generic call-cancel +// token rather than the statement canceller, and so applies to the connect and +// result-fetch calls the statement canceller can't reach. +type ctxWatcher struct { + token *C.kernel_cancel_token_t + done chan struct{} + wg sync.WaitGroup +} + +// newCtxWatcher creates a cancel token and, when ctx is cancellable, starts a +// goroutine that fires the token on ctx.Done(). Returns nil when ctx is nil or +// non-cancellable (ctx.Done() == nil) — the caller then passes a NULL token to +// the plain-equivalent behavior, so there is zero overhead on the common +// background-context path. The caller MUST call stop() to drain the watcher and +// free the token (typically via defer), after the blocking call returns. +// +// A token-creation failure also yields nil (degrade to uncancellable rather +// than fail the operation): cancellation is a robustness improvement, not a +// correctness precondition, so a failure to allocate it must not break connect +// or fetch. +func newCtxWatcher(ctx context.Context) *ctxWatcher { + if ctx == nil || ctx.Done() == nil { + return nil + } + var token *C.kernel_cancel_token_t + if err := call(func() C.KernelStatusCode { return C.kernel_cancel_token_new(&token) }); err != nil { + klog("cancel token_new failed (proceeding uncancellable): %v", err) + return nil + } + w := &ctxWatcher{token: token, done: make(chan struct{})} + w.wg.Add(1) + go func() { + defer w.wg.Done() + select { + case <-ctx.Done(): + klog("ctxWatcher: ctx.Done (%v) → firing cancel token", ctx.Err()) + // Fire is purely local in the kernel (flip an atomic, wake the + // select) — no RPC, so this never blocks. Errors are swallowed: the + // only failure is a null handle, which can't happen here. + _ = call(func() C.KernelStatusCode { return C.kernel_cancel_token_cancel(w.token) }) + case <-w.done: + } + }() + return w +} + +// cancelledErr builds the error returned when a blocking kernel call (connect, +// execute, or a result-stream fetch) was interrupted by the caller's context. It +// is the single home for the subtle dual-%w wrap the three cancellable call sites +// (OpenSession, execute, nextBatch) share, so the contract lives in one place +// rather than being re-derived — and re-risked — at each site. +// +// It prefers the ctx error (database/sql convention: a cancelled call reports the +// cancellation) while keeping the kernel error reachable: BOTH are wrapped with %w +// so errors.Is(err, context.Canceled / context.DeadlineExceeded) still matches AND +// the underlying *KernelError stays reachable via errors.As. Without the second +// wrap, a session-fatal failure racing a cancel would lose its sqlstate / queryId — +// the one handle to what actually went wrong server-side. +// +// op names the interrupted call ("session_open", "execute", "next_batch") for the +// message. kernelErr is the already-classified kernel error (toConnError on the +// connect path, toStatementError on the statement path); ctxErr is ctx.Err(). +// +// It lives here (a tagged file) rather than in the untagged errors_classify.go +// because its only callers are the tagged cgo call sites — under the default +// CGO_ENABLED=0 build the linter would see it as unused. +func cancelledErr(op string, ctxErr, kernelErr error) error { + return fmt.Errorf("kernel: %s cancelled: %w (kernel error: %w)", op, ctxErr, kernelErr) +} + +// tokenPtr returns the underlying token pointer to pass to a *_cancellable +// entry point, or NULL for a nil watcher (uncancellable ctx → the call behaves +// exactly like its plain, non-cancellable variant). +func (w *ctxWatcher) tokenPtr() *C.kernel_cancel_token_t { + if w == nil { + return nil + } + return w.token +} + +// stop drains the watcher goroutine and frees the token. Safe to defer +// immediately after newCtxWatcher. Ordering matters: the watcher may be +// mid-fire (inside kernel_cancel_token_cancel) when the blocking call returns, +// so we close(done) to stop a not-yet-fired watcher, Wait() for any in-flight +// fire to finish, and only then free the token — never free it out from under a +// concurrent cancel (the kernel documents the token as single-owner for +// teardown). +func (w *ctxWatcher) stop() { + if w == nil { + return + } + close(w.done) + w.wg.Wait() + C.kernel_cancel_token_free(w.token) + w.token = nil +} + +// tokenFiredForTest reports whether this watcher's kernel token has been fired +// (kernel_cancel_token_is_cancelled). A test seam so a tagged unit test can +// assert the watcher-goroutine → token-fire wiring end to end through the real +// cgo boundary — without putting cgo in a _test.go file (which Go forbids). Not +// used in production. A nil watcher (uncancellable ctx) reports false. +func (w *ctxWatcher) tokenFiredForTest() bool { + if w == nil { + return false + } + var fired C.bool + if err := call(func() C.KernelStatusCode { + return C.kernel_cancel_token_is_cancelled(w.token, &fired) + }); err != nil { + return false + } + return bool(fired) +} diff --git a/internal/backend/kernel/cancel_connect_test.go b/internal/backend/kernel/cancel_connect_test.go new file mode 100644 index 00000000..08ea3ec5 --- /dev/null +++ b/internal/backend/kernel/cancel_connect_test.go @@ -0,0 +1,125 @@ +//go:build cgo && databricks_kernel + +package kernel + +import ( + "context" + "errors" + "net" + "strings" + "sync" + "testing" + "time" +) + +// blackHoleListener accepts TCP connections and then holds them open forever +// without ever writing a byte, so a client's TLS handshake blocks waiting for the +// ServerHello. It is the deterministic stand-in for a warehouse whose connect +// hangs (a cold-start stall or a network black hole), which is what the +// mid-connect cancel path must be able to interrupt. Returns the "host:port" to +// point the kernel at, and a stop func that closes the listener and every accepted +// connection. Accepted conns are retained so the runtime can't close them early. +func blackHoleListener(t *testing.T) (addr string, stop func()) { + t.Helper() + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatalf("listen: %v", err) + } + done := make(chan struct{}) + // mu guards conns: the accept goroutine appends while stop() ranges over it. + var mu sync.Mutex + var conns []net.Conn + go func() { + for { + c, err := ln.Accept() + if err != nil { + return // listener closed + } + mu.Lock() + conns = append(conns, c) // keep it open (and referenced): never respond + mu.Unlock() + select { + case <-done: + return + default: + } + } + }() + return ln.Addr().String(), func() { + close(done) + _ = ln.Close() // unblocks Accept → the goroutine returns, no more appends + mu.Lock() + defer mu.Unlock() + for _, c := range conns { + _ = c.Close() + } + } +} + +// TestOpenSessionHonorsCtxDeadlineMidConnect is the connect-side counterpart to the +// mid-fetch cancel test: it proves the ctxWatcher fires the kernel cancel token +// WHILE kernel_session_open_cancellable is blocked inside a connect, so a caller's +// deadline is honored mid-connect rather than the connect running until the kernel +// gives up on its own. +// +// It drives the real cgo OpenSession (not database/sql, which can short-circuit an +// already-cancelled ctx before the driver is ever called and can wrap the result in +// its own retry error) against a black-hole server whose TLS handshake never +// completes. A short ctx deadline must then interrupt the blocking connect via the +// token and surface as a "session_open cancelled" error matching +// context.DeadlineExceeded. Without the cancellable wiring the connect would block +// until the kernel's own timeout, so the goroutine guard below fails loudly instead +// of letting a regression hang the suite. +func TestOpenSessionHonorsCtxDeadlineMidConnect(t *testing.T) { + addr, stop := blackHoleListener(t) + defer stop() + + k := New(Config{ + // normalise_host prepends https://, so the kernel dials this host:port over + // TLS and blocks on the handshake the black-hole server never answers. + Host: addr, + WarehouseID: "wh-test", + Auth: Auth{Mode: AuthPAT, Token: "dapi-not-a-real-token"}, + // Relax cert checks so nothing rejects the (never-arriving) server cert before + // the deadline can fire; the block is at the handshake read regardless. + TLSSkipVerify: true, + }) + + const deadline = 750 * time.Millisecond + ctx, cancel := context.WithTimeout(context.Background(), deadline) + defer cancel() + + type result struct { + err error + elapsed time.Duration + } + resCh := make(chan result, 1) + start := time.Now() + go func() { + err := k.OpenSession(ctx) + resCh <- result{err, time.Since(start)} + }() + + select { + case res := <-resCh: + if res.err == nil { + t.Fatal("OpenSession succeeded against a black-hole server, want a cancellation error") + } + // The deadline — not an unrelated connect failure — must be what ended it, and + // the connect-cancel branch must wrap it as "session_open cancelled". + if !errors.Is(res.err, context.DeadlineExceeded) { + t.Errorf("OpenSession err = %v, want it to match context.DeadlineExceeded", res.err) + } + if !strings.Contains(res.err.Error(), "session_open cancelled") { + t.Errorf("OpenSession err = %q, want the mid-connect %q wrap (proves the cancellable path, "+ + "not a pre-connect guard, handled it)", res.err, "session_open cancelled") + } + // It must return promptly after the deadline, not run to some far-off connect + // timeout — the whole point of the mid-connect cancel. + if res.elapsed > 20*time.Second { + t.Errorf("OpenSession took %v after a %v deadline; the mid-connect cancel was not honored promptly", res.elapsed, deadline) + } + case <-time.After(30 * time.Second): + t.Fatal("OpenSession did not return within 30s of a 750ms deadline — the mid-connect cancel token was not fired") + } +} diff --git a/internal/backend/kernel/cancel_test.go b/internal/backend/kernel/cancel_test.go new file mode 100644 index 00000000..93c7a79f --- /dev/null +++ b/internal/backend/kernel/cancel_test.go @@ -0,0 +1,104 @@ +//go:build cgo && databricks_kernel + +package kernel + +import ( + "context" + "testing" + "time" +) + +// A cancellable ctx that fires drives the watcher goroutine to fire the kernel +// cancel token — the wiring the *_cancellable entry points rely on. Exercised +// through the real cgo boundary via the tokenFiredForTest seam (cgo cannot be +// used directly in a _test.go file). This is the safety-critical +// ctx→token bridge; keeping it in the tagged unit test means it runs in the +// build-and-test-kernel CI job, not only in the live e2e. +func TestCtxWatcherFiresTokenOnCancel(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + w := newCtxWatcher(ctx) + if w == nil { + t.Fatal("newCtxWatcher returned nil for a cancellable ctx") + } + defer w.stop() + + if w.tokenFiredForTest() { + t.Fatal("token fired before ctx was cancelled") + } + + cancel() + + // The watcher fires asynchronously; poll briefly for the flip. + fired := false + for i := 0; i < 200; i++ { + if w.tokenFiredForTest() { + fired = true + break + } + time.Sleep(5 * time.Millisecond) + } + if !fired { + t.Fatal("token was not fired after ctx cancellation") + } +} + +// A ctx whose deadline elapses also fires the token (the connect / fetch +// deadline path). +func TestCtxWatcherFiresTokenOnDeadline(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond) + defer cancel() + w := newCtxWatcher(ctx) + if w == nil { + t.Fatal("newCtxWatcher returned nil for a ctx with a deadline") + } + defer w.stop() + + fired := false + for i := 0; i < 200; i++ { + if w.tokenFiredForTest() { + fired = true + break + } + time.Sleep(5 * time.Millisecond) + } + if !fired { + t.Fatal("token was not fired after the ctx deadline elapsed") + } +} + +// A nil / non-cancellable ctx yields a nil watcher (NULL token → the call +// behaves exactly like its plain, non-cancellable variant), and stop()/tokenPtr +// are nil-safe so the call sites need no special-casing. +func TestCtxWatcherNilForUncancellableCtx(t *testing.T) { + if w := newCtxWatcher(nil); w != nil { + t.Error("newCtxWatcher(nil) should return nil") + } + // context.Background() has a nil Done() channel — non-cancellable. + if w := newCtxWatcher(context.Background()); w != nil { + t.Error("newCtxWatcher(background) should return nil (non-cancellable)") + } + // nil watcher is safe to use. + var w *ctxWatcher + if w.tokenPtr() != nil { + t.Error("nil watcher tokenPtr should be NULL") + } + w.stop() // must not panic + if w.tokenFiredForTest() { + t.Error("nil watcher tokenFiredForTest should be false") + } +} + +// A watcher whose ctx never fires cleans up without the token being fired — +// the normal-completion path (stop() drains the idle watcher goroutine). +func TestCtxWatcherCleanUpWithoutFire(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + w := newCtxWatcher(ctx) + if w == nil { + t.Fatal("newCtxWatcher returned nil for a cancellable ctx") + } + if w.tokenFiredForTest() { + t.Fatal("token fired without a cancel") + } + w.stop() // drains the still-waiting watcher; must not hang or fire +} diff --git a/internal/backend/kernel/cgo.go b/internal/backend/kernel/cgo.go index cd604121..34dbd52e 100644 --- a/internal/backend/kernel/cgo.go +++ b/internal/backend/kernel/cgo.go @@ -101,56 +101,82 @@ func klogCtx(ctx context.Context, format string, args ...any) { ).Debug().Msgf("[kernel] "+format, args...) } -// 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 and points their -// verbosity at the driver's log level, so DATABRICKS_LOG_LEVEL drives both the Go -// binding lines and the kernel's Rust lines from one knob. The mapped level is -// passed to kernel_init_logging (Go zerolog level → the kernel's OFF/ERROR/WARN/ -// INFO/DEBUG/TRACE string); DBSQL_KERNEL_DEBUG forces the subscriber on with a -// NULL level so the kernel honors RUST_LOG instead (the advanced override for -// tuning kernel-only verbosity). file_path=NULL sends kernel logs to stderr — the -// kernel ABI has no sink hook, so the Rust lines always go to stderr and are NOT -// routed through logger.SetLogOutput (unlike the Go binding lines). +// initKernelLogging routes the kernel's own Rust (tracing) logs into the driver's +// logger and points their verbosity at the driver's log level, so +// DATABRICKS_LOG_LEVEL drives both the Go binding lines and the kernel's Rust +// lines from one knob. +// +// The kernel exposes exactly one process-global tracing subscriber, reachable +// over the C ABI two mutually-exclusive ways: kernel_init_logging (RUST_LOG → +// stderr) and kernel_set_log_callback (reverse-call → a host sink). We install +// the CALLBACK (see installKernelLogCallback) so kernel-internal events flow into +// logger.Logger — the SAME unified stream as the Go binding lines (klog/klogCtx) +// and honouring logger.SetLogOutput — instead of only reaching stderr. +// +// resolveKernelLogArg maps the driver's zerolog level to the kernel's +// OFF/ERROR/WARN/INFO/DEBUG/TRACE string; DBSQL_KERNEL_DEBUG instead defers to +// RUST_LOG (useNULL → an empty level string → NULL to the kernel, the advanced +// override for tuning kernel-only verbosity). The callback filters at that level +// with the same precedence kernel_init_logging used, so the swap is level-transparent. // -// Best-effort: an Internal return (e.g. the host already installed a global -// subscriber) is a documented, benign outcome — logged at Warn, never fatal to -// connect. The subscriber installs at whatever level is mapped in; a driver left at -// the default Warn level (benchmarks included, and never having set -// DBSQL_KERNEL_DEBUG) installs it at WARN, so the kernel emits nothing below Warn -// and there is no hot-path cost. +// Best-effort: a non-Success install (e.g. the host already installed a global +// subscriber) is a documented, benign outcome — logged, never fatal to connect +// (handled inside installKernelLogCallback). A driver left at the default Warn +// level (benchmarks included, DBSQL_KERNEL_DEBUG unset) installs at WARN, so the +// kernel emits nothing below Warn and there is no hot-path cost. // -// Scope caveat: the kernel subscriber is PROCESS-WIDE, first-call-wins, and never +// Scope caveat: the subscriber is PROCESS-WIDE, first-call-wins, and never // uninstalled — in a long-lived multi-tenant process the first kernel session's -// level/destination applies to ALL subsequent kernel sessions, with no way to -// re-scope or turn it off afterward. That is a kernel-ABI property, not a Go one. -// A direct consequence: the driver level is sampled HERE, once, at the first kernel -// session — a later dbsql.SetLogLevel re-levels the Go binding lines (klog/klogCtx -// re-read GetLevel per call) but NOT the already-installed Rust subscriber. Set the -// level before opening the first kernel connection to govern the Rust logs. +// level applies to ALL subsequent kernel sessions. The driver level is sampled +// HERE, once, at the first kernel session — a later dbsql.SetLogLevel re-levels the +// Go binding lines (klog/klogCtx re-read GetLevel per call) but NOT the +// already-installed Rust subscriber. Set the level before opening the first kernel +// connection to govern the Rust logs. func initKernelLogging() { - initLoggingOnce.Do(func() { - // resolveKernelLogArg decides the level (or NULL for the DBSQL_KERNEL_DEBUG - // override, which lets the kernel honor RUST_LOG). The pure decision lives in - // logging_level.go so it's unit-tested without cgo. - var level cStr - if lvl, useNULL := resolveKernelLogArg(); !useNULL { - level = newCStr(lvl) - defer level.free() - } // else level stays {c: nil} → NULL → kernel honors RUST_LOG - if err := call(func() C.KernelStatusCode { - return C.kernel_init_logging(level.c, nil) - }); err != nil { - // The kernel subscriber didn't install (commonly: the host already - // installed a global tracing subscriber). Non-fatal — surface it through - // the shared logger so it's visible without a separate stderr scrape. - logger.Logger.Warn().Msgf("databricks: kernel_init_logging: %v (kernel logs unavailable; proceeding)", err) - } + // resolveKernelLogArg decides the level string (or "" for the + // DBSQL_KERNEL_DEBUG override, so the kernel honors RUST_LOG). The pure + // decision lives in logging_level.go so it's unit-tested without cgo. + level, _ := resolveKernelLogArg() + // installKernelLogCallback self-guards with logCallbackOnce (process-wide, + // first-call-wins), so no extra sync.Once here. + installKernelLogCallback(level) +} + +// abiCheckOnce ensures the ABI-version handshake runs at most once per process +// (the linked library can't change under us mid-run), and abiCheckErr caches its +// outcome so every OpenSession after the first returns the same verdict cheaply. +var ( + abiCheckOnce sync.Once + abiCheckErr error +) + +// checkABIVersion verifies the C-ABI version compiled into the linked kernel +// library matches the version the driver's header declares. The build-time +// status-code assertions in this file catch header-vs-source drift at compile +// time; this catches the runtime hazard those can't see — a driver built against +// one header linked against a differently-built prebuilt .a (the eventual +// download-a-prebuilt-lib distribution path). A mismatch means the +// KernelStatusCode enum / KernelError struct layout the driver reads may not +// match what the library writes, so we refuse to open rather than silently +// misread status codes / error fields. Runs once; the result is cached. The +// pure got-vs-want verdict lives in compareABI (untagged) so its mismatch path +// is testable without the cgo symbols. +func checkABIVersion() error { + abiCheckOnce.Do(func() { + got, want := abiVersions() + klog("checkABIVersion got=%d want=%d", got, want) + abiCheckErr = compareABI(got, want) }) + return abiCheckErr +} + +// abiVersions returns (library version, header version): the C-ABI version +// compiled into the linked kernel library and the version the driver's header +// declares. A production-side seam so tests can assert the handshake without +// putting cgo in a _test.go file (which Go forbids); not used in production +// beyond checkABIVersion. +func abiVersions() (got, want uint32) { + return uint32(C.kernel_abi_version()), uint32(C.DATABRICKS_KERNEL_ABI_VERSION) } // call runs a fallible kernel entry point and, on a non-Success status, reads diff --git a/internal/backend/kernel/config.go b/internal/backend/kernel/config.go index f0bc8e28..275a8423 100644 --- a/internal/backend/kernel/config.go +++ b/internal/backend/kernel/config.go @@ -35,6 +35,19 @@ type Config struct { TLSTrustedCertsPEM []byte TLSSkipHostnameVerify bool + // TLSClientCertPEM / TLSClientKeyPEM are the mTLS client cert (+ optional + // chain) and matching private key (from WithKernelClientCertificate). They + // travel as a pair — both set or both empty — and are forwarded via the single + // paired kernel_session_config_set_tls_client_certificate. The key is never + // logged. + TLSClientCertPEM []byte + TLSClientKeyPEM []byte + + // CloudFetchEnabled toggles CloudFetch (from WithKernelCloudFetch). Tri-state: + // nil keeps the kernel default (on); a non-nil value is forwarded via + // kernel_session_config_set_cloudfetch_enabled. + CloudFetchEnabled *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 diff --git a/internal/backend/kernel/kernel_test.go b/internal/backend/kernel/kernel_test.go index 4182477c..8d2e8546 100644 --- a/internal/backend/kernel/kernel_test.go +++ b/internal/backend/kernel/kernel_test.go @@ -9,6 +9,7 @@ import ( "encoding/json" "errors" "os" + "runtime/cgo" "testing" "github.com/databricks/databricks-sql-go/driverctx" @@ -58,6 +59,11 @@ func TestSetAuthByMode(t *testing.T) { // drifted. func TestSetKernelTLS(t *testing.T) { ca := []byte("-----BEGIN CERTIFICATE-----\nca\n-----END CERTIFICATE-----\n") + cert := []byte("-----BEGIN CERTIFICATE-----\nleaf\n-----END CERTIFICATE-----\n") + // Assemble the key's PEM marker at runtime so the repo's secret scanner + // doesn't flag a literal BEGIN-PRIVATE-KEY (the bytes are opaque to the + // marshalling path under test). + key := []byte("-----BEGIN " + "PRIVATE" + " KEY-----\nkey\n-----END PRIVATE KEY-----\n") cases := []struct { name string cfg Config @@ -65,6 +71,8 @@ func TestSetKernelTLS(t *testing.T) { {"trusted certs only", Config{TLSTrustedCertsPEM: ca}}, {"skip hostname only", Config{TLSSkipHostnameVerify: true}}, {"both together", Config{TLSTrustedCertsPEM: ca, TLSSkipHostnameVerify: true}}, + {"client certificate (mTLS)", Config{TLSClientCertPEM: cert, TLSClientKeyPEM: key}}, + {"all together", Config{TLSTrustedCertsPEM: ca, TLSSkipHostnameVerify: true, TLSClientCertPEM: cert, TLSClientKeyPEM: key}}, {"none (no-op)", Config{}}, } for _, c := range cases { @@ -76,6 +84,127 @@ func TestSetKernelTLS(t *testing.T) { } } +// TestABIVersionMatches asserts the linked kernel library's C-ABI version +// matches the header the driver compiled against — the same handshake +// checkABIVersion runs at connect. It exercises the real cgo symbols +// (kernel_abi_version + the DATABRICKS_KERNEL_ABI_VERSION macro), so a stale +// .a-vs-header pairing fails here at test time rather than misreading status +// codes at runtime. It also asserts checkABIVersion() returns nil for the +// matched pair the test binary links. +func TestABIVersionMatches(t *testing.T) { + got, want := abiVersions() + if got != want { + t.Fatalf("kernel_abi_version() = %d, header DATABRICKS_KERNEL_ABI_VERSION = %d", got, want) + } + if err := checkABIVersion(); err != nil { + t.Errorf("checkABIVersion() = %v, want nil for the linked (matching) library", err) + } +} + +// TestLogCallbackRoundTrip exercises the full reverse-call path: wrap an +// observable *logSink in a cgo.Handle, drive the C→Go trampoline exactly as the +// kernel's drain thread would (via the test seam), and assert the event arrived +// with the handle correctly unwrapped and the fields intact. This proves the +// cgo.Handle round-trip + recover firewall + sink dispatch that a real kernel log +// event flows through (kernel-side delivery itself is covered by the kernel's own +// integration tests). +func TestLogCallbackRoundTrip(t *testing.T) { + type got struct { + level int + target, message string + } + ch := make(chan got, 1) + sink := &logSink{observe: func(level int, target, message string) { + ch <- got{level, target, message} + }} + h := cgo.NewHandle(sink) + defer h.Delete() + + invokeLogTrampolineForTest(h, kernelLevelWarn, "databricks::sql::kernel", "retrying request") + + select { + case g := <-ch: + if g.level != kernelLevelWarn || g.target != "databricks::sql::kernel" || g.message != "retrying request" { + t.Errorf("delivered event = %+v, want level=WARN target=databricks::sql::kernel msg='retrying request'", g) + } + default: + t.Fatal("the trampoline did not deliver the event to the sink") + } +} + +// captureDriverLog points the global driver logger at a buffer (TraceLevel) for +// the duration of a test, returning the buffer and restoring the logger after. +// Used to assert a trampoline guard branch delivered NOTHING (the branches bail +// before reaching a sink, so there's no observe hook to check — the driver +// logger output is the observable). +func captureDriverLog(t *testing.T) *bytes.Buffer { + t.Helper() + var buf bytes.Buffer + prev := logger.Logger.GetLevel() + logger.SetLogOutput(&buf) + _ = logger.SetLogLevel("trace") + t.Cleanup(func() { + logger.SetLogOutput(os.Stderr) + logger.Logger.Logger = logger.Logger.Level(prev) + }) + return &buf +} + +// A zero ctx must be a no-op (no crash, no delivery). cgo.Handle(0) casts to a +// NULL void* ctx, so the trampoline returns at its `ctx == nil` guard before ever +// calling Value() (calling Value() on the zero handle would itself panic — the +// recover firewall would catch it, but the nil guard means we never reach it). +func TestLogCallbackTrampolineNilCtxSafe(t *testing.T) { + buf := captureDriverLog(t) + invokeLogTrampolineForTest(cgo.Handle(0), kernelLevelError, "databricks::sql::kernel", "should-not-appear") + if buf.Len() != 0 { + t.Errorf("nil-ctx trampoline delivered a line, want none: %q", buf.String()) + } +} + +// TestLogCallbackTrampolineWrongTypeSafe covers the "handle holds a non-*logSink +// value" branch: a live, non-nil handle whose Value() type-asserts to false. The +// trampoline must return without panicking AND without delivering (asserted via +// the captured driver logger staying empty). +func TestLogCallbackTrampolineWrongTypeSafe(t *testing.T) { + buf := captureDriverLog(t) + h := cgo.NewHandle("not a logSink") + defer h.Delete() + invokeLogTrampolineForTest(h, kernelLevelError, "databricks::sql::kernel", "should-not-appear") + if buf.Len() != 0 { + t.Errorf("wrong-type trampoline delivered a line, want none: %q", buf.String()) + } +} + +// TestLogCallbackPanicFirewall is the load-bearing safety test: a sink whose +// routing panics must NOT crash the process — the trampoline's defer recover() +// converts it into a dropped line (a panic unwinding across the cgo boundary +// would abort). We drive a panicking observe hook, assert the call returns, and +// assert a subsequent well-formed sink still delivers (the firewall didn't wedge +// anything process-global). +func TestLogCallbackPanicFirewall(t *testing.T) { + panicSink := &logSink{observe: func(int, string, string) { panic("boom in the sink") }} + ph := cgo.NewHandle(panicSink) + defer ph.Delete() + // Must not crash the test binary. + invokeLogTrampolineForTest(ph, kernelLevelError, "databricks::sql::kernel", "will panic") + + // A subsequent well-formed delivery still works. + ch := make(chan string, 1) + okSink := &logSink{observe: func(_ int, _, message string) { ch <- message }} + oh := cgo.NewHandle(okSink) + defer oh.Delete() + invokeLogTrampolineForTest(oh, kernelLevelWarn, "databricks::sql::kernel", "recovered") + select { + case got := <-ch: + if got != "recovered" { + t.Errorf("post-panic delivery = %q, want 'recovered'", got) + } + default: + t.Fatal("delivery after a panicking callback did not work — firewall wedged") + } +} + // TestKernelLogLevel and TestResolveKernelLogArg — the pure level-resolution tests — // live in the untagged logging_level_test.go so they run under CGO_ENABLED=0. The // tests below exercise klog/klogCtx and so need the cgo build. diff --git a/internal/backend/kernel/log_callback.go b/internal/backend/kernel/log_callback.go new file mode 100644 index 00000000..84f23062 --- /dev/null +++ b/internal/backend/kernel/log_callback.go @@ -0,0 +1,134 @@ +//go:build cgo && databricks_kernel + +package kernel + +/* +#include +#include "databricks_kernel.h" + +// Adapter matching the kernel_log_callback typedef (const char*) that forwards +// to the //export'd Go trampoline. cgo generates the trampoline's prototype in +// _cgo_export.h with non-const char* params, which don't match the typedef, so +// this shim bridges the two and gives kernel_set_log_callback one clean address +// to take. Forward-declared with the SAME (char*) signature cgo generates to +// avoid a conflicting-declaration error; the kernel only ever reads the strings. +void kernelLogTrampoline(void* ctx, int level, char* target, char* message); +// Hand kernel_set_log_callback the trampoline's address as a kernel_log_callback. +// The //export'd trampoline's generated signature uses non-const char* while the +// typedef uses const char*, which are function-pointer compatible (a const-only +// difference); do the cast once here, on the C side, so the Go call site stays +// clean and there is no separately-linked adapter symbol. +static kernel_log_callback kernel_log_cb(void) { + return (kernel_log_callback)kernelLogTrampoline; +} +// Cast a cgo.Handle (an integer token) to the opaque void* ctx on the C side, +// so Go never does unsafe.Pointer(uintptr) directly (which `go vet` flags). +static void* kernel_handle_to_ctx(uintptr_t h) { return (void*)h; } +// Test-only: invoke the trampoline exactly as the kernel drain thread would. +// (A C func-pointer value can't be called directly from Go, so a helper does it.) +static void kernel_invoke_trampoline_for_test(void* ctx, int level, char* target, char* message) { + kernelLogTrampoline(ctx, level, target, message); +} +*/ +import "C" + +import ( + "runtime/cgo" + "sync" + "unsafe" + + "github.com/databricks/databricks-sql-go/logger" +) + +// This file holds the reverse-call machinery for the kernel log callback: +// a cgo-exported Go function the kernel invokes from its log-drain thread, made +// safe with a recover() panic firewall (a panic across the cgo boundary would +// abort the process) and a runtime/cgo.Handle so a Go pointer (the *logSink) can +// round-trip through the C void* ctx under cgo's pointer rules. + +// logCallbackOnce guards the one-time callback registration; logCallbackHandle +// keeps the cgo.Handle alive for the process (the kernel holds its numeric value +// as the opaque ctx and there is no detach in the v0 C ABI, so we never Delete +// it — a deliberate process-lifetime pin, not a leak). +var ( + logCallbackOnce sync.Once + logCallbackHandle cgo.Handle +) + +// invokeLogTrampolineForTest drives the C→Go trampoline exactly as the kernel's +// drain thread would — allocating C strings, casting a cgo.Handle to void*, and +// calling through kernel_log_cb() — so a test can assert the full reverse-call +// round-trip (handle unwrap + recover firewall + sink dispatch) without a live +// kernel event. Not used in production. +func invokeLogTrampolineForTest(h cgo.Handle, level int, target, message string) { + ct := C.CString(target) + defer C.free(unsafe.Pointer(ct)) + cm := C.CString(message) + defer C.free(unsafe.Pointer(cm)) + ctx := C.kernel_handle_to_ctx(C.uintptr_t(h)) + C.kernel_invoke_trampoline_for_test(ctx, C.int(level), ct, cm) +} + +//export kernelLogTrampoline +func kernelLogTrampoline(ctx unsafe.Pointer, level C.int, target, message *C.char) { + // Panic firewall: a Go panic unwinding across the cgo boundary aborts the + // process. Convert any panic in the sink routing into a dropped log line. + defer func() { _ = recover() }() + if ctx == nil { + return + } + sink, ok := cgo.Handle(uintptr(ctx)).Value().(*logSink) + if !ok || sink == nil { + return + } + // Copy the borrowed C strings out immediately — they are valid only for the + // duration of this call. + sink.forward(int(level), C.GoString(target), C.GoString(message)) +} + +// installKernelLogCallback registers the trampoline as the kernel's log sink, +// once per process, wrapping a *logSink in a cgo.Handle passed as the opaque ctx. +// The kernel filters forwarded events at `level` (an OFF/ERROR/WARN/INFO/DEBUG/ +// TRACE string, or "" → NULL to defer to RUST_LOG), which the caller maps from +// the driver's own log level so kernel Rust lines follow the one +// DATABRICKS_LOG_LEVEL knob — the same string kernel_init_logging would have +// received (see initKernelLogging). +// +// Best-effort: a non-Success status (e.g. Internal because a global tracing +// subscriber was already installed) is logged and ignored — kernel logs simply +// won't flow through the callback, which is a documented, benign outcome. Returns +// nothing; the caller does not gate connect on it. +func installKernelLogCallback(level string) { + logCallbackOnce.Do(func() { + // Snapshot the driver logger at install (TraceLevel, immutable) so the + // kernel drain thread never re-gates already-approved events and never + // reads the mutable global logger.Logger (see logSink.log). + logCallbackHandle = cgo.NewHandle(newLogSink()) + // Pass the handle (an integer token) as the opaque ctx via a C-side cast, + // so Go doesn't do unsafe.Pointer(uintptr) directly. + ctx := C.kernel_handle_to_ctx(C.uintptr_t(logCallbackHandle)) + // An empty level maps to a NULL char* (kernel defers to RUST_LOG); a + // non-empty level is passed as a C string the kernel parses with the same + // precedence as kernel_init_logging (explicit wins > RUST_LOG > warn). + clevel := newCStrOrNull(level) + defer clevel.free() + if err := call(func() C.KernelStatusCode { + return C.kernel_set_log_callback(C.kernel_log_cb(), ctx, clevel.c) + }); err != nil { + // On a non-Success return the kernel installed no subscriber and did + // NOT retain ctx (it stores the sink only on the success path), so the + // process-lifetime pin at logCallbackHandle would be a true leak here — + // free it. (On success we deliberately never Delete: the kernel holds + // the handle value as its opaque ctx for the process lifetime.) + logCallbackHandle.Delete() + logCallbackHandle = 0 + // Surface at Warn (visible at the default level), NOT klog: klog is + // Debug-gated and no-ops at the default, and this install runs once so + // a later level change can't re-surface it. The message names the + // benign common cause (a global subscriber already installed) so a + // logging no-op isn't mistaken for a connect failure. + logger.Logger.Warn().Msgf( + "databricks: kernel_set_log_callback: %v (kernel logs not forwarded; proceeding)", err) + } + }) +} diff --git a/internal/backend/kernel/logforward.go b/internal/backend/kernel/logforward.go new file mode 100644 index 00000000..e9258245 --- /dev/null +++ b/internal/backend/kernel/logforward.go @@ -0,0 +1,126 @@ +package kernel + +// This file is intentionally NOT behind the cgo build tag: the level mapping +// and sink-routing logic are pure Go, so their tests run under CGO_ENABLED=0. +// The cgo trampoline that calls into these lives in log_callback.go (tagged). + +import ( + "os" + + "github.com/databricks/databricks-sql-go/logger" + "github.com/rs/zerolog" +) + +// Kernel log levels as delivered over the C ABI (1=ERROR..5=TRACE), mirroring +// the kernel's numeric severity. 0 is unused on the wire (the kernel never +// forwards an OFF event). +const ( + kernelLevelError = 1 + kernelLevelWarn = 2 + kernelLevelInfo = 3 + kernelLevelDebug = 4 + kernelLevelTrace = 5 +) + +// logSink receives forwarded kernel log events and routes them into the driver's +// logger, so kernel-internal diagnostics (HTTP stack, CloudFetch, retry) join the +// unified Go+kernel debug stream instead of only reaching stderr via RUST_LOG. +// A dedicated type (rather than calling the logger directly from the trampoline) +// keeps the routing pure-Go and unit-testable, and gives the cgo.Handle something +// concrete to wrap. +type logSink struct { + // log is the destination for forwarded events. It is a SNAPSHOT of the + // driver logger taken once at install (newLogSink), pinned to TraceLevel, and + // never reassigned. Two reasons: + // - No double gating. The KERNEL already filters events against the level + // the driver passed to kernel_set_log_callback (DATABRICKS_LOG_LEVEL, or + // RUST_LOG under DBSQL_KERNEL_DEBUG); that decision is authoritative. If + // we routed through the live logger.Logger — gated at the driver level — + // a DEBUG kernel event under the DBSQL_KERNEL_DEBUG override would be + // re-dropped by a driver still at Warn, defeating the override. Pinning + // the sink at TraceLevel lets every already-approved event through. + // - No data race. The kernel drain thread (a non-Go OS thread) calls + // forward asynchronously; reading the mutable global logger.Logger here + // would race a concurrent dbsql.SetLogLevel/SetLogOutput. A value copy + // captured at install is immutable, so the drain thread reads only its + // own snapshot. (The writer/output is captured at install; a later + // SetLogOutput does not re-target already-forwarded kernel lines — the + // same "level fixed at first connect" caveat the kernel subscriber has.) + log zerolog.Logger + + // nativeSeverity, when true, emits each forwarded kernel event at its own mapped + // level (kernel ERROR→Error, WARN→Warn). When false (the default) the two + // operationally-loud levels are demoted one step — kernel ERROR→Warn, WARN→Info — + // so kernel-internal transient events (e.g. "retrying request" during a 503 / S3 + // storm) do not inflate the driver's shared Warn/Error stream that log-based + // alerting keys on. INFO/DEBUG/TRACE are unaffected (already below the alerting + // threshold). Set from DBSQL_KERNEL_DEBUG at install: a caller who opted into + // kernel debugging wants the native severities. See newLogSink. + nativeSeverity bool + + // observe, when non-nil, is invoked with each forwarded event before the + // logger routing. A test seam so a round-trip test can assert the trampoline + // unwrapped the handle and delivered the event; nil in production. + observe func(level int, target, message string) +} + +// newLogSink snapshots the current driver logger at TraceLevel for the sink to +// forward through. Called once at install so the drain thread never touches the +// mutable global logger. Kept separate from the zero value so tests can build a +// sink with an explicit writer. +// +// It also decides the severity policy: by default forwarded kernel ERROR/WARN are +// demoted one step (see logSink.nativeSeverity) so routine kernel-internal +// retries/backoff don't page on-call through the driver's own WARN/ERROR rate; +// DBSQL_KERNEL_DEBUG (the same knob that widens kernel verbosity) preserves the +// native severities for a debugging session. +func newLogSink() *logSink { + return &logSink{ + log: logger.Logger.Level(zerolog.TraceLevel), + nativeSeverity: os.Getenv("DBSQL_KERNEL_DEBUG") != "", + } +} + +// forward routes one kernel log event into the sink's logger at the mapped level. +// The sink logger is a TraceLevel snapshot of the driver logger (see logSink.log), +// so events the kernel already approved are emitted without being re-gated by the +// driver level. Kept small and allocation-light: it is on the (debug-only) log +// path, not a hot query path. target and message are already Go strings (copied +// out of the C buffers by the trampoline before this is called). +func (s *logSink) forward(level int, target, message string) { + if s == nil { + return + } + if s.observe != nil { + s.observe(level, target, message) + } + switch level { + case kernelLevelError: + // Demote to Warn by default so a kernel-internal error (e.g. a retried, + // ultimately-recovered request) doesn't inflate the driver's Error rate; + // DBSQL_KERNEL_DEBUG keeps the native Error severity. + if s.nativeSeverity { + s.log.Error().Str("target", target).Msg(message) + } else { + s.log.Warn().Str("target", target).Str("kernelLevel", "error").Msg(message) + } + case kernelLevelWarn: + // Demote to Info by default so transient kernel WARN lines (503/S3 retry, + // backoff) don't page on-call through the driver's WARN rate. + if s.nativeSeverity { + s.log.Warn().Str("target", target).Msg(message) + } else { + s.log.Info().Str("target", target).Str("kernelLevel", "warn").Msg(message) + } + case kernelLevelInfo: + s.log.Info().Str("target", target).Msg(message) + case kernelLevelDebug: + s.log.Debug().Str("target", target).Msg(message) + case kernelLevelTrace: + s.log.Trace().Str("target", target).Msg(message) + default: + // Unknown level: don't drop it — surface at Debug so it's still visible + // without being mistaken for a warning/error. + s.log.Debug().Str("target", target).Int("kernelLevel", level).Msg(message) + } +} diff --git a/internal/backend/kernel/logforward_test.go b/internal/backend/kernel/logforward_test.go new file mode 100644 index 00000000..9d03eb74 --- /dev/null +++ b/internal/backend/kernel/logforward_test.go @@ -0,0 +1,133 @@ +package kernel + +import ( + "bytes" + "encoding/json" + "os" + "strings" + "testing" + + "github.com/databricks/databricks-sql-go/logger" + "github.com/rs/zerolog" +) + +// The level mapping and nil-sink no-op are pure Go, so they run under +// CGO_ENABLED=0. Actual kernel→callback delivery is exercised by the tagged +// TestLogCallbackRoundTrip (which needs the linked kernel). + +// A nil sink must be a safe no-op (the trampoline guards against a missing/typed +// handle by calling forward on a possibly-nil sink). +func TestLogSinkForwardNilIsNoOp(t *testing.T) { + var s *logSink + // Must not panic. + s.forward(kernelLevelError, "databricks::sql::kernel", "boom") +} + +// forward must map each kernel level to the right zerolog level and carry the +// target + message through. A buffer-backed TraceLevel sink (mirroring the +// production snapshot) lets us assert the emitted JSON rather than only "no +// panic" — a bug swapping e.g. ERROR→Debug would otherwise ship green. +// +// Two policies are exercised: the DEFAULT demotes the operationally-loud kernel +// ERROR/WARN one step (→Warn/Info) so kernel-internal retries don't inflate the +// driver's alert-keyed WARN/ERROR rate, tagging the demoted line with a +// kernelLevel="error"/"warn" string so its true origin stays visible; and +// nativeSeverity (DBSQL_KERNEL_DEBUG) preserves the native mapping. +func TestLogSinkForwardMapsLevels(t *testing.T) { + cases := []struct { + kernelLevel int + wantLevel string // zerolog "level" field under the DEFAULT (demoting) policy + wantNative string // zerolog "level" field with nativeSeverity=true + wantKernelStr string // kernelLevel string field on the default line ("" = none expected) + unknown bool // unmapped → Debug WITH a numeric kernelLevel diagnostic field + }{ + {kernelLevelError, "warn", "error", "error", false}, + {kernelLevelWarn, "info", "warn", "warn", false}, + {kernelLevelInfo, "info", "info", "", false}, + {kernelLevelDebug, "debug", "debug", "", false}, + {kernelLevelTrace, "trace", "trace", "", false}, + {0, "debug", "debug", "", true}, + {99, "debug", "debug", "", true}, + {-1, "debug", "debug", "", true}, + } + for _, c := range cases { + // Default (demoting) policy. + var buf bytes.Buffer + s := &logSink{log: zerolog.New(&buf).Level(zerolog.TraceLevel)} + s.forward(c.kernelLevel, "databricks::sql::kernel", "hello") + var rec map[string]any + if err := json.Unmarshal(bytes.TrimSpace(buf.Bytes()), &rec); err != nil { + t.Fatalf("kernelLevel=%d: emitted non-JSON %q: %v", c.kernelLevel, buf.String(), err) + } + if rec["level"] != c.wantLevel { + t.Errorf("kernelLevel=%d (default policy): level=%v, want %s", c.kernelLevel, rec["level"], c.wantLevel) + } + // The unknown-level branch adds a NUMERIC kernelLevel field — the only thing + // distinguishing an unmapped level from a real Debug line. A demoted + // error/warn adds a STRING kernelLevel naming its true origin. Assert both so + // dropping either ships red. + switch { + case c.unknown: + if got, ok := rec["kernelLevel"]; !ok || got != float64(c.kernelLevel) { + t.Errorf("kernelLevel=%d: expected numeric kernelLevel field=%d, got %v (present=%t)", + c.kernelLevel, c.kernelLevel, got, ok) + } + case c.wantKernelStr != "": + if got, ok := rec["kernelLevel"]; !ok || got != c.wantKernelStr { + t.Errorf("kernelLevel=%d: expected demoted line to carry kernelLevel=%q, got %v (present=%t)", + c.kernelLevel, c.wantKernelStr, got, ok) + } + default: + if _, ok := rec["kernelLevel"]; ok { + t.Errorf("kernelLevel=%d: an undemoted mapped level should NOT carry a kernelLevel field", c.kernelLevel) + } + } + if rec["target"] != "databricks::sql::kernel" || rec["message"] != "hello" { + t.Errorf("kernelLevel=%d: target/message not carried through: %v", c.kernelLevel, rec) + } + + // nativeSeverity policy: ERROR/WARN keep their native level and carry no + // demotion tag. + var nbuf bytes.Buffer + ns := &logSink{log: zerolog.New(&nbuf).Level(zerolog.TraceLevel), nativeSeverity: true} + ns.forward(c.kernelLevel, "databricks::sql::kernel", "hello") + var nrec map[string]any + if err := json.Unmarshal(bytes.TrimSpace(nbuf.Bytes()), &nrec); err != nil { + t.Fatalf("kernelLevel=%d (native): emitted non-JSON %q: %v", c.kernelLevel, nbuf.String(), err) + } + if nrec["level"] != c.wantNative { + t.Errorf("kernelLevel=%d (native policy): level=%v, want %s", c.kernelLevel, nrec["level"], c.wantNative) + } + if !c.unknown { + if _, ok := nrec["kernelLevel"]; ok { + t.Errorf("kernelLevel=%d (native): should NOT carry a kernelLevel field", c.kernelLevel) + } + } + } +} + +// A TraceLevel snapshot must emit events the driver's own level would suppress — +// this is the anti-double-gating property (forwarded kernel events the kernel +// already approved must not be re-dropped by a driver still at Warn). Drives the +// PRODUCTION newLogSink() against a Warn-level global logger, so a regression of +// newLogSink to `&logSink{log: logger.Logger}` (routing through the live, +// driver-gated logger) fails here rather than passing against a hand-rolled replica. +func TestLogSinkForwardNotReGatedByDriverLevel(t *testing.T) { + // Point the global driver logger at a buffer and pin it to Warn — the state + // newLogSink() reads at install. Restore afterward so other tests are + // unaffected. + var buf bytes.Buffer + prevLevel := logger.Logger.GetLevel() + logger.SetLogOutput(&buf) + _ = logger.SetLogLevel("warn") + t.Cleanup(func() { + logger.SetLogOutput(os.Stderr) + logger.Logger.Logger = logger.Logger.Level(prevLevel) + }) + + s := newLogSink() // snapshots logger.Logger at TraceLevel — the real path + s.forward(kernelLevelDebug, "databricks::sql::kernel", "debug-line") + if !strings.Contains(buf.String(), "debug-line") { + t.Errorf("a DEBUG kernel event was dropped despite newLogSink's TraceLevel snapshot: %q", buf.String()) + } +} diff --git a/internal/backend/kernel/operation.go b/internal/backend/kernel/operation.go index 4e0228b9..529e5fce 100644 --- a/internal/backend/kernel/operation.go +++ b/internal/backend/kernel/operation.go @@ -173,16 +173,13 @@ func (k *KernelBackend) execute(ctx context.Context, req backend.ExecRequest) (b // EXECUTE_STATEMENT on StatementID() != "", so without this a failed // kernel query records no metric while the same Thrift failure does. op.statementID = statementIDFromError(execErr) - // Prefer the caller's ctx error when the ctx was cancelled (database/sql - // convention). Wrap BOTH the ctx error and the kernel error so - // errors.Is(err, context.Canceled/DeadlineExceeded) still matches (the - // cancellation contract) AND the *KernelError stays reachable via errors.As — - // otherwise a session-fatal failure racing a cancel would lose its - // sqlstate/queryId, the one handle to what actually went wrong server-side. + // Prefer the caller's ctx error when the ctx was cancelled; cancelledErr + // holds the shared dual-%w wrap (see its doc) so errors.Is still matches the + // ctx error AND the *KernelError stays reachable via errors.As. if ctx.Err() != nil { klogCtx(ctx, "Execute failed under cancelled ctx: kernelErr=%v ctxErr=%v", execErr, ctx.Err()) op.close() - return op, fmt.Errorf("kernel: execute cancelled: %w (kernel error: %w)", ctx.Err(), toStatementError(execErr)) + return op, cancelledErr("execute", ctx.Err(), toStatementError(execErr)) } klogCtx(ctx, "Execute failed: %v", execErr) // We still return the PLAIN error (toStatementError, never ErrBadConn), so the diff --git a/internal/backend/kernel/rows.go b/internal/backend/kernel/rows.go index 44557bf1..10887ab7 100644 --- a/internal/backend/kernel/rows.go +++ b/internal/backend/kernel/rows.go @@ -41,6 +41,16 @@ type kernelRows struct { op *kernelOp stream *C.kernel_result_stream_t callbacks *dbsqlrows.TelemetryCallbacks + // watcher bridges r.ctx onto ONE kernel cancel token for the whole result set, + // created once here rather than per nextBatch: a cancellable request-scoped ctx + // is the common service case, and a multi-GB CloudFetch stream is thousands of + // batches, so a fresh token + goroutine + teardown per batch would be pure churn + // on the driver's known-sensitive large-result path. Amortizing to one watcher + // matches the execute path's once-per-statement canceller. Firing the token + // aborts the in-flight fetch and, since it stays fired, any subsequent one — which + // is correct: iteration stops on the returned error. nil for a non-cancellable ctx + // (NULL token → the plain fetch path, zero overhead). Stopped in Close. + watcher *ctxWatcher cols []string cur arrow.Record // current batch (nil until first Next) @@ -69,6 +79,11 @@ func newKernelRows(ctx context.Context, op *kernelOp, stream *C.kernel_result_st // 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()} + // One cancel-token watcher for the whole result set (see kernelRows.watcher), + // nil on a non-cancellable ctx. Created before the fallible schema fetch so the + // two construction-failure r.Close() calls below tear it down too; Close is + // idempotent and stop() is nil-safe. + r.watcher = newCtxWatcher(ctx) var csch C.struct_ArrowSchema if err := call(func() C.KernelStatusCode { @@ -106,6 +121,10 @@ func (r *kernelRows) Close() error { } r.closed = true closeStart := time.Now() + // Drain the ctx watcher and free its token first: it may be mid-fire (inside + // kernel_cancel_token_cancel) and its token must not be freed out from under a + // concurrent fire. stop() is nil-safe (non-cancellable ctx). + r.watcher.stop() if r.cur != nil { r.cur.Release() r.cur = nil @@ -168,14 +187,19 @@ func (r *kernelRows) next(dest []driver.Value) error { // nextBatch pulls the next Arrow batch. A released array (release==NULL) is the // kernel's end-of-stream sentinel. func (r *kernelRows) nextBatch() error { - // Honor cancellation at batch boundaries: check ctx before entering the - // blocking C fetch (which cannot itself observe ctx). This does NOT interrupt - // a fetch already in flight — and database/sql's own cancel watcher can't - // either: its Rows.Close takes rs.closemu.Lock(), which blocks until the - // in-progress Next (holding the RLock) returns, so the stream close waits for - // the C call to finish on its own. A single hung CloudFetch batch is therefore - // uninterruptible (the kernel exposes no per-download timeout); mid-fetch - // cancellation would need the execute path's watcher/canceller applied here. + // Honor cancellation both at the batch boundary AND mid-fetch. The pre-fetch + // ctx check fast-fails an already-cancelled ctx without dialing; the result-set + // cancel token (r.watcher, created once in newKernelRows) then bridges the + // deadline into the kernel so a fetch already in flight — a hung CloudFetch chunk + // (a wedged S3 / pre-signed-URL GET) — no longer blocks Next past the deadline. + // Firing the token unblocks this call promptly (it stops waiting on the batch); + // the kernel's background download of that chunk continues to completion or its + // own read-timeout (~60s) rather than being torn down mid-flight, so the caller's + // deadline is honored while the socket is reclaimed shortly after — see the scope + // note on kernel_result_stream_next_batch_cancellable in databricks_kernel.h. A + // NULL token (uncancellable ctx) makes the cancellable fetch behave exactly like + // the plain kernel_result_stream_next_batch, so there is no watcher overhead on + // the common background-context path. if r.ctx != nil { if err := r.ctx.Err(); err != nil { return err @@ -188,9 +212,21 @@ func (r *kernelRows) nextBatch() error { 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) + return C.kernel_result_stream_next_batch_cancellable(r.stream, &carr, &csch, r.watcher.tokenPtr()) }); err != nil { + // Evict a session-fatal conn BEFORE the ctx-cancelled branch below: a + // session-fatal fetch failure (expired token, dropped/unavailable session) + // racing a ctx deadline/cancel is still session-fatal, so returning the ctx + // error first would leave a dead conn marked valid in the pool. Mirrors the + // execute path's evict-before-ctx ordering in operation.go. r.op.backend.evictIfSessionFatal(err) + // Prefer the caller's ctx error when the fetch was interrupted; cancelledErr + // holds the shared dual-%w wrap (see its doc) so errors.Is still matches the + // ctx error AND the *KernelError stays reachable via errors.As. + if r.ctx != nil && r.ctx.Err() != nil { + klogCtx(r.ctx, "nextBatch interrupted by ctx: kernelErr=%v ctxErr=%v", err, r.ctx.Err()) + return cancelledErr("next_batch", r.ctx.Err(), toStatementError(err)) + } return fmt.Errorf("kernel: next_batch: %w", toStatementError(err)) } if carr.release == nil { diff --git a/internal/config/config.go b/internal/config/config.go index f689f414..b22e13a1 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -76,6 +76,29 @@ type KernelExperimentalConfig struct { // of the blanket InsecureSkipVerify (which relaxes both chain and hostname). // Maps to kernel_session_config_set_tls_skip_hostname_verification. TLSSkipHostnameVerify bool + // TLSClientCertPEM / TLSClientKeyPEM are the client leaf cert (+ optional + // chain) and matching private key for mutual TLS (mTLS). They travel as a + // pair — both set or both nil (WithKernelClientCertificate is the only setter + // and takes both) — and map to the single paired + // kernel_session_config_set_tls_client_certificate. The key is never logged. + TLSClientCertPEM []byte + TLSClientKeyPEM []byte + // TLSClientCertConfigured records that WithKernelClientCertificate was + // invoked, independent of whether the caller passed non-empty bytes. It + // exists because an empty cert+key pair is otherwise indistinguishable from + // the option never being called (KernelExperimental can be non-nil from any + // other WithKernel* option), which would let a caller who explicitly asked + // for mTLS — e.g. after a failed/empty PEM load — silently connect with no + // client identity (applyKernelTLS gates the setter on a non-empty cert). + // validateKernelConfig uses this marker to reject an incomplete mTLS request + // loudly instead of failing open. Never forwarded to the kernel C ABI. + TLSClientCertConfigured bool + // CloudFetchEnabled toggles CloudFetch on the kernel path. Tri-state: nil + // keeps the kernel default (on); a non-nil *false forces the inline-Arrow + // path. A *bool (not a plain bool) so "unset" is distinguishable from + // "explicitly false" — maps to kernel_session_config_set_cloudfetch_enabled. + // Deliberately NOT the plain-bool WithCloudFetch (which can't express unset). + CloudFetchEnabled *bool } // DeepCopy returns a deep copy of the experimental config, or nil for a nil @@ -85,10 +108,23 @@ func (k *KernelExperimentalConfig) DeepCopy() *KernelExperimentalConfig { if k == nil { return nil } - cp := &KernelExperimentalConfig{TLSSkipHostnameVerify: k.TLSSkipHostnameVerify} + cp := &KernelExperimentalConfig{ + TLSSkipHostnameVerify: k.TLSSkipHostnameVerify, + TLSClientCertConfigured: k.TLSClientCertConfigured, + } if k.TLSTrustedCertsPEM != nil { cp.TLSTrustedCertsPEM = append([]byte(nil), k.TLSTrustedCertsPEM...) } + if k.TLSClientCertPEM != nil { + cp.TLSClientCertPEM = append([]byte(nil), k.TLSClientCertPEM...) + } + if k.TLSClientKeyPEM != nil { + cp.TLSClientKeyPEM = append([]byte(nil), k.TLSClientKeyPEM...) + } + if k.CloudFetchEnabled != nil { + v := *k.CloudFetchEnabled + cp.CloudFetchEnabled = &v + } return cp } diff --git a/kernel_backend.go b/kernel_backend.go index e956058f..4a9345da 100644 --- a/kernel_backend.go +++ b/kernel_backend.go @@ -26,9 +26,10 @@ func newKernelBackend(_ context.Context, cfg *config.Config) (backend.Backend, e } // Assemble the flat kernel.Config from the driver config. The field-by-field - // mapping (including the experimental TLS forwarding) is the pure, cgo-free - // buildKernelConfig in kernel_config.go, unit-tested under CGO_ENABLED=0. - // SPOG org routing rides in HTTPPath's ?o= and is parsed kernel-side. + // mapping (including the experimental TLS / mTLS / CloudFetch forwarding) is the + // pure, cgo-free buildKernelConfig in kernel_config.go, unit-tested under + // CGO_ENABLED=0. SPOG org routing rides in HTTPPath's ?o= and is parsed + // kernel-side. kc := buildKernelConfig(cfg, kauth) // Proxy: the Thrift path uses http.ProxyFromEnvironment; mirror it by reading // the same HTTP(S)_PROXY / NO_PROXY environment for the kernel. Resolved here diff --git a/kernel_config.go b/kernel_config.go index ecf4cc84..9a256dcc 100644 --- a/kernel_config.go +++ b/kernel_config.go @@ -86,6 +86,42 @@ func validateKernelConfig(cfg *config.Config) (kernel.Auth, error) { return kernel.Auth{}, fmt.Errorf("databricks: disabling retries via WithRetries is %w "+ "(the kernel retries internally); omit it or use the default (Thrift) backend", dbsqlerr.ErrNotSupportedByKernel) } + // CloudFetch on the kernel path is controlled by the tri-state WithKernelCloudFetch + // (forwarded to kernel_session_config_set_cloudfetch_enabled), NOT the + // backend-neutral plain-bool WithCloudFetch. UseCloudFetch defaults to true, so an + // explicit WithCloudFetch(false) is distinguishable from unset — but the kernel + // path does not read it (it is not forwarded), so honoring it silently is + // impossible. Rather than let a disable request be dropped (CloudFetch would stay + // on), reject it loudly and steer the caller to the kernel option, keeping the + // "nothing silently ignored" contract literally true. When WithKernelCloudFetch + // WAS used (CloudFetchEnabled != nil) that explicit kernel knob is authoritative, + // so a redundant/contradictory WithCloudFetch is not consulted and not rejected. + kernelCloudFetchSet := cfg.KernelExperimental != nil && cfg.KernelExperimental.CloudFetchEnabled != nil + if !cfg.UseCloudFetch && !kernelCloudFetchSet { + return kernel.Auth{}, fmt.Errorf("databricks: disabling CloudFetch via WithCloudFetch(false) is %w "+ + "on the kernel path; use WithKernelCloudFetch(false) instead, or the default (Thrift) backend", dbsqlerr.ErrNotSupportedByKernel) + } + // mTLS client identity (WithKernelClientCertificate) must be a complete pair: + // mTLS needs both a non-empty client cert and its private key. Reject any + // incomplete request loudly at connect — otherwise applyKernelTLS (which gates + // the mTLS forward on the cert being non-empty) would silently skip the setter + // and connect with no client identity, a fail-open that contradicts the "both + // halves required together / never silently dropped" contract. + // + // "Incomplete" covers three cases: cert-without-key, key-without-cert, and — + // critically — the option being invoked with both halves empty/nil (e.g. from + // a failed PEM load), which the marker distinguishes from the option never + // being called. We also treat a bare cert/key present as a request even + // without the marker, so a config assembled without the option can't fail open + // either. + if ke := cfg.KernelExperimental; ke != nil { + certLen, keyLen := len(ke.TLSClientCertPEM), len(ke.TLSClientKeyPEM) + mtlsRequested := ke.TLSClientCertConfigured || certLen > 0 || keyLen > 0 + if mtlsRequested && (certLen == 0 || keyLen == 0) { + return kernel.Auth{}, fmt.Errorf("databricks: WithKernelClientCertificate requires " + + "both a non-empty client certificate and a non-empty private key") + } + } return kauth, nil } @@ -117,13 +153,17 @@ func buildKernelConfig(cfg *config.Config, kauth kernel.Auth) kernel.Config { if cfg.TLSConfig != nil && cfg.TLSConfig.InsecureSkipVerify { kc.TLSSkipVerify = true } - // Experimental kernel-only TLS knobs (WithKernelTrustedCerts / - // WithKernelSkipHostnameVerify), if any. These have no Thrift-path equivalent - // (the connector rejects them on that path) and are forwarded verbatim to the + // Experimental kernel-only knobs (WithKernelTrustedCerts / + // WithKernelSkipHostnameVerify / WithKernelClientCertificate / + // WithKernelCloudFetch), if any. These have no Thrift-path equivalent (the + // connector rejects them on that path) and are forwarded verbatim to the // kernel C ABI in OpenSession. if ke := cfg.KernelExperimental; ke != nil { kc.TLSTrustedCertsPEM = ke.TLSTrustedCertsPEM kc.TLSSkipHostnameVerify = ke.TLSSkipHostnameVerify + kc.TLSClientCertPEM = ke.TLSClientCertPEM + kc.TLSClientKeyPEM = ke.TLSClientKeyPEM + kc.CloudFetchEnabled = ke.CloudFetchEnabled } return kc } diff --git a/kernel_config_test.go b/kernel_config_test.go index 4da8efcb..0eddae02 100644 --- a/kernel_config_test.go +++ b/kernel_config_test.go @@ -127,6 +127,98 @@ func TestValidateKernelConfig(t *testing.T) { }) } + // mTLS client identity must be a complete pair. An unpaired credential (cert + // or key alone) is rejected loudly so a lone key can't be silently dropped by + // applyKernelTLS; a complete pair (and the no-mTLS case) validates. + t.Run("mTLS unpaired cert rejected", func(t *testing.T) { + c := baseKernelConfig() + c.KernelExperimental = &config.KernelExperimentalConfig{TLSClientCertPEM: []byte("cert")} + if _, err := validateKernelConfig(c); err == nil { + t.Fatal("expected an error for a cert without a key") + } + }) + t.Run("mTLS unpaired key rejected", func(t *testing.T) { + c := baseKernelConfig() + c.KernelExperimental = &config.KernelExperimentalConfig{TLSClientKeyPEM: []byte("key")} + if _, err := validateKernelConfig(c); err == nil { + t.Fatal("expected an error for a key without a cert") + } + }) + t.Run("mTLS complete pair validates", func(t *testing.T) { + c := baseKernelConfig() + c.KernelExperimental = &config.KernelExperimentalConfig{ + TLSClientCertPEM: []byte("cert"), + TLSClientKeyPEM: []byte("key"), + } + if _, err := validateKernelConfig(c); err != nil { + t.Errorf("a complete mTLS pair should validate, got %v", err) + } + }) + // Fail-open guard: WithKernelClientCertificate invoked with empty bytes (e.g. + // a failed PEM load) sets the configured marker with no cert/key. Validation + // must reject it rather than let applyKernelTLS silently skip the setter and + // connect with no client identity. + t.Run("mTLS configured but empty rejected", func(t *testing.T) { + c := baseKernelConfig() + c.KernelExperimental = &config.KernelExperimentalConfig{TLSClientCertConfigured: true} + if _, err := validateKernelConfig(c); err == nil { + t.Fatal("expected an error when WithKernelClientCertificate was used with empty cert/key") + } + }) + // The same request expressed through the public option (nil, nil): it must be + // rejected all the way through, not just when the config is hand-built. + t.Run("mTLS via option with empty args rejected", func(t *testing.T) { + c := baseKernelConfig() + WithKernelClientCertificate(nil, nil)(c) + if _, err := validateKernelConfig(c); err == nil { + t.Fatal("expected an error when WithKernelClientCertificate(nil, nil) is used") + } + }) + // A non-mTLS experimental config (only some other WithKernel* option set) must + // still validate — the marker, not merely a non-nil KernelExperimental, gates + // the mTLS completeness check. + t.Run("non-mTLS experimental config validates", func(t *testing.T) { + c := baseKernelConfig() + c.KernelExperimental = &config.KernelExperimentalConfig{TLSSkipHostnameVerify: true} + if _, err := validateKernelConfig(c); err != nil { + t.Errorf("a non-mTLS experimental config should validate, got %v", err) + } + }) + + // CloudFetch: a bare WithCloudFetch(false) can't be honored on the kernel path + // (UseCloudFetch is not forwarded), so it must be rejected loudly — not silently + // dropped, leaving CloudFetch on — and steer the caller to WithKernelCloudFetch. + // When the kernel-specific knob IS set it is authoritative, so a redundant or + // contradictory WithCloudFetch is not consulted and not rejected. + t.Run("bare WithCloudFetch(false) rejected", func(t *testing.T) { + c := baseKernelConfig() + c.UseCloudFetch = false // what WithCloudFetch(false) sets + _, err := validateKernelConfig(c) + if err == nil { + t.Fatal("expected WithCloudFetch(false) to be rejected on the kernel path") + } + if !errors.Is(err, dbsqlerr.ErrNotSupportedByKernel) { + t.Errorf("WithCloudFetch(false) rejection should wrap ErrNotSupportedByKernel, got %v", err) + } + }) + t.Run("default UseCloudFetch(true) accepted", func(t *testing.T) { + c := baseKernelConfig() // WithDefaults sets UseCloudFetch=true + if _, err := validateKernelConfig(c); err != nil { + t.Errorf("the CloudFetch-on default should validate, got %v", err) + } + }) + t.Run("WithCloudFetch(false) + WithKernelCloudFetch authoritative, accepted", func(t *testing.T) { + for _, kernelEnabled := range []bool{false, true} { + c := baseKernelConfig() + c.UseCloudFetch = false // the neutral bool says off... + WithKernelCloudFetch(kernelEnabled)(c) // ...but the explicit kernel knob wins + if _, err := validateKernelConfig(c); err != nil { + t.Errorf("WithKernelCloudFetch(%v) should make the config valid despite WithCloudFetch(false), got %v", + kernelEnabled, err) + } + } + }) + t.Run("PAT via WithAuthenticator resolves the token", func(t *testing.T) { c := baseKernelConfig() c.AccessToken = "" @@ -296,9 +388,12 @@ var kernelConfigFieldDisposition = map[string]string{ "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", + // CloudFetch internally (below the C ABI), so the tuning knobs are inert — but + // each is classified individually so a new CloudFetch option can't slip the guard. + // UseCloudFetch is the exception: a bare WithCloudFetch(false) is rejected on the + // kernel path (steering to WithKernelCloudFetch), since silently keeping CloudFetch + // on would violate the "nothing silently ignored" contract — see validateKernelConfig. + "UseCloudFetch": "rejected", "MaxDownloadThreads": "inert", "MaxFilesInMemory": "inert", "MinTimeToExpiry": "inert", @@ -366,6 +461,45 @@ func TestBuildKernelConfig(t *testing.T) { } }) + // mTLS client cert/key and the tri-state CloudFetch toggle are the other three + // fields buildKernelConfig forwards; each is copied ONLY here, so a dropped line + // (e.g. deleting kc.CloudFetchEnabled = ke.CloudFetchEnabled, silently leaving + // CloudFetch on for WithKernelCloudFetch(false)) would pass every other test. + t.Run("experimental mTLS + CloudFetch fields forwarded", func(t *testing.T) { + c := baseKernelConfig() + disabled := false + c.KernelExperimental = &config.KernelExperimentalConfig{ + TLSClientCertPEM: []byte("client-cert"), + TLSClientKeyPEM: []byte("client-key"), + CloudFetchEnabled: &disabled, + } + kc := buildKernelConfig(c, kernel.Auth{Mode: kernel.AuthPAT, Token: "dapi-x"}) + if got := string(kc.TLSClientCertPEM); got != "client-cert" { + t.Errorf("TLSClientCertPEM = %q, want %q (WithKernelClientCertificate cert not forwarded)", got, "client-cert") + } + if got := string(kc.TLSClientKeyPEM); got != "client-key" { + t.Errorf("TLSClientKeyPEM = %q, want %q (WithKernelClientCertificate key not forwarded)", got, "client-key") + } + if kc.CloudFetchEnabled == nil { + t.Fatal("CloudFetchEnabled = nil, want a non-nil false (WithKernelCloudFetch not forwarded)") + } + if *kc.CloudFetchEnabled { + t.Error("CloudFetchEnabled = true, want false (WithKernelCloudFetch(false) not forwarded)") + } + }) + + t.Run("unset CloudFetchEnabled stays nil (kernel default)", func(t *testing.T) { + c := baseKernelConfig() + // An experimental block with mTLS but no CloudFetch toggle: CloudFetchEnabled + // must remain nil so OpenSession leaves the kernel default (CloudFetch on) + // rather than forwarding a false. + c.KernelExperimental = &config.KernelExperimentalConfig{TLSSkipHostnameVerify: true} + kc := buildKernelConfig(c, kernel.Auth{Mode: kernel.AuthPAT, Token: "dapi-x"}) + if kc.CloudFetchEnabled != nil { + t.Errorf("CloudFetchEnabled = %v, want nil (unset must not forward a value)", *kc.CloudFetchEnabled) + } + }) + t.Run("nil KernelExperimental leaves TLS fields zero", func(t *testing.T) { c := baseKernelConfig() // KernelExperimental nil kc := buildKernelConfig(c, kernel.Auth{Mode: kernel.AuthPAT, Token: "dapi-x"}) @@ -373,6 +507,10 @@ func TestBuildKernelConfig(t *testing.T) { t.Errorf("expected zero experimental TLS fields with nil KernelExperimental, got certs=%v skipHost=%v", kc.TLSTrustedCertsPEM, kc.TLSSkipHostnameVerify) } + if kc.TLSClientCertPEM != nil || kc.TLSClientKeyPEM != nil || kc.CloudFetchEnabled != nil { + t.Errorf("expected zero experimental mTLS/CloudFetch fields with nil KernelExperimental, got cert=%v key=%v cloudfetch=%v", + kc.TLSClientCertPEM, kc.TLSClientKeyPEM, kc.CloudFetchEnabled) + } }) t.Run("InsecureSkipVerify maps to blanket TLSSkipVerify", func(t *testing.T) { diff --git a/kernel_e2e_test.go b/kernel_e2e_test.go index cc92d091..d513914e 100644 --- a/kernel_e2e_test.go +++ b/kernel_e2e_test.go @@ -9,6 +9,7 @@ import ( "database/sql/driver" "errors" "os" + "strings" "sync" "sync/atomic" "testing" @@ -406,6 +407,214 @@ func TestKernelE2ECancellation(t *testing.T) { t.Logf("cancelled after %v with err=%v", elapsed, err) } +// TestKernelE2EConnectHonorsCancelledCtx proves the connect-side pre-connect guard +// against a live warehouse: opening a session under an already-cancelled ctx must +// fail fast, not connect. It exercises the OpenSession entry guard (ctx.Err() before +// any dialing) — NOT the mid-connect cancellable path, which fires the cancel token +// WHILE kernel_session_open_cancellable is blocked; that path is proven +// deterministically and hermetically (no live warehouse) by +// TestOpenSessionHonorsCtxDeadlineMidConnect in the kernel package's tagged unit +// tests. (The execute-path TestKernelE2ECancellation does not reach connect, which +// happens before any statement runs.) +func TestKernelE2EConnectHonorsCancelledCtx(t *testing.T) { + db := kernelTestDB(t) + defer db.Close() + + ctx, cancel := context.WithCancel(context.Background()) + cancel() // already cancelled before the first connection is opened + + // The first statement forces a connect; with a pre-cancelled ctx it must + // return a context error, not succeed. + var got int64 + err := db.QueryRowContext(ctx, "SELECT 1").Scan(&got) + if err == nil { + t.Fatal("expected a context error connecting with an already-cancelled ctx, got nil") + } + if !errors.Is(err, context.Canceled) { + t.Logf("connect under cancelled ctx returned %v (want context.Canceled in the chain)", err) + // database/sql may surface its own driver.ErrBadConn retry wrapper; the + // key property is that it FAILED rather than silently connecting. + } +} + +// TestKernelE2EQueryCancelDuringExecution: a ctx cancelled while a long-running +// query is in flight is honored — the query returns promptly with a context +// error rather than blocking until the server finishes. This drives the +// server-side statement cancel on the execute path AND, once results begin +// streaming, the mid-fetch cancel token on the read path (nextBatch → +// kernel_result_stream_next_batch_cancellable). Uses a query that would otherwise +// run for many seconds. +func TestKernelE2EQueryCancelDuringExecution(t *testing.T) { + db := kernelTestDB(t) + defer db.Close() + + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + + start := time.Now() + // A deliberately slow, HIGH-CARDINALITY query: an un-aggregated cross join + // streams billions of rows, so the result keeps flowing far past the deadline + // (unlike a COUNT(*), which collapses to one row that returns instantly and + // exercises no mid-fetch cancel). An honored cancel — on the execute path or + // mid-fetch on the read path — must return well before it could complete. + var err error + rows, qerr := db.QueryContext(ctx, + "SELECT a.id, b.id FROM range(1000000000) a CROSS JOIN range(1000000) b") + if qerr != nil { + // Deadline fired during execute/connect (the common case for a slow query). + err = qerr + } else { + // QueryContext returned before the deadline; the deadline then fires while + // we drain. Drive rows.Next() (NOT a bare Scan, which would return a + // usage error unrelated to cancellation) until it stops, then read the + // terminal error — this exercises the mid-fetch path. + for rows.Next() { + } + err = rows.Err() + _ = rows.Close() + } + elapsed := time.Since(start) + + if err == nil { + t.Fatal("expected a context deadline error on a long query past its deadline, got nil") + } + // It must be the deadline firing, not an unrelated failure — otherwise the + // test would pass without proving cancellation works. (Mirrors the assertion + // in TestKernelE2ECancellation.) + if !errors.Is(err, context.DeadlineExceeded) { + t.Errorf("expected context.DeadlineExceeded, got %v", err) + } + if elapsed > 30*time.Second { + t.Errorf("cancel took %v — the deadline was not honored promptly", elapsed) + } + t.Logf("long query cancelled after %v with err=%v", elapsed, err) +} + +// TestKernelE2ECancelDuringFetch proves the mid-fetch cancel path: a ctx +// cancelled WHILE a nextBatch fetch is blocked inside +// kernel_result_stream_next_batch_cancellable is honored via the kernel cancel +// token, not merely by nextBatch's pre-fetch ctx.Err() guard. +// +// That guard short-circuits an already-cancelled ctx at the top of nextBatch and +// returns a bare context error, so a cancel landing *between* batches never +// reaches the cancellable C call. To hit the in-flight path we cancel from a +// separate goroutine while actively draining a huge, slow, multi-batch stream — +// where wall-clock time is dominated by the blocking fetch. Which path handled +// the cancel is observable in the terminal error: the cancellable call wraps its +// error as "next_batch cancelled", while the pre-fetch guard returns a bare +// "context canceled". Because landing mid-fetch is a timing window, we retry a +// bounded number of times and require at least one run to provably reach the +// cancellable call. +func TestKernelE2ECancelDuringFetch(t *testing.T) { + db := kernelTestDB(t) + defer db.Close() + + const attempts = 6 + var last string + for i := 1; i <= attempts; i++ { + midFetch, errStr := cancelDuringFetchOnce(t, db) + last = errStr + if midFetch { + t.Logf("attempt %d reached the mid-fetch cancellable path: %s", i, errStr) + return + } + t.Logf("attempt %d cancelled via the pre-fetch guard (retrying): %s", i, errStr) + } + t.Fatalf("no attempt reached kernel_result_stream_next_batch_cancellable in %d tries; last err=%q", attempts, last) +} + +// cancelDuringFetchOnce runs one attempt of the mid-fetch cancel: it starts a +// large slow stream, proves it is streaming (reads a row after QueryContext +// succeeds), then cancels from a separate goroutine while the drain loop is +// blocked in a fetch. It asserts the drain stopped promptly with a +// context.Canceled error, and reports whether that error came from the in-flight +// cancellable call (the "next_batch cancelled" wrapper) so the caller can retry +// until the mid-fetch path is provably exercised. +func cancelDuringFetchOnce(t *testing.T, db *sql.DB) (midFetch bool, errStr string) { + t.Helper() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + // A narrow, multi-chunk stream (~2.4 GB, under the external-links byte limit): + // small enough that execute/QueryContext returns promptly, but spanning many + // CloudFetch chunks so the drain re-enters a blocking network fetch many times. + rows, err := db.QueryContext(ctx, "SELECT id FROM range(0, 300000000)") + if err != nil { + t.Fatalf("QueryContext should succeed before any cancel; got %v", err) + } + defer rows.Close() + + // The pre-fetch ctx.Err() guard means a cancel landing BETWEEN batches never + // reaches the cancellable C call, so we can't just cancel on a fixed timer + // (that lands mid-batch-consume). Instead a watcher fires the cancel only once + // the main goroutine has been blocked inside a single rows.Next() longer than + // any buffered-row return could take — i.e. it is provably inside the blocking + // kernel_result_stream_next_batch_cancellable fetch, not iterating cached rows. + // nextInStart holds the UnixNano when the current Next() began (0 when not in + // Next); the watcher polls it and cancels when that dwell exceeds the fetch + // threshold. + const ( + fetchThreshold = 15 * time.Millisecond // a buffered-row return is sub-ms; a network fetch is far longer + giveUp = 8 * time.Second // fallback so a run never hangs if fetches stay hot + ) + var nextInStart atomic.Int64 + done := make(chan struct{}) + var wg sync.WaitGroup + wg.Add(1) + go func() { + defer wg.Done() + ticker := time.NewTicker(2 * time.Millisecond) + defer ticker.Stop() + deadline := time.Now().Add(giveUp) + for { + select { + case <-done: + return + case <-ticker.C: + if start := nextInStart.Load(); start != 0 && + time.Now().UnixNano()-start > int64(fetchThreshold) { + cancel() // in-flight fetch: fire the token mid-call + return + } + if time.Now().After(deadline) { + cancel() // fallback; this attempt likely won't prove mid-fetch + return + } + } + } + }() + + start := time.Now() + for { + nextInStart.Store(time.Now().UnixNano()) + ok := rows.Next() + nextInStart.Store(0) + if !ok { + break + } + } + err = rows.Err() + elapsed := time.Since(start) + close(done) + wg.Wait() + + if err == nil { + t.Fatal("expected a context error while fetching after cancel, got nil") + } + if !errors.Is(err, context.Canceled) { + t.Errorf("expected context.Canceled during fetch, got %v", err) + } + // A no-op cancel that only returned when the whole stream drained would blow + // past this; a real abort returns promptly. + if elapsed > 30*time.Second { + t.Errorf("fetch cancel took %v; expected prompt abort", elapsed) + } + // The cancellable C call wraps its error as "next_batch cancelled"; the + // pre-fetch guard returns a bare context error. That distinguishes which path + // handled the cancel. + return strings.Contains(err.Error(), "next_batch cancelled"), err.Error() +} + // TestKernelE2EInitialNamespace proves WithInitialNamespace selects the initial // catalog/schema on the kernel session — applied post-connect via USE CATALOG / // USE SCHEMA, since the kernel C ABI has no namespace setter. current_catalog() / diff --git a/kernel_experimental_test.go b/kernel_experimental_test.go index 1c3534d2..f131ead1 100644 --- a/kernel_experimental_test.go +++ b/kernel_experimental_test.go @@ -17,16 +17,22 @@ import ( // kernel_config tests. // kernelExperimentalFieldDisposition records how each experimental (kernel-only) -// option is handled. Every experimental field is forwarded to the kernel C ABI -// (there is no "inert" experimental knob — they exist precisely because the kernel -// supports them) and rejected on the Thrift path (the connector fails loud when -// KernelExperimental is non-nil). A new field on config.KernelExperimentalConfig -// without an entry here fails TestKernelExperimentalFieldsClassified, forcing a -// deliberate decision and a setter in KernelBackend.OpenSession so it can't be +// option is handled. Most experimental fields are "forwarded" to the kernel C ABI +// (they exist precisely because the kernel supports them) and rejected on the +// Thrift path (the connector fails loud when KernelExperimental is non-nil). The +// lone exception is "validation-only" metadata (TLSClientCertConfigured), which is +// consumed by validateKernelConfig and deliberately never forwarded. A new field +// on config.KernelExperimentalConfig without an entry here fails +// TestKernelExperimentalFieldsClassified, forcing a deliberate decision (and, for +// a forwarded field, a setter in KernelBackend.OpenSession) so it can't be // silently dropped. var kernelExperimentalFieldDisposition = map[string]string{ - "TLSTrustedCertsPEM": "forwarded", // set_tls_trusted_certs - "TLSSkipHostnameVerify": "forwarded", // set_tls_skip_hostname_verification + "TLSTrustedCertsPEM": "forwarded", // set_tls_trusted_certs + "TLSSkipHostnameVerify": "forwarded", // set_tls_skip_hostname_verification + "TLSClientCertPEM": "forwarded", // set_tls_client_certificate (cert half) + "TLSClientKeyPEM": "forwarded", // set_tls_client_certificate (key half) + "TLSClientCertConfigured": "validation-only", // consumed by validateKernelConfig; never forwarded + "CloudFetchEnabled": "forwarded", // set_cloudfetch_enabled } func TestKernelExperimentalFieldsClassified(t *testing.T) { @@ -66,6 +72,21 @@ func TestWithKernelTLSOptionsSetExperimental(t *testing.T) { {"skip hostname", WithKernelSkipHostnameVerify(), func(k *config.KernelExperimentalConfig) bool { return k.TLSSkipHostnameVerify }}, + {"client certificate", WithKernelClientCertificate([]byte("cert"), []byte("key")), func(k *config.KernelExperimentalConfig) bool { + return string(k.TLSClientCertPEM) == "cert" && string(k.TLSClientKeyPEM) == "key" && k.TLSClientCertConfigured + }}, + {"client certificate empty still marks configured", WithKernelClientCertificate(nil, nil), func(k *config.KernelExperimentalConfig) bool { + // Even with empty bytes the option must mark itself configured, so + // validateKernelConfig can reject the incomplete mTLS request rather + // than fail open (connect with no client identity). + return k.TLSClientCertConfigured && len(k.TLSClientCertPEM) == 0 && len(k.TLSClientKeyPEM) == 0 + }}, + {"cloudfetch off", WithKernelCloudFetch(false), func(k *config.KernelExperimentalConfig) bool { + return k.CloudFetchEnabled != nil && !*k.CloudFetchEnabled + }}, + {"cloudfetch on", WithKernelCloudFetch(true), func(k *config.KernelExperimentalConfig) bool { + return k.CloudFetchEnabled != nil && *k.CloudFetchEnabled + }}, } for _, c := range cases { t.Run(c.name, func(t *testing.T) { @@ -95,6 +116,8 @@ func TestWithKernelOptionsRejectedOnThriftPath(t *testing.T) { }{ {"trusted certs", WithKernelTrustedCerts([]byte("ca"))}, {"skip hostname", WithKernelSkipHostnameVerify()}, + {"client certificate", WithKernelClientCertificate([]byte("cert"), []byte("key"))}, + {"cloudfetch", WithKernelCloudFetch(false)}, } for _, tc := range cases { t.Run(tc.name, func(t *testing.T) { @@ -142,18 +165,41 @@ func TestWithKernelTrustedCertsCopiesPEM(t *testing.T) { // the whole Config per conn, and a shared backing array would let one conn's // mutation reach another. func TestKernelExperimentalDeepCopy(t *testing.T) { + cf := false orig := &config.KernelExperimentalConfig{ - TLSTrustedCertsPEM: []byte("ca-bundle"), - TLSSkipHostnameVerify: true, + TLSTrustedCertsPEM: []byte("ca-bundle"), + TLSSkipHostnameVerify: true, + TLSClientCertPEM: []byte("cert-pem"), + TLSClientKeyPEM: []byte("key-pem"), + TLSClientCertConfigured: true, + CloudFetchEnabled: &cf, } cp := orig.DeepCopy() if cp == nil || string(cp.TLSTrustedCertsPEM) != "ca-bundle" || !cp.TLSSkipHostnameVerify { t.Fatalf("DeepCopy lost data: %+v", cp) } + if string(cp.TLSClientCertPEM) != "cert-pem" || string(cp.TLSClientKeyPEM) != "key-pem" { + t.Fatalf("DeepCopy lost the mTLS cert/key: %+v", cp) + } + if !cp.TLSClientCertConfigured { + t.Fatalf("DeepCopy lost the TLSClientCertConfigured marker: %+v", cp) + } + if cp.CloudFetchEnabled == nil || *cp.CloudFetchEnabled != false { + t.Fatalf("DeepCopy lost CloudFetchEnabled: %+v", cp) + } cp.TLSTrustedCertsPEM[0] = 'X' if orig.TLSTrustedCertsPEM[0] == 'X' { t.Error("DeepCopy aliased the CA byte slice; a copy mutation reached the original") } + cp.TLSClientCertPEM[0] = 'X' + cp.TLSClientKeyPEM[0] = 'X' + if orig.TLSClientCertPEM[0] == 'X' || orig.TLSClientKeyPEM[0] == 'X' { + t.Error("DeepCopy aliased an mTLS byte slice; a copy mutation reached the original") + } + *cp.CloudFetchEnabled = true + if *orig.CloudFetchEnabled { + t.Error("DeepCopy aliased the CloudFetchEnabled pointer; a copy mutation reached the original") + } if (*config.KernelExperimentalConfig)(nil).DeepCopy() != nil { t.Error("nil.DeepCopy() should be nil") } diff --git a/kernel_mtls_test.go b/kernel_mtls_test.go new file mode 100644 index 00000000..a7314562 --- /dev/null +++ b/kernel_mtls_test.go @@ -0,0 +1,294 @@ +//go:build cgo && databricks_kernel + +package dbsql + +// Hermetic TLS-feature coverage for the kernel backend: mutual TLS +// (WithKernelClientCertificate), custom-CA trust (WithKernelTrustedCerts), and +// hostname-skip (WithKernelSkipHostnameVerify), exercised end-to-end through the +// real cgo → kernel → rustls stack against a local httptest TLS server — no +// warehouse, so this runs in the tagged build-and-test-kernel CI job (and nightly). +// +// The kernel speaks SEA (JSON over HTTPS); a bare httptest server is not a SEA +// endpoint, so the *query* always fails at the application layer. That is fine: the +// property under test is whether the TLS/mTLS HANDSHAKE completes with the forwarded +// material. The server handler records, per connection, whether it was reached and +// whether a client certificate was presented — reaching the handler AT ALL means the +// handshake (including client-cert verification when demanded) succeeded. Each case +// asserts the reached/​saw-cert signal, then tears the connect down immediately +// rather than waiting out the kernel's connect-retry budget, so the whole file runs +// in a few seconds. + +import ( + "context" + "crypto/ecdsa" + "crypto/elliptic" + "crypto/rand" + "crypto/tls" + "crypto/x509" + "crypto/x509/pkix" + "database/sql" + "encoding/pem" + "math/big" + "net" + "net/http" + "net/http/httptest" + "sync/atomic" + "testing" + "time" +) + +// testCertPair is a leaf certificate + its PKCS#8 private key, both PEM-encoded. +type testCertPair struct { + certPEM []byte + keyPEM []byte + tlsCert tls.Certificate +} + +// mkTestCA returns a self-signed CA cert (parsed, for signing leaves) plus its key +// and PEM encoding (for WithKernelTrustedCerts). +func mkTestCA(t *testing.T) (*x509.Certificate, *ecdsa.PrivateKey, []byte) { + t.Helper() + key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader) + if err != nil { + t.Fatalf("CA key: %v", err) + } + // Fixed validity window (no time.Now) so the cert is deterministic; well inside + // any realistic test clock. + tmpl := &x509.Certificate{ + SerialNumber: big.NewInt(1), + Subject: pkix.Name{CommonName: "databricks-sql-go test CA"}, + NotBefore: time.Unix(1_600_000_000, 0), + NotAfter: time.Unix(4_000_000_000, 0), + IsCA: true, + KeyUsage: x509.KeyUsageCertSign | x509.KeyUsageDigitalSignature, + BasicConstraintsValid: true, + } + der, err := x509.CreateCertificate(rand.Reader, tmpl, tmpl, &key.PublicKey, key) + if err != nil { + t.Fatalf("CA cert: %v", err) + } + ca, err := x509.ParseCertificate(der) + if err != nil { + t.Fatalf("parse CA: %v", err) + } + caPEM := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: der}) + return ca, key, caPEM +} + +// mkTestLeaf mints a CA-signed leaf. serverSANs (IPs or DNS names) are set for a +// server leaf; an empty slice yields a client-auth leaf. The key is emitted as +// PKCS#8 — the form the kernel's tls-rustls Identity::from_pem accepts. +func mkTestLeaf(t *testing.T, ca *x509.Certificate, caKey *ecdsa.PrivateKey, cn string, serverSANs []string) testCertPair { + t.Helper() + key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader) + if err != nil { + t.Fatalf("leaf key: %v", err) + } + tmpl := &x509.Certificate{ + SerialNumber: big.NewInt(2), + Subject: pkix.Name{CommonName: cn}, + NotBefore: time.Unix(1_600_000_000, 0), + NotAfter: time.Unix(4_000_000_000, 0), + KeyUsage: x509.KeyUsageDigitalSignature, + } + if len(serverSANs) > 0 { + tmpl.ExtKeyUsage = []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth} + for _, s := range serverSANs { + if ip := net.ParseIP(s); ip != nil { + tmpl.IPAddresses = append(tmpl.IPAddresses, ip) + } else { + tmpl.DNSNames = append(tmpl.DNSNames, s) + } + } + } else { + tmpl.ExtKeyUsage = []x509.ExtKeyUsage{x509.ExtKeyUsageClientAuth} + } + der, err := x509.CreateCertificate(rand.Reader, tmpl, ca, &key.PublicKey, caKey) + if err != nil { + t.Fatalf("leaf cert: %v", err) + } + certPEM := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: der}) + keyDER, err := x509.MarshalPKCS8PrivateKey(key) + if err != nil { + t.Fatalf("marshal key: %v", err) + } + keyPEM := pem.EncodeToMemory(&pem.Block{Type: "PRIVATE KEY", Bytes: keyDER}) + tlsCert, err := tls.X509KeyPair(certPEM, keyPEM) + if err != nil { + t.Fatalf("x509 key pair: %v", err) + } + return testCertPair{certPEM: certPEM, keyPEM: keyPEM, tlsCert: tlsCert} +} + +// tlsProbe is a local httptest TLS server that records, across connections, +// whether its handler was reached and whether a client cert was presented. +type tlsProbe struct { + srv *httptest.Server + reached atomic.Bool + sawCert atomic.Bool +} + +func (p *tlsProbe) host() string { return p.srv.Listener.Addr().String() } + +// startTLSProbe serves HTTPS with the given server cert and client-auth policy. +// The handler always returns 503 (not a real SEA endpoint); the point is that +// reaching it proves the handshake completed. +func startTLSProbe(t *testing.T, serverCert tls.Certificate, clientCAs *x509.CertPool, clientAuth tls.ClientAuthType) *tlsProbe { + t.Helper() + p := &tlsProbe{} + p.srv = httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + p.reached.Store(true) + if r.TLS != nil && len(r.TLS.PeerCertificates) > 0 { + p.sawCert.Store(true) + } + http.Error(w, "not a SEA endpoint", http.StatusServiceUnavailable) + })) + p.srv.TLS = &tls.Config{ + Certificates: []tls.Certificate{serverCert}, + ClientAuth: clientAuth, + ClientCAs: clientCAs, + MinVersion: tls.VersionTLS12, + } + p.srv.StartTLS() + t.Cleanup(p.srv.Close) + return p +} + +// runKernelConnect opens a kernel-backed connection to host and drives a query in a +// goroutine to force the TLS handshake. It returns as soon as onReached() reports +// the outcome is decided (the handshake succeeded, observed via the probe) OR the +// query returns on its own OR settle elapses — whichever first — then cancels the +// query. This avoids waiting out the kernel's connect-retry budget on the negative +// (handshake-should-fail) cases. It does not assert; the caller inspects the probe. +func runKernelConnect(t *testing.T, host string, settle time.Duration, onReached func() bool, extra ...ConnOption) { + t.Helper() + opts := append([]ConnOption{ + WithServerHostname(host), + WithHTTPPath("/sql/1.0/warehouses/hermetic"), + WithAccessToken("dapi-hermetic-placeholder"), + WithUseKernel(true), + }, extra...) + connector, err := NewConnector(opts...) + if err != nil { + t.Fatalf("NewConnector: %v", err) + } + db := sql.OpenDB(connector) + defer db.Close() + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + done := make(chan struct{}) + go func() { + defer close(done) + var x int64 + _ = db.QueryRowContext(ctx, "SELECT 1").Scan(&x) + }() + + // Poll the decision signal frequently; bail (cancel) the instant it flips so a + // positive case finishes sub-second instead of riding the retry budget. On a + // negative case the signal never flips, so we wait out `settle` — long enough + // for at least one handshake attempt to complete-or-fail — then cancel. + deadline := time.Now().Add(settle) + for time.Now().Before(deadline) { + if onReached() { + break + } + select { + case <-done: + return + case <-time.After(20 * time.Millisecond): + } + } + cancel() + <-done +} + +// TestKernelMTLSHandshake proves WithKernelClientCertificate is forwarded and +// enforced: with a client cert the mTLS handshake completes (handler reached, cert +// seen); without one, a RequireAndVerifyClientCert server rejects the handshake +// (handler never reached). +func TestKernelMTLSHandshake(t *testing.T) { + ca, caKey, caPEM := mkTestCA(t) + server := mkTestLeaf(t, ca, caKey, "127.0.0.1", []string{"127.0.0.1"}) + client := mkTestLeaf(t, ca, caKey, "hermetic-client", nil) + clientCAs := x509.NewCertPool() + clientCAs.AddCert(ca) + + t.Run("client cert forwarded -> handshake succeeds", func(t *testing.T) { + p := startTLSProbe(t, server.tlsCert, clientCAs, tls.RequireAndVerifyClientCert) + runKernelConnect(t, p.host(), 10*time.Second, p.reached.Load, + WithKernelTrustedCerts(caPEM), + WithKernelClientCertificate(client.certPEM, client.keyPEM), + ) + if !p.reached.Load() { + t.Fatal("handler never reached: the mTLS handshake failed — client cert not forwarded/accepted") + } + if !p.sawCert.Load() { + t.Error("handler reached but no client cert was presented — mTLS identity not forwarded") + } + }) + + t.Run("no client cert -> handshake rejected", func(t *testing.T) { + p := startTLSProbe(t, server.tlsCert, clientCAs, tls.RequireAndVerifyClientCert) + // Trust the CA so the SERVER cert validates; present no client identity. The + // server demands one, so the handshake must fail before the handler runs. + runKernelConnect(t, p.host(), 3*time.Second, p.reached.Load, + WithKernelTrustedCerts(caPEM), + ) + if p.reached.Load() { + t.Error("handler was reached without a client cert — mTLS was not enforced") + } + }) +} + +// TestKernelCustomCATrust proves WithKernelTrustedCerts is forwarded: a +// privately-signed server cert validates only when its CA is trusted. +func TestKernelCustomCATrust(t *testing.T) { + ca, caKey, caPEM := mkTestCA(t) + server := mkTestLeaf(t, ca, caKey, "127.0.0.1", []string{"127.0.0.1"}) + + t.Run("CA trusted -> TLS validates", func(t *testing.T) { + p := startTLSProbe(t, server.tlsCert, nil, tls.NoClientCert) + runKernelConnect(t, p.host(), 10*time.Second, p.reached.Load, + WithKernelTrustedCerts(caPEM)) + if !p.reached.Load() { + t.Fatal("handler not reached with the CA trusted — WithKernelTrustedCerts not forwarded") + } + }) + + t.Run("CA untrusted -> TLS fails", func(t *testing.T) { + p := startTLSProbe(t, server.tlsCert, nil, tls.NoClientCert) + // No trusted certs: the private CA is unknown, so the chain must not validate. + runKernelConnect(t, p.host(), 3*time.Second, p.reached.Load) + if p.reached.Load() { + t.Error("handler reached without trusting the CA — an unknown issuer was accepted") + } + }) +} + +// TestKernelSkipHostnameVerify proves WithKernelSkipHostnameVerify is forwarded: a +// server cert whose SAN does not cover the dialed host is rejected on hostname +// grounds unless the skip is set (chain validation stays on either way). +func TestKernelSkipHostnameVerify(t *testing.T) { + ca, caKey, caPEM := mkTestCA(t) + // SAN deliberately does NOT include 127.0.0.1 (the dialed host). + wrongHost := mkTestLeaf(t, ca, caKey, "wrong.example", []string{"not-the-host.example"}) + + t.Run("wrong SAN, no skip -> hostname check fails", func(t *testing.T) { + p := startTLSProbe(t, wrongHost.tlsCert, nil, tls.NoClientCert) + runKernelConnect(t, p.host(), 3*time.Second, p.reached.Load, + WithKernelTrustedCerts(caPEM)) // chain trusted, but hostname won't match + if p.reached.Load() { + t.Error("handler reached despite a hostname mismatch and no skip — hostname verification not enforced") + } + }) + + t.Run("wrong SAN, skip on -> handshake succeeds", func(t *testing.T) { + p := startTLSProbe(t, wrongHost.tlsCert, nil, tls.NoClientCert) + runKernelConnect(t, p.host(), 10*time.Second, p.reached.Load, + WithKernelTrustedCerts(caPEM), WithKernelSkipHostnameVerify()) + if !p.reached.Load() { + t.Fatal("handler not reached with hostname-skip on — WithKernelSkipHostnameVerify not forwarded") + } + }) +}