From 7a717bfeb4767865247a85b32e4f1a574ea64e53 Mon Sep 17 00:00:00 2001 From: Srinivas Lingampalli Date: Sat, 23 May 2026 04:35:07 +0100 Subject: [PATCH 1/2] feat: added publish feature --- dataconnect/__init__.py | 11 +- dataconnect/client.py | 67 +- dataconnect/models.py | 15 + dataconnect/service/base.py | 26 +- dataconnect/service/default.py | 73 ++- dataconnect/service/mappers.py | 36 +- dataconnect/transport/__init__.py | 2 + .../transport/arrow_flight/transport.py | 65 ++ dataconnect/transport/base.py | 9 + dataconnect/transport/models.py | 29 + tests/test_dry_publish.py | 8 +- tests/test_publish.py | 580 ++++++++++++++++++ 12 files changed, 912 insertions(+), 9 deletions(-) create mode 100644 tests/test_publish.py diff --git a/dataconnect/__init__.py b/dataconnect/__init__.py index e676ca1..a13e675 100644 --- a/dataconnect/__init__.py +++ b/dataconnect/__init__.py @@ -11,7 +11,15 @@ ServerError, ValidationError, ) -from dataconnect.models import DatasetVersion, DryPublishResult, PaginatedResponse, Pagination, Study, StudyEnvironment +from dataconnect.models import ( + DatasetVersion, + DryPublishResult, + PaginatedResponse, + Pagination, + PublishResult, + Study, + StudyEnvironment, +) __all__ = [ # Client @@ -23,6 +31,7 @@ "DryPublishResult", "PaginatedResponse", "Pagination", + "PublishResult", # Exceptions — catch these in user application code "DataConnectError", "AuthenticationError", diff --git a/dataconnect/client.py b/dataconnect/client.py index 2a59da5..a7acd2e 100644 --- a/dataconnect/client.py +++ b/dataconnect/client.py @@ -12,7 +12,14 @@ import pandas as pd -from dataconnect.models import Dataset, DatasetVersion, DryPublishResult, PaginatedResponse, StudiesResult +from dataconnect.models import ( + Dataset, + DatasetVersion, + DryPublishResult, + PaginatedResponse, + PublishResult, + StudiesResult, +) from dataconnect.service import DataConnectService, DefaultDataConnectService _DEFAULT_HOST = "enodia-gateway.platform.imedidata.com" @@ -96,7 +103,27 @@ def dry_publish( data: pd.DataFrame, datetime_formats: dict[str, str] | None = None, ) -> DryPublishResult: + """Validate a dataset against the server without persisting any changes. + Delegates directly to :meth:`DataConnectService.dry_publish`. + + Args: + project_token: Base64-encoded project token identifying the target + study, study environment, and project. + dataset_name: Name of the dataset to validate. + key_columns: Column names that form the unique key for the dataset. + source_datasets: UUIDs of the source datasets the published dataset + is derived from. + data: The dataset to validate as a ``pd.DataFrame``. + datetime_formats: Optional mapping of column name → datetime format + string (e.g. ``{"visit_date": "yyyy-MM-dd"}``). Defaults to + an empty dict when omitted. + + Returns: + A :class:`DryPublishResult` containing the server's validation + outcome, including per-field validity flags, error messages, and + an optional ``invalid_records`` DataFrame. + """ return self._service.dry_publish( project_token=project_token, dataset_name=dataset_name, @@ -106,6 +133,44 @@ def dry_publish( datetime_formats=datetime_formats, ) + def publish( + self, + project_token: str, + dataset_name: str, + key_columns: list[str], + source_datasets: list[UUID], + data: pd.DataFrame, + datetime_formats: dict[str, str] | None = None, + ) -> PublishResult: + """Publish a dataset to the server. + + Delegates directly to :meth:`DataConnectService.publish`. + + Args: + project_token: Base64-encoded project token identifying the target + study, study environment, and project. + dataset_name: Name of the dataset to publish. + key_columns: Column names that form the unique key for the dataset. + source_datasets: UUIDs of the source datasets the published dataset + is derived from. + data: The dataset to publish as a ``pd.DataFrame``. + datetime_formats: Optional mapping of column name → datetime format + string (e.g. ``{"visit_date": "yyyy-MM-dd"}``). Defaults to + an empty dict when omitted. + + Returns: + A :class:`PublishResult` containing the server's publish outcome, + including dataset UUID, version, and record counts. + """ + return self._service.publish( + project_token=project_token, + dataset_name=dataset_name, + key_columns=key_columns, + source_datasets=source_datasets, + data=data, + datetime_formats=datetime_formats, + ) + # Lifecycle def close(self) -> None: diff --git a/dataconnect/models.py b/dataconnect/models.py index 8677946..e6bbe60 100644 --- a/dataconnect/models.py +++ b/dataconnect/models.py @@ -82,3 +82,18 @@ class DryPublishResult: duplicate_record_count: int | None = None invalid_record_count: int | None = None invalid_records: pd.DataFrame | None = None + + +@dataclass +class PublishResult: + """Result of a publish operation, including status and details.""" + + status: bool + dataset_name: str | None = None + dataset_uuid: str | None = None + dataset_version: int | None = None + dataset_batch_number: int | None = None + valid_record_count: int | None = None + duplicate_record_count: int | None = None + invalid_record_count: int | None = None + invalid_records: pd.DataFrame | None = None diff --git a/dataconnect/service/base.py b/dataconnect/service/base.py index c315cd4..261eb48 100644 --- a/dataconnect/service/base.py +++ b/dataconnect/service/base.py @@ -7,7 +7,14 @@ import pandas as pd -from dataconnect.models import Dataset, DatasetVersion, DryPublishResult, PaginatedResponse, StudiesResult +from dataconnect.models import ( + Dataset, + DatasetVersion, + DryPublishResult, + PaginatedResponse, + PublishResult, + StudiesResult, +) class DataConnectService(ABC): @@ -44,7 +51,22 @@ def dry_publish( source_datasets: list[UUID], data: pd.DataFrame, datetime_formats: dict[str, str] | None = None, - ) -> DryPublishResult: ... + ) -> DryPublishResult: + """Validate a dataset against the server without persisting any changes.""" + ... + + @abstractmethod + def publish( + self, + project_token: str, + dataset_name: str, + key_columns: list[str], + source_datasets: list[UUID], + data: pd.DataFrame, + datetime_formats: dict[str, str] | None = None, + ) -> PublishResult: + """Publish a dataset to the server and return the result.""" + ... @abstractmethod def close(self) -> None: ... diff --git a/dataconnect/service/default.py b/dataconnect/service/default.py index 0e96f55..38b1fda 100644 --- a/dataconnect/service/default.py +++ b/dataconnect/service/default.py @@ -7,11 +7,20 @@ import pandas as pd -from dataconnect.models import Dataset, DatasetVersion, DryPublishResult, PaginatedResponse, Pagination, StudiesResult +from dataconnect.models import ( + Dataset, + DatasetVersion, + DryPublishResult, + PaginatedResponse, + Pagination, + PublishResult, + StudiesResult, +) from dataconnect.service.base import DataConnectService from dataconnect.service.error_handler import translate_error from dataconnect.service.mappers import ( dry_publish_response_to_domain, + publish_response_to_domain, resource_to_dataset, resource_to_dataset_version, resource_to_fetched_data, @@ -205,6 +214,68 @@ def dry_publish( except TransportError as ex: raise translate_error(ex) from ex + def publish( + self, + project_token: str, + dataset_name: str, + key_columns: list[str], + source_datasets: list[UUID], + data: pd.DataFrame, + datetime_formats: dict[str, str] | None = None, + ) -> PublishResult: + """Publish a dataset to the server and return the result. + + Encodes the publish configuration as a JSON ``input_config`` string, + sets ``is_dry_publish: False`` so the server treats the call as a + live publish run, then delegates to the transport layer. + + The server response is mapped to a :class:`PublishResult` via + :func:`publish_response_to_domain`. + + Args: + project_token: Base64-encoded project token identifying the target + study, study environment, and project. + dataset_name: Name of the dataset to publish. + key_columns: Column names that form the unique key for the dataset. + source_datasets: UUIDs of the source datasets the published dataset + is derived from. + data: The dataset to publish as a ``pd.DataFrame``. + datetime_formats: Optional mapping of column name → datetime format + string (e.g. ``{"visit_date": "yyyy-MM-dd"}``). Defaults to + an empty dict when omitted. + + Returns: + A :class:`PublishResult` containing the server's publish outcome, + including dataset UUID, version, and record counts. + + Raises: + DataConnectError: Any :class:`TransportError` from the transport + layer is translated by :func:`translate_error` into the public + API's :class:`DataConnectError` hierarchy before propagating. + """ + request = PublishRequest( + input_config=json.dumps( + { + "is_dry_publish": False, + "project_token": project_token, + "dataset_name": dataset_name, + "dataset_description": dataset_name, + "key_columns": key_columns, + "source_datasets": [str(uuid) for uuid in source_datasets], + "datetime_formats": datetime_formats or {}, + } + ), + data=data, + ) + + try: + publish_result = self._transport.publish_dataset(request) + + return publish_response_to_domain(publish_result) + + except TransportError as ex: + raise translate_error(ex) from ex + def close(self) -> None: """Close the underlying transport connection.""" diff --git a/dataconnect/service/mappers.py b/dataconnect/service/mappers.py index 0a8513e..7ca8d0e 100644 --- a/dataconnect/service/mappers.py +++ b/dataconnect/service/mappers.py @@ -15,8 +15,8 @@ import pyarrow as pa from dataconnect.exceptions import NotFoundError -from dataconnect.models import Dataset, DatasetVersion, DryPublishResult, Study, StudyEnvironment -from dataconnect.transport.models import DataTable, DryPublishResponse, ResourceInfo +from dataconnect.models import Dataset, DatasetVersion, DryPublishResult, PublishResult, Study, StudyEnvironment +from dataconnect.transport.models import DataTable, DryPublishResponse, PublishResponse, ResourceInfo def resource_to_study(resource: ResourceInfo) -> Study: @@ -124,3 +124,35 @@ def dry_publish_response_to_domain(result: DryPublishResponse | None) -> DryPubl invalid_record_count=result.invalid_record_count, invalid_records=result.invalid_records, ) + + +def publish_response_to_domain(result: PublishResponse | None) -> PublishResult: + """Map a transport-layer ``PublishResponse`` to a ``PublishResult`` domain object. + + ``PublishResponse`` carries flat, typed fields returned by the server after a + publish call. The mapping is direct for all shared fields with one + exception: + + Args: + result: The transport-layer result returned by + :meth:`Transport.publish_dataset`. Pass ``None`` to obtain a + default :class:`PublishResult` with ``status=False`` and all + other fields at their zero values. + + Returns: + A :class:`PublishResult` suitable for returning to the caller. + """ + if result is None: + return PublishResult(status=False) + + return PublishResult( + status=result.status, + dataset_name=result.dataset_name, + dataset_uuid=result.dataset_uuid, + dataset_version=result.dataset_version, + dataset_batch_number=result.dataset_batch_number, + valid_record_count=result.valid_record_count, + duplicate_record_count=result.duplicate_record_count, + invalid_record_count=result.invalid_record_count, + invalid_records=result.invalid_records, + ) diff --git a/dataconnect/transport/__init__.py b/dataconnect/transport/__init__.py index d3024f2..19b7658 100644 --- a/dataconnect/transport/__init__.py +++ b/dataconnect/transport/__init__.py @@ -12,6 +12,7 @@ DataTable, DryPublishResponse, PublishRequest, + PublishResponse, ResourceInfo, ResourceQuery, ) @@ -24,5 +25,6 @@ "DatasetTicket", "DataTable", "PublishRequest", + "PublishResponse", "DryPublishResponse", ] diff --git a/dataconnect/transport/arrow_flight/transport.py b/dataconnect/transport/arrow_flight/transport.py index d33b562..e52d32b 100644 --- a/dataconnect/transport/arrow_flight/transport.py +++ b/dataconnect/transport/arrow_flight/transport.py @@ -28,6 +28,7 @@ DataTable, DryPublishResponse, PublishRequest, + PublishResponse, ResourceInfo, ResourceQuery, ) @@ -321,5 +322,69 @@ def dry_publish_dataset(self, publish_request: PublishRequest) -> DryPublishResp except Exception as ex: raise parse_dataconnect_error(ex) from ex + def publish_dataset(self, publish_request: PublishRequest) -> PublishResponse: + """Send a dataset to the server via Arrow Flight ``do_put`` and return the publish result. + + Follows the same two-phase metadata protocol as :meth:`dry_publish_dataset`: + the server first returns a JSON result buffer, then an optional Arrow IPC + buffer containing any invalid records. + + Args: + publish_request: A :class:`PublishRequest` carrying the encoded server + config (``input_config``) and the dataset as a ``pd.DataFrame``. + + Returns: + A :class:`PublishResponse` populated from the server's JSON response, + including the assigned dataset UUID, version, and record counts. + + Raises: + TransportError: Any Arrow Flight or gRPC error is translated by + :func:`parse_dataconnect_error` before propagating. + """ + descriptor_bytes = publish_request.input_config.encode("utf-8") + flight_descriptor = flight.FlightDescriptor.for_path(descriptor_bytes) + + arrow_table = pa.Table.from_pandas(publish_request.data, preserve_index=False) + + # pa.Table.from_pandas() widens string→large_string, binary→large_binary, + # and list→large_list. Normalize back so the schema matches the server. + arrow_table = arrow_table.cast( + pa.schema([f.with_type(_normalize_arrow_type(f.type)) for f in arrow_table.schema]) + ) + + try: + writer, reader = self._client.do_put(flight_descriptor, arrow_table.schema, self._options()) + + try: + for batch in arrow_table.to_batches(): + writer.write_batch(batch) + writer.done_writing() # only signal completion when all batches succeeded + + # Read while the RPC call is still open (before writer.close()) + # The server first writes a JSON result, then the Arrow table as IPC bytes + _json_buf = reader.read() + json_result = json.loads(_json_buf.to_pybytes()) + + # Read the Arrow table returned by the server upon successful publishing + metadata_buf = reader.read() + result_table = pa.ipc.open_stream(pa.BufferReader(metadata_buf)).read_all() if metadata_buf else None + finally: + writer.close() # terminates the RPC call — must happen after all reads + + return PublishResponse( + status=json_result.get("status", False), + dataset_name=json_result.get("dataset_name", None), + dataset_uuid=json_result.get("dataset_uuid", None), + dataset_version=json_result.get("dataset_version", None), + dataset_batch_number=json_result.get("dataset_batch_number", None), + valid_record_count=json_result.get("valid_record_count", None), + duplicate_record_count=json_result.get("duplicate_record_count", None), + invalid_record_count=json_result.get("invalid_record_count", None), + invalid_records=result_table.to_pandas() if result_table else None, + ) + + except Exception as ex: + raise parse_dataconnect_error(ex) from ex + def close(self) -> None: self._client.close() diff --git a/dataconnect/transport/base.py b/dataconnect/transport/base.py index da668ed..7898d3e 100644 --- a/dataconnect/transport/base.py +++ b/dataconnect/transport/base.py @@ -14,6 +14,7 @@ DataTable, DryPublishResponse, PublishRequest, + PublishResponse, ResourceInfo, ResourceQuery, ) @@ -46,6 +47,14 @@ def dry_publish_dataset(self, publish_request: PublishRequest) -> DryPublishResp ``DryPublishResponse`` containing the complete result set. """ + @abstractmethod + def publish_dataset(self, publish_request: PublishRequest) -> PublishResponse: + """Publish a dataset described by ``publish_request``. + + All record batches from the stream are read and returned as a single + ``PublishResponse`` 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 348e7b6..4111798 100644 --- a/dataconnect/transport/models.py +++ b/dataconnect/transport/models.py @@ -67,12 +67,22 @@ class DataTable: @dataclass(frozen=True) class PublishRequest: + """A publish request containing the input configuration and the dataset to be published.""" + input_config: str + """JSON-encoded server configuration string, including ``is_dry_publish`` flag.""" data: pd.DataFrame + """The dataset to send to the server.""" @dataclass(frozen=True) class DryPublishResponse: + """Transport-layer response from a dry-publish call. + + Carries the server's validation outcome for all rows and schema checks + without any data being persisted. + """ + status: bool is_schema_valid: bool is_config_valid: bool @@ -86,3 +96,22 @@ class DryPublishResponse: duplicate_record_count: int invalid_record_count: int = 0 invalid_records: pd.DataFrame | None = None + + +@dataclass(frozen=True) +class PublishResponse: + """Transport-layer response from a live publish call. + + Carries the server's outcome after persisting the submitted dataset, + including the assigned dataset UUID and version number. + """ + + status: bool + dataset_name: str | None = None + dataset_uuid: str | None = None + dataset_version: int | None = None + dataset_batch_number: int | None = None + valid_record_count: int | None = None + duplicate_record_count: int | None = None + invalid_record_count: int | None = None + invalid_records: pd.DataFrame | None = None diff --git a/tests/test_dry_publish.py b/tests/test_dry_publish.py index 316525b..4f2ec8f 100644 --- a/tests/test_dry_publish.py +++ b/tests/test_dry_publish.py @@ -32,6 +32,7 @@ DataTable, DryPublishResponse, PublishRequest, + PublishResponse, ResourceInfo, ResourceQuery, ) @@ -223,12 +224,15 @@ def list_resources(self, request: ResourceQuery) -> list[ResourceInfo]: def get_ticket(self, ticket: DatasetTicket) -> DataTable: raise NotImplementedError - def dry_publish_dataset(self, request: PublishRequest) -> DryPublishResponse: - self.last_request = request + def dry_publish_dataset(self, publish_request: PublishRequest) -> DryPublishResponse: + self.last_request = publish_request if self._raise is not None: raise self._raise return self._return # type: ignore[return-value] + def publish_dataset(self, publish_request: PublishRequest) -> PublishResponse: + raise NotImplementedError + def close(self) -> None: pass diff --git a/tests/test_publish.py b/tests/test_publish.py new file mode 100644 index 0000000..04384bc --- /dev/null +++ b/tests/test_publish.py @@ -0,0 +1,580 @@ +"""Unit tests for the publish feature. + +Covers: +- ``publish_response_to_domain`` (service mapper) +- ``DefaultDataConnectService.publish`` (service layer) +- ``ArrowFlightTransport.publish_dataset`` (transport layer) +- ``DataConnectClient.dry_publish`` (client façade) +- ``DataConnectClient.publish`` (client façade) +""" + +from __future__ import annotations + +import json +from unittest.mock import MagicMock, patch +from uuid import UUID + +import pandas as pd +import pyarrow as pa +import pytest + +from dataconnect.client import DataConnectClient +from dataconnect.exceptions import ValidationError +from dataconnect.models import DryPublishResult, PublishResult +from dataconnect.service.default import DefaultDataConnectService +from dataconnect.service.mappers import publish_response_to_domain +from dataconnect.transport.arrow_flight.transport import ArrowFlightTransport +from dataconnect.transport.base import Transport +from dataconnect.transport.errors import TransportValidationError +from dataconnect.transport.models import ( + DatasetTicket, + DataTable, + DryPublishResponse, + PublishRequest, + PublishResponse, + ResourceInfo, + ResourceQuery, +) + +# --------------------------------------------------------------------------- +# Helpers shared across test suites +# --------------------------------------------------------------------------- + + +def _make_publish_response(**overrides: object) -> PublishResponse: + """Return a fully-populated ``PublishResponse`` with sensible defaults.""" + defaults: dict = dict( + status=True, + dataset_name="demo_dataset", + dataset_uuid="0158ea12-4004-3817-899b-2de6becbc0f9", + dataset_version=1, + dataset_batch_number=1, + valid_record_count=10, + duplicate_record_count=0, + invalid_record_count=0, + invalid_records=None, + ) + return PublishResponse(**{**defaults, **overrides}) + + +def _make_json_buf(d: dict) -> pa.Buffer: + return pa.py_buffer(json.dumps(d).encode("utf-8")) + + +def _make_ipc_buf(df: pd.DataFrame) -> pa.Buffer: + table = pa.Table.from_pandas(df) + sink = pa.BufferOutputStream() + w = pa.ipc.new_stream(sink, table.schema) + w.write_table(table) + w.close() + return pa.py_buffer(sink.getvalue().to_pybytes()) + + +# --------------------------------------------------------------------------- +# publish_response_to_domain +# --------------------------------------------------------------------------- + + +class TestPublishResponseToDomain: + """Every field of ``PublishResponse`` must land on the correct attribute of + ``PublishResult``.""" + + def test_status_true_is_mapped(self) -> None: + assert publish_response_to_domain(_make_publish_response(status=True)).status is True + + def test_status_false_is_mapped(self) -> None: + assert publish_response_to_domain(_make_publish_response(status=False)).status is False + + def test_dataset_name_mapped(self) -> None: + result = publish_response_to_domain(_make_publish_response(dataset_name="my_ds")) + assert result.dataset_name == "my_ds" + + def test_dataset_uuid_mapped(self) -> None: + uid = "0158ea12-4004-3817-899b-2de6becbc0f9" + result = publish_response_to_domain(_make_publish_response(dataset_uuid=uid)) + assert result.dataset_uuid == uid + + def test_dataset_version_mapped(self) -> None: + result = publish_response_to_domain(_make_publish_response(dataset_version=3)) + assert result.dataset_version == 3 + + def test_dataset_batch_number_mapped(self) -> None: + result = publish_response_to_domain(_make_publish_response(dataset_batch_number=2)) + assert result.dataset_batch_number == 2 + + def test_valid_record_count_mapped(self) -> None: + result = publish_response_to_domain(_make_publish_response(valid_record_count=50)) + assert result.valid_record_count == 50 + + def test_duplicate_record_count_mapped(self) -> None: + result = publish_response_to_domain(_make_publish_response(duplicate_record_count=5)) + assert result.duplicate_record_count == 5 + + def test_invalid_record_count_mapped(self) -> None: + result = publish_response_to_domain(_make_publish_response(invalid_record_count=2)) + assert result.invalid_record_count == 2 + + def test_invalid_records_none_preserved(self) -> None: + result = publish_response_to_domain(_make_publish_response(invalid_records=None)) + assert result.invalid_records is None + + def test_invalid_records_dataframe_preserved(self) -> None: + df = pd.DataFrame({"col": [1, 2, 3]}) + result = publish_response_to_domain(_make_publish_response(invalid_records=df)) + assert result.invalid_records is df + + def test_none_result_returns_status_false(self) -> None: + result = publish_response_to_domain(None) + assert isinstance(result, PublishResult) + assert result.status is False + + def test_returns_publish_result_instance(self) -> None: + result = publish_response_to_domain(_make_publish_response()) + assert isinstance(result, PublishResult) + + +# --------------------------------------------------------------------------- +# DefaultDataConnectService.publish +# --------------------------------------------------------------------------- + + +class _StubTransport(Transport): + """Minimal stub that satisfies the Transport ABC for publish tests.""" + + def __init__( + self, + publish_return: PublishResponse | None = None, + raise_error: Exception | None = None, + ) -> None: + self._return = publish_return + self._raise = raise_error + self.last_request: PublishRequest | None = None + + def list_resources(self, request: ResourceQuery) -> list[ResourceInfo]: + return [] + + def get_ticket(self, ticket: DatasetTicket) -> DataTable: + raise NotImplementedError + + def dry_publish_dataset(self, publish_request: PublishRequest) -> DryPublishResponse: + raise NotImplementedError + + def publish_dataset(self, publish_request: PublishRequest) -> PublishResponse: + self.last_request = publish_request + if self._raise is not None: + raise self._raise + return self._return # type: ignore[return-value] + + def close(self) -> None: + pass + + +def _make_service( + publish_return: PublishResponse | None = None, + raise_error: Exception | None = None, +) -> tuple[DefaultDataConnectService, _StubTransport]: + transport = _StubTransport(publish_return=publish_return, raise_error=raise_error) + return DefaultDataConnectService(transport), transport + + +def _default_publish_args() -> dict: + return dict( + project_token="tok_abc123", + dataset_name="my_dataset", + key_columns=["site_id"], + source_datasets=[UUID("0158ea12-4004-3817-899b-2de6becbc0f9")], + data=pd.DataFrame({"site_id": ["S1"], "value": [1]}), + ) + + +class TestPublishService: + """``DefaultDataConnectService.publish`` must build the correct request, + delegate to the transport, map the response, and translate errors.""" + + def test_returns_publish_result_instance(self) -> None: + service, _ = _make_service(publish_return=_make_publish_response()) + result = service.publish(**_default_publish_args()) + assert isinstance(result, PublishResult) + + def test_successful_status_is_propagated(self) -> None: + service, _ = _make_service(publish_return=_make_publish_response(status=True)) + assert service.publish(**_default_publish_args()).status is True + + # --- input_config JSON content --- + + def test_input_config_sets_is_dry_publish_false(self) -> None: + service, transport = _make_service(publish_return=_make_publish_response()) + service.publish(**_default_publish_args()) + assert transport.last_request is not None + config = json.loads(transport.last_request.input_config) + assert config["is_dry_publish"] is False + + def test_input_config_contains_project_token(self) -> None: + service, transport = _make_service(publish_return=_make_publish_response()) + service.publish(**_default_publish_args()) + assert transport.last_request is not None + assert json.loads(transport.last_request.input_config)["project_token"] == "tok_abc123" + + def test_input_config_contains_dataset_name(self) -> None: + service, transport = _make_service(publish_return=_make_publish_response()) + service.publish(**_default_publish_args()) + assert transport.last_request is not None + assert json.loads(transport.last_request.input_config)["dataset_name"] == "my_dataset" + + def test_input_config_contains_key_columns(self) -> None: + service, transport = _make_service(publish_return=_make_publish_response()) + service.publish(**_default_publish_args()) + assert transport.last_request is not None + assert json.loads(transport.last_request.input_config)["key_columns"] == ["site_id"] + + def test_input_config_source_datasets_serialized_as_strings(self) -> None: + service, transport = _make_service(publish_return=_make_publish_response()) + service.publish(**_default_publish_args()) + assert transport.last_request is not None + config = json.loads(transport.last_request.input_config) + assert config["source_datasets"] == ["0158ea12-4004-3817-899b-2de6becbc0f9"] + + def test_datetime_formats_defaults_to_empty_dict(self) -> None: + service, transport = _make_service(publish_return=_make_publish_response()) + service.publish(**_default_publish_args()) # no datetime_formats kwarg + assert transport.last_request is not None + assert json.loads(transport.last_request.input_config)["datetime_formats"] == {} + + def test_datetime_formats_passed_through_when_provided(self) -> None: + service, transport = _make_service(publish_return=_make_publish_response()) + args = _default_publish_args() + args["datetime_formats"] = {"visit_date": "yyyy-MM-dd"} + service.publish(**args) + assert transport.last_request is not None + assert json.loads(transport.last_request.input_config)["datetime_formats"] == {"visit_date": "yyyy-MM-dd"} + + def test_data_dataframe_passed_to_transport(self) -> None: + service, transport = _make_service(publish_return=_make_publish_response()) + df = pd.DataFrame({"a": [1, 2, 3]}) + args = _default_publish_args() + args["data"] = df + service.publish(**args) + assert transport.last_request is not None + assert transport.last_request.data is df + + # --- falsy / None result --- + + def test_none_transport_result_returns_status_false(self) -> None: + service, _ = _make_service(publish_return=None) + result = service.publish(**_default_publish_args()) + assert isinstance(result, PublishResult) + assert result.status is False + + # --- error translation --- + + def test_transport_validation_error_is_translated_to_service_error(self) -> None: + err = TransportValidationError( + error_code="VAL_001", + message="schema mismatch", + timestamp="2024-01-01T00:00:00Z", + ) + service, _ = _make_service(raise_error=err) + with pytest.raises(ValidationError): + service.publish(**_default_publish_args()) + + +# --------------------------------------------------------------------------- +# ArrowFlightTransport.publish_dataset +# --------------------------------------------------------------------------- + + +def _make_flight_transport() -> ArrowFlightTransport: + """Create an ``ArrowFlightTransport`` with a mocked ``FlightClient``.""" + with patch.object(ArrowFlightTransport, "_get_client", return_value=MagicMock()): + return ArrowFlightTransport(host="localhost", port=5005, use_tls=False) + + +def _wire_do_put( + transport: ArrowFlightTransport, + json_resp: dict, + invalid_records_df: pd.DataFrame | None = None, +) -> tuple[MagicMock, MagicMock]: + """Configure ``transport._client.do_put`` to return controlled writer/reader mocks.""" + json_buf = _make_json_buf(json_resp) + ipc_buf = _make_ipc_buf(invalid_records_df) if invalid_records_df is not None else None + + reader_mock = MagicMock() + reader_mock.read.side_effect = [json_buf, ipc_buf] + writer_mock = MagicMock() + transport._client.do_put.return_value = (writer_mock, reader_mock) + return writer_mock, reader_mock + + +# A minimal valid JSON response the server would return for a live publish call. +_VALID_JSON_RESP: dict = { + "status": True, + "dataset_name": "demo_dataset", + "dataset_uuid": "0158ea12-4004-3817-899b-2de6becbc0f9", + "dataset_version": 1, + "dataset_batch_number": 1, + "valid_record_count": 1, + "duplicate_record_count": 0, + "invalid_record_count": 0, +} + + +class TestPublishDatasetTransport: + """``ArrowFlightTransport.publish_dataset`` must normalise the Arrow schema, + manage the writer lifecycle correctly, and build ``PublishResponse`` from + the server's two-phase metadata response.""" + + # --- schema normalisation --- + + def test_large_string_columns_cast_to_string_before_send(self) -> None: + transport = _make_flight_transport() + _wire_do_put(transport, _VALID_JSON_RESP) + + df = pd.DataFrame({"name": ["Alice", "Bob"], "age": [30, 25]}) + transport.publish_dataset(PublishRequest(input_config="{}", data=df)) + + sent_schema = transport._client.do_put.call_args[0][1] + assert sent_schema.field("name").type == pa.utf8(), "large_string should have been normalised to string" + + def test_integer_and_float_columns_not_affected_by_normalisation(self) -> None: + transport = _make_flight_transport() + _wire_do_put(transport, _VALID_JSON_RESP) + + df = pd.DataFrame({"count": [1, 2], "score": [1.1, 2.2]}) + transport.publish_dataset(PublishRequest(input_config="{}", data=df)) + + sent_schema = transport._client.do_put.call_args[0][1] + assert pa.types.is_integer(sent_schema.field("count").type) + assert pa.types.is_floating(sent_schema.field("score").type) + + # --- writer lifecycle --- + + def test_done_writing_is_called_once_on_success(self) -> None: + transport = _make_flight_transport() + writer_mock, _ = _wire_do_put(transport, _VALID_JSON_RESP) + + transport.publish_dataset(PublishRequest(input_config="{}", data=pd.DataFrame({"x": [1]}))) + writer_mock.done_writing.assert_called_once() + + def test_close_is_called_once_on_success(self) -> None: + transport = _make_flight_transport() + writer_mock, _ = _wire_do_put(transport, _VALID_JSON_RESP) + + transport.publish_dataset(PublishRequest(input_config="{}", data=pd.DataFrame({"x": [1]}))) + writer_mock.close.assert_called_once() + + def test_done_writing_precedes_close(self) -> None: + transport = _make_flight_transport() + writer_mock, _ = _wire_do_put(transport, _VALID_JSON_RESP) + + transport.publish_dataset(PublishRequest(input_config="{}", data=pd.DataFrame({"x": [1]}))) + method_names = [str(c) for c in writer_mock.method_calls] + done_idx = next(i for i, n in enumerate(method_names) if "done_writing" in n) + close_idx = next(i for i, n in enumerate(method_names) if "close" in n) + assert done_idx < close_idx, "done_writing() must be called before close()" + + def test_close_called_even_when_write_batch_raises(self) -> None: + transport = _make_flight_transport() + writer_mock = MagicMock() + writer_mock.write_batch.side_effect = RuntimeError("network error") + reader_mock = MagicMock() + transport._client.do_put.return_value = (writer_mock, reader_mock) + + from dataconnect.transport.errors import TransportError + + with pytest.raises(TransportError): + transport.publish_dataset(PublishRequest(input_config="{}", data=pd.DataFrame({"x": [1]}))) + + writer_mock.close.assert_called_once() + + def test_done_writing_not_called_when_write_batch_raises(self) -> None: + transport = _make_flight_transport() + writer_mock = MagicMock() + writer_mock.write_batch.side_effect = RuntimeError("network error") + transport._client.do_put.return_value = (writer_mock, MagicMock()) + + from dataconnect.transport.errors import TransportError + + with pytest.raises(TransportError): + transport.publish_dataset(PublishRequest(input_config="{}", data=pd.DataFrame({"x": [1]}))) + + writer_mock.done_writing.assert_not_called() + + # --- response parsing --- + + def test_returns_publish_response_instance(self) -> None: + transport = _make_flight_transport() + _wire_do_put(transport, _VALID_JSON_RESP) + + result = transport.publish_dataset(PublishRequest(input_config="{}", data=pd.DataFrame({"x": [1]}))) + assert isinstance(result, PublishResponse) + + def test_status_parsed_from_json(self) -> None: + transport = _make_flight_transport() + _wire_do_put(transport, {**_VALID_JSON_RESP, "status": False}) + + result = transport.publish_dataset(PublishRequest(input_config="{}", data=pd.DataFrame({"x": [1]}))) + assert result.status is False + + def test_dataset_uuid_parsed_from_json(self) -> None: + transport = _make_flight_transport() + uid = "aaaabbbb-1111-2222-3333-ccccddddeeee" + _wire_do_put(transport, {**_VALID_JSON_RESP, "dataset_uuid": uid}) + + result = transport.publish_dataset(PublishRequest(input_config="{}", data=pd.DataFrame({"x": [1]}))) + assert result.dataset_uuid == uid + + def test_dataset_version_parsed_from_json(self) -> None: + transport = _make_flight_transport() + _wire_do_put(transport, {**_VALID_JSON_RESP, "dataset_version": 7}) + + result = transport.publish_dataset(PublishRequest(input_config="{}", data=pd.DataFrame({"x": [1]}))) + assert result.dataset_version == 7 + + def test_none_ipc_buf_yields_no_invalid_records(self) -> None: + transport = _make_flight_transport() + json_buf = _make_json_buf(_VALID_JSON_RESP) + reader_mock = MagicMock() + reader_mock.read.side_effect = [json_buf, None] + writer_mock = MagicMock() + transport._client.do_put.return_value = (writer_mock, reader_mock) + + result = transport.publish_dataset(PublishRequest(input_config="{}", data=pd.DataFrame({"x": [1]}))) + assert result.invalid_records is None + + def test_ipc_buf_is_decoded_to_dataframe(self) -> None: + transport = _make_flight_transport() + invalid_df = pd.DataFrame({"row_id": [10, 20], "reason": ["bad type", "null value"]}) + _wire_do_put(transport, _VALID_JSON_RESP, invalid_records_df=invalid_df) + + result = transport.publish_dataset(PublishRequest(input_config="{}", data=pd.DataFrame({"x": [1]}))) + assert result.invalid_records is not None + assert list(result.invalid_records.columns) == ["row_id", "reason"] + assert len(result.invalid_records) == 2 + + +# --------------------------------------------------------------------------- +# DataConnectClient.dry_publish and DataConnectClient.publish (façade) +# --------------------------------------------------------------------------- + + +class _FakeService: + """Minimal fake service to verify the client delegates correctly.""" + + def __init__( + self, + dry_publish_return: DryPublishResult | None = None, + publish_return: PublishResult | None = None, + ) -> None: + self._dry_publish_return = dry_publish_return or DryPublishResult(status=True) + self._publish_return = publish_return or PublishResult(status=True) + self.dry_publish_calls: list[dict] = [] + self.publish_calls: list[dict] = [] + + def dry_publish(self, **kwargs: object) -> DryPublishResult: + self.dry_publish_calls.append(kwargs) + return self._dry_publish_return + + def publish(self, **kwargs: object) -> PublishResult: + self.publish_calls.append(kwargs) + return self._publish_return + + def get_studies(self, **kwargs: object) -> None: # type: ignore[return] + pass + + def get_datasets(self, **kwargs: object) -> None: # type: ignore[return] + pass + + def get_dataset_versions(self, **kwargs: object) -> None: # type: ignore[return] + pass + + def fetch_data(self, **kwargs: object) -> None: # type: ignore[return] + pass + + def close(self) -> None: + pass + + +def _default_args() -> dict: + return dict( + project_token="tok_xyz", + dataset_name="labs", + key_columns=["subject_id"], + source_datasets=[UUID("0158ea12-4004-3817-899b-2de6becbc0f9")], + data=pd.DataFrame({"subject_id": ["S01"]}), + ) + + +class TestClientDryPublish: + """``DataConnectClient.dry_publish`` must delegate all arguments to the service.""" + + def test_returns_dry_publish_result(self) -> None: + service = _FakeService(dry_publish_return=DryPublishResult(status=True)) + client = DataConnectClient(service) # type: ignore[arg-type] + result = client.dry_publish(**_default_args()) + assert isinstance(result, DryPublishResult) + + def test_delegates_all_args_to_service(self) -> None: + service = _FakeService() + client = DataConnectClient(service) # type: ignore[arg-type] + args = _default_args() + args["datetime_formats"] = {"visit_date": "yyyy-MM-dd"} + client.dry_publish(**args) + + assert len(service.dry_publish_calls) == 1 + call = service.dry_publish_calls[0] + assert call["project_token"] == "tok_xyz" + assert call["dataset_name"] == "labs" + assert call["key_columns"] == ["subject_id"] + assert call["datetime_formats"] == {"visit_date": "yyyy-MM-dd"} + + def test_datetime_formats_none_forwarded_as_none(self) -> None: + service = _FakeService() + client = DataConnectClient(service) # type: ignore[arg-type] + client.dry_publish(**_default_args()) # no datetime_formats kwarg + call = service.dry_publish_calls[0] + assert call["datetime_formats"] is None + + def test_result_from_service_is_returned_unchanged(self) -> None: + expected = DryPublishResult(status=False, errors=["schema mismatch"]) + service = _FakeService(dry_publish_return=expected) + client = DataConnectClient(service) # type: ignore[arg-type] + result = client.dry_publish(**_default_args()) + assert result is expected + + +class TestClientPublish: + """``DataConnectClient.publish`` must delegate all arguments to the service.""" + + def test_returns_publish_result(self) -> None: + service = _FakeService(publish_return=PublishResult(status=True)) + client = DataConnectClient(service) # type: ignore[arg-type] + result = client.publish(**_default_args()) + assert isinstance(result, PublishResult) + + def test_delegates_all_args_to_service(self) -> None: + service = _FakeService() + client = DataConnectClient(service) # type: ignore[arg-type] + args = _default_args() + args["datetime_formats"] = {"collected_dt": "yyyy-MM-dd"} + client.publish(**args) + + assert len(service.publish_calls) == 1 + call = service.publish_calls[0] + assert call["project_token"] == "tok_xyz" + assert call["dataset_name"] == "labs" + assert call["key_columns"] == ["subject_id"] + assert call["datetime_formats"] == {"collected_dt": "yyyy-MM-dd"} + + def test_datetime_formats_none_forwarded_as_none(self) -> None: + service = _FakeService() + client = DataConnectClient(service) # type: ignore[arg-type] + client.publish(**_default_args()) # no datetime_formats kwarg + call = service.publish_calls[0] + assert call["datetime_formats"] is None + + def test_result_from_service_is_returned_unchanged(self) -> None: + expected = PublishResult(status=True, dataset_uuid="abc-123", dataset_version=2) + service = _FakeService(publish_return=expected) + client = DataConnectClient(service) # type: ignore[arg-type] + result = client.publish(**_default_args()) + assert result is expected From 304221f7e8e4e00eec7eeb3ef26ad63449a674a3 Mon Sep 17 00:00:00 2001 From: Srinivas Lingampalli Date: Sat, 23 May 2026 04:49:20 +0100 Subject: [PATCH 2/2] fix: fixed review comments --- dataconnect/client.py | 6 ++---- dataconnect/service/mappers.py | 7 +++---- dataconnect/transport/models.py | 9 ++++++++- 3 files changed, 13 insertions(+), 9 deletions(-) diff --git a/dataconnect/client.py b/dataconnect/client.py index a7acd2e..577c6b4 100644 --- a/dataconnect/client.py +++ b/dataconnect/client.py @@ -116,8 +116,7 @@ def dry_publish( is derived from. data: The dataset to validate as a ``pd.DataFrame``. datetime_formats: Optional mapping of column name → datetime format - string (e.g. ``{"visit_date": "yyyy-MM-dd"}``). Defaults to - an empty dict when omitted. + string (e.g. ``{"visit_date": "yyyy-MM-dd"}``). Returns: A :class:`DryPublishResult` containing the server's validation @@ -155,8 +154,7 @@ def publish( is derived from. data: The dataset to publish as a ``pd.DataFrame``. datetime_formats: Optional mapping of column name → datetime format - string (e.g. ``{"visit_date": "yyyy-MM-dd"}``). Defaults to - an empty dict when omitted. + string (e.g. ``{"visit_date": "yyyy-MM-dd"}``). Returns: A :class:`PublishResult` containing the server's publish outcome, diff --git a/dataconnect/service/mappers.py b/dataconnect/service/mappers.py index 7ca8d0e..f65e747 100644 --- a/dataconnect/service/mappers.py +++ b/dataconnect/service/mappers.py @@ -130,14 +130,13 @@ def publish_response_to_domain(result: PublishResponse | None) -> PublishResult: """Map a transport-layer ``PublishResponse`` to a ``PublishResult`` domain object. ``PublishResponse`` carries flat, typed fields returned by the server after a - publish call. The mapping is direct for all shared fields with one - exception: + publish call. The mapping is direct for all shared fields. Args: result: The transport-layer result returned by - :meth:`Transport.publish_dataset`. Pass ``None`` to obtain a + :meth:`Transport.publish_dataset`. Pass ``None`` to obtain a default :class:`PublishResult` with ``status=False`` and all - other fields at their zero values. + other fields left at their default values. Returns: A :class:`PublishResult` suitable for returning to the caller. diff --git a/dataconnect/transport/models.py b/dataconnect/transport/models.py index 4111798..a6f87d9 100644 --- a/dataconnect/transport/models.py +++ b/dataconnect/transport/models.py @@ -67,10 +67,17 @@ class DataTable: @dataclass(frozen=True) class PublishRequest: - """A publish request containing the input configuration and the dataset to be published.""" + """A publish request containing the input configuration and the dataset to be published. + + Attributes: + input_config: JSON-encoded server configuration string, including + ``is_dry_publish`` flag. + data: The dataset to send to the server. + """ input_config: str """JSON-encoded server configuration string, including ``is_dry_publish`` flag.""" + data: pd.DataFrame """The dataset to send to the server."""