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>

Author: rahuls-db

src/databricks/sql/backend/thrift_backend.py

@@ -284,17 +284,22 @@ 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. Surface a distinct marker so the connection
-            # layer can transparently re-open on the kernel backend. host_url is
-            # deliberately omitted: this is a recoverable signal, not a terminal
-            # failure, so it must not emit a failure-telemetry event here.
-            if response.status.sqlState == ReydenThriftUnsupportedError.SQL_STATE:
+            # 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,
@@ -527,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

tests/unit/test_thrift_backend.py

@@ -595,6 +595,47 @@ def test_make_request_checks_status_code(self):
             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,
@@ -604,27 +645,25 @@ def test_reyden_sqlstate_raises_distinct_marker(self):
             ssl_options=SSLOptions(),
             http_client=MagicMock(),
         )
-
-        # KP001 on an ERROR_STATUS → the recoverable Reyden marker.
         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(lambda _: reyden_resp, Mock())
+            thrift_backend.make_request(OpenSession, Mock())
 
-        # Any other sqlState on an ERROR_STATUS → the generic DatabaseError, and
-        # explicitly NOT the Reyden marker.
-        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:
-            thrift_backend.make_request(lambda _: other_resp, Mock())
+            thrift_backend.make_request(ExecuteStatement, Mock())
         self.assertNotIsInstance(cm.exception, ReydenThriftUnsupportedError)
 
     def test_handle_execute_response_checks_operation_state_in_direct_results(self):