Improve error handling in connection and cursor context managers to ensure proper closure during exceptions, including KeyboardInterrupt. Add tests for nested cursor management and verify operation closure on server-side errors.

Author: madhav-db

src/databricks/sql/client.py

@@ -315,7 +315,13 @@ def __enter__(self) -> "Connection":
         return self
 
     def __exit__(self, exc_type, exc_value, traceback):
-        self.close()
+        try:
+            self.close()
+        except BaseException as e:
+            logger.warning(f"Exception during connection close in __exit__: {e}")
+            if exc_type is None:
+                raise
+        return False
 
     def __del__(self):
         if self.open:
@@ -459,11 +465,9 @@ def __exit__(self, exc_type, exc_value, traceback):
         try:
             logger.debug("Cursor context manager exiting, calling close()")
             self.close()
-        except Exception as e:
+        except BaseException as e:
             logger.warning(f"Exception during cursor close in __exit__: {e}")
-            # Don't suppress the original exception if there was one
             if exc_type is None:
-                # Only raise our new exception if there wasn't already one in progress
                 raise
         return False
 

tests/e2e/test_driver.py

@@ -50,7 +50,7 @@
 
 from tests.e2e.common.uc_volume_tests import PySQLUCVolumeTestSuiteMixin
 
-from databricks.sql.exc import SessionAlreadyClosedError
+from databricks.sql.exc import SessionAlreadyClosedError, CursorAlreadyClosedError
 
 log = logging.getLogger(__name__)
 
@@ -813,7 +813,6 @@ def test_close_connection_closes_cursors(self):
             ars = cursor.active_result_set
 
             # We must manually run this check because thrift_backend always forces `has_been_closed_server_side` to True
-
             # Cursor op state should be open before connection is closed
             status_request = ttypes.TGetOperationStatusReq(
                 operationHandle=ars.command_id, getProgressUpdate=False
@@ -840,9 +839,106 @@ def test_closing_a_closed_connection_doesnt_fail(self, caplog):
         with self.connection() as conn:
             # First .close() call is explicit here
             conn.close()
-
         assert "Session appears to have been closed already" in caplog.text
 
+        # --- Integrated KeyboardInterrupt test ---
+        conn = None
+        try:
+            with pytest.raises(KeyboardInterrupt):
+                with self.connection() as c:
+                    conn = c
+                    raise KeyboardInterrupt("Simulated interrupt")
+        finally:
+            if conn is not None:
+                assert not conn.open, "Connection should be closed after KeyboardInterrupt"
+
+    def test_cursor_close_properly_closes_operation(self):
+        """Test that Cursor.close() properly closes the active operation handle on the server."""
+        with self.connection() as conn:
+            cursor = conn.cursor()
+            try:
+                cursor.execute("SELECT 1 AS test")
+                assert cursor.active_op_handle is not None
+                cursor.close()
+                assert cursor.active_op_handle is None
+                assert not cursor.open
+            finally:
+                if cursor.open:
+                    cursor.close()
+
+        # --- Integrated KeyboardInterrupt test ---
+        conn = None
+        cursor = None
+        try:
+            with self.connection() as c:
+                conn = c
+                with pytest.raises(KeyboardInterrupt):
+                    with conn.cursor() as cur:
+                        cursor = cur
+                        raise KeyboardInterrupt("Simulated interrupt")
+        finally:
+            if cursor is not None:
+                assert not cursor.open, "Cursor should be closed after KeyboardInterrupt"
+
+    def test_nested_cursor_context_managers(self):
+        """Test that nested cursor context managers properly close operations on the server."""
+        with self.connection() as conn:
+            with conn.cursor() as cursor1:
+                cursor1.execute("SELECT 1 AS test1")
+                assert cursor1.active_op_handle is not None
+
+                with conn.cursor() as cursor2:
+                    cursor2.execute("SELECT 2 AS test2")
+                    assert cursor2.active_op_handle is not None
+
+                # After inner context manager exit, cursor2 should be not open
+                assert not cursor2.open
+                assert cursor2.active_op_handle is None
+
+            # After outer context manager exit, cursor1 should be not open
+            assert not cursor1.open
+            assert cursor1.active_op_handle is None
+
+    def test_cursor_error_handling(self):
+        """Test that cursor close handles errors properly to prevent orphaned operations."""
+        with self.connection() as conn:
+            cursor = conn.cursor()
+
+            cursor.execute("SELECT 1 AS test")
+
+            op_handle = cursor.active_op_handle
+
+            assert op_handle is not None
+
+            # Manually close the operation to simulate server-side closure
+            conn.thrift_backend.close_command(op_handle)
+
+            cursor.close()
+
+            assert not cursor.open
+
+    def test_result_set_close(self):
+        """Test that ResultSet.close() properly closes operations on the server and handles state correctly."""
+        with self.connection() as conn:
+            cursor = conn.cursor()
+            try:
+                cursor.execute("SELECT * FROM RANGE(10)")
+
+                result_set = cursor.active_result_set
+                assert result_set is not None
+
+                initial_op_state = result_set.op_state
+
+                result_set.close()
+
+                assert result_set.op_state == result_set.thrift_backend.CLOSED_OP_STATE
+                assert result_set.op_state != initial_op_state
+
+                # Closing the result set again should be a no-op and not raise exceptions
+                result_set.close()
+            finally:
+                cursor.close()
+
 
 # use a RetrySuite to encapsulate these tests which we'll typically want to run together; however keep
 # the 429/503 subsuites separate since they execute under different circumstances.
@@ -875,3 +971,32 @@ def test_initial_namespace(self):
             assert cursor.fetchone()[0] == self.arguments["catalog"]
             cursor.execute("select current_database()")
             assert cursor.fetchone()[0] == table_name
+
+
+class TestContextManagerInterrupts(PySQLPytestTestCase):
+    def test_connection_context_manager_handles_keyboard_interrupt(self):
+        # This test ensures that a KeyboardInterrupt inside a connection context propagates and closes the connection
+        conn = None
+        try:
+            with pytest.raises(KeyboardInterrupt):
+                with self.connection() as c:
+                    conn = c
+                    raise KeyboardInterrupt("Simulated interrupt")
+        finally:
+            if conn is not None:
+                assert not conn.open, "Connection should be closed after KeyboardInterrupt"
+
+    def test_cursor_context_manager_handles_keyboard_interrupt(self):
+        # This test ensures that a KeyboardInterrupt inside a cursor context propagates and closes the cursor
+        conn = None
+        cursor = None
+        try:
+            with self.connection() as c:
+                conn = c
+                with pytest.raises(KeyboardInterrupt):
+                    with conn.cursor() as cur:
+                        cursor = cur
+                        raise KeyboardInterrupt("Simulated interrupt")
+        finally:
+            if cursor is not None:
+                assert not cursor.open, "Cursor should be closed after KeyboardInterrupt"

tests/unit/test_client.py

@@ -284,6 +284,15 @@ def test_context_manager_closes_cursor(self):
             cursor.close = mock_close
         mock_close.assert_called_once_with()
 
+        cursor = client.Cursor(Mock(), Mock())
+        cursor.close = Mock()
+        try:
+            with self.assertRaises(KeyboardInterrupt):
+                with cursor:
+                    raise KeyboardInterrupt("Simulated interrupt")
+        finally:
+            cursor.close.assert_called()
+
     @patch("%s.client.ThriftBackend" % PACKAGE_NAME)
     def test_context_manager_closes_connection(self, mock_client_class):
         instance = mock_client_class.return_value
@@ -299,6 +308,15 @@ def test_context_manager_closes_connection(self, mock_client_class):
         close_session_id = instance.close_session.call_args[0][0].sessionId
         self.assertEqual(close_session_id, b"\x22")
 
+        connection = databricks.sql.connect(**self.DUMMY_CONNECTION_ARGS)
+        connection.close = Mock()
+        try:
+            with self.assertRaises(KeyboardInterrupt):
+                with connection:
+                    raise KeyboardInterrupt("Simulated interrupt")
+        finally:
+            connection.close.assert_called()
+
     def dict_product(self, dicts):
         """
         Generate cartesion product of values in input dictionary, outputting a dictionary