5050
5151from tests .e2e .common .uc_volume_tests import PySQLUCVolumeTestSuiteMixin
5252
53- from databricks .sql .exc import SessionAlreadyClosedError
53+ from databricks .sql .exc import SessionAlreadyClosedError , CursorAlreadyClosedError
5454
5555log = logging .getLogger (__name__ )
5656
@@ -820,7 +820,6 @@ def test_close_connection_closes_cursors(self):
820820 ars = cursor .active_result_set
821821
822822 # We must manually run this check because thrift_backend always forces `has_been_closed_server_side` to True
823-
824823 # Cursor op state should be open before connection is closed
825824 status_request = ttypes .TGetOperationStatusReq (
826825 operationHandle = ars .command_id , getProgressUpdate = False
@@ -847,9 +846,104 @@ def test_closing_a_closed_connection_doesnt_fail(self, caplog):
847846 with self .connection () as conn :
848847 # First .close() call is explicit here
849848 conn .close ()
850-
851849 assert "Session appears to have been closed already" in caplog .text
852850
851+ conn = None
852+ try :
853+ with pytest .raises (KeyboardInterrupt ):
854+ with self .connection () as c :
855+ conn = c
856+ raise KeyboardInterrupt ("Simulated interrupt" )
857+ finally :
858+ if conn is not None :
859+ assert not conn .open , "Connection should be closed after KeyboardInterrupt"
860+
861+ def test_cursor_close_properly_closes_operation (self ):
862+ """Test that Cursor.close() properly closes the active operation handle on the server."""
863+ with self .connection () as conn :
864+ cursor = conn .cursor ()
865+ try :
866+ cursor .execute ("SELECT 1 AS test" )
867+ assert cursor .active_op_handle is not None
868+ cursor .close ()
869+ assert cursor .active_op_handle is None
870+ assert not cursor .open
871+ finally :
872+ if cursor .open :
873+ cursor .close ()
874+
875+ conn = None
876+ cursor = None
877+ try :
878+ with self .connection () as c :
879+ conn = c
880+ with pytest .raises (KeyboardInterrupt ):
881+ with conn .cursor () as cur :
882+ cursor = cur
883+ raise KeyboardInterrupt ("Simulated interrupt" )
884+ finally :
885+ if cursor is not None :
886+ assert not cursor .open , "Cursor should be closed after KeyboardInterrupt"
887+
888+ def test_nested_cursor_context_managers (self ):
889+ """Test that nested cursor context managers properly close operations on the server."""
890+ with self .connection () as conn :
891+ with conn .cursor () as cursor1 :
892+ cursor1 .execute ("SELECT 1 AS test1" )
893+ assert cursor1 .active_op_handle is not None
894+
895+ with conn .cursor () as cursor2 :
896+ cursor2 .execute ("SELECT 2 AS test2" )
897+ assert cursor2 .active_op_handle is not None
898+
899+ # After inner context manager exit, cursor2 should be not open
900+ assert not cursor2 .open
901+ assert cursor2 .active_op_handle is None
902+
903+ # After outer context manager exit, cursor1 should be not open
904+ assert not cursor1 .open
905+ assert cursor1 .active_op_handle is None
906+
907+ def test_cursor_error_handling (self ):
908+ """Test that cursor close handles errors properly to prevent orphaned operations."""
909+ with self .connection () as conn :
910+ cursor = conn .cursor ()
911+
912+ cursor .execute ("SELECT 1 AS test" )
913+
914+ op_handle = cursor .active_op_handle
915+
916+ assert op_handle is not None
917+
918+ # Manually close the operation to simulate server-side closure
919+ conn .thrift_backend .close_command (op_handle )
920+
921+ cursor .close ()
922+
923+ assert not cursor .open
924+
925+ def test_result_set_close (self ):
926+ """Test that ResultSet.close() properly closes operations on the server and handles state correctly."""
927+ with self .connection () as conn :
928+ cursor = conn .cursor ()
929+ try :
930+ cursor .execute ("SELECT * FROM RANGE(10)" )
931+
932+ result_set = cursor .active_result_set
933+ assert result_set is not None
934+
935+ initial_op_state = result_set .op_state
936+
937+ result_set .close ()
938+
939+ assert result_set .op_state == result_set .thrift_backend .CLOSED_OP_STATE
940+ assert result_set .op_state != initial_op_state
941+
942+ # Closing the result set again should be a no-op and not raise exceptions
943+ result_set .close ()
944+ finally :
945+ cursor .close ()
946+
853947
854948# use a RetrySuite to encapsulate these tests which we'll typically want to run together; however keep
855949# the 429/503 subsuites separate since they execute under different circumstances.
0 commit comments