diff --git a/packages/gooddata-sdk/src/gooddata_sdk/__init__.py b/packages/gooddata-sdk/src/gooddata_sdk/__init__.py index 1be061a11..199f2fe99 100644 --- a/packages/gooddata-sdk/src/gooddata_sdk/__init__.py +++ b/packages/gooddata-sdk/src/gooddata_sdk/__init__.py @@ -259,6 +259,7 @@ from gooddata_sdk.compute.model.attribute import Attribute from gooddata_sdk.compute.model.base import ExecModelEntity, ObjId from gooddata_sdk.compute.model.execution import ( + ArrowFormat, BareExecutionResponse, Execution, ExecutionDefinition, diff --git a/packages/gooddata-sdk/src/gooddata_sdk/compute/model/execution.py b/packages/gooddata-sdk/src/gooddata_sdk/compute/model/execution.py index a81b807ac..0fd9a2da5 100644 --- a/packages/gooddata-sdk/src/gooddata_sdk/compute/model/execution.py +++ b/packages/gooddata-sdk/src/gooddata_sdk/compute/model/execution.py @@ -2,7 +2,7 @@ from __future__ import annotations import logging -from typing import Any, Union +from typing import Any, Literal, Union from attrs import define, field from attrs.setters import frozen as frozen_attr @@ -18,6 +18,8 @@ logger = logging.getLogger(__name__) +ArrowFormat = Literal["application/vnd.apache.arrow.file", "application/vnd.apache.arrow.stream"] + @define class TotalDimension: @@ -372,6 +374,29 @@ def read_result( ) return ExecutionResult(execution_result) + def read_result_binary( + self, + accept: ArrowFormat = "application/vnd.apache.arrow.file", + ) -> bytes: + """ + Reads the execution result in Apache Arrow IPC binary format. + + Args: + accept: Arrow format to request; either 'application/vnd.apache.arrow.file' + (Arrow IPC File format, default) or 'application/vnd.apache.arrow.stream' + (Arrow IPC Stream format). + Returns: + Raw bytes of the Arrow IPC response from the /binary endpoint. + """ + response = self._actions_api.retrieve_result_binary( + workspace_id=self._workspace_id, + result_id=self.result_id, + accept_content_types=[accept], + _check_return_type=False, + _preload_content=False, + ) + return response.data + def cancel(self) -> None: """ Cancels the execution backing this execution result. @@ -464,6 +489,22 @@ def read_result( ) -> ExecutionResult: return self.bare_exec_response.read_result(limit, offset, timeout) + def read_result_binary( + self, + accept: ArrowFormat = "application/vnd.apache.arrow.file", + ) -> bytes: + """ + Reads the execution result in Apache Arrow IPC binary format. + + Args: + accept: Arrow format to request; either 'application/vnd.apache.arrow.file' + (Arrow IPC File format, default) or 'application/vnd.apache.arrow.stream' + (Arrow IPC Stream format). + Returns: + Raw bytes of the Arrow IPC response. + """ + return self.bare_exec_response.read_result_binary(accept) + def cancel(self) -> None: """ Cancels the execution. diff --git a/packages/gooddata-sdk/src/gooddata_sdk/compute/service.py b/packages/gooddata-sdk/src/gooddata_sdk/compute/service.py index 6163798b9..0d5bb5bc3 100644 --- a/packages/gooddata-sdk/src/gooddata_sdk/compute/service.py +++ b/packages/gooddata-sdk/src/gooddata_sdk/compute/service.py @@ -18,6 +18,7 @@ from gooddata_sdk.client import GoodDataApiClient from gooddata_sdk.compute.model.execution import ( + ArrowFormat, Execution, ExecutionDefinition, ResultCacheMetadata, @@ -73,6 +74,33 @@ def for_exec_def( else None, ) + def retrieve_result_binary( + self, + workspace_id: str, + result_id: str, + accept: ArrowFormat = "application/vnd.apache.arrow.file", + ) -> bytes: + """ + Gets execution result in Apache Arrow IPC binary format from GoodData workspace. + + Args: + workspace_id (str): workspace identifier + result_id (str): execution result ID + accept (ArrowFormat): Arrow format to request; either + 'application/vnd.apache.arrow.file' (default) or + 'application/vnd.apache.arrow.stream'. + Returns: + bytes: Raw Arrow IPC bytes. + """ + response = self._actions_api.retrieve_result_binary( + workspace_id=workspace_id, + result_id=result_id, + accept_content_types=[accept], + _check_return_type=False, + _preload_content=False, + ) + return response.data + def retrieve_result_cache_metadata(self, workspace_id: str, result_id: str) -> ResultCacheMetadata: """ Gets execution result's metadata from GoodData.CN workspace for given execution result ID. diff --git a/packages/gooddata-sdk/tests/compute/test_read_result_binary.py b/packages/gooddata-sdk/tests/compute/test_read_result_binary.py new file mode 100644 index 000000000..ef283e82d --- /dev/null +++ b/packages/gooddata-sdk/tests/compute/test_read_result_binary.py @@ -0,0 +1,102 @@ +# (C) 2025 GoodData Corporation +from __future__ import annotations + +from unittest.mock import MagicMock + +import pytest +from gooddata_sdk.compute.model.execution import ArrowFormat, BareExecutionResponse, Execution + + +def _make_bare_exec_response(mock_response_data: bytes = b"arrow-data") -> tuple[BareExecutionResponse, MagicMock]: + """Build a BareExecutionResponse with a mocked actions API.""" + api_client = MagicMock() + exec_response = MagicMock() + exec_response.__getitem__ = MagicMock( + side_effect=lambda k: {"executionResult": "result-id-123"} if k == "links" else MagicMock() + ) + afm_exec_response = MagicMock() + afm_exec_response.__getitem__ = MagicMock( + side_effect=lambda k: exec_response if k == "execution_response" else MagicMock() + ) + + mock_http_response = MagicMock() + mock_http_response.data = mock_response_data + api_client.actions_api.retrieve_result_binary.return_value = mock_http_response + + bare = BareExecutionResponse( + api_client=api_client, + workspace_id="ws-1", + execution_response=afm_exec_response, + ) + return bare, api_client.actions_api + + +@pytest.mark.parametrize( + "accept", + [ + "application/vnd.apache.arrow.file", + "application/vnd.apache.arrow.stream", + ], +) +def test_bare_execution_response_read_result_binary_formats(accept: ArrowFormat) -> None: + """BareExecutionResponse.read_result_binary() passes the correct accept type to the API.""" + expected_bytes = b"\x00\x00\x00arrow" + bare, actions_api = _make_bare_exec_response(mock_response_data=expected_bytes) + + result = bare.read_result_binary(accept=accept) + + assert result == expected_bytes + actions_api.retrieve_result_binary.assert_called_once_with( + workspace_id="ws-1", + result_id="result-id-123", + accept_content_types=[accept], + _check_return_type=False, + _preload_content=False, + ) + + +def test_bare_execution_response_read_result_binary_default_format() -> None: + """BareExecutionResponse.read_result_binary() defaults to arrow.file format.""" + bare, actions_api = _make_bare_exec_response() + + bare.read_result_binary() + + call_kwargs = actions_api.retrieve_result_binary.call_args[1] + assert call_kwargs["accept_content_types"] == ["application/vnd.apache.arrow.file"] + + +def test_execution_read_result_binary_delegates() -> None: + """Execution.read_result_binary() delegates to BareExecutionResponse.read_result_binary().""" + api_client = MagicMock() + exec_response = MagicMock() + exec_response.__getitem__ = MagicMock( + side_effect=lambda k: {"executionResult": "result-id-456"} if k == "links" else [] + ) + afm_exec_response = MagicMock() + afm_exec_response.__getitem__ = MagicMock( + side_effect=lambda k: exec_response if k == "execution_response" else MagicMock() + ) + + expected_bytes = b"stream-data" + mock_http_response = MagicMock() + mock_http_response.data = expected_bytes + api_client.actions_api.retrieve_result_binary.return_value = mock_http_response + + exec_def = MagicMock() + execution = Execution( + api_client=api_client, + workspace_id="ws-2", + exec_def=exec_def, + response=afm_exec_response, + ) + + result = execution.read_result_binary(accept="application/vnd.apache.arrow.stream") + + assert result == expected_bytes + api_client.actions_api.retrieve_result_binary.assert_called_once_with( + workspace_id="ws-2", + result_id="result-id-456", + accept_content_types=["application/vnd.apache.arrow.stream"], + _check_return_type=False, + _preload_content=False, + )