feat: auto-recover Reyden Thrift connections onto the kernel - #948
Conversation
An unconfigured connect() 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 and transparently re-open the session on the kernel backend, and remember the warehouse (process-wide cache keyed by (host, warehouse_id), ~6h TTL) so subsequent connects skip the doomed Thrift attempt. Only the default path auto-recovers; an explicit use_kernel/use_sea is always honored. 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 · 2 Low
Solid, well-tested feature — the KP001 marker, cache, and recovery logic are correct and the name resolution via from databricks.sql import * works. One medium concern: the connection-failure telemetry suppression reads the original kwargs (never updated on the Reyden→kernel retry), so a recovered kernel connection's failure gets logged despite the "kernel owns telemetry" design. Two low notes on detection scope and dropping the rejected Thrift session without close().
… telemetry The connection-failure telemetry suppression read the original connect() kwargs to decide whether the failed connection was a kernel connection. On the Reyden auto-recovery path the kernel retry uses a kwargs copy, so the original still said Thrift — a kernel open-failure was logged by the wrapper despite the kernel owning telemetry for kernel connections. Decide from the session that actually failed (self.session.use_kernel) instead. If the kernel was never constructed (e.g. its wheel is missing), self.session stays Thrift and the wrapper still logs, so that otherwise-invisible failure is still recorded. 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-guarded feature with strong unit coverage (cache behavior, reactive recovery, cache pre-check, explicit-backend guardrail, error chaining, and telemetry suppression are all exercised). Name resolution (ReydenThriftUnsupportedError via from databricks.sql import *), the self.session.use_kernel telemetry-suppression read, and the exception-chaining logic all check out against the surrounding code. One low-severity note about the KP001 detection being global to all Thrift responses rather than scoped to OpenSession.
_check_response_for_error runs 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 ReydenThriftUnsupportedError with no handler. Gate the mapping behind a detect_reyden flag that make_request sets only for the OpenSession method (mirroring the existing method.__name__ discrimination). Every other RPC now surfaces a KP001 as a generic DatabaseError, 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: 1 Nit
Looks good — solid, well-tested feature. Detection is correctly scoped to OpenSession, the cache is host-keyed and thread-safe with sensible eviction, explicit-backend choices are honored, and the kernel-retry error chaining is preserved. One low note: the telemetry-suppression change also alters behavior for explicit use_kernel connections when kernel construction fails — flagged inline to confirm intent.
Other findings
- ⚪ Nit — The telemetry-suppression source changed from the caller's kwargs to the actually-failed session, which is the correct fix for the Reyden auto-recovery path. Note this also changes an unrelated case: for an explicit
use_kernel=Trueconnection whose kernelSession.__init__raises (e.g. the kernel wheel is missing, so_create_backendfails beforeself.sessionis ever assigned),getattr(self, "session", None)isNone, soattempted_kernelisFalseand the wrapper-side failure log now emits — whereas the oldnot kwargs.get("use_kernel", False)would have suppressed it.
This is arguably an improvement (a missing-wheel failure produces no kernel telemetry, so the wrapper event is the only signal), but it's a behavior change beyond the PR's stated scope. Worth confirming it's intended; if so, no action needed.
The connection-failure telemetry suppression reads session.use_kernel, but these tests mock Session, so use_kernel was a truthy MagicMock — which suppressed the wrapper failure log and broke test_connection_failure_sends_correct_telemetry_payload. Set use_kernel explicitly on the mock (False for the default Thrift case, True for the kernel case) to reflect production. 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
Solid, well-tested feature — the KP001→marker detection is correctly scoped to OpenSession, the marker propagates without triggering retries, the process-wide cache is thread-safe and host-scoped, and error chaining/guardrails are sound. One low-severity behavior change: the telemetry-suppression refactor now emits the wrapper failure log for an explicit use_kernel=True connection whose construction fails (missing wheel), which the previous kwargs-based check suppressed and which the tests don't cover.
On auto-recovery the Thrift default of auth_type=None (implicitly databricks-oauth) was forwarded unchanged to the kernel, which has no such fallback and rejects auth_type=None unless a PAT/M2M credential shape is present — so a bare OAuth-U2M connection failed to recover. Inject auth_type=databricks-oauth on the default-path kernel retry when auth_type is unset and no credential shape is present, mirroring Thrift; skip it when a credential shape exists so kernel routing is unchanged. Also correct the warehouse-cache docstring: warehouse ids are globally unique, so there is no cross-workspace collision on a shared SPOG host; the host component is an optimization, not a correctness requirement. 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 careful, well-tested feature with only one low/style note. Verified the recovery control flow (Thrift ReydenThriftUnsupportedError propagates cleanly out of make_request → session.open() and is caught in _open_session_with_reyden_fallback), the AuthType.DATABRICKS_OAUTH value used for kernel-recovery kwargs, the Session.use_kernel attribute the new telemetry-suppression logic reads, and the explicit-backend guardrail; each behavior has corresponding unit coverage. Nit (not filed inline): CHANGELOG adds a bare # Unreleased h1 rather than the existing # X.Y.Z (date) heading form — fine, just noting the divergence.
vuanhphung
left a comment
There was a problem hiding this comment.
btw, you can also add skip-coverage label for this feature, since it's expected that there are multiple PRs to be landed and we can add one final test at the end
What type of PR is this?
Description
An unconfigured
connect()to a Reyden / Real-Time SQL warehouse defaults to the Thrift backend, which the SQL Gateway proxy rejects (SQLSTATEKP001). This change detects that rejection atOpenSessionand transparently re-opens the session on the kernel (SEA) backend, so no customer configuration change is needed.ThriftDatabricksClient._check_response_for_errorraises a distinctReydenThriftUnsupportedErrorwhen the OpenSessionTStatuscarries SQLSTATEKP001(matched on the SQLSTATE only — no message-string matching).Connectioncatches the marker and re-opens the session once withuse_kernel=True. If the kernel open also fails, the kernel error is surfaced with the Thrift rejection preserved in the exception chain.(host, warehouse_id)(~6h TTL, host-scoped for multi-tenant safety, opportunistic eviction of expired entries) remembers warehouses that reject Thrift, so subsequent connects skip the doomed Thrift round-trip and open the kernel directly.use_kernel/use_seais always honored. There is no client kill switch — the server-side SAFE flag gates whether the rejection is emitted at all.Note: the kernel backend is currently an optional extra, so on a base install a rejection surfaces a clear "install
databricks-sql-connector[kernel]" error. Making the kernel a hard dependency for fully silent recovery is a planned follow-up (packaging change, likely a major version bump).How is this tested?
Unit: new
tests/unit/test_reyden_warehouse_cache.py(cache behavior + warehouse-id extraction + expiry sweep);TestReydenThriftFallbackintests/unit/test_session.py(reactive recovery, cache pre-check, cache marking, explicit-backend guardrail, non-Reyden error not recovered, kernel-retry error chaining);test_reyden_sqlstate_raises_distinct_markerintests/unit/test_thrift_backend.py.E2E (against a real prod Reyden warehouse): an unconfigured
connect()transparently recovered onto the kernel andSELECT * from range(10)returned all rows; a secondconnect()in the same process took the cache pre-check path and opened the kernel without any Thrift round-trip.Related Tickets & Documents
Design: "Simplifying Reyden Onboarding on Drivers" (Option A). Mirrors the ADBC driver change adbc-drivers/databricks#670.
This PR was created with GitHub MCP.