diff --git a/src/databricks/sql/backend/kernel/client.py b/src/databricks/sql/backend/kernel/client.py index 0ca7ba6a8..012710540 100644 --- a/src/databricks/sql/backend/kernel/client.py +++ b/src/databricks/sql/backend/kernel/client.py @@ -966,7 +966,9 @@ def get_execution_result( # telemetry — an accepted, narrow gap limited to the case where # result-set construction fails after a successful await. try: - return self._make_result_set(stream, cursor, command_id) + result_set = self._make_result_set(stream, cursor, command_id) + result_set.num_modified_rows = getattr(handle, "num_modified_rows", None) + return result_set except Exception as exc: raise _wrap_kernel_exception("get_execution_result", exc) from exc diff --git a/src/databricks/sql/backend/thrift_backend.py b/src/databricks/sql/backend/thrift_backend.py index 19e378292..9439861e1 100644 --- a/src/databricks/sql/backend/thrift_backend.py +++ b/src/databricks/sql/backend/thrift_backend.py @@ -931,6 +931,11 @@ def get_execution_result( if not thrift_handle: raise ValueError("Not a valid Thrift command ID") + operation_status = self._poll_for_status(thrift_handle) + self._check_command_not_in_error_or_closed_state( + thrift_handle, operation_status + ) + req = ttypes.TFetchResultsReq( operationHandle=ttypes.TOperationHandle( thrift_handle.operationId, @@ -981,6 +986,7 @@ def get_execution_result( is_staging_operation=is_staging_operation, arrow_schema_bytes=schema_bytes, result_format=t_result_set_metadata_resp.resultFormat, + num_modified_rows=operation_status.numModifiedRows, ) return ThriftResultSet( diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py index 6e03ed227..0ae0afea1 100755 --- a/src/databricks/sql/client.py +++ b/src/databricks/sql/client.py @@ -1701,6 +1701,9 @@ def get_async_execution_result(self): self.active_command_id, self ) + num_modified_rows = self.active_result_set.num_modified_rows + self.rowcount = num_modified_rows if num_modified_rows is not None else -1 + if self.active_result_set and self.active_result_set.is_staging_operation: self._handle_staging_operation( staging_allowed_local_path=self.connection.staging_allowed_local_path diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index 047252488..b9c42dd7c 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -303,6 +303,29 @@ def test_rowcount_resets_between_statements(self): cursor.execute("SELECT 1") self.assertEqual(cursor.rowcount, -1) + def test_async_rowcount_on_completion_and_reset(self): + for count in (5, 0, None): + with self.subTest(count=count): + result_set = Mock(num_modified_rows=count, is_staging_operation=False) + backend = ThriftDatabricksClientMockFactory.new() + backend.get_execution_result.return_value = result_set + backend.get_query_state.return_value = CommandState.SUCCEEDED + cursor = client.Cursor(connection=Mock(), backend=backend) + cursor.rowcount = 9 + + cursor.execute_async("UPDATE t SET x = 1") + self.assertEqual(cursor.rowcount, -1) + cursor.active_command_id = Mock(spec=CommandId) + self.assertIs(cursor.get_async_execution_result(), cursor) + self.assertEqual(cursor.rowcount, count if count is not None else -1) + result_set.fetchall.assert_not_called() + + result_set.num_modified_rows = None + cursor.execute_async("SELECT 1") + self.assertEqual(cursor.rowcount, -1) + cursor.get_async_execution_result() + self.assertEqual(cursor.rowcount, -1) + def test_closed_cursor_doesnt_allow_operations(self): cursor = client.Cursor(Mock(), Mock()) cursor.close() diff --git a/tests/unit/test_kernel_client.py b/tests/unit/test_kernel_client.py index b285ce6e1..d8dacdbc9 100644 --- a/tests/unit/test_kernel_client.py +++ b/tests/unit/test_kernel_client.py @@ -1233,6 +1233,48 @@ def test_get_query_state_uses_reader_after_result_stream_started(): reader.status.assert_called_once_with() +@pytest.mark.parametrize("count", [7, 0, None, "unavailable"]) +@pytest.mark.parametrize("source", ["owner", "reader", "attached"]) +def test_async_result_count_after_await(source, count): + c = _make_client() + c._kernel_session = MagicMock() + cid = CommandId.from_sea_statement_id("async-count") + owner = MagicMock() + handle = owner + if source == "attached": + c._kernel_session.attach_async_statement.return_value = handle + else: + c._async_handles[cid.guid] = owner + if source == "reader": + c._async_result_stream_started.add(cid.guid) + handle = owner.reader.return_value + + del handle.num_modified_rows + stream = MagicMock() + stream.arrow_schema.return_value = pa.schema([("num_affected_rows", pa.int64())]) + stream.fetch_next_batch.side_effect = [ + pa.record_batch( + [[count if isinstance(count, int) else 1]], names=["num_affected_rows"] + ), + None, + ] + + def await_result(): + if count != "unavailable": + handle.num_modified_rows = count + return stream + + handle.await_result.side_effect = await_result + cursor = MagicMock(arraysize=100, buffer_size_bytes=1024) + rs = c.get_execution_result(cid, cursor) + + assert rs.num_modified_rows == (None if count == "unavailable" else count) + stream.fetch_next_batch.assert_not_called() + assert [list(row) for row in rs.fetchall()] == [ + [count if isinstance(count, int) else 1] + ] + + def test_get_execution_result_uses_retained_owning_handle_first(): """The first in-process result fetch uses the retained submitting handle so the kernel finalizes the original async statement telemetry.""" diff --git a/tests/unit/test_thrift_backend.py b/tests/unit/test_thrift_backend.py index dedadaeb2..a51c0a9d5 100644 --- a/tests/unit/test_thrift_backend.py +++ b/tests/unit/test_thrift_backend.py @@ -3,7 +3,7 @@ import itertools import unittest import pytest -from unittest.mock import patch, MagicMock, Mock +from unittest.mock import patch, MagicMock, Mock, sentinel from ssl import CERT_NONE, CERT_REQUIRED from urllib3 import HTTPSConnectionPool @@ -134,6 +134,51 @@ def test_metadata_catalog_escaping( assert getattr(request, thrift_field) == value +@pytest.mark.skipif(pyarrow is None, reason="PyArrow is not installed") +@pytest.mark.parametrize("count", [7, 0, None]) +def test_async_result_reads_status_count_and_preserves_rows(count): + with ( + patch( + "databricks.sql.backend.thrift_backend.TCLIService.Client", autospec=True + ), + patch("databricks.sql.backend.thrift_backend.ThriftResultSet") as result_set, + ): + backend = ThriftDatabricksClient( + "foobar", + 443, + "path", + [], + auth_provider=AuthProvider(), + ssl_options=SSLOptions(), + http_client=MagicMock(), + ) + operation_status = ttypes.TGetOperationStatusResp( + status=ThriftBackendTestSuite.okay_status, + operationState=ttypes.TOperationState.FINISHED_STATE, + numModifiedRows=count, + ) + backend._poll_for_status = Mock(return_value=operation_status) + metadata = ttypes.TGetResultSetMetadataResp( + schema=ttypes.TTableSchema(columns=[]), + arrowSchema=pyarrow.schema([]).serialize().to_pybytes(), + resultFormat=ttypes.TSparkRowSetType.ARROW_BASED_SET, + ) + backend.make_request = Mock( + return_value=Mock( + resultSetMetadata=metadata, + results=sentinel.counter_rows, + hasMoreRows=False, + ) + ) + cid = CommandId.from_thrift_handle(ThriftBackendTestSuite.operation_handle) + backend.get_execution_result(cid, Mock()) + + backend._poll_for_status.assert_called_once_with(cid.to_thrift_handle()) + response = result_set.call_args.kwargs["execute_response"] + assert response.num_modified_rows == count + assert result_set.call_args.kwargs["t_row_set"] is sentinel.counter_rows + + @pytest.mark.skipif(pyarrow is None, reason="PyArrow is not installed") class ThriftBackendTestSuite(unittest.TestCase): okay_status = ttypes.TStatus(statusCode=ttypes.TStatusCode.SUCCESS_STATUS)