Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -504,6 +504,15 @@ def read_result(
) -> ExecutionResult:
return self.bare_exec_response.read_result(limit, offset, timeout)

def read_result_arrow(self) -> pyarrow.Table:
"""
Reads the full execution result as a pyarrow Table via the binary endpoint.

The binary endpoint returns the complete result in one shot (no paging).
Requires pyarrow to be installed (pip install gooddata-sdk[arrow]).
"""
return self.bare_exec_response.read_result_arrow()

def cancel(self) -> None:
"""
Cancels the execution.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,31 @@ def readinto(self, b: bytearray) -> int:
return n


def _make_execution(ipc_bytes: bytes):
"""Return an Execution backed by a mock API client."""
from gooddata_sdk.compute.model.execution import Execution

mock_api_client = MagicMock()
mock_response = _FakeResponse(ipc_bytes)
mock_api_client.actions_api.api_client.call_api.return_value = mock_response

afm_exec_response = {
"execution_response": {
"links": {"executionResult": "result-id-123"},
"dimensions": [],
}
}
mock_exec_def = MagicMock()

execution = Execution(
api_client=mock_api_client,
workspace_id="ws-id",
exec_def=mock_exec_def,
response=afm_exec_response,
)
return execution, mock_response


def _make_bare(ipc_bytes: bytes):
"""Return a BareExecutionResponse backed by a mock API client."""
from gooddata_sdk.compute.model.execution import BareExecutionResponse
Expand Down Expand Up @@ -108,3 +133,33 @@ def test_read_result_arrow_no_pyarrow_raises() -> None:

with patch.object(_exec_mod, "_ipc", None), pytest.raises(ImportError, match="pyarrow is required"):
bare.read_result_arrow()


# ---------------------------------------------------------------------------
# Execution.read_result_arrow – delegates to BareExecutionResponse
# ---------------------------------------------------------------------------


def test_execution_read_result_arrow_returns_table() -> None:
"""Execution.read_result_arrow() delegates to BareExecutionResponse and returns a pa.Table."""
import pyarrow as pa

ipc_bytes = _make_ipc_stream_bytes()
execution, mock_response = _make_execution(ipc_bytes)

result = execution.read_result_arrow()

assert isinstance(result, pa.Table)
mock_response.release_conn.assert_called_once()


def test_execution_read_result_arrow_calls_binary_endpoint() -> None:
"""Execution.read_result_arrow() hits the /binary endpoint with the correct Accept header."""
ipc_bytes = _make_ipc_stream_bytes()
execution, _ = _make_execution(ipc_bytes)

execution.read_result_arrow()

call_kwargs = execution.bare_exec_response._actions_api.api_client.call_api.call_args.kwargs
assert call_kwargs["header_params"]["Accept"] == "application/vnd.apache.arrow.stream"
assert "/binary" in call_kwargs["resource_path"]
Loading