diff --git a/dataconnect/client.py b/dataconnect/client.py index 83e7ef1..6702a97 100644 --- a/dataconnect/client.py +++ b/dataconnect/client.py @@ -10,6 +10,8 @@ from types import TracebackType from uuid import UUID +import pandas as pd + from dataconnect.models import Dataset, DatasetVersion, Study from dataconnect.service import DataConnectService, DefaultDataConnectService @@ -51,6 +53,14 @@ def get_dataset_versions(self, dataset_uuid: UUID) -> list[DatasetVersion]: """List the dataset versions the client is authorized to access.""" return self._service.get_dataset_versions(dataset_uuid) + def fetch_data( + self, + dataset_uuid: UUID, + first_n_rows: int | None = None, + ) -> pd.DataFrame: + """Fetch data frames for a given dataset UUID.""" + return self._service.fetch_data(dataset_uuid, first_n_rows) + def get_datasets( self, study_environment_uuid: UUID, diff --git a/dataconnect/service/base.py b/dataconnect/service/base.py index ab0a224..6709ed9 100644 --- a/dataconnect/service/base.py +++ b/dataconnect/service/base.py @@ -5,6 +5,8 @@ from abc import ABC, abstractmethod from uuid import UUID +import pandas as pd + from dataconnect.models import Dataset, DatasetVersion, Study @@ -26,5 +28,12 @@ def get_datasets( @abstractmethod def get_dataset_versions(self, dataset_uuid: UUID) -> list[DatasetVersion]: ... + @abstractmethod + def fetch_data( + self, + dataset_uuid: UUID, + first_n_rows: int | None = None, + ) -> pd.DataFrame: ... + @abstractmethod def close(self) -> None: ... diff --git a/dataconnect/service/default.py b/dataconnect/service/default.py index c67fac3..d009bb1 100644 --- a/dataconnect/service/default.py +++ b/dataconnect/service/default.py @@ -4,6 +4,8 @@ from uuid import UUID +import pandas as pd + from dataconnect.exceptions import ( AuthenticationError, AuthorizationError, @@ -16,7 +18,12 @@ ) from dataconnect.models import Dataset, DatasetVersion, Study from dataconnect.service.base import DataConnectService -from dataconnect.service.mappers import resource_to_dataset, resource_to_dataset_version, resource_to_study +from dataconnect.service.mappers import ( + resource_to_dataset, + resource_to_dataset_version, + resource_to_fetched_data, + resource_to_study, +) from dataconnect.service.validators import validate_search_study_name from dataconnect.transport.base import Transport from dataconnect.transport.errors import ( @@ -34,6 +41,7 @@ _ACTION_LIST_STUDIES = "studies.list" _ACTION_LIST_DATASETS = "datasets.list" _ACTION_LIST_DATASET_VERSIONS = "dataset_versions.list" +_ACTION_FETCH_TICKET = "data.fetch_ticket" def _translate_error(ex: TransportError) -> DataConnectError: @@ -101,6 +109,32 @@ def get_dataset_versions(self, dataset_uuid: UUID) -> list[DatasetVersion]: except (IndexError, KeyError, TypeError, ValueError) as ex: raise ValidationError(f"Unexpected dataset versions response format: {ex}") from ex + def fetch_data(self, dataset_uuid: UUID, first_n_rows: int | None = None) -> pd.DataFrame: + + if not dataset_uuid or not str(dataset_uuid).strip(): + raise ValueError("dataset_uuid must be provided.") + + if dataset_uuid.int == 0: + raise ValueError("dataset_uuid must not be an empty UUID.") + + if first_n_rows is not None and (not isinstance(first_n_rows, int) or first_n_rows <= 0): + raise ValueError("first_n_rows must be a positive integer when provided.") + + request = ResourceQuery(action=_ACTION_FETCH_TICKET).append_body( + { + "study_env_uuid": None, + "dataset_name": None, + "dataset_uuid": str(dataset_uuid), + "limit": first_n_rows, + } + ) + + try: + table = self._transport.do_get(request) + return resource_to_fetched_data(table) + except TransportError as ex: + raise _translate_error(ex) from ex + def get_datasets( self, study_environment_uuid: UUID, diff --git a/dataconnect/service/mappers.py b/dataconnect/service/mappers.py index 00d93d2..97dfcf2 100644 --- a/dataconnect/service/mappers.py +++ b/dataconnect/service/mappers.py @@ -10,9 +10,12 @@ import json from uuid import UUID +import pandas as pd +import pyarrow as pa + from dataconnect.exceptions import NotFoundError from dataconnect.models import Dataset, DatasetVersion, Study, StudyEnvironment -from dataconnect.transport.models import ResourceInfo +from dataconnect.transport.models import DataTable, ResourceInfo def resource_to_study(resource: ResourceInfo) -> Study: @@ -47,6 +50,14 @@ def resource_to_dataset_version(resource: ResourceInfo) -> DatasetVersion: ) +def resource_to_fetched_data(table: DataTable) -> pd.DataFrame: + """Convert a transport-layer ``DataTable`` into a ``pandas.DataFrame``.""" + + ipc_buffer = pa.BufferReader(table.ipc_bytes) + with pa.ipc.open_stream(ipc_buffer) as reader: + return reader.read_all().to_pandas() + + def resource_to_dataset(resource: ResourceInfo) -> Dataset: """Parse a transport-layer ``ResourceInfo`` into a ``Dataset`` domain object.""" diff --git a/dataconnect/transport/arrow_flight/transport.py b/dataconnect/transport/arrow_flight/transport.py index 824163b..57cc4e7 100644 --- a/dataconnect/transport/arrow_flight/transport.py +++ b/dataconnect/transport/arrow_flight/transport.py @@ -12,6 +12,7 @@ import platform import subprocess +import pyarrow as pa import pyarrow.flight as flight from dataconnect.transport.base import Transport @@ -21,7 +22,7 @@ TransportConnectionError, TransportStatusError, ) -from dataconnect.transport.models import DataRef, ResourceInfo, ResourceQuery +from dataconnect.transport.models import DataRef, DataTable, ResourceInfo, ResourceQuery def _to_resource_info(info: flight.FlightInfo) -> ResourceInfo: @@ -39,11 +40,31 @@ def _to_resource_info(info: flight.FlightInfo) -> ResourceInfo: ) +def _to_bytes(table: pa.Table) -> DataTable: + """Serialize a ``pa.Table`` to a technology-agnostic ``DataTable``. + + Each record batch is serialized individually as Arrow IPC bytes. + The schema is serialized separately so it can be recovered without + the data batches. + """ + schema_bytes = table.schema.serialize().to_pybytes() + + sink = pa.BufferOutputStream() + writer = pa.ipc.new_stream(sink, table.schema) + for batch in table.to_batches(): + writer.write_batch(batch) + writer.close() + ipc_bytes = sink.getvalue().to_pybytes() + + return DataTable(schema_bytes=schema_bytes, ipc_bytes=ipc_bytes) + + # Maps service-layer action names to the flight_type value the Arrow Flight server expects. _ACTION_FLIGHT_TYPE: dict[str, str] = { "studies.list": "STUDIES", "datasets.list": "DATASETS", "dataset_versions.list": "VERSIONS", + "data.fetch_ticket": "DATA_FETCH_TICKET", } @@ -141,5 +162,46 @@ def list_resources(self, request: ResourceQuery) -> list[ResourceInfo]: except Exception as ex: raise TransportConnectionError(f"Unexpected error during list_resources: {ex}") from ex + def do_get(self, request: ResourceQuery) -> DataTable: + """Call FlightClient.do_get and read all chunks into a single pa.Table.""" + + flight_type = _ACTION_FLIGHT_TYPE.get(request.action) + + if flight_type is None: + raise TransportStatusError( + f"Unknown action: {request.action!r}", status_code=3, grpc_status="INVALID_ARGUMENT" + ) + + body = json.loads(request.body) if request.body else {} + ticket_bytes = json.dumps({**body, "flight_type": flight_type}, separators=(",", ":")).encode("utf-8") + ticket = flight.Ticket(ticket_bytes) + + try: + table = self._client.do_get(ticket, self._options()) + batches: list[pa.RecordBatch] = [] + while True: + try: + chunk, _metadata = table.read_chunk() + batches.append(chunk) + except StopIteration: + break + except flight.FlightError as ex: + raise TransportConnectionError(f"Error reading stream: {ex}") from ex + + return _to_bytes(pa.Table.from_batches(batches, schema=table.schema)) + + except flight.FlightUnauthenticatedError as ex: + raise TransportAuthenticationError(str(ex)) from ex + except flight.FlightUnauthorizedError as ex: + raise TransportAuthorizationError(str(ex)) from ex + except flight.FlightUnavailableError as ex: + raise TransportConnectionError(str(ex)) from ex + except flight.FlightInternalError as ex: + raise TransportStatusError(str(ex), status_code=13, grpc_status="INTERNAL") from ex + except flight.FlightError as ex: + raise TransportConnectionError(str(ex)) from ex + except Exception as ex: + raise TransportConnectionError(f"Unexpected error during do_get: {ex}") from ex + def close(self) -> None: self._client.close() diff --git a/dataconnect/transport/base.py b/dataconnect/transport/base.py index df12c50..3c06161 100644 --- a/dataconnect/transport/base.py +++ b/dataconnect/transport/base.py @@ -9,7 +9,7 @@ from abc import ABC, abstractmethod -from dataconnect.transport.models import ResourceInfo, ResourceQuery +from dataconnect.transport.models import DataTable, ResourceInfo, ResourceQuery class Transport(ABC): @@ -23,6 +23,14 @@ def list_resources(self, request: ResourceQuery) -> list[ResourceInfo]: service layer's responsibility. """ + @abstractmethod + def do_get(self, request: ResourceQuery) -> DataTable: + """Fetch the full dataset described by ``request``. + + All record batches from the stream are read and returned as a single + ``DataTable`` containing the complete result set. + """ + @abstractmethod def close(self) -> None: """Close the transport connection.""" diff --git a/dataconnect/transport/models.py b/dataconnect/transport/models.py index 65ea576..6f3027b 100644 --- a/dataconnect/transport/models.py +++ b/dataconnect/transport/models.py @@ -38,3 +38,15 @@ class ResourceInfo: endpoints: list[DataRef] total_records: int schema_bytes: bytes + + +@dataclass(frozen=True) +class DataTable: + """Technology-agnostic representation of a fetched data result. + + ``schema_bytes`` holds the Arrow IPC-serialized schema. + ``ipc_bytes`` holds the full Arrow IPC stream (schema + all batches). + """ + + schema_bytes: bytes + ipc_bytes: bytes diff --git a/poetry.lock b/poetry.lock index d4db17a..d8d3bb9 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1,4 +1,4 @@ -# This file is automatically @generated by Poetry 2.2.1 and should not be changed by hand. +# This file is automatically @generated by Poetry 2.3.1 and should not be changed by hand. [[package]] name = "cfgv" @@ -498,7 +498,7 @@ version = "2.4.4" description = "Fundamental package for array computing in Python" optional = false python-versions = ">=3.11" -groups = ["ml"] +groups = ["main"] files = [ {file = "numpy-2.4.4-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:f983334aea213c99992053ede6168500e5f086ce74fbc4acc3f2b00f5762e9db"}, {file = "numpy-2.4.4-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:72944b19f2324114e9dc86a159787333b77874143efcf89a5167ef83cfee8af0"}, @@ -602,7 +602,7 @@ version = "2.3.3" description = "Powerful data structures for data analysis, time series, and statistics" optional = false python-versions = ">=3.9" -groups = ["ml"] +groups = ["main"] files = [ {file = "pandas-2.3.3-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:376c6446ae31770764215a6c937f72d917f214b43560603cd60da6408f183b6c"}, {file = "pandas-2.3.3-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:e19d192383eab2f4ceb30b412b22ea30690c9e618f78870357ae1d682912015a"}, @@ -988,7 +988,7 @@ version = "2.9.0.post0" description = "Extensions to the standard Python datetime module" optional = false python-versions = "!=3.0.*,!=3.1.*,!=3.2.*,>=2.7" -groups = ["ml"] +groups = ["main"] files = [ {file = "python-dateutil-2.9.0.post0.tar.gz", hash = "sha256:37dd54208da7e1cd875388217d5e00ebd4179249f90fb72437e91a35459a0ad3"}, {file = "python_dateutil-2.9.0.post0-py2.py3-none-any.whl", hash = "sha256:a8b2bc7bffae282281c8140a97d3aa9c14da0b136dfe83f850eea9a5f7470427"}, @@ -1033,7 +1033,7 @@ version = "2026.1.post1" description = "World timezone definitions, modern and historical" optional = false python-versions = "*" -groups = ["ml"] +groups = ["main"] files = [ {file = "pytz-2026.1.post1-py2.py3-none-any.whl", hash = "sha256:f2fd16142fda348286a75e1a524be810bb05d444e5a081f37f7affc635035f7a"}, {file = "pytz-2026.1.post1.tar.gz", hash = "sha256:3378dde6a0c3d26719182142c56e60c7f9af7e968076f31aae569d72a0358ee1"}, @@ -1197,7 +1197,7 @@ version = "1.17.0" description = "Python 2 and 3 compatibility utilities" optional = false python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*" -groups = ["ml"] +groups = ["main"] files = [ {file = "six-1.17.0-py2.py3-none-any.whl", hash = "sha256:4721f391ed90541fddacab5acf947aa0d3dc7d27b2e1e8eda2be8970586c3274"}, {file = "six-1.17.0.tar.gz", hash = "sha256:ff70335d468e7eb6ec65b95b99d3a2836546063f63acc5171de367e834932a81"}, @@ -1231,7 +1231,7 @@ version = "2026.1" description = "Provider of IANA time zone data" optional = false python-versions = ">=2" -groups = ["ml"] +groups = ["main"] files = [ {file = "tzdata-2026.1-py2.py3-none-any.whl", hash = "sha256:4b1d2be7ac37ceafd7327b961aa3a54e467efbdb563a23655fbfe0d39cfc42a9"}, {file = "tzdata-2026.1.tar.gz", hash = "sha256:67658a1903c75917309e753fdc349ac0efd8c27db7a0cb406a25be4840f87f98"}, @@ -1268,4 +1268,4 @@ reference = "mdsol" [metadata] lock-version = "2.1" python-versions = "^3.13" -content-hash = "82ed9c8a49cc8f08ca7870effa1ffbecd1999728135df3ff8c86bb704ac34013" +content-hash = "b6f7ded000698b328661a0002d12b90abacab17b6dcfa87b5601f3bbe038ef4a" diff --git a/pyproject.toml b/pyproject.toml index ae7990a..768f12e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -19,6 +19,7 @@ url = "https://mdsol.jfrog.io/artifactory/api/pypi/pypi-prod-virtual/simple" [tool.poetry.dependencies] python = "^3.13" pyarrow = "^19.0.0" +pandas = "^2.0.2" # SERVICE [tool.poetry.group.service] @@ -32,7 +33,6 @@ gunicorn ="^20.1.0" optional = true [tool.poetry.group.ml.dependencies] -pandas = "^2.0.2" # DEV [tool.poetry.group.dev] diff --git a/tests/test_client.py b/tests/test_client.py index 760bd20..615cd81 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -4,24 +4,34 @@ from types import ModuleType from uuid import UUID +import pandas as pd import pytest from dataconnect.client import DataConnectClient from dataconnect.models import Dataset, DatasetVersion, Study +def _make_dataframe() -> pd.DataFrame: + """Return a small DataFrame for use in fetch_data wiring tests.""" + return pd.DataFrame({"subject_id": ["S001"], "age": [30]}) + + class _FakeService: def __init__( self, studies: list[Study] | None = None, versions: list[DatasetVersion] | None = None, + fetch_data_result: pd.DataFrame | None = None, datasets: list[Dataset] | None = None, ) -> None: self._studies = studies or [] self._versions = versions or [] + self._fetch_data_result = fetch_data_result if fetch_data_result is not None else pd.DataFrame() self._datasets = datasets or [] self.closed = 0 self.last_dataset_uuid: UUID | None = None + self.last_fetch_data_uuid: UUID | None = None + self.last_first_n_rows: int | None = None self.last_get_datasets_kwargs: dict[str, object] | None = None def get_studies(self) -> list[Study]: @@ -31,6 +41,11 @@ def get_dataset_versions(self, dataset_uuid: UUID) -> list[DatasetVersion]: self.last_dataset_uuid = dataset_uuid return self._versions + def fetch_data(self, dataset_uuid: UUID, first_n_rows: int | None = None) -> pd.DataFrame: + self.last_fetch_data_uuid = dataset_uuid + self.last_first_n_rows = first_n_rows + return self._fetch_data_result + def get_datasets(self, **kwargs: object) -> list[Dataset]: self.last_get_datasets_kwargs = kwargs return self._datasets @@ -109,6 +124,31 @@ def close(self) -> None: assert isinstance(captured["transport"], FakeArrowFlightTransport) +def test_fetch_data_forwards_args_to_service() -> None: + """Client.fetch_data must delegate uuid and first_n_rows to the service unchanged.""" + dataset_uuid = UUID("073410b6-79be-3e7d-ae37-92f6e054013e") + expected_df = _make_dataframe() + service = _FakeService(fetch_data_result=expected_df) + client = DataConnectClient(service) + + result = client.fetch_data(dataset_uuid, first_n_rows=5) + + assert result is expected_df + assert service.last_fetch_data_uuid == dataset_uuid + assert service.last_first_n_rows == 5 + + +def test_fetch_data_forwards_no_limit_to_service() -> None: + """When first_n_rows is omitted, None must be forwarded to the service.""" + dataset_uuid = UUID("073410b6-79be-3e7d-ae37-92f6e054013e") + service = _FakeService(fetch_data_result=_make_dataframe()) + client = DataConnectClient(service) + + client.fetch_data(dataset_uuid) + + assert service.last_first_n_rows is None + + @pytest.mark.benchmark def test_dummy_benchmark() -> None: # Dummy benchmark test to satisfy CI diff --git a/tests/test_service.py b/tests/test_service.py index 65c375a..c16c45b 100644 --- a/tests/test_service.py +++ b/tests/test_service.py @@ -3,13 +3,30 @@ import json from uuid import UUID +import pandas as pd +import pyarrow as pa import pytest -from dataconnect.exceptions import ConnectionError, ValidationError +from dataconnect.exceptions import ( + AuthenticationError, + AuthorizationError, + ConnectionError, + NotFoundError, + QueryError, + ServerError, + ValidationError, +) from dataconnect.models import Dataset, DatasetVersion from dataconnect.service.default import DefaultDataConnectService -from dataconnect.transport.errors import TransportConnectionError -from dataconnect.transport.models import DataRef, ResourceInfo, ResourceQuery +from dataconnect.transport.errors import ( + TransportAuthenticationError, + TransportAuthorizationError, + TransportConnectionError, + TransportIOError, + TransportNotFoundError, + TransportStatusError, +) +from dataconnect.transport.models import DataRef, DataTable, ResourceInfo, ResourceQuery class _FakeTransport: @@ -17,10 +34,15 @@ def __init__( self, resources: list[ResourceInfo] | None = None, error: Exception | None = None, + data_table: DataTable | None = None, + do_get_error: Exception | None = None, ) -> None: self._resources = resources or [] self._error = error + self._data_table = data_table + self._do_get_error = do_get_error self.last_request: ResourceQuery | None = None + self.last_do_get_request: ResourceQuery | None = None def list_resources(self, request: ResourceQuery) -> list[ResourceInfo]: self.last_request = request @@ -28,10 +50,29 @@ def list_resources(self, request: ResourceQuery) -> list[ResourceInfo]: raise self._error return self._resources + def do_get(self, request: ResourceQuery) -> DataTable: + self.last_do_get_request = request + if self._do_get_error is not None: + raise self._do_get_error + assert self._data_table is not None, "_FakeTransport: no data_table configured" + return self._data_table + def close(self) -> None: return None +def _make_ipc_table(data: dict[str, list[object]]) -> DataTable: + """Serialise a dict of columns to an Arrow IPC stream wrapped in DataTable.""" + arrow_table = pa.table(data) + sink = pa.BufferOutputStream() + writer = pa.ipc.new_stream(sink, arrow_table.schema) + writer.write_table(arrow_table) + writer.close() + ipc_bytes = sink.getvalue().to_pybytes() + schema_bytes = arrow_table.schema.serialize().to_pybytes() + return DataTable(schema_bytes=schema_bytes, ipc_bytes=ipc_bytes) + + def _resource_with_ticket_json(payload: dict[str, object]) -> ResourceInfo: return ResourceInfo( descriptor=b"", @@ -149,6 +190,165 @@ def test_get_dataset_versions_raises_validation_error_on_zero_input() -> None: assert "dataset_uuid must not be empty" in str(excinfo.value) +# --------------------------------------------------------------------------- +# fetch_data tests +# --------------------------------------------------------------------------- + + +def test_fetch_data_returns_dataframe_with_correct_values() -> None: + """Happy path: fetched IPC bytes are mapped to a DataFrame matching the source data.""" + dataset_uuid = UUID("073410b6-79be-3e7d-ae37-92f6e054013e") + source = {"subject_id": ["S001", "S002"], "age": [30, 45]} + transport = _FakeTransport(data_table=_make_ipc_table(source)) + service = DefaultDataConnectService(transport) + + result = service.fetch_data(dataset_uuid) + + assert isinstance(result, pd.DataFrame) + assert result["subject_id"].tolist() == source["subject_id"] + assert result["age"].tolist() == source["age"] + + +def test_fetch_data_builds_correct_request() -> None: + """The request sent to the transport must carry the right action, dataset_uuid, and limit.""" + dataset_uuid = UUID("073410b6-79be-3e7d-ae37-92f6e054013e") + transport = _FakeTransport(data_table=_make_ipc_table({"x": [1]})) + service = DefaultDataConnectService(transport) + + service.fetch_data(dataset_uuid, first_n_rows=10) + + assert transport.last_do_get_request is not None + assert transport.last_do_get_request.action == "data.fetch_ticket" + body = json.loads(transport.last_do_get_request.body) + assert body["dataset_uuid"] == str(dataset_uuid) + assert body["limit"] == 10 + + +def test_fetch_data_no_limit_sends_none_in_body() -> None: + """When first_n_rows is omitted, limit must be None in the request body.""" + dataset_uuid = UUID("073410b6-79be-3e7d-ae37-92f6e054013e") + transport = _FakeTransport(data_table=_make_ipc_table({"x": [1]})) + service = DefaultDataConnectService(transport) + + service.fetch_data(dataset_uuid) + + body = json.loads(transport.last_do_get_request.body) # type: ignore[union-attr] + assert body["limit"] is None + + +def test_fetch_data_returns_empty_dataframe_for_empty_table() -> None: + """An IPC stream with no rows should produce an empty DataFrame.""" + dataset_uuid = UUID("073410b6-79be-3e7d-ae37-92f6e054013e") + transport = _FakeTransport(data_table=_make_ipc_table({"col": pa.array([], type=pa.int32())})) + service = DefaultDataConnectService(transport) + + result = service.fetch_data(dataset_uuid) + + assert isinstance(result, pd.DataFrame) + assert len(result) == 0 + assert "col" in result.columns + + +def test_fetch_data_raises_value_error_on_zero_first_n_rows() -> None: + dataset_uuid = UUID("073410b6-79be-3e7d-ae37-92f6e054013e") + transport = _FakeTransport(data_table=_make_ipc_table({"x": [1]})) + service = DefaultDataConnectService(transport) + + with pytest.raises(ValueError, match="first_n_rows must be a positive integer"): + service.fetch_data(dataset_uuid, first_n_rows=0) + + +def test_fetch_data_raises_value_error_on_negative_first_n_rows() -> None: + dataset_uuid = UUID("073410b6-79be-3e7d-ae37-92f6e054013e") + transport = _FakeTransport(data_table=_make_ipc_table({"x": [1]})) + service = DefaultDataConnectService(transport) + + with pytest.raises(ValueError, match="first_n_rows must be a positive integer"): + service.fetch_data(dataset_uuid, first_n_rows=-5) + + +def test_fetch_data_raises_value_error_on_non_int_first_n_rows() -> None: + dataset_uuid = UUID("073410b6-79be-3e7d-ae37-92f6e054013e") + transport = _FakeTransport(data_table=_make_ipc_table({"x": [1]})) + service = DefaultDataConnectService(transport) + + with pytest.raises(ValueError, match="first_n_rows must be a positive integer"): + service.fetch_data(dataset_uuid, first_n_rows="abc") # type: ignore[arg-type] + + +def test_fetch_data_raises_value_error_on_none_uuid() -> None: + transport = _FakeTransport(data_table=_make_ipc_table({"x": [1]})) + service = DefaultDataConnectService(transport) + + with pytest.raises((ValueError, TypeError)): + service.fetch_data(None) # type: ignore[arg-type] + + +def test_fetch_data_raises_value_error_on_empty_uuid() -> None: + transport = _FakeTransport(data_table=_make_ipc_table({"x": [1]})) + service = DefaultDataConnectService(transport) + + with pytest.raises(ValueError, match="must not be an empty UUID"): + service.fetch_data(UUID(int=0)) + + +# Transport-error translation tests + + +def test_fetch_data_translates_connection_error() -> None: + dataset_uuid = UUID("073410b6-79be-3e7d-ae37-92f6e054013e") + transport = _FakeTransport(do_get_error=TransportConnectionError("timeout")) + service = DefaultDataConnectService(transport) + + with pytest.raises(ConnectionError, match="timeout"): + service.fetch_data(dataset_uuid) + + +def test_fetch_data_translates_authentication_error() -> None: + dataset_uuid = UUID("073410b6-79be-3e7d-ae37-92f6e054013e") + transport = _FakeTransport(do_get_error=TransportAuthenticationError("bad token")) + service = DefaultDataConnectService(transport) + + with pytest.raises(AuthenticationError, match="bad token"): + service.fetch_data(dataset_uuid) + + +def test_fetch_data_translates_authorization_error() -> None: + dataset_uuid = UUID("073410b6-79be-3e7d-ae37-92f6e054013e") + transport = _FakeTransport(do_get_error=TransportAuthorizationError("forbidden")) + service = DefaultDataConnectService(transport) + + with pytest.raises(AuthorizationError, match="forbidden"): + service.fetch_data(dataset_uuid) + + +def test_fetch_data_translates_not_found_error() -> None: + dataset_uuid = UUID("073410b6-79be-3e7d-ae37-92f6e054013e") + transport = _FakeTransport(do_get_error=TransportNotFoundError("dataset not found")) + service = DefaultDataConnectService(transport) + + with pytest.raises(NotFoundError, match="dataset not found"): + service.fetch_data(dataset_uuid) + + +def test_fetch_data_translates_io_error() -> None: + dataset_uuid = UUID("073410b6-79be-3e7d-ae37-92f6e054013e") + transport = _FakeTransport(do_get_error=TransportIOError("stream broken")) + service = DefaultDataConnectService(transport) + + with pytest.raises(QueryError, match="stream broken"): + service.fetch_data(dataset_uuid) + + +def test_fetch_data_translates_status_error() -> None: + dataset_uuid = UUID("073410b6-79be-3e7d-ae37-92f6e054013e") + transport = _FakeTransport(do_get_error=TransportStatusError("internal", status_code=13)) + service = DefaultDataConnectService(transport) + + with pytest.raises(ServerError, match="internal"): + service.fetch_data(dataset_uuid) + + # --- get_datasets tests ---