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

* feat: auto-recover Reyden Thrift connections onto the kernel

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>

* fix: attribute recovered-kernel connection failures to the kernel for 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>

* refactor: scope KP001 Reyden detection to OpenSession only

_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>

* Fix telemetry test mocks to set session.use_kernel

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>

* Mirror the Thrift auth default onto the Reyden kernel recovery path

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>

---------

Signed-off-by: Rahul Singhal <rahul.singhal@databricks.com>
Co-authored-by: Isaac <no-reply@databricks.com>

Author: rahuls-db

CHANGELOG.md

@@ -1,5 +1,8 @@
 # Release History
 
+# Unreleased
+- Transparently auto-recover Thrift connections to Reyden / Real-Time warehouses: when a warehouse rejects the default Thrift protocol (SQLSTATE `KP001`), the session is re-opened on the kernel backend and the warehouse is remembered so later connections skip Thrift. Applies only when no backend was chosen explicitly.
+
 # 4.5.0 (2026-09-01)
 - Upgrade Databricks SQL Kernel to 1.0.0.
 - Add JWT private-key M2M and Azure Entra authentication for kernel connections.

src/databricks/sql/backend/reyden_warehouse_cache.py

@@ -0,0 +1,101 @@
+"""Process-wide cache of warehouses known to reject the legacy Thrift protocol.
+
+A Reyden / Real-Time SQL warehouse rejects a Thrift ``OpenSession`` — the SQL
+Gateway proxy stamps SQLSTATE ``KP001`` on the rejection. When the driver
+auto-recovers by re-opening on the kernel backend, it records the warehouse
+here so later connections to the same warehouse skip the doomed Thrift attempt
+and open on the kernel directly.
+
+Keyed by ``(host, warehouse_id)``. Warehouse ids are globally unique, so the
+warehouse id alone identifies the warehouse — even on a SPOG host shared by many
+workspaces (where only the ``?o=<workspace-id>`` path param distinguishes them),
+there is no cross-workspace collision. The host is kept in the key only as a
+cheap optimization (scoping lookups) and defense-in-depth, not for correctness.
+Entries expire after ``_TTL_SECONDS`` so a warehouse later reconfigured to accept
+Thrift is eventually retried.
+"""
+
+import re
+import threading
+import time
+from typing import Dict, Optional, Tuple
+
+# A warehouse's Reyden membership can change (an id may be recreated on a
+# Thrift-capable endpoint), so cached entries are re-validated after this long.
+# Matches the ADBC driver's 6-hour horizon.
+_TTL_SECONDS = 6 * 60 * 60
+
+# Warehouse paths look like ``/sql/1.0/warehouses/<id>`` or
+# ``.../endpoints/<id>``; the id stops at the next ``/``, ``?`` or ``&`` (e.g. a
+# ``?o=`` SPOG routing param). All-purpose-compute cluster paths carry no
+# warehouse id and never match — they are never Reyden warehouses.
+_WAREHOUSE_PATH_RE = re.compile(r".*/(?:warehouses|endpoints)/([^?&/]+)")
+
+
+def extract_warehouse_id(http_path: Optional[str]) -> Optional[str]:
+    """Return the warehouse/endpoint id embedded in ``http_path``, or ``None``."""
+    if not http_path:
+        return None
+    match = _WAREHOUSE_PATH_RE.match(http_path)
+    return match.group(1) if match else None
+
+
+class _ReydenWarehouseCache:
+    def __init__(self, ttl_seconds: float = _TTL_SECONDS) -> None:
+        self._ttl_seconds = ttl_seconds
+        self._lock = threading.Lock()
+        # (host_lowercased, warehouse_id) -> monotonic expiry deadline
+        self._expiry: Dict[Tuple[str, str], float] = {}
+
+    @staticmethod
+    def _key(host: str, warehouse_id: str) -> Tuple[str, str]:
+        return (host.lower(), warehouse_id)
+
+    def mark_reyden(self, host: str, warehouse_id: str) -> None:
+        now = time.monotonic()
+        with self._lock:
+            # Opportunistic sweep: mark_reyden only runs on an actual Thrift
+            # rejection (rare), so purging every expired entry here is near-free
+            # and bounds the cache to warehouses seen within the TTL window
+            # rather than every warehouse ever seen (the per-key lazy eviction
+            # in is_known_reyden never reclaims a warehouse that is not looked
+            # up again).
+            for key in [k for k, deadline in self._expiry.items() if deadline <= now]:
+                del self._expiry[key]
+            self._expiry[self._key(host, warehouse_id)] = now + self._ttl_seconds
+
+    def is_known_reyden(self, host: str, warehouse_id: str) -> bool:
+        key = self._key(host, warehouse_id)
+        now = time.monotonic()
+        with self._lock:
+            deadline = self._expiry.get(key)
+            if deadline is None:
+                return False
+            if deadline <= now:
+                # Lazily evict so a reconfigured warehouse is retried over Thrift.
+                del self._expiry[key]
+                return False
+            return True
+
+    def clear(self) -> None:
+        with self._lock:
+            self._expiry.clear()
+
+
+# Process-wide singleton; multi-tenant safe via the host component of the key.
+_CACHE = _ReydenWarehouseCache()
+
+
+def mark_reyden(host: str, warehouse_id: str) -> None:
+    """Record that ``warehouse_id`` on ``host`` rejects the Thrift protocol."""
+    _CACHE.mark_reyden(host, warehouse_id)
+
+
+def is_known_reyden(host: str, warehouse_id: str) -> bool:
+    """Whether ``warehouse_id`` on ``host`` is known (unexpired) to reject Thrift."""
+    return _CACHE.is_known_reyden(host, warehouse_id)
+
+
+def clear_cache() -> None:
+    """Reset the cache. Intended for tests."""
+    _CACHE.clear()

src/databricks/sql/backend/thrift_backend.py

@@ -284,11 +284,23 @@ def _initialize_retry_args(self, kwargs):
             )
 
     @staticmethod
-    def _check_response_for_error(response, host_url=None):
+    def _check_response_for_error(response, host_url=None, detect_reyden=False):
         if response.status and response.status.statusCode in [
             ttypes.TStatusCode.ERROR_STATUS,
             ttypes.TStatusCode.INVALID_HANDLE_STATUS,
         ]:
+            # A Reyden / Real-Time warehouse rejects the legacy Thrift protocol
+            # with SQLSTATE KP001, but only at OpenSession. `detect_reyden` gates
+            # the marker to that call so a stray KP001 on any other RPC surfaces
+            # as a normal DatabaseError (the connection-layer recovery only wraps
+            # session open). host_url is deliberately omitted on the marker: it is
+            # a recoverable signal, not a terminal failure, so it must not emit a
+            # failure-telemetry event here.
+            if (
+                detect_reyden
+                and response.status.sqlState == ReydenThriftUnsupportedError.SQL_STATE
+            ):
+                raise ReydenThriftUnsupportedError(response.status.errorMessage)
             raise DatabaseError(
                 response.status.errorMessage,
                 host_url=host_url,
@@ -520,7 +532,14 @@ def attempt_request(attempt):
             if not isinstance(response_or_error_info, RequestErrorInfo):
                 # log nothing here, presume that main request logging covers
                 response = response_or_error_info
-                ThriftDatabricksClient._check_response_for_error(response, self._host)
+                # Only OpenSession opts into KP001→Reyden-marker detection (the
+                # rejection is stamped only there). Mirrors the method.__name__
+                # discrimination already used above for GetOperationStatus.
+                ThriftDatabricksClient._check_response_for_error(
+                    response,
+                    self._host,
+                    detect_reyden=getattr(method, "__name__", None) == "OpenSession",
+                )
                 return response
 
             error_info = response_or_error_info

src/databricks/sql/client.py

@@ -35,6 +35,12 @@
     ProgrammingError,
     TransactionError,
     DatabaseError,
+    ReydenThriftUnsupportedError,
+)
+from databricks.sql.backend.reyden_warehouse_cache import (
+    extract_warehouse_id,
+    is_known_reyden,
+    mark_reyden,
 )
 
 from databricks.sql.backend.databricks_client import DatabricksClient
@@ -66,7 +72,7 @@
 from databricks.sql.session import Session
 from databricks.sql.backend.types import CommandId, BackendType, CommandState, SessionId
 
-from databricks.sql.auth.common import ClientContext
+from databricks.sql.auth.common import AuthType, ClientContext
 from databricks.sql.common.unified_http_client import UnifiedHttpClient
 from databricks.sql.common.http import HttpMethod
 
@@ -399,24 +405,30 @@ def read(self) -> Optional[OAuthToken]:
         self.http_client = UnifiedHttpClient(client_context)
 
         try:
-            self.session = Session(
+            self.session = self._open_session_with_reyden_fallback(
                 server_hostname,
                 http_path,
-                self.http_client,
                 http_headers,
                 session_configuration,
                 catalog,
                 schema,
                 _use_arrow_native_complex_types,
-                **kwargs,
+                kwargs,
             )
-            self.session.open()
         except Exception as e:
             # Respect user's telemetry preference even during connection failure.
-            # For use_kernel connections the kernel owns telemetry, so suppress
-            # the wrapper-side failure log to avoid wrapper-vs-kernel duplication.
-            enable_telemetry = kwargs.get("enable_telemetry", True) and not kwargs.get(
-                "use_kernel", False
+            # For a kernel connection the kernel owns telemetry, so suppress the
+            # wrapper-side failure log to avoid wrapper-vs-kernel duplication.
+            # Read the backend from the session that actually failed rather than
+            # the caller's kwargs: on the Reyden auto-recovery path we retry on
+            # the kernel via a kwargs copy, so the original kwargs still says
+            # Thrift. If the kernel never got constructed (e.g. its wheel is
+            # missing), self.session is the Thrift session and we still log.
+            attempted_kernel = getattr(
+                getattr(self, "session", None), "use_kernel", False
+            )
+            enable_telemetry = (
+                kwargs.get("enable_telemetry", True) and not attempted_kernel
             )
             TelemetryClientFactory.connection_failure_log(
                 error_name="Exception",
@@ -512,6 +524,106 @@ def read(self) -> Optional[OAuthToken]:
             session_id=self.get_session_id_hex(),
         )
 
+    def _open_session_with_reyden_fallback(
+        self,
+        server_hostname: str,
+        http_path: str,
+        http_headers,
+        session_configuration,
+        catalog,
+        schema,
+        _use_arrow_native_complex_types,
+        kwargs: dict,
+    ) -> Session:
+        """Open a ``Session``, transparently recovering onto the kernel backend
+        when a Reyden / Real-Time warehouse rejects the default Thrift protocol.
+
+        Reyden warehouses reject a Thrift ``OpenSession`` (SQLSTATE ``KP001``);
+        the kernel (SEA) backend is the supported path. Auto-recovery applies
+        only when the caller did not pick a backend explicitly (neither
+        ``use_kernel`` nor ``use_sea``). On a rejection the warehouse is
+        remembered so later connections skip the doomed Thrift attempt.
+        """
+
+        def build_session(session_kwargs: dict) -> Session:
+            # Assign self.session before open() so a failed open still leaves the
+            # attempted session on the connection — __del__ and the failure
+            # telemetry log both rely on self.session being present.
+            self.session = Session(
+                server_hostname,
+                http_path,
+                self.http_client,
+                http_headers,
+                session_configuration,
+                catalog,
+                schema,
+                _use_arrow_native_complex_types,
+                **session_kwargs,
+            )
+            self.session.open()
+            return self.session
+
+        def kernel_recovery_kwargs() -> dict:
+            # Kwargs for re-opening on the kernel. The Thrift path treats an
+            # unset auth_type as databricks-oauth (see get_auth_provider); the
+            # kernel path has no such fallback and rejects auth_type=None unless
+            # a credential shape (PAT / OAuth M2M) is present. Mirror the Thrift
+            # default so a bare OAuth-U2M connection recovers instead of failing
+            # with NotSupportedError. Skip the injection when a credential shape
+            # is already present — the kernel routes on it regardless of
+            # auth_type, and forcing databricks-oauth alongside an M2M secret or
+            # a credentials_provider would change that routing.
+            recovery_kwargs = {**kwargs, "use_kernel": True}
+            has_credential_shape = (
+                recovery_kwargs.get("access_token")
+                or recovery_kwargs.get("oauth_client_secret")
+                or recovery_kwargs.get("oauth_jwt_key_file")
+                or recovery_kwargs.get("credentials_provider")
+            )
+            if recovery_kwargs.get("auth_type") is None and not has_credential_shape:
+                recovery_kwargs["auth_type"] = AuthType.DATABRICKS_OAUTH.value
+            return recovery_kwargs
+
+        # An explicit backend choice is always honored — auto-recovery engages
+        # only on the default (Thrift) path.
+        explicit_backend = kwargs.get("use_kernel", False) or kwargs.get(
+            "use_sea", False
+        )
+        if explicit_backend:
+            return build_session(kwargs)
+
+        warehouse_id = extract_warehouse_id(http_path)
+
+        # Pre-check: a warehouse already seen to reject Thrift opens straight on
+        # the kernel, skipping the doomed Thrift OpenSession round-trip.
+        if warehouse_id and is_known_reyden(server_hostname, warehouse_id):
+            logger.info(
+                "Warehouse %s on %s is known to require the kernel backend; "
+                "opening on the kernel and skipping Thrift.",
+                warehouse_id,
+                server_hostname,
+            )
+            return build_session(kernel_recovery_kwargs())
+
+        try:
+            return build_session(kwargs)
+        except ReydenThriftUnsupportedError as thrift_ex:
+            logger.info(
+                "Thrift is not supported for this Reyden/Real-Time warehouse; "
+                "transparently re-opening the session on the kernel backend."
+            )
+            # Remember the rejection regardless of the retry's outcome — the
+            # warehouse is Reyden either way, so future connects should skip
+            # Thrift; a kernel failure below is a separate, orthogonal problem.
+            if warehouse_id:
+                mark_reyden(server_hostname, warehouse_id)
+            try:
+                return build_session(kernel_recovery_kwargs())
+            except Exception as kernel_ex:
+                # Surface the kernel failure (the actionable one) while keeping
+                # the original Thrift rejection in the chain for diagnosis.
+                raise kernel_ex from thrift_ex
+
     def _set_use_inline_params_with_warning(self, value: Union[bool, str]):
         """Valid values are True, False, and "silent"
 

src/databricks/sql/exc.py

@@ -114,6 +114,20 @@ class ServerOperationError(DatabaseError):
     pass
 
 
+class ReydenThriftUnsupportedError(DatabaseError):
+    """Marker for a Reyden / Real-Time warehouse rejecting the legacy Thrift
+    protocol at OpenSession (the SQL Gateway proxy stamps SQLSTATE ``KP001``).
+
+    It signals the connection layer to transparently re-open the session on the
+    kernel backend. Subclassing ``DatabaseError`` means that when auto-recovery
+    does not apply (an explicit backend was chosen) or the kernel retry also
+    fails, callers catching ``DatabaseError`` still observe it.
+    """
+
+    # SQLSTATE the SQL Gateway proxy stamps on the Reyden Thrift rejection.
+    SQL_STATE = "KP001"
+
+
 class RequestError(OperationalError):
     """Thrown if there was a error during request to the server.
     Its context will have the following keys: