feat: auto-recover Reyden Thrift connections onto the kernel

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,98 @@
+"""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)`` — the host is part of the key so the same
+warehouse id observed on two different workspaces never collides. 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
@@ -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,85 @@ 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
+
+        # 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({**kwargs, "use_kernel": True})
+
+        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({**kwargs, "use_kernel": True})
+            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:

tests/unit/test_reyden_warehouse_cache.py

@@ -0,0 +1,85 @@
+import pytest
+
+from databricks.sql.backend import reyden_warehouse_cache
+from databricks.sql.backend.reyden_warehouse_cache import (
+    _ReydenWarehouseCache,
+    extract_warehouse_id,
+)
+
+
+class TestExtractWarehouseId:
+    @pytest.mark.parametrize(
+        "path, expected",
+        [
+            ("/sql/1.0/warehouses/abc123", "abc123"),
+            ("/sql/1.0/endpoints/def456", "def456"),
+            ("/sql/1.0/warehouses/abc123?o=42", "abc123"),
+            ("sql/1.0/warehouses/wh?param=1&o=2", "wh"),
+            # All-purpose-compute cluster path — no warehouse id.
+            ("/sql/protocolv1/o/1234567890/0101-cluster", None),
+            ("", None),
+            (None, None),
+        ],
+    )
+    def test_extract(self, path, expected):
+        assert extract_warehouse_id(path) == expected
+
+
+class TestReydenWarehouseCacheClass:
+    def test_mark_then_known(self):
+        cache = _ReydenWarehouseCache()
+        assert cache.is_known_reyden("host", "wh") is False
+        cache.mark_reyden("host", "wh")
+        assert cache.is_known_reyden("host", "wh") is True
+
+    def test_host_case_insensitive(self):
+        cache = _ReydenWarehouseCache()
+        cache.mark_reyden("Host.Example.COM", "wh")
+        assert cache.is_known_reyden("host.example.com", "wh") is True
+
+    def test_distinct_hosts_do_not_collide(self):
+        cache = _ReydenWarehouseCache()
+        cache.mark_reyden("host-a", "wh")
+        # Same warehouse id on a different host must not be treated as Reyden.
+        assert cache.is_known_reyden("host-b", "wh") is False
+
+    def test_distinct_warehouses_do_not_collide(self):
+        cache = _ReydenWarehouseCache()
+        cache.mark_reyden("host", "wh-a")
+        assert cache.is_known_reyden("host", "wh-b") is False
+
+    def test_entry_expires_and_is_evicted(self):
+        cache = _ReydenWarehouseCache(ttl_seconds=0)
+        cache.mark_reyden("host", "wh")
+        # A zero TTL means the deadline is already in the past on read.
+        assert cache.is_known_reyden("host", "wh") is False
+        # Expired entry is evicted, not just reported absent.
+        assert cache._expiry == {}
+
+    def test_mark_sweeps_expired_entries(self):
+        # A zero TTL makes the first entry expired by the time the second mark
+        # runs, so the opportunistic sweep must purge it even though it was
+        # never read back.
+        cache = _ReydenWarehouseCache(ttl_seconds=0)
+        cache.mark_reyden("host", "old")
+        assert ("host", "old") in cache._expiry
+        cache.mark_reyden("host", "new")
+        assert ("host", "old") not in cache._expiry
+        assert ("host", "new") in cache._expiry
+
+    def test_mark_keeps_live_entries(self):
+        # Live (unexpired) entries survive the sweep on a subsequent mark.
+        cache = _ReydenWarehouseCache(ttl_seconds=3600)
+        cache.mark_reyden("host", "a")
+        cache.mark_reyden("host", "b")
+        assert ("host", "a") in cache._expiry
+        assert ("host", "b") in cache._expiry
+
+
+class TestModuleSingleton:
+    def test_mark_and_clear(self):
+        reyden_warehouse_cache.clear_cache()
+        reyden_warehouse_cache.mark_reyden("host", "wh")
+        assert reyden_warehouse_cache.is_known_reyden("host", "wh") is True
+        reyden_warehouse_cache.clear_cache()
+        assert reyden_warehouse_cache.is_known_reyden("host", "wh") is False