From efb6e5bfd6f43dd2217219f03a85a6cbf6dc8b65 Mon Sep 17 00:00:00 2001 From: Srinivas Lingampalli Date: Wed, 6 May 2026 21:47:12 +0100 Subject: [PATCH 1/2] fix: fixed copilot review comments --- README.md | 3 +-- dataconnect/service/default.py | 2 +- dataconnect/service/mappers.py | 5 ++++- dataconnect/transport/arrow_flight/transport.py | 6 +++++- 4 files changed, 11 insertions(+), 5 deletions(-) diff --git a/README.md b/README.md index 186a77f..4d76d30 100644 --- a/README.md +++ b/README.md @@ -38,8 +38,7 @@ with DataConnectClient.connect( token="your-bearer-token", ) as client: - studies = client.studies(search_study_name="ACME", page=1, page_size=10) - study = studies[0] + studies = client.get_studies(search_study_name="ACME") ``` ## Development diff --git a/dataconnect/service/default.py b/dataconnect/service/default.py index 2135b89..dd287cf 100644 --- a/dataconnect/service/default.py +++ b/dataconnect/service/default.py @@ -69,7 +69,7 @@ def get_studies(self) -> list[Study]: try: return [resource_to_study(r) for r in resources] - except (KeyError, TypeError, ValueError) as ex: + except (IndexError, KeyError, TypeError, ValueError) as ex: raise ValidationError(f"Unexpected studies response format: {ex}") from ex def close(self) -> None: diff --git a/dataconnect/service/mappers.py b/dataconnect/service/mappers.py index 00eb9c2..46a307d 100644 --- a/dataconnect/service/mappers.py +++ b/dataconnect/service/mappers.py @@ -12,11 +12,14 @@ from dataconnect.models import Study, StudyEnvironment from dataconnect.transport.models import ResourceInfo - +from dataconnect.exceptions import NotFoundError def resource_to_study(resource: ResourceInfo) -> Study: """Parse a transport-layer ``ResourceInfo`` into a ``Study`` domain object.""" + if not resource or not resource.endpoints or not resource.endpoints[0].ticket: + raise NotFoundError("Invalid resource: missing endpoints or ticket") + data = json.loads(resource.endpoints[0].ticket.decode("utf-8")) return Study( diff --git a/dataconnect/transport/arrow_flight/transport.py b/dataconnect/transport/arrow_flight/transport.py index 8fae243..1476686 100644 --- a/dataconnect/transport/arrow_flight/transport.py +++ b/dataconnect/transport/arrow_flight/transport.py @@ -77,7 +77,11 @@ def list_resources(self, request: ResourceQuery) -> list[ResourceInfo]: flight_type = _ACTION_FLIGHT_TYPE.get(request.action) if flight_type is None: - raise TransportConnectionError(f"Unknown action: {request.action!r}") + raise TransportStatusError( + f"Unknown action: {request.action!r}", + status_code=3, + grpc_status="INVALID_ARGUMENT" + ) body = json.loads(request.body) if request.body else {} criteria = json.dumps({**body, "flight_type": flight_type}, separators=(",", ":")).encode("utf-8") From c9a0992ea211a61475964cc7bfffae87c2c7e62f Mon Sep 17 00:00:00 2001 From: Srinivas Lingampalli Date: Wed, 6 May 2026 21:56:29 +0100 Subject: [PATCH 2/2] refactor: removing obsolete code --- dataconnect/_encoding.py | 16 --- dataconnect/auth.py | 16 --- dataconnect/framework/__init__.py | 0 dataconnect/framework/pyarrow_transport.py | 99 ------------------- dataconnect/framework/transport.py | 44 --------- dataconnect/service/mappers.py | 5 +- .../transport/arrow_flight/transport.py | 4 +- 7 files changed, 4 insertions(+), 180 deletions(-) delete mode 100644 dataconnect/_encoding.py delete mode 100644 dataconnect/auth.py delete mode 100644 dataconnect/framework/__init__.py delete mode 100644 dataconnect/framework/pyarrow_transport.py delete mode 100644 dataconnect/framework/transport.py diff --git a/dataconnect/_encoding.py b/dataconnect/_encoding.py deleted file mode 100644 index bd6bdcb..0000000 --- a/dataconnect/_encoding.py +++ /dev/null @@ -1,16 +0,0 @@ -"""Encoding utilities for DataConnect.""" - -from __future__ import annotations - -import json -from typing import Any - - -def dumps(obj: Any) -> bytes: - """Serialize *obj* to JSON (bytes).""" - return json.dumps(obj, separators=(",", ":")).encode("utf-8") - - -def loads(data: bytes) -> Any: - """Deserialize JSON from *data* (bytes).""" - return json.loads(data.decode("utf-8")) diff --git a/dataconnect/auth.py b/dataconnect/auth.py deleted file mode 100644 index 61f722f..0000000 --- a/dataconnect/auth.py +++ /dev/null @@ -1,16 +0,0 @@ -"""Authentication and authorization utilities for DataConnect.""" - -from __future__ import annotations - -from dataclasses import dataclass - - -class Credentials: - """Marker base class for all credentials types.""" - - -@dataclass(frozen=True) -class BearerTokenAuth(Credentials): - """Bearer token authentication credentials (OAuth access token).""" - - token: str diff --git a/dataconnect/framework/__init__.py b/dataconnect/framework/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/dataconnect/framework/pyarrow_transport.py b/dataconnect/framework/pyarrow_transport.py deleted file mode 100644 index df7d49a..0000000 --- a/dataconnect/framework/pyarrow_transport.py +++ /dev/null @@ -1,99 +0,0 @@ -"""Default Flight transport implementation using pyarrow.""" - -from __future__ import annotations - -from collections.abc import Iterator - -import pyarrow as pa -from pyarrow import flight - -from dataconnect.auth import BearerTokenAuth, Credentials -from dataconnect.exceptions import AuthenticationError, DataConnectError, QueryError -from dataconnect.framework.transport import FlightTransport, RecordBatchStream - - -class PyArrowFlightTransport(FlightTransport): - """Flight transport implementation using pyarrow.""" - - def __init__( - self, - location: str, - *, - credentials: Credentials | None = None, - tls_root_certs: bytes | None = None, - headers: dict[str, str] | None = None, - ) -> None: - try: - self._client = flight.FlightClient( - location, - tls_root_certs=tls_root_certs, - ) - except Exception as exc: - raise ConnectionError(f"Failed to connect to {location}: {exc}") from exc - - self._call_headers: list[tuple[bytes, bytes]] = [ - (k.lower().encode("ascii"), v.encode("utf-8")) for k, v in (headers or {}).items() - ] - - if credentials is not None: - self._apply_credentials(credentials) - - # Auth - def _apply_credentials(self, credentials: Credentials) -> None: - try: - if isinstance(credentials, BearerTokenAuth): - self._call_headers.append((b"authorization", f"Bearer {credentials.token}".encode())) - else: - raise DataConnectError(f"Unsupported credentials type: {type(credentials)}") - except flight.FlightUnauthenticatedError as exc: - raise AuthenticationError(f"Authentication failed: {exc}") from exc - except flight.FlightError as exc: - raise ConnectionError(str(exc)) from exc - - def __options(self) -> flight.FlightCallOptions: - return flight.FlightCallOptions(headers=self._call_headers) - - # FlightTransport - def do_get(self, ticket: bytes) -> RecordBatchStream: - try: - reader = self._client.do_get(flight.Ticket(ticket), self.__options()) - except flight.FlightUnauthenticatedError as exc: - raise AuthenticationError(f"Authentication failed: {exc}") from exc - except flight.FlightError as exc: - raise QueryError(str(exc)) from exc - return _PyArrowRecordBatchStream(reader) - - def do_put(self, command: bytes, table: pa.Table) -> bytes | None: - """Upload a table via DoPut with a command descriptor.""" - return b"" - - def do_action(self, action: str, body: bytes = b"") -> bytes: - """Invoke a Flight action and return a response.""" - return b"" - - def close(self) -> None: - self._client.close() - - -class _PyArrowRecordBatchStream(RecordBatchStream): - """Adapter for pyarrow RecordBatchReader to implement RecordBatchStream.""" - - def __init__(self, reader: flight.FlightStreamReader) -> None: - self._reader = reader - - def read_all(self) -> pa.Table: - """Read all record batches into a single PyArrow Table.""" - try: - return self._reader.read_all() - except flight.FlightError as exc: - raise QueryError(str(exc)) from exc - - def __iter__(self) -> Iterator[pa.RecordBatch]: - try: - while True: - chunk = self._reader.read_chunk() - yield chunk.data - except StopIteration: - return - except flight.FlightError as exc: - raise QueryError(str(exc)) from exc diff --git a/dataconnect/framework/transport.py b/dataconnect/framework/transport.py deleted file mode 100644 index 09d8986..0000000 --- a/dataconnect/framework/transport.py +++ /dev/null @@ -1,44 +0,0 @@ -"""Abstract transport interface.""" - -from __future__ import annotations - -from abc import ABC, abstractmethod -from collections.abc import Iterator - -import pyarrow as pa - - -class FlightTransport(ABC): - """Minimal abstract transport interface for DataConnect Flight operations.""" - - @abstractmethod - def do_action(self, action: str, body: bytes = b"") -> bytes: - """Invoke a Flight action and return a response.""" - - @abstractmethod - def do_get(self, ticket: bytes) -> RecordBatchStream: - """Open a DoGet stream for ticket.""" - - @abstractmethod - def do_put( - self, - command: bytes, - table: pa.Table, - ) -> bytes | None: - """Upload a table via DoPut with a command descriptor.""" - - @abstractmethod - def close(self) -> None: - """Close the transport connection.""" - - -class RecordBatchStream(ABC): - """Minimal abstract stream interface for Flight record batches.""" - - @abstractmethod - def read_all(self) -> pa.Table: - """Read all record batches into a single PyArrow Table.""" - - @abstractmethod - def __iter__(self) -> Iterator[pa.RecordBatch]: - """Iterate over record batches.""" diff --git a/dataconnect/service/mappers.py b/dataconnect/service/mappers.py index 46a307d..a4a6407 100644 --- a/dataconnect/service/mappers.py +++ b/dataconnect/service/mappers.py @@ -10,16 +10,17 @@ import json from uuid import UUID +from dataconnect.exceptions import NotFoundError from dataconnect.models import Study, StudyEnvironment from dataconnect.transport.models import ResourceInfo -from dataconnect.exceptions import NotFoundError + def resource_to_study(resource: ResourceInfo) -> Study: """Parse a transport-layer ``ResourceInfo`` into a ``Study`` domain object.""" if not resource or not resource.endpoints or not resource.endpoints[0].ticket: raise NotFoundError("Invalid resource: missing endpoints or ticket") - + data = json.loads(resource.endpoints[0].ticket.decode("utf-8")) return Study( diff --git a/dataconnect/transport/arrow_flight/transport.py b/dataconnect/transport/arrow_flight/transport.py index 1476686..cd714ee 100644 --- a/dataconnect/transport/arrow_flight/transport.py +++ b/dataconnect/transport/arrow_flight/transport.py @@ -78,9 +78,7 @@ def list_resources(self, request: ResourceQuery) -> list[ResourceInfo]: if flight_type is None: raise TransportStatusError( - f"Unknown action: {request.action!r}", - status_code=3, - grpc_status="INVALID_ARGUMENT" + f"Unknown action: {request.action!r}", status_code=3, grpc_status="INVALID_ARGUMENT" ) body = json.loads(request.body) if request.body else {}