WIP: this gets me a magicmock returned but doesn't honor my setting

Freezing this code here so I can come back to it after making adjustments

Signed-off-by: Jesse Whitehouse <[email protected]>

src/databricks/sql/client.py

@@ -8,7 +8,7 @@
 
 from databricks.sql import __version__
 from databricks.sql import *
-from databricks.sql.exc import OperationalError, SessionAlreadyClosedError
+from databricks.sql.exc import OperationalError, SessionAlreadyClosedError, CursorAlreadyClosedError
 from databricks.sql.thrift_backend import ThriftBackend
 from databricks.sql.utils import ExecuteResponse, ParamEscaper, inject_parameters
 from databricks.sql.types import Row
@@ -961,6 +961,9 @@ def close(self) -> None:
                 and self.connection.open
             ):
                 self.thrift_backend.close_command(self.command_id)
+        except RequestError as e:
+            if isinstance(e.args[1], CursorAlreadyClosedError):
+                logger.info("Operation was canceled by a prior request")
         finally:
             self.has_been_closed_server_side = True
             self.op_state = self.thrift_backend.CLOSED_OP_STATE

src/databricks/sql/exc.py

@@ -107,4 +107,7 @@ class UnsafeToRetryError(RequestError):
     """Thrown if ExecuteStatement request receives a code other than 200, 429, or 503"""
 
 class SessionAlreadyClosedError(RequestError):
-    """Thrown if CloseSession receives a code 404. ThriftBackend should gracefully proceed as this is expected."""
+    """Thrown if CloseSession receives a code 404. ThriftBackend should gracefully proceed as this is expected."""
+
+class CursorAlreadyClosedError(RequestError):
+    """Thrown if CancelOperation receives a code 404. ThriftBackend should gracefully proceed as this is expected."""

src/databricks/sql/thrift_backend.py

@@ -362,6 +362,9 @@ def attempt_request(attempt):
                     self._transport.setAllowRetries(False)
                 elif this_method_name == "CloseSession":
                     self._transport.set_retry_command_type(databricks.sql.auth.thrift_http_client.CommandType.CLOSE_SESSION)
+                elif this_method_name == "CancelOperation":
+                    self._transport.set_retry_command_type(databricks.sql.auth.thrift_http_client.CommandType.CANCEL_OPERATION)
+
                 else:
                     self._transport.set_retry_command_type(databricks.sql.auth.thrift_http_client.CommandType.OTHER)
 

tests/e2e/common/retry_test_mixins.py

@@ -42,7 +42,7 @@ def _test_retry_disabled_with_message(self, error_msg_substring, exception_type)
 
 from typing import List
 import pytest
-from unittest.mock import patch, MagicMock
+from unittest.mock import patch, MagicMock, PropertyMock
 from urllib3.exceptions import MaxRetryError
 from databricks.sql.exc import (
     RequestError,
@@ -209,6 +209,35 @@ def test_retry_abort_close_session_on_404(self):
                         expected_message_was_found = target in log
             self.assertTrue(expected_message_was_found, "Did not find expected log messages")
 
+    def test_retry_abort_cancel_operation_on_404(self):
+        """ GIVEN the connector sends a CancelOperation command
+            WHEN server sends a 404 (which is normally retried)
+            THEN nothing is retried because 404 means the operation was already canceled
+        """
+
+        # First response is a Bad Gateway -> Result is the command actually goes through
+        # Second response is a 404 because the session is no longer found
+        responses = [
+            {"status": 502, "headers": {"Retry-After": "1"}},
+            {"status": 404, "headers": {"Retry-After": None}},
+        ]
+
+        with self.connection(extra_params={**self._retry_policy}) as conn:
+            with conn.cursor() as curs:
+                with patch("databricks.sql.utils.ExecuteResponse.has_been_closed_server_side", new_callable=PropertyMock) as mock_prop:
+                    curs.execute("SELECT 1")
+                    mock_prop.return_value = False
+                    with mock_sequential_server_responses(responses):
+                        with self.assertLogs("databricks.sql", level="INFO",) as cm:
+                            curs.close()
+                        expected_message_was_found = False
+                        for log in cm.output:
+                            if expected_message_was_found:
+                                break
+                            target = "Operation was canceled by a prior request"
+                            expected_message_was_found = target in log
+                self.assertTrue(expected_message_was_found, "Did not find expected log messages")
+