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/_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/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..a4a6407 100644 --- a/dataconnect/service/mappers.py +++ b/dataconnect/service/mappers.py @@ -10,6 +10,7 @@ import json from uuid import UUID +from dataconnect.exceptions import NotFoundError from dataconnect.models import Study, StudyEnvironment from dataconnect.transport.models import ResourceInfo @@ -17,6 +18,9 @@ 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..cd714ee 100644 --- a/dataconnect/transport/arrow_flight/transport.py +++ b/dataconnect/transport/arrow_flight/transport.py @@ -77,7 +77,9 @@ 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")