diff --git a/CHANGELOG.md b/CHANGELOG.md index 2f021ed09..37a508355 100644 --- a/CHANGELOG.md +++ b/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. diff --git a/src/databricks/sql/backend/reyden_warehouse_cache.py b/src/databricks/sql/backend/reyden_warehouse_cache.py new file mode 100644 index 000000000..6f443a7ea --- /dev/null +++ b/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/`` or +# ``.../endpoints/``; 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() diff --git a/src/databricks/sql/backend/thrift_backend.py b/src/databricks/sql/backend/thrift_backend.py index e047aedf6..fcf363cf0 100644 --- a/src/databricks/sql/backend/thrift_backend.py +++ b/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 diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py index 4d9ca0327..009080e3e 100755 --- a/src/databricks/sql/client.py +++ b/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" diff --git a/src/databricks/sql/exc.py b/src/databricks/sql/exc.py index 9e918a936..ae648fd10 100644 --- a/src/databricks/sql/exc.py +++ b/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: diff --git a/tests/unit/test_reyden_warehouse_cache.py b/tests/unit/test_reyden_warehouse_cache.py new file mode 100644 index 000000000..dff154ef8 --- /dev/null +++ b/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 diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py index dc08470d9..8b14a5c52 100644 --- a/tests/unit/test_session.py +++ b/tests/unit/test_session.py @@ -1,5 +1,6 @@ import pytest from unittest.mock import patch, MagicMock, Mock, PropertyMock +from contextlib import contextmanager import gc from databricks.sql.thrift_api.TCLIService.ttypes import ( @@ -9,6 +10,11 @@ ) from databricks.sql.backend.types import SessionId, BackendType from databricks.sql.common.agent import KNOWN_AGENTS +from databricks.sql.exc import ( + ReydenThriftUnsupportedError, + DatabaseError, + OperationalError, +) from databricks.sql.session import Session import databricks.sql @@ -781,3 +787,146 @@ def test_connect_use_kernel_instantiates_real_kernel_backend(self): ) finally: conn.close() + + +class TestReydenThriftFallback: + """Transparent auto-recovery from a Reyden / Real-Time warehouse rejecting + the legacy Thrift protocol (SQLSTATE KP001) onto the kernel backend.""" + + PACKAGE = "databricks.sql" + HOST = "reyden.example.com" + WAREHOUSE_PATH = "/sql/1.0/warehouses/wh-reyden" + + @pytest.fixture(autouse=True) + def _clear_cache(self): + from databricks.sql.backend import reyden_warehouse_cache + + reyden_warehouse_cache.clear_cache() + yield + reyden_warehouse_cache.clear_cache() + + @contextmanager + def _fake_kernel(self): + """Fake the Rust wheel and patch KernelDatabricksClient so open_session + returns a valid SessionId; yields the mock for call assertions.""" + import sys + import types + + pytest.importorskip( + "pyarrow", reason="kernel client module imports pyarrow at load" + ) + fake = types.ModuleType("databricks_sql_kernel") + fake.KernelError = type("KernelError", (Exception,), {}) + fake.Session = MagicMock() + with patch.dict(sys.modules, {"databricks_sql_kernel": fake}), patch( + "databricks.sql.backend.kernel.client.KernelDatabricksClient" + ) as mock_kernel: + mock_kernel.return_value.open_session.return_value = SessionId( + BackendType.SEA, "sess-id", None + ) + yield mock_kernel + + @staticmethod + def _reject(): + return ReydenThriftUnsupportedError( + "Lakehouse/RT is not supported for Thrift protocol" + ) + + def _connect(self, **overrides): + args = dict( + server_hostname=self.HOST, + http_path=self.WAREHOUSE_PATH, + access_token="tok", + enable_telemetry=False, + ) + args.update(overrides) + return databricks.sql.connect(**args) + + @patch("%s.session.ThriftDatabricksClient" % PACKAGE) + def test_thrift_rejection_recovers_onto_kernel(self, mock_thrift): + mock_thrift.return_value.open_session.side_effect = self._reject() + with self._fake_kernel() as mock_kernel: + conn = self._connect() + try: + assert conn.session.use_kernel is True + mock_kernel.return_value.open_session.assert_called_once() + finally: + conn.close() + + @patch("%s.session.ThriftDatabricksClient" % PACKAGE) + def test_known_reyden_warehouse_skips_thrift(self, mock_thrift): + from databricks.sql.backend import reyden_warehouse_cache + + # A previously observed rejection is remembered — casing of the host in + # the cache key must not matter. + reyden_warehouse_cache.mark_reyden(self.HOST.upper(), "wh-reyden") + with self._fake_kernel() as mock_kernel: + conn = self._connect() + try: + assert conn.session.use_kernel is True + mock_kernel.return_value.open_session.assert_called_once() + # Pre-check must short-circuit before any Thrift OpenSession. + mock_thrift.return_value.open_session.assert_not_called() + finally: + conn.close() + + @patch("%s.session.ThriftDatabricksClient" % PACKAGE) + def test_rejection_marks_cache_for_next_connect(self, mock_thrift): + from databricks.sql.backend import reyden_warehouse_cache + + mock_thrift.return_value.open_session.side_effect = self._reject() + with self._fake_kernel(): + conn = self._connect() + conn.close() + assert reyden_warehouse_cache.is_known_reyden(self.HOST, "wh-reyden") + + @patch("%s.session.ThriftDatabricksClient" % PACKAGE) + def test_non_reyden_thrift_error_not_recovered(self, mock_thrift): + mock_thrift.return_value.open_session.side_effect = DatabaseError( + "some unrelated server error" + ) + with pytest.raises(DatabaseError) as excinfo: + self._connect() + assert not isinstance(excinfo.value, ReydenThriftUnsupportedError) + + @patch("%s.session.SeaDatabricksClient" % PACKAGE) + def test_explicit_use_sea_not_recovered(self, mock_sea): + # An explicit backend choice is always honored — even the (contrived) + # case of SEA surfacing the marker must not trigger Thrift→kernel + # recovery. + mock_sea.return_value.open_session.side_effect = self._reject() + with pytest.raises(ReydenThriftUnsupportedError): + self._connect(use_sea=True) + + @patch("%s.session.ThriftDatabricksClient" % PACKAGE) + def test_kernel_retry_failure_chains_both_errors(self, mock_thrift): + mock_thrift.return_value.open_session.side_effect = self._reject() + with self._fake_kernel() as mock_kernel: + mock_kernel.return_value.open_session.side_effect = OperationalError( + "kernel could not open session" + ) + with pytest.raises(OperationalError) as excinfo: + self._connect() + # Kernel failure is surfaced as primary; the Thrift rejection is + # preserved in the chain for diagnosis. + assert isinstance(excinfo.value.__cause__, ReydenThriftUnsupportedError) + + @patch("%s.session.ThriftDatabricksClient" % PACKAGE) + def test_recovered_kernel_failure_suppresses_wrapper_telemetry(self, mock_thrift): + # A connection that recovered onto the kernel and then failed there must + # NOT emit the wrapper's connection-failure log — the kernel owns + # telemetry for kernel connections. Guards against reading the original + # (Thrift) kwargs instead of the session that actually failed. + mock_thrift.return_value.open_session.side_effect = self._reject() + with self._fake_kernel() as mock_kernel, patch( + "databricks.sql.client.TelemetryClientFactory.connection_failure_log" + ) as mock_fail_log: + mock_kernel.return_value.open_session.side_effect = OperationalError( + "kernel boom" + ) + # enable_telemetry=True so only the kernel-suppression logic can flip + # it off — proving the fix rather than the user's opt-out. + with pytest.raises(OperationalError): + self._connect(enable_telemetry=True) + mock_fail_log.assert_called_once() + assert mock_fail_log.call_args.kwargs["enable_telemetry"] is False diff --git a/tests/unit/test_telemetry.py b/tests/unit/test_telemetry.py index d2d69b9f9..5663fa119 100644 --- a/tests/unit/test_telemetry.py +++ b/tests/unit/test_telemetry.py @@ -482,6 +482,9 @@ def test_connection_failure_sends_correct_telemetry_payload( # Set up the mock to create a session instance first, then make open() fail mock_session_instance = MagicMock() mock_session_instance.is_open = False # Ensure cleanup is safe + # Default (Thrift) session: the failure-telemetry suppression reads + # session.use_kernel, so the mock must expose a real bool, not a truthy Mock. + mock_session_instance.use_kernel = False mock_session_instance.open.side_effect = Exception(error_message) mock_session.return_value = mock_session_instance @@ -511,6 +514,8 @@ def test_connection_failure_does_not_send_telemetry_for_kernel( error_message = "Could not connect to host" mock_session_instance = MagicMock() mock_session_instance.is_open = False + # Kernel session: use_kernel is read to suppress the wrapper-side failure log. + mock_session_instance.use_kernel = True mock_session_instance.open.side_effect = Exception(error_message) mock_session.return_value = mock_session_instance diff --git a/tests/unit/test_thrift_backend.py b/tests/unit/test_thrift_backend.py index 4746b18ff..1dff470dc 100644 --- a/tests/unit/test_thrift_backend.py +++ b/tests/unit/test_thrift_backend.py @@ -594,6 +594,78 @@ def test_make_request_checks_status_code(self): mock_response.status.statusCode = code thrift_backend.make_request(lambda _: mock_response, Mock()) + def test_reyden_sqlstate_raises_distinct_marker(self): + reyden_resp = Mock() + reyden_resp.status = ttypes.TStatus( + statusCode=ttypes.TStatusCode.ERROR_STATUS, + sqlState=ReydenThriftUnsupportedError.SQL_STATE, + errorMessage="Lakehouse/RT is not supported for Thrift protocol", + ) + + # KP001 on an ERROR_STATUS → the recoverable Reyden marker, but ONLY when + # detection is enabled (i.e. the OpenSession path). + with self.assertRaises(ReydenThriftUnsupportedError): + ThriftDatabricksClient._check_response_for_error( + reyden_resp, detect_reyden=True + ) + + # The same KP001 on any other RPC (detect_reyden=False, the default) is a + # generic DatabaseError — the marker is scoped to OpenSession so recovery + # never has to handle it elsewhere. + with self.assertRaises(DatabaseError) as cm: + ThriftDatabricksClient._check_response_for_error( + reyden_resp, detect_reyden=False + ) + self.assertNotIsInstance(cm.exception, ReydenThriftUnsupportedError) + + # Any other sqlState on an ERROR_STATUS is never the marker, even on the + # OpenSession path. + other_resp = Mock() + other_resp.status = ttypes.TStatus( + statusCode=ttypes.TStatusCode.ERROR_STATUS, + sqlState="42000", + errorMessage="a syntax error", + ) + with self.assertRaises(DatabaseError) as cm: + ThriftDatabricksClient._check_response_for_error( + other_resp, detect_reyden=True + ) + self.assertNotIsInstance(cm.exception, ReydenThriftUnsupportedError) + + def test_reyden_detection_wired_only_for_open_session(self): + # make_request enables Reyden detection based on the RPC method name, so + # a KP001 from OpenSession maps to the marker while the same status from + # any other RPC stays a generic DatabaseError. + thrift_backend = ThriftDatabricksClient( + "foobar", + 443, + "path", + [], + auth_provider=AuthProvider(), + ssl_options=SSLOptions(), + http_client=MagicMock(), + ) + reyden_resp = Mock() + reyden_resp.status = ttypes.TStatus( + statusCode=ttypes.TStatusCode.ERROR_STATUS, + sqlState=ReydenThriftUnsupportedError.SQL_STATE, + errorMessage="Lakehouse/RT is not supported for Thrift protocol", + ) + + # make_request keys detection off method.__name__. + def OpenSession(_): + return reyden_resp + + def ExecuteStatement(_): + return reyden_resp + + with self.assertRaises(ReydenThriftUnsupportedError): + thrift_backend.make_request(OpenSession, Mock()) + + with self.assertRaises(DatabaseError) as cm: + thrift_backend.make_request(ExecuteStatement, Mock()) + self.assertNotIsInstance(cm.exception, ReydenThriftUnsupportedError) + def test_handle_execute_response_checks_operation_state_in_direct_results(self): for resp_type in self.execute_response_types: with self.subTest(resp_type=resp_type):