Skip to content

feat: auto-recover Reyden Thrift connections onto the kernel - #948

Merged
rahuls-db merged 5 commits into
mainfrom
feat/reyden-thrift-auto-recovery
Sep 10, 2026
Merged

feat: auto-recover Reyden Thrift connections onto the kernel#948
rahuls-db merged 5 commits into
mainfrom
feat/reyden-thrift-auto-recovery

Conversation

@rahuls-db

Copy link
Copy Markdown
Collaborator

What type of PR is this?

  • Feature

Description

An unconfigured connect() to a Reyden / Real-Time SQL warehouse defaults to the Thrift backend, which the SQL Gateway proxy rejects (SQLSTATE KP001). This change detects that rejection at OpenSession and transparently re-opens the session on the kernel (SEA) backend, so no customer configuration change is needed.

  • Detection: ThriftDatabricksClient._check_response_for_error raises a distinct ReydenThriftUnsupportedError when the OpenSession TStatus carries SQLSTATE KP001 (matched on the SQLSTATE only — no message-string matching).
  • Recovery: Connection catches the marker and re-opens the session once with use_kernel=True. If the kernel open also fails, the kernel error is surfaced with the Thrift rejection preserved in the exception chain.
  • Pinning: a process-wide cache keyed by (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.
  • Scope / guardrails: only the default (unspecified) backend path auto-recovers; an explicit use_kernel/use_sea is 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 tests
  • E2E Tests
  • Manually

Unit: new tests/unit/test_reyden_warehouse_cache.py (cache behavior + warehouse-id extraction + expiry sweep); TestReydenThriftFallback in tests/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_marker in tests/unit/test_thrift_backend.py.

E2E (against a real prod Reyden warehouse): an unconfigured connect() transparently recovered onto the kernel and SELECT * from range(10) returned all rows; a second connect() 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.

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>
@rahuls-db
rahuls-db marked this pull request as ready for review September 9, 2026 20:41

@peco-review-bot peco-review-bot Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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().

Comment thread src/databricks/sql/client.py
Comment thread src/databricks/sql/backend/thrift_backend.py Outdated
Comment thread src/databricks/sql/client.py Outdated
… 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>

@peco-review-bot peco-review-bot Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread src/databricks/sql/backend/thrift_backend.py Outdated
_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>

@peco-review-bot peco-review-bot Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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=True connection whose kernel Session.__init__ raises (e.g. the kernel wheel is missing, so _create_backend fails before self.session is ever assigned), getattr(self, "session", None) is None, so attempted_kernel is False and the wrapper-side failure log now emits — whereas the old not 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>

@peco-review-bot peco-review-bot Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread src/databricks/sql/client.py
Comment thread src/databricks/sql/client.py
Comment thread src/databricks/sql/client.py
Comment thread src/databricks/sql/client.py
Comment thread src/databricks/sql/client.py
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>

@peco-review-bot peco-review-bot Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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_requestsession.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.

Comment thread src/databricks/sql/backend/reyden_warehouse_cache.py

@vuanhphung vuanhphung left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

@rahuls-db rahuls-db added the skip-coverage Skip the coverage fan-out for this PR (no tracking issue opened in databricks-driver-test) label Sep 10, 2026
@rahuls-db
rahuls-db enabled auto-merge September 10, 2026 18:54
@rahuls-db
rahuls-db added this pull request to the merge queue Sep 10, 2026
Merged via the queue into main with commit 696e780 Sep 10, 2026
63 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ai-assisted skip-coverage Skip the coverage fan-out for this PR (no tracking issue opened in databricks-driver-test)

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants