Merge branch 'main' into backend-refactors

Author: varun-edachali-dbx

src/databricks/sql/client.py

@@ -282,13 +282,7 @@ def __enter__(self) -> "Connection":
         return self
 
     def __exit__(self, exc_type, exc_value, traceback):
-        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
+        self.close()
 
     def __del__(self):
         if self.open:
@@ -418,14 +412,7 @@ def __enter__(self) -> "Cursor":
         return self
 
     def __exit__(self, exc_type, exc_value, traceback):
-        try:
-            logger.debug("Cursor context manager exiting, calling close()")
-            self.close()
-        except BaseException as e:
-            logger.warning(f"Exception during cursor close in __exit__: {e}")
-            if exc_type is None:
-                raise
-        return False
+        self.close()
 
     def __iter__(self):
         if self.active_result_set:
@@ -1092,21 +1079,7 @@ def cancel(self) -> None:
     def close(self) -> None:
         """Close cursor"""
         self.open = False
-
-        # Close active operation handle if it exists
-        if self.active_command_id:
-            try:
-                self.backend.close_command(self.active_command_id)
-            except RequestError as e:
-                if isinstance(e.args[1], CursorAlreadyClosedError):
-                    logger.info("Operation was canceled by a prior request")
-                else:
-                    logging.warning(f"Error closing operation handle: {e}")
-            except Exception as e:
-                logging.warning(f"Error closing operation handle: {e}")
-            finally:
-                self.active_command_id = None
-
+        self.active_op_handle = None
         if self.active_result_set:
             self._close_and_clear_active_result_set()
 

tests/e2e/common/large_queries_mixin.py

@@ -83,11 +83,11 @@ def test_query_with_large_narrow_result_set(self):
                 assert row[0] == row_id
 
     def test_long_running_query(self):
-        """Incrementally increase query size until it takes at least 5 minutes,
+        """Incrementally increase query size until it takes at least 3 minutes,
         and asserts that the query completes successfully.
         """
         minutes = 60
-        min_duration = 5 * minutes
+        min_duration = 3 * minutes
 
         duration = -1
         scale0 = 10000
@@ -113,5 +113,5 @@ def test_long_running_query(self):
                 duration = time.time() - start
                 current_fraction = duration / min_duration
                 print("Took {} s with scale factor={}".format(duration, scale_factor))
-                # Extrapolate linearly to reach 5 min and add 50% padding to push over the limit
+                # Extrapolate linearly to reach 3 min and add 50% padding to push over the limit
                 scale_factor = math.ceil(1.5 * scale_factor / current_fraction)

tests/e2e/common/staging_ingestion_tests.py

@@ -46,7 +46,7 @@ def test_staging_ingestion_life_cycle(self, ingestion_user):
         ) as conn:
 
             cursor = conn.cursor()
-            query = f"PUT '{temp_path}' INTO 'stage://tmp/{ingestion_user}/tmp/11/15/file1.csv' OVERWRITE"
+            query = f"PUT '{temp_path}' INTO 'stage://tmp/{ingestion_user}/tmp/11/16/file1.csv' OVERWRITE"
             cursor.execute(query)
 
         # GET should succeed
@@ -57,7 +57,7 @@ def test_staging_ingestion_life_cycle(self, ingestion_user):
             extra_params={"staging_allowed_local_path": new_temp_path}
         ) as conn:
             cursor = conn.cursor()
-            query = f"GET 'stage://tmp/{ingestion_user}/tmp/11/15/file1.csv' TO '{new_temp_path}'"
+            query = f"GET 'stage://tmp/{ingestion_user}/tmp/11/16/file1.csv' TO '{new_temp_path}'"
             cursor.execute(query)
 
         with open(new_fh, "rb") as fp:
@@ -67,7 +67,7 @@ def test_staging_ingestion_life_cycle(self, ingestion_user):
 
         # REMOVE should succeed
 
-        remove_query = f"REMOVE 'stage://tmp/{ingestion_user}/tmp/11/15/file1.csv'"
+        remove_query = f"REMOVE 'stage://tmp/{ingestion_user}/tmp/11/16/file1.csv'"
 
         with self.connection(extra_params={"staging_allowed_local_path": "/"}) as conn:
             cursor = conn.cursor()
@@ -79,7 +79,7 @@ def test_staging_ingestion_life_cycle(self, ingestion_user):
                 Error, match="Staging operation over HTTP was unsuccessful: 404"
             ):
                 cursor = conn.cursor()
-                query = f"GET 'stage://tmp/{ingestion_user}/tmp/11/15/file1.csv' TO '{new_temp_path}'"
+                query = f"GET 'stage://tmp/{ingestion_user}/tmp/11/16/file1.csv' TO '{new_temp_path}'"
                 cursor.execute(query)
 
         os.remove(temp_path)

tests/e2e/test_driver.py

@@ -51,7 +51,7 @@
 
 from tests.e2e.common.uc_volume_tests import PySQLUCVolumeTestSuiteMixin
 
-from databricks.sql.exc import SessionAlreadyClosedError, CursorAlreadyClosedError
+from databricks.sql.exc import SessionAlreadyClosedError
 
 log = logging.getLogger(__name__)
 
@@ -809,142 +809,6 @@ def test_catalogs_returns_arrow_table(self):
             results = cursor.fetchall_arrow()
             assert isinstance(results, pyarrow.Table)
 
-    def test_close_connection_closes_cursors(self):
-
-        from databricks.sql.thrift_api.TCLIService import ttypes
-
-        with self.connection() as conn:
-            cursor = conn.cursor()
-            cursor.execute(
-                "SELECT id, id `id2`, id `id3` FROM RANGE(1000000) order by RANDOM()"
-            )
-            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.to_thrift_handle(),
-                getProgressUpdate=False,
-            )
-            op_status_at_server = ars.backend._client.GetOperationStatus(status_request)
-            assert op_status_at_server.operationState != CommandState.CLOSED
-
-            conn.close()
-
-            # When connection closes, any cursor operations should no longer exist at the server
-            with pytest.raises(SessionAlreadyClosedError) as cm:
-                op_status_at_server = ars.backend._client.GetOperationStatus(
-                    status_request
-                )
-
-    def test_closing_a_closed_connection_doesnt_fail(self, caplog):
-        caplog.set_level(logging.DEBUG)
-        # Second .close() call is when this context manager exits
-        with self.connection() as conn:
-            # First .close() call is explicit here
-            conn.close()
-        assert "Session appears to have been closed already" in caplog.text
-
-        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_command_id is not None
-                cursor.close()
-                assert cursor.active_command_id is None
-                assert not cursor.open
-            finally:
-                if cursor.open:
-                    cursor.close()
-
-        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_command_id is not None
-
-                with conn.cursor() as cursor2:
-                    cursor2.execute("SELECT 2 AS test2")
-                    assert cursor2.active_command_id is not None
-
-                # After inner context manager exit, cursor2 should be not open
-                assert not cursor2.open
-                assert cursor2.active_command_id is None
-
-            # After outer context manager exit, cursor1 should be not open
-            assert not cursor1.open
-            assert cursor1.active_command_id 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_command_id
-
-            assert op_handle is not None
-
-            # Manually close the operation to simulate server-side closure
-            conn.session.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 == CommandState.CLOSED
-                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.

tests/unit/test_client.py

@@ -23,7 +23,6 @@
 import databricks.sql
 import databricks.sql.client as client
 from databricks.sql import InterfaceError, DatabaseError, Error, NotSupportedError
-from databricks.sql.exc import RequestError, CursorAlreadyClosedError
 from databricks.sql.types import Row
 from databricks.sql.result_set import ResultSet, ThriftResultSet
 from databricks.sql.backend.types import CommandId, CommandState
@@ -261,15 +260,20 @@ 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()
+    @patch("%s.client.ThriftBackend" % PACKAGE_NAME)
+    def test_context_manager_closes_connection(self, mock_client_class):
+        instance = mock_client_class.return_value
 
-        try:
-            with self.assertRaises(KeyboardInterrupt):
-                with cursor:
-                    raise KeyboardInterrupt("Simulated interrupt")
-        finally:
-            cursor.close.assert_called()
+        mock_open_session_resp = MagicMock(spec=TOpenSessionResp)()
+        mock_open_session_resp.sessionHandle.sessionId = b"\x22"
+        instance.open_session.return_value = mock_open_session_resp
+
+        with databricks.sql.connect(**self.DUMMY_CONNECTION_ARGS) as connection:
+            pass
+
+        # Check the close session request has an id of x22
+        close_session_id = instance.close_session.call_args[0][0].sessionId
+        self.assertEqual(close_session_id, b"\x22")
 
     def dict_product(self, dicts):
         """
@@ -599,42 +603,6 @@ def test_access_current_query_id(self):
         cursor.close()
         self.assertIsNone(cursor.query_id)
 
-    def test_cursor_close_handles_exception(self):
-        """Test that Cursor.close() handles exceptions from close_command properly."""
-        mock_backend = Mock()
-        mock_connection = Mock()
-        mock_command_id = Mock()
-
-        mock_backend.close_command.side_effect = Exception("Test error")
-
-        cursor = client.Cursor(mock_connection, mock_backend)
-        cursor.active_command_id = mock_command_id
-
-        cursor.close()
-
-        mock_backend.close_command.assert_called_once_with(mock_command_id)
-
-        self.assertIsNone(cursor.active_command_id)
-
-        self.assertFalse(cursor.open)
-
-    def test_cursor_context_manager_handles_exit_exception(self):
-        """Test that cursor's context manager handles exceptions during __exit__."""
-        mock_backend = Mock()
-        mock_connection = Mock()
-
-        cursor = client.Cursor(mock_connection, mock_backend)
-        original_close = cursor.close
-        cursor.close = Mock(side_effect=Exception("Test error during close"))
-
-        try:
-            with cursor:
-                raise ValueError("Test error inside context")
-        except ValueError:
-            pass
-
-        cursor.close.assert_called_once()
-
     def test_connection_close_handles_cursor_close_exception(self):
         """Test that _close handles exceptions from cursor.close() properly."""
         cursors_closed = []
@@ -670,49 +638,6 @@ def mock_close_normal():
             cursors_closed, [1, 2], "Both cursors should have close called"
         )
 
-    def test_resultset_close_handles_cursor_already_closed_error(self):
-        """Test that ResultSet.close() handles CursorAlreadyClosedError properly."""
-        result_set = client.ThriftResultSet.__new__(client.ThriftResultSet)
-        result_set.backend = Mock()
-        result_set.backend.CLOSED_OP_STATE = "CLOSED"
-        result_set.connection = Mock()
-        result_set.connection.open = True
-        result_set.op_state = "RUNNING"
-        result_set.has_been_closed_server_side = False
-        result_set.command_id = Mock()
-
-        class MockRequestError(Exception):
-            def __init__(self):
-                self.args = ["Error message", CursorAlreadyClosedError()]
-
-        result_set.backend.close_command.side_effect = MockRequestError()
-
-        original_close = client.ResultSet.close
-        try:
-            try:
-                if (
-                    result_set.op_state != result_set.backend.CLOSED_OP_STATE
-                    and not result_set.has_been_closed_server_side
-                    and result_set.connection.open
-                ):
-                    result_set.backend.close_command(result_set.command_id)
-            except MockRequestError as e:
-                if isinstance(e.args[1], CursorAlreadyClosedError):
-                    pass
-            finally:
-                result_set.has_been_closed_server_side = True
-                result_set.op_state = result_set.backend.CLOSED_OP_STATE
-
-            result_set.backend.close_command.assert_called_once_with(
-                result_set.command_id
-            )
-
-            assert result_set.has_been_closed_server_side is True
-
-            assert result_set.op_state == result_set.backend.CLOSED_OP_STATE
-        finally:
-            pass
-
 
 if __name__ == "__main__":
     suite = unittest.TestLoader().loadTestsFromModule(sys.modules[__name__])