@@ -811,142 +811,6 @@ def test_catalogs_returns_arrow_table(self):
811811 results = cursor .fetchall_arrow ()
812812 assert isinstance (results , pyarrow .Table )
813813
814- def test_close_connection_closes_cursors (self ):
815-
816- from databricks .sql .thrift_api .TCLIService import ttypes
817-
818- with self .connection () as conn :
819- cursor = conn .cursor ()
820- cursor .execute (
821- "SELECT id, id `id2`, id `id3` FROM RANGE(1000000) order by RANDOM()"
822- )
823- ars = cursor .active_result_set
824-
825- # We must manually run this check because thrift_backend always forces `has_been_closed_server_side` to True
826- # Cursor op state should be open before connection is closed
827- status_request = ttypes .TGetOperationStatusReq (
828- operationHandle = ars .command_id .to_thrift_handle (),
829- getProgressUpdate = False ,
830- )
831- op_status_at_server = ars .backend ._client .GetOperationStatus (status_request )
832- assert op_status_at_server .operationState != CommandState .CLOSED
833-
834- conn .close ()
835-
836- # When connection closes, any cursor operations should no longer exist at the server
837- with pytest .raises (SessionAlreadyClosedError ) as cm :
838- op_status_at_server = ars .backend ._client .GetOperationStatus (
839- status_request
840- )
841-
842- def test_closing_a_closed_connection_doesnt_fail (self , caplog ):
843- caplog .set_level (logging .DEBUG )
844- # Second .close() call is when this context manager exits
845- with self .connection () as conn :
846- # First .close() call is explicit here
847- conn .close ()
848- assert "Session appears to have been closed already" in caplog .text
849-
850- conn = None
851- try :
852- with pytest .raises (KeyboardInterrupt ):
853- with self .connection () as c :
854- conn = c
855- raise KeyboardInterrupt ("Simulated interrupt" )
856- finally :
857- if conn is not None :
858- assert (
859- not conn .open
860- ), "Connection should be closed after KeyboardInterrupt"
861-
862- def test_cursor_close_properly_closes_operation (self ):
863- """Test that Cursor.close() properly closes the active operation handle on the server."""
864- with self .connection () as conn :
865- cursor = conn .cursor ()
866- try :
867- cursor .execute ("SELECT 1 AS test" )
868- assert cursor .active_command_id is not None
869- cursor .close ()
870- assert cursor .active_command_id is None
871- assert not cursor .open
872- finally :
873- if cursor .open :
874- cursor .close ()
875-
876- conn = None
877- cursor = None
878- try :
879- with self .connection () as c :
880- conn = c
881- with pytest .raises (KeyboardInterrupt ):
882- with conn .cursor () as cur :
883- cursor = cur
884- raise KeyboardInterrupt ("Simulated interrupt" )
885- finally :
886- if cursor is not None :
887- assert (
888- not cursor .open
889- ), "Cursor should be closed after KeyboardInterrupt"
890-
891- def test_nested_cursor_context_managers (self ):
892- """Test that nested cursor context managers properly close operations on the server."""
893- with self .connection () as conn :
894- with conn .cursor () as cursor1 :
895- cursor1 .execute ("SELECT 1 AS test1" )
896- assert cursor1 .active_command_id is not None
897-
898- with conn .cursor () as cursor2 :
899- cursor2 .execute ("SELECT 2 AS test2" )
900- assert cursor2 .active_command_id is not None
901-
902- # After inner context manager exit, cursor2 should be not open
903- assert not cursor2 .open
904- assert cursor2 .active_command_id is None
905-
906- # After outer context manager exit, cursor1 should be not open
907- assert not cursor1 .open
908- assert cursor1 .active_command_id is None
909-
910- def test_cursor_error_handling (self ):
911- """Test that cursor close handles errors properly to prevent orphaned operations."""
912- with self .connection () as conn :
913- cursor = conn .cursor ()
914-
915- cursor .execute ("SELECT 1 AS test" )
916-
917- op_handle = cursor .active_command_id
918-
919- assert op_handle is not None
920-
921- # Manually close the operation to simulate server-side closure
922- conn .session .backend .close_command (op_handle )
923-
924- cursor .close ()
925-
926- assert not cursor .open
927-
928- def test_result_set_close (self ):
929- """Test that ResultSet.close() properly closes operations on the server and handles state correctly."""
930- with self .connection () as conn :
931- cursor = conn .cursor ()
932- try :
933- cursor .execute ("SELECT * FROM RANGE(10)" )
934-
935- result_set = cursor .active_result_set
936- assert result_set is not None
937-
938- initial_op_state = result_set .status
939-
940- result_set .close ()
941-
942- assert result_set .status == CommandState .CLOSED
943- assert result_set .status != initial_op_state
944-
945- # Closing the result set again should be a no-op and not raise exceptions
946- result_set .close ()
947- finally :
948- cursor .close ()
949-
950814 def test_row_limit_with_larger_result (self ):
951815 """Test that row_limit properly constrains results when query would return more rows"""
952816 row_limit = 1000
0 commit comments