feat: auto-recover Reyden Thrift connections onto the kernel backend - #479
feat: auto-recover Reyden Thrift connections onto the kernel backend#479rahuls-db wants to merge 6 commits into
Conversation
An unconfigured connection to a Reyden / Real-Time SQL warehouse defaults to the Thrift backend, which the SQL Gateway proxy rejects with SQLSTATE KP001. Detect that rejection at OpenSession (CheckStatus) and transparently re-open the session on the SEA/kernel backend, remembering the warehouse in a process-wide cache keyed by (host, warehouse_id) with a ~6h TTL so later connects skip the doomed Thrift attempt. Only the default path auto-recovers; an explicit WithUseKernel is always honored. On a double failure the kernel error is surfaced with the Thrift rejection preserved via errors.Join. Recovery requires a databricks_kernel build, since the kernel backend is otherwise not linked in (the fallback then surfaces the not-compiled error joined with the Reyden rejection). Co-authored-by: Isaac <no-reply@databricks.com> Signed-off-by: Rahul Singhal <rahul.singhal@databricks.com>
skipDriverTelemetry read cfg.UseKernel, which stays false on the Reyden auto-recovery path (the fallback opens the kernel without mutating cfg). A recovered-kernel connection therefore kept the Go driver's telemetry active, duplicating the kernel's own telemetry. Derive the skip decision from the backend that actually opened instead. Mirrors the analogous fix in the Python driver. Co-authored-by: Isaac <no-reply@databricks.com> Signed-off-by: Rahul Singhal <rahul.singhal@databricks.com>
![peco-review-bot[bot]](https://avatars.githubusercontent.com/in/3801177?s=60&v=4){kind=small}
There was a problem hiding this comment.
Verdict: 1 Medium · 1 Low
Solid, well-tested port of the Reyden Thrift→kernel auto-recovery feature; the error-marker plumbing, cache, and guardrails are coherent and the production errors.Is chain checks out. Two non-blocking notes: a coverage gap where the recovery tests bypass the real Thrift error-wrapping path (Medium), and a cache-dependent inconsistency in the WithKernel*-without-WithUseKernel guardrail (Low). Nit: skipDriverTelemetry (connector.go:48) is now vestigial production code — replaced by shouldSkipDriverTelemetry in Connect and referenced only by connector_kernel_u2m_test.go; consider removing it and repointing that test.
CheckStatus is the shared status checker for every Thrift RPC, so mapping KP001 to the recoverable marker there gave it a wider blast radius than the recovery logic (which only wraps session open): a stray KP001 on any other RPC would have surfaced as ErrReydenThriftUnsupported with no handler. Keep CheckStatus generic and add an OpenSession-scoped CheckOpenSessionStatus that does the KP001 mapping; only the OpenSession wrapper uses it. Every other RPC now surfaces a KP001 as a plain error, unchanged from before. Co-authored-by: Isaac <no-reply@databricks.com> Signed-off-by: Rahul Singhal <rahul.singhal@databricks.com>
There was a problem hiding this comment.
Verdict: 3 Low
Solid, well-tested port of the Reyden Thrift→kernel auto-recovery. The KP001 detection is correctly scoped to OpenSession, the marker's errors.Is chaining works through the production requestError wrapper, the cache is properly RWMutex-guarded and host-keyed, and no new third-party imports/leaks are introduced. Only 3 low-severity notes: a test-fidelity gap (tests inject the raw marker, not the wrapped form production emits), a now-test-only stale skipDriverTelemetry, and a guardrail bypass in the cache pre-check.
| assert.True(t, kb.SessionValid(), "kernel session should be open") | ||
|
|
||
| // The rejection is remembered for future connects. | ||
| assert.True(t, warehouse_cache.IsKnownReyden(tt.host, tt.warehouseID), |
There was a problem hiding this comment.
🔵 Low — The reactive-recovery test injects the raw reydenThriftUnsupportedError as openSessionErr, but in production thrift.Backend.OpenSession wraps every OpenSession failure in NewRequestError(ctx, "error connecting: ...", err) before it reaches openSessionWithReydenFallback. So the connector actually calls errors.Is(err, ErrReydenThriftUnsupported) against a *requestError → withStack → withMessage → marker chain, not the bare marker.
The detection does still work (I traced requestError.Unwrap → withStack.Unwrap → withMessage.Cause → reydenThriftUnsupportedError.Is), so this is not a live bug. But no test exercises the wrapped form the connector sees in production — a future change to the unwrap chain or to how the Thrift backend wraps the error would silently break auto-recovery while all these tests stay green. Consider having the fake Thrift backend return dbsqlerrint.NewRequestError(context.Background(), "error connecting", dbsqlerrint.NewReydenThriftUnsupportedError(...)) so the test matches the production error shape.
| cfg *config.Config | ||
| client *http.Client | ||
| kernelBackendFactory backendFactory // Seam for testing; nil in production | ||
| thriftBackendFactory thriftBackendFactory // Seam for testing; nil in production |
There was a problem hiding this comment.
🔵 Low — skipDriverTelemetry(cfg *config.Config) is no longer used by production code — Connect now derives the skip decision from the active backend via shouldSkipDriverTelemetry(be) (connector.go:101). The only remaining references are in connector_kernel_u2m_test.go, so the function survives unused/staticcheck purely as test-only code that duplicates the now-superseded cfg.UseKernel logic.
Since the real decision moved to shouldSkipDriverTelemetry, this stale helper (and its test) can drift from actual behavior — e.g. it returns false for a Reyden auto-recovery onto the kernel, which is exactly the case the new logic was written to fix. Recommend deleting skipDriverTelemetry and repointing that test at shouldSkipDriverTelemetry.
(Anchored to the nearest changed line — see the description for the exact location.)
| // returns the backend, session latency, and error. | ||
| func (c *connector) openSessionWithReydenFallback(ctx context.Context) (backend.Backend, int64, error) { | ||
| // Extract warehouse ID from HTTPPath for cache lookups. | ||
| warehouseID := warehouse_cache.ExtractWarehouseID(c.cfg.HTTPPath) |
There was a problem hiding this comment.
🔵 Low — The known-Reyden pre-check opens the kernel backend directly and never runs the KernelExperimental != nil guardrail that the default (non-UseKernel) path enforces at connector.go:846. As a result, the same config (no WithUseKernel, but a WithKernel* option set) produces two different outcomes depending on process-global cache state: a clear ErrRequiresKernelBackend error on a cache miss, but a silent kernel open on a cache hit.
This isn't harmful (the pre-check path does construct the kernel backend, so the experimental option is honored rather than ignored), but the nondeterministic branch-on-cache behavior is surprising and undocumented. Worth either applying the same guardrail in the pre-check branch, or adding a comment noting the intentional divergence.
…apped-marker tests - connector: hoist the WithKernel*-without-WithUseKernel guardrail above the cache pre-check so the misconfiguration is rejected deterministically, regardless of process-global cache state (previously a warm cache let the pre-check open the kernel and silently bypass the guardrail). - connector: delete the now-unused skipDriverTelemetry(cfg); the skip decision is derived from the active backend via shouldSkipDriverTelemetry(be). Remove its redundant test (shouldSkipDriverTelemetry is already covered). - connector_reyden_test: inject the Reyden marker WRAPPED in NewRequestError, as thrift.Backend.OpenSession does in production, so the errors.Is unwrap chain the recovery relies on is pinned; add a test that the guardrail fires on a warm cache. Co-authored-by: Isaac <no-reply@databricks.com> Signed-off-by: Rahul Singhal <rahul.singhal@databricks.com>
There was a problem hiding this comment.
Verdict: 1 Medium
Solid, well-tested port of the Reyden Thrift→kernel auto-recovery. The recovery logic, error-marker plumbing (errors.Is chain through NewRequestError), guardrail ordering, and cache keying all look correct. One medium concern: a stub-only test is not build-tagged and will fail under the databricks_kernel build. Nit: Cache.MarkReyden's sweep re-indexes the map (c.expiry[key].Before(now)) instead of ranging over key, deadline — harmless but slightly wasteful.
| // In the default pure-Go build, newKernelBackend returns an error indicating | ||
| // the kernel backend is not compiled in. This naturally exercises the | ||
| // error-chaining path for double-failure scenarios. | ||
| cfg := config.WithDefaults() |
There was a problem hiding this comment.
🟡 Medium — TestReydenDefaultBuildKernelNotCompiled asserts that newKernelBackend returns an error wrapping ErrKernelNotCompiled (assert.Nil(be), assert.Error(err), errors.Is(err, ErrKernelNotCompiled)). That behavior is specific to the default-build stub (kernel_stub.go, //go:build !cgo || !databricks_kernel). This test file has no build constraint, so it is also compiled and run under -tags databricks_kernel + CGO_ENABLED=1, where newKernelBackend resolves to the real implementation in kernel_backend.go and will not return ErrKernelNotCompiled — the three assertions fail and make test under the kernel build breaks.
Every other kernel-behavior test in this package (kernel_stub_test.go, kernel_backend_test.go, etc.) is build-tagged accordingly. This test asserts stub-only behavior, so it should carry //go:build !cgo || !databricks_kernel (matching kernel_stub_test.go). The other tests in this file are backend-agnostic (they use the injected fakes), so only this one test needs to move to a tagged file or gain the constraint.
- warehouse_cache/cache.go: drop the redundant `match != nil` (staticcheck S1009: len(nil) is 0) and fix struct field gofmt alignment. Both failed golangci-lint. - Move TestReydenDefaultBuildKernelNotCompiled into a //go:build !databricks_kernel file. It asserts newKernelBackend reports "not compiled", which only holds in the default build; under -tags databricks_kernel the real backend is linked in, so the test was failing the "Test (kernel backend)" CI job. Co-authored-by: Isaac <no-reply@databricks.com> Signed-off-by: Rahul Singhal <rahul.singhal@databricks.com>
There was a problem hiding this comment.
Verdict: 1 Low
Looks good — solid, well-tested port of the Reyden auto-recovery feature. The errors.Is unwrap chain through NewRequestError is verified correct, the WithKernel*-without-WithUseKernel guardrail is preserved (now checked before the cache pre-check), and telemetry attribution correctly derives from the active backend rather than cfg.UseKernel. One low-severity behavioral note on the cache pre-check path (no Thrift fallback / stripped context on a kernel failure) is posted inline.
- connector: on the cache pre-check path, wrap a kernel backend-create/OpenSession failure with context noting the warehouse was cached as Reyden and Thrift was skipped. Uses %w so the underlying error (including a default-build ErrKernelNotCompiled) stays reachable via errors.Is; adds a test asserting both. - warehouse_cache: range the sweep over (key, deadline) instead of re-indexing the map per key. Co-authored-by: Isaac <no-reply@databricks.com> Signed-off-by: Rahul Singhal <rahul.singhal@databricks.com>
There was a problem hiding this comment.
Verdict: 1 Low
Looks good — a clean, well-tested port. The Reyden marker's errors.Is chain survives the production NewRequestError wrapping, shouldSkipDriverTelemetry correctly derives from the active backend (fixing recovered-onto-kernel attribution), and the warehouse_cache package is stdlib-only with sound locking/TTL semantics. One low-severity note inline about the pre-check not being gated on UseKernel. Nit (summary-only): the warehouse_cache package name uses an underscore, which is non-idiomatic Go — it won't fail CI since ST1003 is disabled in .golangci.yml, but warehousecache would match convention.
| // than a stale subset that drifts as options are added. | ||
| if !c.cfg.UseKernel && c.cfg.KernelExperimental != nil { | ||
| return nil, 0, fmt.Errorf("databricks: a WithKernel* option %w; "+ | ||
| "add WithUseKernel(true) or remove it", dbsqlerr.ErrRequiresKernelBackend) |
There was a problem hiding this comment.
🔵 Low — The cache pre-check fires regardless of c.cfg.UseKernel. When a caller explicitly sets WithUseKernel(true) for a warehouse that happens to be cached as Reyden, control enters this pre-check branch instead of the normal if c.cfg.UseKernel path below. Functionally the outcome is identical (kernel either way), but on a kernel failure the error is wrapped with a misleading message — "warehouse %s is cached as Reyden so Thrift was skipped, but the kernel backend could not be created" — even though Thrift was never in play for an explicit-kernel connection.
Consider gating the pre-check on !c.cfg.UseKernel so the auto-recovery messaging only applies to the default (Thrift) path, and explicit-kernel connections fall through to the plain getKernelBackend branch with its normal error surface.
1 participant
Description
Ports the Reyden Thrift auto-recovery feature (already in the Python driver, databricks/databricks-sql-python#948) to the Go driver.
An unconfigured connection to a Reyden / Real-Time SQL warehouse defaults to the Thrift backend, which the SQL Gateway proxy rejects with SQLSTATE
KP001. This change detects that rejection atOpenSessionand transparently re-opens the session on the SEA/kernel backend, so no connection-parameter change is needed.client.CheckStatusreturns a distinctErrReydenThriftUnsupported(viaerrors.Is) when the OpenSessionTStatuscarries SQLSTATEKP001(matched on the SQLSTATE only).connector.openSessionWithReydenFallbackcatches the marker and re-opens once on the kernel backend. On a double failure the kernel error is surfaced with the original Thrift rejection preserved viaerrors.Join.internal/warehouse_cache) keyed by(host_lowercased, warehouse_id), ~6h TTL,sync.RWMutex-guarded, with opportunistic eviction; a pre-check skips the Thrift round-trip for a known-Reyden warehouse.WithUseKernelis always honored.Build-tag caveat: the kernel backend is only linked in under
-tags databricks_kernel+CGO_ENABLED=1. In a default build the fallback surfaces the not-compiled error joined with the Reyden rejection, rather than silently recovering — the same kernel-availability constraint the Python driver has with its optional[kernel]extra.Testing
Unit tests drive the real
openSessionWithReydenFallbackvia injected backend-factory seams (connector_reyden_test.go): KP001 detection, reactive recovery onto the kernel, cache pre-check skipping Thrift, cache marking, explicit-UseKernelguardrail, non-Reyden error pass-through, and double-failureerrors.Joinchaining; plusinternal/warehouse_cachecache/extraction/expiry tests.go test(root + changed internal packages),gofmt, andgo vetare clean.Related
Design: "Simplifying Reyden Onboarding on Drivers" (Option A). Sibling PRs: Python databricks/databricks-sql-python#948; a Node port is in flight.
This PR was created with GitHub MCP.