From d74d30a47ae2bea298b133c0a7b3cc567d6f0c5e Mon Sep 17 00:00:00 2001 From: Chris Knight Date: Sat, 13 Jun 2026 21:35:28 +0200 Subject: [PATCH 01/11] chore: add event field coverage to discovery script --- plugboard/cli/server/__init__.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/plugboard/cli/server/__init__.py b/plugboard/cli/server/__init__.py index 02d8a696..519e4355 100644 --- a/plugboard/cli/server/__init__.py +++ b/plugboard/cli/server/__init__.py @@ -86,12 +86,14 @@ async def _discover_components(api_url: str, base_cls: type) -> None: outputs = [] input_events = [] output_events = [] + event_field_coverage = {} if io: inputs = list(io.inputs) outputs = list(io.outputs) input_events = [getattr(e, "type", str(e)) for e in io.input_events] output_events = [getattr(e, "type", str(e)) for e in io.output_events] + event_field_coverage = getattr(io, "event_field_coverage", {}) data = { "id": f"{c.__module__}.{c.__qualname__}", @@ -102,6 +104,7 @@ async def _discover_components(api_url: str, base_cls: type) -> None: "outputs": outputs, "input_events": input_events, "output_events": output_events, + "event_field_coverage": event_field_coverage, } await _post_to_api(f"{api_url}/types/component", data) From d75fe798b19ecc8a318cfa5c52fe72181d84c7d7 Mon Sep 17 00:00:00 2001 From: Chris Knight Date: Sat, 15 Aug 2026 16:52:32 +0200 Subject: [PATCH 02/11] fix: add type annotation for event_field_coverage --- plugboard/cli/server/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/plugboard/cli/server/__init__.py b/plugboard/cli/server/__init__.py index 519e4355..5571f1d9 100644 --- a/plugboard/cli/server/__init__.py +++ b/plugboard/cli/server/__init__.py @@ -86,7 +86,7 @@ async def _discover_components(api_url: str, base_cls: type) -> None: outputs = [] input_events = [] output_events = [] - event_field_coverage = {} + event_field_coverage: dict[str, list[str]] = {} if io: inputs = list(io.inputs) From 81460bc22ab1aacf0fd1ab36535f4df0c8142589 Mon Sep 17 00:00:00 2001 From: Chris Knight Date: Sat, 15 Aug 2026 17:06:27 +0200 Subject: [PATCH 03/11] fix: ignore fsspec in license check fsspec license metadata is not detectable by licensecheck, causing CI failures. Add to ignore list since it's BSD-3-Clause licensed. --- pyproject.toml | 3 +++ 1 file changed, 3 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index 777c9af6..50bdf483 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -193,3 +193,6 @@ lines-after-imports = 2 exclude = "tests/*" cc_min = "C" mi_min = "B" + +[tool.licensecheck] +ignore_packages = ["fsspec"] From e14a0fae6b301197b551761442357d49fd5450a1 Mon Sep 17 00:00:00 2001 From: Chris Knight Date: Sat, 15 Aug 2026 18:11:40 +0200 Subject: [PATCH 04/11] feat: Add MessageDataReader and MessageDataWriter base classes Implements issue #102: Base component for external communication via pub/sub message broker infrastructure. New base classes: - MessageDataReader: Abstract base for reading data from message brokers with connection management, reconnection with exponential backoff, retry logic, message acknowledgment, and chunked/buffered reading. - MessageDataWriter: Abstract base for writing data to message brokers with connection management, reconnection, retry logic, and chunked/buffered writing. Concrete implementations: - GCPPubSubDataReader/Writer: Google Cloud PubSub - AWSSQSDataReader/AWSSNSDataWriter: AWS SQS/SNS - KafkaDataReader/Writer: Apache Kafka Also includes: - Message broker exceptions (ConnectionError, TransientError, PermanentError) - Settings for GCP PubSub, AWS, and Kafka - Optional dependencies in pyproject.toml - Proposal document with design rationale - Comprehensive unit tests (72 new tests) --- docs/message-data-reader-writer-proposal.md | 592 ++++++++++++++++++++ plugboard/exceptions/__init__.py | 24 + plugboard/library/__init__.py | 4 + plugboard/library/aws_messaging_io.py | 243 ++++++++ plugboard/library/gcp_pubsub_io.py | 232 ++++++++ plugboard/library/kafka_io.py | 224 ++++++++ plugboard/library/message_reader.py | 266 +++++++++ plugboard/library/message_writer.py | 253 +++++++++ plugboard/utils/settings.py | 42 ++ pyproject.toml | 3 + tests/unit/test_aws_messaging_io.py | 331 +++++++++++ tests/unit/test_gcp_pubsub_io.py | 385 +++++++++++++ tests/unit/test_kafka_io.py | 377 +++++++++++++ tests/unit/test_message_data_reader.py | 366 ++++++++++++ tests/unit/test_message_data_writer.py | 364 ++++++++++++ 15 files changed, 3706 insertions(+) create mode 100644 docs/message-data-reader-writer-proposal.md create mode 100644 plugboard/library/aws_messaging_io.py create mode 100644 plugboard/library/gcp_pubsub_io.py create mode 100644 plugboard/library/kafka_io.py create mode 100644 plugboard/library/message_reader.py create mode 100644 plugboard/library/message_writer.py create mode 100644 tests/unit/test_aws_messaging_io.py create mode 100644 tests/unit/test_gcp_pubsub_io.py create mode 100644 tests/unit/test_kafka_io.py create mode 100644 tests/unit/test_message_data_reader.py create mode 100644 tests/unit/test_message_data_writer.py diff --git a/docs/message-data-reader-writer-proposal.md b/docs/message-data-reader-writer-proposal.md new file mode 100644 index 00000000..db4180ab --- /dev/null +++ b/docs/message-data-reader-writer-proposal.md @@ -0,0 +1,592 @@ +# Proposal: MessageDataReader and MessageDataWriter Base Classes + +## Issue Reference + +[Issue #102: feat: Base component for external communication](https://github.com/plugboard-dev/plugboard/issues/102) + +## Summary + +Develop `MessageDataReader` and `MessageDataWriter` abstract base classes that provide common logic for reading from and writing to pub/sub message broker infrastructure. These are analogous to the existing `DataReader` and `DataWriter` components (which handle chunking/transforming for file access), but focused on message broker communication — including connection management, reconnection, retries, and message acknowledgment. + +Three concrete implementations will be provided: +1. **Google Cloud PubSub** (`GCPPubSubDataReader` / `GCPPubSubDataWriter`) +2. **AWS SNS/SQS** (`AWSSNSQSDataReader` / `AWSSQSDataWriter`) +3. **Apache Kafka** (`KafkaDataReader` / `KafkaDataWriter`) + +--- + +## Design Rationale + +### Why not extend `DataReader`/`DataWriter`? + +The existing `DataReader`/`DataWriter` classes are designed for finite data sources (files, databases) where: +- `_fetch()` raises `NoMoreDataException` when data is exhausted +- Data is read in chunks until the source is depleted +- No connection lifecycle management is needed (connections are per-query) + +Message brokers have fundamentally different semantics: +- Data arrives continuously (no natural "end of data") +- Connections are long-lived and must be managed (connect, reconnect, disconnect) +- Messages require acknowledgment after processing +- Transient failures require retry with exponential backoff + +Therefore, `MessageDataReader`/`MessageDataWriter` will be standalone `Component` subclasses that follow a *similar* pattern to `DataReader`/`DataWriter` (field-based IO, chunking, buffering) but with message-broker-specific lifecycle management. + +### Relationship to existing patterns + +| Pattern | Base Class | Handles | Subclasses implement | +|---------|-----------|---------|---------------------| +| File I/O | `DataReader`/`DataWriter` | Chunking, buffering, field IO | `_fetch()`, `_convert()`, `_save()` | +| WebSocket | `WebsocketBase` | Connection lifecycle, reconnection | `step()` for read/write | +| **Message Broker** | `MessageDataReader`/`MessageDataWriter` | Connection lifecycle, reconnection, retry, chunking, buffering, acknowledgment | `_connect()`, `_disconnect()`, `_receive()`/`_send()`, `_convert()`, `_ack()` | + +--- + +## Interface Design + +### `MessageDataReader` + +```python +class MessageDataReader(Component, ABC): + """Abstract base class for reading data from a pub/sub message broker. + + Provides connection management, reconnection with exponential backoff, + retry logic, message acknowledgment, and chunked/buffered reading + analogous to `DataReader`. + + Subclasses must implement broker-specific methods for connecting, + receiving messages, converting messages to field buffers, and + acknowledging processed messages. + """ + + io = IOController() + + def __init__( + self, + field_names: list[str], + topic: str, + subscription_id: str | None = None, + chunk_size: int | None = None, + max_retries: int = 3, + retry_base_delay: float = 1.0, + retry_max_delay: float = 60.0, + **kwargs: Unpack[ComponentArgsDict], + ) -> None: + """Instantiate the `MessageDataReader`. + + Args: + field_names: The names of the fields to extract from messages. + topic: The topic/queue to read from. + subscription_id: Optional; A subscription ID (required for some brokers like GCP PubSub). + chunk_size: Optional; Number of messages to fetch per batch. + max_retries: Maximum number of retry attempts for transient failures. + retry_base_delay: Base delay in seconds for exponential backoff. + retry_max_delay: Maximum delay in seconds for exponential backoff. + **kwargs: Additional keyword arguments for `Component`. + """ +``` + +#### Abstract methods (implemented by subclasses): + +| Method | Signature | Description | +|--------|-----------|-------------| +| `_connect` | `async def _connect(self) -> None` | Establish connection to the message broker. | +| `_disconnect` | `async def _disconnect(self) -> None` | Close the connection to the message broker. | +| `_receive` | `async def _receive(self) -> list[Any]` | Receive a batch of raw messages from the broker. Should block until at least one message is available or a timeout occurs. Return empty list on timeout. | +| `_convert` | `async def _convert(self, messages: list[Any]) -> dict[str, deque]` | Convert raw messages into a `dict[str, deque]` field buffer. | +| `_ack` | `async def _ack(self, messages: list[Any]) -> None` | Acknowledge successful processing of messages. | + +#### Concrete methods (provided by base class): + +| Method | Description | +|--------|-------------| +| `init()` | Calls `_connect()` with retry logic. Pre-fetches first batch. | +| `step()` | Consumes one record from the buffer. Fetches next batch if buffer empty. Calls `_ack()` on processed messages. | +| `destroy()` | Calls `_disconnect()` to clean up broker connection. | +| `_receive_with_retry()` | Wraps `_receive()` with exponential backoff retry and automatic reconnection. | + +### `MessageDataWriter` + +```python +class MessageDataWriter(Component, ABC): + """Abstract base class for writing data to a pub/sub message broker. + + Provides connection management, reconnection with exponential backoff, + retry logic, and chunked/buffered writing analogous to `DataWriter`. + + Subclasses must implement broker-specific methods for connecting, + sending messages, converting field data to messages, and + broker-specific message formatting. + """ + + io = IOController() + + def __init__( + self, + field_names: list[str], + topic: str, + chunk_size: int | None = None, + max_retries: int = 3, + retry_base_delay: float = 1.0, + retry_max_delay: float = 60.0, + **kwargs: Unpack[ComponentArgsDict], + ) -> None: + """Instantiate the `MessageDataWriter`. + + Args: + field_names: The names of the fields to include in messages. + topic: The topic/queue to write to. + chunk_size: Optional; Number of records to batch into a single message. + max_retries: Maximum number of retry attempts for transient failures. + retry_base_delay: Base delay in seconds for exponential backoff. + retry_max_delay: Maximum delay in seconds for exponential backoff. + **kwargs: Additional keyword arguments for `Component`. + """ +``` + +#### Abstract methods (implemented by subclasses): + +| Method | Signature | Description | +|--------|-----------|-------------| +| `_connect` | `async def _connect(self) -> None` | Establish connection to the message broker. | +| `_disconnect` | `async def _disconnect(self) -> None` | Close the connection to the message broker. | +| `_send` | `async def _send(self, messages: list[Any]) -> None` | Send a batch of messages to the broker. | +| `_convert` | `async def _convert(self, data: dict[str, deque]) -> list[Any]` | Convert field buffer data into broker-specific message format. | + +#### Concrete methods (provided by base class): + +| Method | Description | +|--------|-------------| +| `init()` | Calls `_connect()` with retry logic. | +| `step()` | Buffers input fields. Triggers `_send()` when `chunk_size` reached. | +| `run()` | Runs step loop to completion, then flushes remaining buffered data. | +| `destroy()` | Calls `_disconnect()` to clean up broker connection. | +| `_send_with_retry()` | Wraps `_send()` with exponential backoff retry and automatic reconnection. | + +--- + +## Connection Management & Retry Strategy + +The base classes provide robust connection management: + +### Connection lifecycle + +``` +init() → _connect() [with retry] → ready for step() +step() → _receive_with_retry() / _send_with_retry() → process messages +destroy() → _disconnect() +``` + +### Reconnection with exponential backoff + +```python +async def _receive_with_retry(self) -> list[Any]: + """Receives messages with retry and exponential backoff.""" + last_exception = None + for attempt in range(self._max_retries + 1): + try: + return await self._receive() + except TransientError as e: + last_exception = e + if attempt < self._max_retries: + delay = min( + self._retry_base_delay * (2 ** attempt), + self._retry_max_delay, + ) + self._logger.warning( + "Transient error receiving messages, retrying", + attempt=attempt + 1, + delay=delay, + error=str(e), + ) + await asyncio.sleep(delay) + # Attempt reconnection before retry + await self._reconnect() + raise last_exception # type: ignore[misc] +``` + +### Reconnection strategy + +```python +async def _reconnect(self) -> None: + """Attempts to reconnect to the message broker.""" + self._logger.info("Attempting reconnection to message broker") + try: + await self._disconnect() + except Exception: + pass # Best-effort disconnect + await self._connect() + self._logger.info("Reconnected to message broker") +``` + +--- + +## Concrete Implementations + +### 1. Google Cloud PubSub + +**Dependencies**: `google-cloud-pubsub` (added as optional dependency `gcp-pubsub`) + +#### `GCPPubSubDataReader` + +```python +class GCPPubSubDataReader(MessageDataReader): + """Reads data from Google Cloud PubSub subscription.""" + + def __init__( + self, + project_id: str, + subscription_id: str, + parse_json: bool = True, + **kwargs: Unpack[MessageDataReaderArgsSpec], + ) -> None: + ... + + async def _connect(self) -> None: + # Create AsyncSubscriberClient + # Subscribe to subscription + + async def _disconnect(self) -> None: + # Close subscriber client + + async def _receive(self) -> list[Any]: + # Pull batch of messages (up to chunk_size) + # Return list of PubSubMessage + + async def _convert(self, messages: list[Any]) -> dict[str, deque]: + # Parse message data (JSON or raw bytes) + # Extract fields into dict[str, deque] + + async def _ack(self, messages: list[Any]) -> None: + # Acknowledge messages via subscriber +``` + +#### `GCPPubSubDataWriter` + +```python +class GCPPubSubDataWriter(MessageDataWriter): + """Writes data to Google Cloud PubSub topic.""" + + def __init__( + self, + project_id: str, + topic_id: str, + parse_json: bool = True, + **kwargs: Unpack[MessageDataWriterArgsSpec], + ) -> None: + ... + + async def _connect(self) -> None: + # Create AsyncPublisherClient + + async def _disconnect(self) -> None: + # Close publisher client + + async def _send(self, messages: list[Any]) -> None: + # Publish messages to topic + + async def _convert(self, data: dict[str, deque]) -> list[Any]: + # Convert field data to JSON-encoded bytes +``` + +### 2. AWS SNS/SQS + +**Dependencies**: `aioboto3` or `aws-sdk-pandas` (added as optional dependency `aws-messaging`) + +> **Note**: AWS uses SQS for receiving (queue-based) and SNS for publishing (topic-based). The reader uses SQS; the writer can use either SNS (pub/sub) or SQS (queue). We'll implement both. + +#### `AWSSQSDataReader` + +```python +class AWSSQSDataReader(MessageDataReader): + """Reads data from AWS SQS queue.""" + + def __init__( + self, + queue_url: str, + region: str, + parse_json: bool = True, + wait_time_seconds: int = 20, # Long polling + **kwargs: Unpack[MessageDataReaderArgsSpec], + ) -> None: + ... + + async def _connect(self) -> None: + # Create aioboto3 SQS client + + async def _disconnect(self) -> None: + # Close session + + async def _receive(self) -> list[Any]: + # ReceiveMessage with MaxNumberOfMessages=chunk_size + # Long-polling with WaitTimeSeconds + + async def _convert(self, messages: list[Any]) -> dict[str, deque]: + # Parse message body (JSON) + # Extract fields + + async def _ack(self, messages: list[Any]) -> None: + # DeleteMessage for each processed message +``` + +#### `AWSSNSDataWriter` + +```python +class AWSSNSDataWriter(MessageDataWriter): + """Writes data to AWS SNS topic.""" + + def __init__( + self, + topic_arn: str, + region: str, + parse_json: bool = True, + **kwargs: Unpack[MessageDataWriterArgsSpec], + ) -> None: + ... + + async def _connect(self) -> None: + # Create aioboto3 SNS client + + async def _disconnect(self) -> None: + # Close session + + async def _send(self, messages: list[Any]) -> None: + # Publish each message to SNS topic + + async def _convert(self, data: dict[str, deque]) -> list[Any]: + # Convert field data to JSON strings +``` + +### 3. Apache Kafka + +**Dependencies**: `aiokafka` (added as optional dependency `kafka`) + +#### `KafkaDataReader` + +```python +class KafkaDataReader(MessageDataReader): + """Reads data from Apache Kafka topic.""" + + def __init__( + self, + bootstrap_servers: str | list[str], + topic: str, + group_id: str, + parse_json: bool = True, + **kwargs: Unpack[MessageDataReaderArgsSpec], + ) -> None: + ... + + async def _connect(self) -> None: + # Create AIOKafkaConsumer + # Subscribe to topic + + async def _disconnect(self) -> None: + # Stop consumer + + async def _receive(self) -> list[Any]: + # getmany() with timeout to fetch batch of messages + + async def _convert(self, messages: list[Any]) -> dict[str, deque]: + # Parse message value (JSON or raw bytes) + # Extract fields + + async def _ack(self, messages: list[Any]) -> None: + # Commit offsets for processed messages +``` + +#### `KafkaDataWriter` + +```python +class KafkaDataWriter(MessageDataWriter): + """Writes data to Apache Kafka topic.""" + + def __init__( + self, + bootstrap_servers: str | list[str], + topic: str, + parse_json: bool = True, + **kwargs: Unpack[MessageDataWriterArgsSpec], + ) -> None: + ... + + async def _connect(self) -> None: + # Create AIOKafkaProducer + + async def _disconnect(self) -> None: + # Stop producer + + async def _send(self, messages: list[Any]) -> None: + # send_and_wait for each message + + async def _convert(self, data: dict[str, deque]) -> list[Any]: + # Convert field data to JSON-encoded bytes +``` + +--- + +## Module Structure + +``` +plugboard/library/ +├── __init__.py # Updated exports +├── data_reader.py # Existing DataReader +├── data_writer.py # Existing DataWriter +├── file_io.py # Existing FileReader/FileWriter +├── sql_io.py # Existing SQLReader/SQLWriter +├── websocket_io.py # Existing WebsocketBase/Reader/Writer +├── message_reader.py # NEW: MessageDataReader base class +├── message_writer.py # NEW: MessageDataWriter base class +├── gcp_pubsub_io.py # NEW: GCPPubSubDataReader/Writer +├── aws_messaging_io.py # NEW: AWSSQSDataReader/Writer, AWSSNSDataWriter +└── kafka_io.py # NEW: KafkaDataReader/Writer +``` + +--- + +## Settings & Dependency Injection + +### Settings additions (`utils/settings.py`) + +```python +class _GCPPubSubSettings(BaseSettings): + project_id: str | None = None + model_config = SettingsConfigDict(env_prefix="GCP_PUBSUB_") + +class _AWSSettings(BaseSettings): + region: str | None = None + access_key_id: str | None = None + secret_access_key: str | None = None + model_config = SettingsConfigDict(env_prefix="AWS_") + +class _KafkaSettings(BaseSettings): + bootstrap_servers: str | list[str] | None = None + model_config = SettingsConfigDict(env_prefix="KAFKA_") +``` + +### DI additions (`utils/di.py`) + +No new DI resources are needed initially — each concrete implementation manages its own client lifecycle via `_connect()`/`_disconnect()`. DI resources can be added later when integrating against real infrastructure. + +--- + +## Optional Dependencies (`pyproject.toml`) + +```toml +[project.optional-dependencies] +gcp-pubsub = ["google-cloud-pubsub>=2.25,<3"] +aws-messaging = ["aioboto3>=13.0,<15"] +kafka = ["aiokafka>=0.11,<1"] +``` + +--- + +## Testing Strategy + +### Unit Tests (no cloud infrastructure required) + +For each base class and concrete implementation, we'll create unit tests using mocks: + +1. **`tests/unit/test_message_data_reader.py`**: + - Test `MessageDataReader` base class behavior with a mock implementation + - Test connection lifecycle (init → connect, destroy → disconnect) + - Test retry logic with simulated transient failures + - Test reconnection behavior + - Test chunked reading and buffering + - Test message acknowledgment + - Test field extraction from messages + +2. **`tests/unit/test_message_data_writer.py`**: + - Test `MessageDataWriter` base class behavior with a mock implementation + - Test connection lifecycle + - Test retry logic + - Test chunked writing and buffering + - Test flush on `run()` completion + - Test field data conversion to messages + +3. **`tests/unit/test_gcp_pubsub_io.py`**: + - Test `GCPPubSubDataReader`/`Writer` with mocked `google.cloud.pubsub` clients + - Test connection setup/teardown + - Test message receive/convert/ack + - Test message send/convert + +4. **`tests/unit/test_aws_messaging_io.py`**: + - Test `AWSSQSDataReader`/`AWSSNSDataWriter` with mocked `aioboto3` clients + - Test SQS receive/ack (delete) + - Test SNS publish + - Test long-polling configuration + +5. **`tests/unit/test_kafka_io.py`**: + - Test `KafkaDataReader`/`Writer` with mocked `aiokafka` clients + - Test consumer/producer lifecycle + - Test message receive/convert/commit + - Test message send/convert + +### Integration Tests (require cloud infrastructure — for later) + +Integration tests will be added in `tests/integration/` once cloud infrastructure is set up: +- `tests/integration/test_gcp_pubsub_io.py` +- `tests/integration/test_aws_messaging_io.py` +- `tests/integration/test_kafka_io.py` + +### Test patterns + +Following existing patterns: +- `pytest.mark.asyncio` for async tests +- Mock classes extending the abstract base (like `MockDataReader` in existing tests) +- `pytest.fixture` for test data +- Parametrized tests for chunk_size variations +- `structlog` for test logging + +--- + +## Implementation Order + +1. **Phase 1**: Base classes (`message_reader.py`, `message_writer.py`) + unit tests +2. **Phase 2**: Google Cloud PubSub implementation + unit tests +3. **Phase 3**: AWS SNS/SQS implementation + unit tests +4. **Phase 4**: Kafka implementation + unit tests +5. **Phase 5**: Update `__init__.py` exports, settings, pyproject.toml dependencies +6. **Phase 6**: Integration tests (when cloud infrastructure is available) + +--- + +## Error Handling + +### Custom exceptions + +```python +class MessageBrokerConnectionError(Exception): + """Raised when connection to message broker fails.""" + +class MessageBrokerTransientError(Exception): + """Raised on transient broker errors (eligible for retry).""" + +class MessageBrokerPermanentError(Exception): + """Raised on permanent broker errors (not eligible for retry).""" +``` + +### Error classification + +Each concrete implementation is responsible for classifying broker-specific errors into these categories. The base class handles retry logic based on these classifications. + +--- + +## Serialization + +Messages will be serialized as JSON by default (configurable via `parse_json` flag). This follows the pattern established by `WebsocketReader`/`WebsocketWriter` and ensures interoperability across different broker implementations. + +For the `_convert()` method: +- **Reader**: Parse JSON message data → extract named fields → `dict[str, deque]` +- **Writer**: Take `dict[str, deque]` → combine into records → serialize as JSON + +--- + +## Future Enhancements + +- Dead-letter queue handling +- Message filtering / schema validation +- Metrics collection (message rates, latencies) +- Schema registry integration (Avro, Protobuf) +- DI-managed broker connections (for connection pooling across components) +- Batch acknowledgment optimizations diff --git a/plugboard/exceptions/__init__.py b/plugboard/exceptions/__init__.py index 499a7a05..815778c8 100644 --- a/plugboard/exceptions/__init__.py +++ b/plugboard/exceptions/__init__.py @@ -116,3 +116,27 @@ class ProcessStatusError(Exception): """Raised when a `Process` is in an invalid state for the requested operation.""" pass + + +class MessageBrokerError(Exception): + """Base exception for message broker errors.""" + + pass + + +class MessageBrokerConnectionError(MessageBrokerError): + """Raised when connection to a message broker fails.""" + + pass + + +class MessageBrokerTransientError(MessageBrokerError): + """Raised on transient message broker errors (eligible for retry).""" + + pass + + +class MessageBrokerPermanentError(MessageBrokerError): + """Raised on permanent message broker errors (not eligible for retry).""" + + pass diff --git a/plugboard/library/__init__.py b/plugboard/library/__init__.py index b6909900..51b77ab1 100644 --- a/plugboard/library/__init__.py +++ b/plugboard/library/__init__.py @@ -4,6 +4,8 @@ from .data_writer import DataWriter from .file_io import FileReader, FileWriter from .llm import LLMChat, LLMImageProcessor +from .message_reader import MessageDataReader +from .message_writer import MessageDataWriter from .sql_io import SQLReader, SQLWriter from .websocket_io import WebsocketBase, WebsocketReader, WebsocketWriter @@ -15,6 +17,8 @@ "LLMImageProcessor", "FileReader", "FileWriter", + "MessageDataReader", + "MessageDataWriter", "SQLReader", "SQLWriter", "WebsocketBase", diff --git a/plugboard/library/aws_messaging_io.py b/plugboard/library/aws_messaging_io.py new file mode 100644 index 00000000..cd30d39d --- /dev/null +++ b/plugboard/library/aws_messaging_io.py @@ -0,0 +1,243 @@ +"""Provides `AWSSQSDataReader` and `AWSSNSDataWriter` for AWS SQS/SNS messaging.""" + +from __future__ import annotations + +from collections import deque +import json +import typing as _t + +from plugboard.exceptions import NoMoreDataException +from plugboard.library.message_reader import MessageDataReader, MessageDataReaderArgsDict +from plugboard.library.message_writer import MessageDataWriter, MessageDataWriterArgsDict +from plugboard.utils import depends_on_optional + + +try: + import aioboto3 +except ImportError: # pragma: no cover + pass + + +class AWSSQSDataReaderArgsDict(MessageDataReaderArgsDict): + """Specification of the `AWSSQSDataReader` constructor arguments. + + Attributes: + queue_url: The SQS queue URL. + region: The AWS region. + parse_json: Whether to parse message bodies as JSON. + wait_time_seconds: Long-polling wait time in seconds. + """ + + queue_url: str + region: str + parse_json: _t.NotRequired[bool] + wait_time_seconds: _t.NotRequired[int] + + +class AWSSNSDataWriterArgsDict(MessageDataWriterArgsDict): + """Specification of the `AWSSNSDataWriter` constructor arguments. + + Attributes: + topic_arn: The SNS topic ARN. + region: The AWS region. + parse_json: Whether to encode message data as JSON. + """ + + topic_arn: str + region: str + parse_json: _t.NotRequired[bool] + + +class AWSSQSDataReader(MessageDataReader): + """Reads data from an AWS SQS queue. + + Messages are received from the queue using long-polling and converted + to field values. Messages are deleted from the queue after processing + (acknowledgment). + """ + + @depends_on_optional("aioboto3", extra="aws-messaging") + def __init__( + self, + queue_url: str, + region: str, + parse_json: bool = True, + wait_time_seconds: int = 20, + **kwargs: _t.Unpack[AWSSQSDataReaderArgsDict], + ) -> None: + """Instantiates the `AWSSQSDataReader`. + + Args: + queue_url: The SQS queue URL. + region: The AWS region. + parse_json: Whether to parse message bodies as JSON. + wait_time_seconds: Long-polling wait time in seconds (max 20). + **kwargs: Additional keyword arguments for + [`MessageDataReader`][plugboard.library.MessageDataReader]. + """ + topic = kwargs.pop("topic", queue_url) + super().__init__(topic=topic, **kwargs) + self._queue_url = queue_url + self._region = region + self._parse_json = parse_json + self._wait_time_seconds = wait_time_seconds + self._session: _t.Any = None + self._client: _t.Any = None + + async def _connect(self) -> None: + """Creates an SQS client session.""" + self._session = aioboto3.Session() + self._client_ctx = self._session.client("sqs", region_name=self._region) + self._client = await self._client_ctx.__aenter__() + + async def _disconnect(self) -> None: + """Closes the SQS client session.""" + if self._client is not None: + try: + await self._client_ctx.__aexit__(None, None, None) + except Exception: # noqa: S102 + pass + self._client = None + self._session = None + + async def _receive(self) -> list[_t.Any]: + """Receives a batch of messages from the SQS queue. + + Returns: + A list of SQS message dicts. + + Raises: + NoMoreDataException: If the queue does not exist. + """ + if self._client is None: + raise RuntimeError("SQS client not initialized") + max_messages = min(self._chunk_size or 10, 10) # SQS max is 10 + try: + response = await self._client.receive_message( + QueueUrl=self._queue_url, + MaxNumberOfMessages=max_messages, + WaitTimeSeconds=self._wait_time_seconds, + ) + except Exception as e: + if "QueueDoesNotExist" in str(type(e).__name__) or "NonExistentQueue" in str(e): + raise NoMoreDataException from e + raise + return response.get("Messages", []) + + async def _convert(self, messages: list[_t.Any]) -> dict[str, deque]: + """Converts SQS messages to a field buffer. + + Args: + messages: A list of SQS message dicts. + + Returns: + A dictionary mapping field names to deques of field values. + """ + converted: dict[str, deque] = {field: deque() for field in self.io.outputs} + for msg in messages: + body = msg.get("Body", "") + if self._parse_json: + record = json.loads(body) + else: + record = {"data": body} + for field in self.io.outputs: + converted[field].append(record.get(field)) + return converted + + async def _ack(self, messages: list[_t.Any]) -> None: + """Deletes processed messages from the SQS queue. + + Args: + messages: The SQS message dicts to delete. + """ + if self._client is None: + raise RuntimeError("SQS client not initialized") + for msg in messages: + receipt_handle = msg.get("ReceiptHandle") + if receipt_handle: + await self._client.delete_message( + QueueUrl=self._queue_url, ReceiptHandle=receipt_handle + ) + + +class AWSSNSDataWriter(MessageDataWriter): + """Writes data to an AWS SNS topic. + + Field data is converted to JSON-encoded messages and published + to the specified SNS topic. + """ + + @depends_on_optional("aioboto3", extra="aws-messaging") + def __init__( + self, + topic_arn: str, + region: str, + parse_json: bool = True, + **kwargs: _t.Unpack[AWSSNSDataWriterArgsDict], + ) -> None: + """Instantiates the `AWSSNSDataWriter`. + + Args: + topic_arn: The SNS topic ARN. + region: The AWS region. + parse_json: Whether to encode message data as JSON. + **kwargs: Additional keyword arguments for + [`MessageDataWriter`][plugboard.library.MessageDataWriter]. + """ + topic = kwargs.pop("topic", topic_arn) + super().__init__(topic=topic, **kwargs) + self._topic_arn = topic_arn + self._region = region + self._parse_json = parse_json + self._session: _t.Any = None + self._client: _t.Any = None + + async def _connect(self) -> None: + """Creates an SNS client session.""" + self._session = aioboto3.Session() + self._client_ctx = self._session.client("sns", region_name=self._region) + self._client = await self._client_ctx.__aenter__() + + async def _disconnect(self) -> None: + """Closes the SNS client session.""" + if self._client is not None: + try: + await self._client_ctx.__aexit__(None, None, None) + except Exception: # noqa: S102 + pass + self._client = None + self._session = None + + async def _send(self, messages: list[_t.Any]) -> None: + """Publishes messages to the SNS topic. + + Args: + messages: A list of message strings to publish. + """ + if self._client is None: + raise RuntimeError("SNS client not initialized") + for msg_data in messages: + await self._client.publish( + TopicArn=self._topic_arn, + Message=msg_data, + ) + + async def _convert(self, data: dict[str, deque]) -> list[_t.Any]: + """Converts field buffer data to JSON-encoded message strings. + + Args: + data: A dictionary mapping field names to deques of field values. + + Returns: + A list of message strings ready to publish. + """ + completed_rows = min(len(d) for d in data.values()) if data else 0 + messages: list[str] = [] + for i in range(completed_rows): + record = {field: data[field][i] for field in data} + if self._parse_json: + messages.append(json.dumps(record)) + else: + first_field = next(iter(record.values())) + messages.append(str(first_field)) + return messages diff --git a/plugboard/library/gcp_pubsub_io.py b/plugboard/library/gcp_pubsub_io.py new file mode 100644 index 00000000..57b97efd --- /dev/null +++ b/plugboard/library/gcp_pubsub_io.py @@ -0,0 +1,232 @@ +"""Provides `GCPPubSubDataReader` and `GCPPubSubDataWriter` for Google Cloud PubSub.""" + +from __future__ import annotations + +from collections import deque +import json +import typing as _t + +from plugboard.exceptions import NoMoreDataException +from plugboard.library.message_reader import MessageDataReader, MessageDataReaderArgsDict +from plugboard.library.message_writer import MessageDataWriter, MessageDataWriterArgsDict +from plugboard.utils import depends_on_optional + + +try: + from google.cloud import pubsub_v1 + from google.cloud.pubsub_v1.subscriber.message import Message as PubSubMessage +except ImportError: # pragma: no cover + pass + + +class GCPPubSubDataReaderArgsDict(MessageDataReaderArgsDict): + """Specification of the `GCPPubSubDataReader` constructor arguments. + + Attributes: + project_id: The GCP project ID. + subscription_id: The PubSub subscription ID. + parse_json: Whether to parse message data as JSON. + """ + + project_id: str + subscription_id: str + parse_json: _t.NotRequired[bool] + + +class GCPPubSubDataWriterArgsDict(MessageDataWriterArgsDict): + """Specification of the `GCPPubSubDataWriter` constructor arguments. + + Attributes: + project_id: The GCP project ID. + topic_id: The PubSub topic ID. + parse_json: Whether to encode message data as JSON. + """ + + project_id: str + topic_id: str + parse_json: _t.NotRequired[bool] + + +class GCPPubSubDataReader(MessageDataReader): + """Reads data from a Google Cloud PubSub subscription. + + Messages are pulled from the subscription in batches and converted + to field values. Messages are acknowledged after processing. + """ + + @depends_on_optional("google.cloud.pubsub_v1", extra="gcp-pubsub") + def __init__( + self, + project_id: str, + subscription_id: str, + parse_json: bool = True, + **kwargs: _t.Unpack[GCPPubSubDataReaderArgsDict], + ) -> None: + """Instantiates the `GCPPubSubDataReader`. + + Args: + project_id: The GCP project ID. + subscription_id: The PubSub subscription ID. + parse_json: Whether to parse message data as JSON. + **kwargs: Additional keyword arguments for + [`MessageDataReader`][plugboard.library.MessageDataReader]. + """ + topic = kwargs.pop("topic", f"{project_id}/{subscription_id}") + super().__init__(topic=topic, **kwargs) + self._project_id = project_id + self._subscription_id = subscription_id + self._subscription_path = ( + f"projects/{project_id}/subscriptions/{subscription_id}" + ) + self._parse_json = parse_json + self._subscriber: _t.Optional[pubsub_v1.SubscriberClient] = None + + async def _connect(self) -> None: + """Creates a PubSub subscriber client.""" + self._subscriber = pubsub_v1.SubscriberClient() + + async def _disconnect(self) -> None: + """Closes the PubSub subscriber client.""" + if self._subscriber is not None: + self._subscriber.close() + self._subscriber = None + + async def _receive(self) -> list[_t.Any]: + """Pulls a batch of messages from the PubSub subscription. + + Returns: + A list of PubSub `Message` objects. + + Raises: + NoMoreDataException: If the subscription is deleted or unreachable. + """ + if self._subscriber is None: + raise RuntimeError("Subscriber client not initialized") + max_messages = self._chunk_size or 10 + try: + response = self._subscriber.pull( + request={"subscription": self._subscription_path, "max_messages": max_messages}, + timeout=30.0, + ) + except Exception as e: + if "NOT_FOUND" in str(e) or "Subscription not found" in str(e): + raise NoMoreDataException from e + raise + if not response.received_messages: + return [] + return list(response.received_messages) + + async def _convert(self, messages: list[_t.Any]) -> dict[str, deque]: + """Converts PubSub messages to a field buffer. + + Args: + messages: A list of `ReceivedMessage` objects. + + Returns: + A dictionary mapping field names to deques of field values. + """ + converted: dict[str, deque] = {field: deque() for field in self.io.outputs} + for msg_wrapper in messages: + data = msg_wrapper.message.data + if self._parse_json: + record = json.loads(data.decode("utf-8")) + else: + record = {"data": data} + for field in self.io.outputs: + converted[field].append(record.get(field)) + return converted + + async def _ack(self, messages: list[_t.Any]) -> None: + """Acknowledges processed PubSub messages. + + Args: + messages: The `ReceivedMessage` objects to acknowledge. + """ + if self._subscriber is None: + raise RuntimeError("Subscriber client not initialized") + ack_ids = [msg_wrapper.ack_id for msg_wrapper in messages] + self._subscriber.acknowledge( + request={"subscription": self._subscription_path, "ack_ids": ack_ids} + ) + + +class GCPPubSubDataWriter(MessageDataWriter): + """Writes data to a Google Cloud PubSub topic. + + Field data is converted to JSON-encoded messages and published + to the specified topic. + """ + + @depends_on_optional("google.cloud.pubsub_v1", extra="gcp-pubsub") + def __init__( + self, + project_id: str, + topic_id: str, + parse_json: bool = True, + **kwargs: _t.Unpack[GCPPubSubDataWriterArgsDict], + ) -> None: + """Instantiates the `GCPPubSubDataWriter`. + + Args: + project_id: The GCP project ID. + topic_id: The PubSub topic ID. + parse_json: Whether to encode message data as JSON. + **kwargs: Additional keyword arguments for + [`MessageDataWriter`][plugboard.library.MessageDataWriter]. + """ + topic = kwargs.pop("topic", f"{project_id}/{topic_id}") + super().__init__(topic=topic, **kwargs) + self._project_id = project_id + self._topic_id = topic_id + self._topic_path = f"projects/{project_id}/topics/{topic_id}" + self._parse_json = parse_json + self._publisher: _t.Optional[pubsub_v1.PublisherClient] = None + + async def _connect(self) -> None: + """Creates a PubSub publisher client.""" + self._publisher = pubsub_v1.PublisherClient() + + async def _disconnect(self) -> None: + """Closes the PubSub publisher client.""" + if self._publisher is not None: + self._publisher.close() # type: ignore[no-untyped-call] + self._publisher = None + + async def _send(self, messages: list[_t.Any]) -> None: + """Publishes messages to the PubSub topic. + + Args: + messages: A list of bytes objects to publish. + """ + if self._publisher is None: + raise RuntimeError("Publisher client not initialized") + futures = [] + for msg_data in messages: + future = self._publisher.publish(self._topic_path, data=msg_data) + futures.append(future) + # Wait for all publishes to complete + for future in futures: + future.result(timeout=60.0) + + async def _convert(self, data: dict[str, deque]) -> list[_t.Any]: + """Converts field buffer data to JSON-encoded bytes messages. + + Args: + data: A dictionary mapping field names to deques of field values. + + Returns: + A list of bytes objects ready to publish. + """ + completed_rows = min(len(d) for d in data.values()) if data else 0 + messages: list[bytes] = [] + for i in range(completed_rows): + record = {field: data[field][i] for field in data} + if self._parse_json: + messages.append(json.dumps(record).encode("utf-8")) + else: + # Send raw data from the first field + first_field = next(iter(record.values())) + messages.append( + first_field if isinstance(first_field, bytes) else str(first_field).encode() + ) + return messages diff --git a/plugboard/library/kafka_io.py b/plugboard/library/kafka_io.py new file mode 100644 index 00000000..c8ef5646 --- /dev/null +++ b/plugboard/library/kafka_io.py @@ -0,0 +1,224 @@ +"""Provides `KafkaDataReader` and `KafkaDataWriter` for Apache Kafka messaging.""" + +from __future__ import annotations + +from collections import deque +import json +import typing as _t + +from plugboard.exceptions import NoMoreDataException +from plugboard.library.message_reader import MessageDataReader, MessageDataReaderArgsDict +from plugboard.library.message_writer import MessageDataWriter, MessageDataWriterArgsDict +from plugboard.utils import depends_on_optional + + +try: + from aiokafka import AIOKafkaConsumer, AIOKafkaProducer +except ImportError: # pragma: no cover + pass + + +class KafkaDataReaderArgsDict(MessageDataReaderArgsDict): + """Specification of the `KafkaDataReader` constructor arguments. + + Attributes: + bootstrap_servers: Kafka broker address(es). + group_id: Consumer group ID. + parse_json: Whether to parse message values as JSON. + """ + + bootstrap_servers: _t.Union[str, list[str]] + group_id: str + parse_json: _t.NotRequired[bool] + + +class KafkaDataWriterArgsDict(MessageDataWriterArgsDict): + """Specification of the `KafkaDataWriter` constructor arguments. + + Attributes: + bootstrap_servers: Kafka broker address(es). + parse_json: Whether to encode message values as JSON. + """ + + bootstrap_servers: _t.Union[str, list[str]] + parse_json: _t.NotRequired[bool] + + +class KafkaDataReader(MessageDataReader): + """Reads data from an Apache Kafka topic. + + Messages are consumed from the topic using a consumer group and converted + to field values. Offsets are committed after processing (acknowledgment). + """ + + @depends_on_optional("aiokafka", extra="kafka") + def __init__( + self, + bootstrap_servers: str | list[str], + group_id: str, + parse_json: bool = True, + **kwargs: _t.Unpack[KafkaDataReaderArgsDict], + ) -> None: + """Instantiates the `KafkaDataReader`. + + Args: + bootstrap_servers: Kafka broker address(es) (e.g. `"localhost:9092"`). + group_id: Consumer group ID. + parse_json: Whether to parse message values as JSON. + **kwargs: Additional keyword arguments for + [`MessageDataReader`][plugboard.library.MessageDataReader]. + """ + super().__init__(**kwargs) + if isinstance(bootstrap_servers, str): + bootstrap_servers = [bootstrap_servers] + self._bootstrap_servers = bootstrap_servers + self._group_id = group_id + self._parse_json = parse_json + self._consumer: _t.Optional[AIOKafkaConsumer] = None + + async def _connect(self) -> None: + """Creates and starts a Kafka consumer.""" + self._consumer = AIOKafkaConsumer( + self._topic, + bootstrap_servers=self._bootstrap_servers, + group_id=self._group_id, + auto_offset_reset="earliest", + enable_auto_commit=False, + max_poll_records=self._chunk_size or 10, + ) + await self._consumer.start() + + async def _disconnect(self) -> None: + """Stops and closes the Kafka consumer.""" + if self._consumer is not None: + await self._consumer.stop() + self._consumer = None + + async def _receive(self) -> list[_t.Any]: + """Receives a batch of messages from the Kafka topic. + + Returns: + A list of Kafka `ConsumerRecord` objects. + + Raises: + NoMoreDataException: If the consumer has been closed. + """ + if self._consumer is None: + raise RuntimeError("Kafka consumer not initialized") + max_messages = self._chunk_size or 10 + # Use getmany to fetch a batch with timeout + data = await self._consumer.getmany(timeout_ms=30000, max_records=max_messages) + messages: list[_t.Any] = [] + for _tp, records in data.items(): + messages.extend(records) + if not messages: + raise NoMoreDataException + return messages[:max_messages] + + async def _convert(self, messages: list[_t.Any]) -> dict[str, deque]: + """Converts Kafka messages to a field buffer. + + Args: + messages: A list of `ConsumerRecord` objects. + + Returns: + A dictionary mapping field names to deques of field values. + """ + converted: dict[str, deque] = {field: deque() for field in self.io.outputs} + for record in messages: + value = record.value + if isinstance(value, bytes): + value = value.decode("utf-8") + if self._parse_json: + record_data = json.loads(value) + else: + record_data = {"data": value} + for field in self.io.outputs: + converted[field].append(record_data.get(field)) + return converted + + async def _ack(self, messages: list[_t.Any]) -> None: + """Commits offsets for processed Kafka messages. + + Args: + messages: The `ConsumerRecord` objects to acknowledge. + """ + if self._consumer is None: + raise RuntimeError("Kafka consumer not initialized") + await self._consumer.commit() + + +class KafkaDataWriter(MessageDataWriter): + """Writes data to an Apache Kafka topic. + + Field data is converted to JSON-encoded messages and produced + to the specified Kafka topic. + """ + + @depends_on_optional("aiokafka", extra="kafka") + def __init__( + self, + bootstrap_servers: str | list[str], + parse_json: bool = True, + **kwargs: _t.Unpack[KafkaDataWriterArgsDict], + ) -> None: + """Instantiates the `KafkaDataWriter`. + + Args: + bootstrap_servers: Kafka broker address(es) (e.g. `"localhost:9092"`). + parse_json: Whether to encode message values as JSON. + **kwargs: Additional keyword arguments for + [`MessageDataWriter`][plugboard.library.MessageDataWriter]. + """ + super().__init__(**kwargs) + if isinstance(bootstrap_servers, str): + bootstrap_servers = [bootstrap_servers] + self._bootstrap_servers = bootstrap_servers + self._parse_json = parse_json + self._producer: _t.Optional[AIOKafkaProducer] = None + + async def _connect(self) -> None: + """Creates and starts a Kafka producer.""" + self._producer = AIOKafkaProducer( + bootstrap_servers=self._bootstrap_servers, + ) + await self._producer.start() + + async def _disconnect(self) -> None: + """Stops and closes the Kafka producer.""" + if self._producer is not None: + await self._producer.stop() + self._producer = None + + async def _send(self, messages: list[_t.Any]) -> None: + """Sends messages to the Kafka topic. + + Args: + messages: A list of bytes objects to send. + """ + if self._producer is None: + raise RuntimeError("Kafka producer not initialized") + for msg_data in messages: + await self._producer.send_and_wait(self._topic, value=msg_data) + + async def _convert(self, data: dict[str, deque]) -> list[_t.Any]: + """Converts field buffer data to JSON-encoded bytes messages. + + Args: + data: A dictionary mapping field names to deques of field values. + + Returns: + A list of bytes objects ready to send. + """ + completed_rows = min(len(d) for d in data.values()) if data else 0 + messages: list[bytes] = [] + for i in range(completed_rows): + record = {field: data[field][i] for field in data} + if self._parse_json: + messages.append(json.dumps(record).encode("utf-8")) + else: + first_field = next(iter(record.values())) + messages.append( + first_field if isinstance(first_field, bytes) else str(first_field).encode() + ) + return messages diff --git a/plugboard/library/message_reader.py b/plugboard/library/message_reader.py new file mode 100644 index 00000000..eafd611c --- /dev/null +++ b/plugboard/library/message_reader.py @@ -0,0 +1,266 @@ +"""Provides `MessageDataReader` base class for reading data from pub/sub message brokers.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +import asyncio +from collections import deque +from asyncio.tasks import Task +import typing as _t + +from plugboard.component import Component, IOController +from plugboard.exceptions import IOSetupError, IOStreamClosedError, NoMoreDataException +from plugboard.schemas import ComponentArgsDict + + +class MessageDataReaderArgsDict(ComponentArgsDict): + """Specification of the `MessageDataReader` constructor arguments. + + Attributes: + field_names: The names of the fields to read from messages. + topic: The topic/queue to read from. + chunk_size: Optional; The number of messages to fetch per batch. + max_retries: Maximum number of retry attempts for transient failures. + retry_base_delay: Base delay in seconds for exponential backoff. + retry_max_delay: Maximum delay in seconds for exponential backoff. + """ + + field_names: list[str] + topic: str + chunk_size: _t.NotRequired[int | None] + max_retries: _t.NotRequired[int] + retry_base_delay: _t.NotRequired[float] + retry_max_delay: _t.NotRequired[float] + + +class MessageDataReader(Component, ABC): + """Abstract base class for reading data from a pub/sub message broker. + + Provides connection management, reconnection with exponential backoff, + retry logic, message acknowledgment, and chunked/buffered reading + analogous to [`DataReader`][plugboard.library.DataReader]. + + Subclasses must implement broker-specific methods for connecting, + receiving messages, converting messages to field buffers, and + acknowledging processed messages. + """ + + io = IOController() + + def __init__( + self, + field_names: list[str], + topic: str, + chunk_size: _t.Optional[int] = None, + max_retries: int = 3, + retry_base_delay: float = 1.0, + retry_max_delay: float = 60.0, + **kwargs: _t.Unpack[ComponentArgsDict], + ) -> None: + """Instantiates the `MessageDataReader`. + + Args: + field_names: The names of the fields to extract from messages. + topic: The topic/queue to read from. + chunk_size: Optional; The number of messages to fetch per batch. + max_retries: Maximum number of retry attempts for transient failures. + retry_base_delay: Base delay in seconds for exponential backoff. + retry_max_delay: Maximum delay in seconds for exponential backoff. + **kwargs: Additional keyword arguments for [`Component`][plugboard.component.Component]. + """ + super().__init__(**kwargs) + self._topic = topic + self._buffer: dict[str, deque] = dict() + self._chunk_size = chunk_size + self._max_retries = max_retries + self._retry_base_delay = retry_base_delay + self._retry_max_delay = retry_max_delay + self._pending_ack: list[_t.Any] = [] + self._task: _t.Optional[Task] = None + self.io = IOController( + inputs=None, + outputs=field_names, + input_events=self.__class__.io.input_events, + output_events=self.__class__.io.output_events, + namespace=self.name, + component=self, + ) + + def __init_subclass__(cls, *args: _t.Any, **kwargs: _t.Any) -> None: + try: + return super().__init_subclass__(*args, **kwargs) + except IOSetupError: + # Concrete subclasses of the abstract data io classes represent a special case for io + # setup. They receive io args at run time, not declaration time, so skip error. + pass + + @abstractmethod + async def _connect(self) -> None: + """Establishes connection to the message broker. + + Raises: + MessageBrokerConnectionError: If connection cannot be established. + """ + pass + + @abstractmethod + async def _disconnect(self) -> None: + """Closes the connection to the message broker.""" + pass + + @abstractmethod + async def _receive(self) -> list[_t.Any]: + """Receives a batch of raw messages from the broker. + + Should block until at least one message is available or a timeout occurs. + Returns an empty list on timeout. + + Returns: + A list of raw broker-specific message objects. + + Raises: + NoMoreDataException: If the subscription/source is exhausted and no + more messages will arrive. + """ + pass + + @abstractmethod + async def _convert(self, messages: list[_t.Any]) -> dict[str, deque]: + """Converts raw messages into a `dict[str, deque]` field buffer. + + Args: + messages: Raw broker-specific message objects. + + Returns: + A dictionary mapping field names to deques of field values. + """ + pass + + @abstractmethod + async def _ack(self, messages: list[_t.Any]) -> None: + """Acknowledges successful processing of messages. + + Args: + messages: The raw messages to acknowledge. + """ + pass + + async def _receive_with_retry(self) -> list[_t.Any]: + """Receives messages with exponential backoff retry and reconnection. + + Returns: + A list of raw broker-specific message objects. + + Raises: + NoMoreDataException: If the source is exhausted. + MessageBrokerConnectionError: If all retries are exhausted. + """ + last_exception: _t.Optional[Exception] = None + for attempt in range(self._max_retries + 1): + try: + return await self._receive() + except NoMoreDataException: + raise + except Exception as e: + last_exception = e + if attempt < self._max_retries: + delay = min( + self._retry_base_delay * (2**attempt), + self._retry_max_delay, + ) + self._logger.warning( + "Transient error receiving messages, retrying", + attempt=attempt + 1, + delay=delay, + error=str(e), + ) + await asyncio.sleep(delay) + await self._reconnect() + raise last_exception # type: ignore[misc] + + async def _reconnect(self) -> None: + """Attempts to reconnect to the message broker.""" + self._logger.info("Attempting reconnection to message broker", topic=self._topic) + try: + await self._disconnect() + except Exception: # noqa: S102 + self._logger.warning("Error during disconnect in reconnection", exc_info=True) + await self._connect() + self._logger.info("Reconnected to message broker", topic=self._topic) + + async def _fetch_batch(self) -> None: + """Fetches a batch of messages and updates the internal buffer.""" + if self._task is None: + self._task = asyncio.create_task(self._receive_with_retry()) + messages = await self._task + # Start fetching next batch concurrently + self._task = asyncio.create_task(self._receive_with_retry()) + if len(messages) == 0: + raise NoMoreDataException + new_buffer = await self._convert(messages) + self._buffer = {field_name: new_buffer[field_name] for field_name in self.io.outputs} + self._pending_ack = messages + + def _consume_record(self) -> None: + """Consumes one record from the buffer and sets field attributes.""" + for field in self.io.outputs: + setattr(self, field, self._buffer[field].popleft()) + + async def _ack_pending(self) -> None: + """Acknowledges all pending messages.""" + if self._pending_ack: + await self._ack(self._pending_ack) + self._pending_ack = [] + + async def init(self) -> None: + """Initialises the `MessageDataReader`. + + Connects to the message broker and pre-fetches the first batch of messages. + If no messages are available, the reader will raise `IOStreamClosedError` + on the first `step()` call. + """ + await self._connect() + self._logger.info("Connected to message broker", topic=self._topic) + try: + await self._fetch_batch() + except NoMoreDataException: + # No messages available at init time; step() will raise IOStreamClosedError + pass + + async def step(self) -> None: + """Reads data from the message broker and updates outputs. + + Consumes one record from the buffer. If the buffer is empty, + fetches the next batch. Acknowledges processed messages. + + Raises: + IOStreamClosedError: If there is no more data to read. + """ + if not self._buffer: + # Buffer was never populated (e.g. empty source at init) + await self.io.close() + raise IOStreamClosedError("No more messages from broker") + try: + self._consume_record() + await self._ack_pending() + except IndexError: + try: + await self._fetch_batch() + self._consume_record() + await self._ack_pending() + except NoMoreDataException: + await self.io.close() + raise IOStreamClosedError("No more messages from broker") + + async def destroy(self) -> None: + """Destroys the `MessageDataReader` and disconnects from the broker.""" + if self._task is not None: + self._task.cancel() + try: + await self._task + except (asyncio.CancelledError, Exception): + pass + self._task = None + await self._disconnect() + self._logger.info("Disconnected from message broker", topic=self._topic) + await super().destroy() diff --git a/plugboard/library/message_writer.py b/plugboard/library/message_writer.py new file mode 100644 index 00000000..6e1ee7a3 --- /dev/null +++ b/plugboard/library/message_writer.py @@ -0,0 +1,253 @@ +"""Provides `MessageDataWriter` base class for writing data to pub/sub message brokers.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +import asyncio +from collections import defaultdict, deque +from asyncio.tasks import Task +import typing as _t + +from plugboard.component import Component, IOController +from plugboard.exceptions import IOSetupError +from plugboard.schemas import ComponentArgsDict + + +class MessageDataWriterArgsDict(ComponentArgsDict): + """Specification of the `MessageDataWriter` constructor arguments. + + Attributes: + field_names: The names of the fields to include in messages. + topic: The topic/queue to write to. + chunk_size: Optional; The number of records to batch into messages. + max_retries: Maximum number of retry attempts for transient failures. + retry_base_delay: Base delay in seconds for exponential backoff. + retry_max_delay: Maximum delay in seconds for exponential backoff. + """ + + field_names: list[str] + topic: str + chunk_size: _t.NotRequired[int | None] + max_retries: _t.NotRequired[int] + retry_base_delay: _t.NotRequired[float] + retry_max_delay: _t.NotRequired[float] + + +class MessageDataWriter(Component, ABC): + """Abstract base class for writing data to a pub/sub message broker. + + Provides connection management, reconnection with exponential backoff, + retry logic, and chunked/buffered writing analogous to + [`DataWriter`][plugboard.library.DataWriter]. + + Subclasses must implement broker-specific methods for connecting, + sending messages, and converting field data to broker-specific + message format. + """ + + io = IOController() + + def __init__( + self, + field_names: list[str], + topic: str, + chunk_size: _t.Optional[int] = None, + max_retries: int = 3, + retry_base_delay: float = 1.0, + retry_max_delay: float = 60.0, + **kwargs: _t.Unpack[ComponentArgsDict], + ) -> None: + """Instantiates the `MessageDataWriter`. + + Args: + field_names: The names of the fields to include in messages. + topic: The topic/queue to write to. + chunk_size: Optional; The number of records to batch into a single send operation. + max_retries: Maximum number of retry attempts for transient failures. + retry_base_delay: Base delay in seconds for exponential backoff. + retry_max_delay: Maximum delay in seconds for exponential backoff. + **kwargs: Additional keyword arguments for [`Component`][plugboard.component.Component]. + """ + super().__init__(**kwargs) + self._topic = topic + self._buffer: dict[str, deque] = defaultdict(deque) + self._chunk_size = chunk_size + self._max_retries = max_retries + self._retry_base_delay = retry_base_delay + self._retry_max_delay = retry_max_delay + self._task: _t.Optional[Task] = None + self.io = IOController( + inputs=field_names, + outputs=None, + input_events=self.__class__.io.input_events, + output_events=self.__class__.io.output_events, + event_field_coverage=self.__class__.io.event_field_coverage, + namespace=self.name, + component=self, + ) + + def __init_subclass__(cls, *args: _t.Any, **kwargs: _t.Any) -> None: + try: + return super().__init_subclass__(*args, **kwargs) + except IOSetupError: + # Concrete subclasses of the abstract data io classes represent a special case for io + # setup. They receive io args at run time, not declaration time, so skip error. + pass + + @abstractmethod + async def _connect(self) -> None: + """Establishes connection to the message broker. + + Raises: + MessageBrokerConnectionError: If connection cannot be established. + """ + pass + + @abstractmethod + async def _disconnect(self) -> None: + """Closes the connection to the message broker.""" + pass + + @abstractmethod + async def _send(self, messages: list[_t.Any]) -> None: + """Sends a batch of messages to the broker. + + Args: + messages: A list of broker-specific message objects to send. + + Raises: + MessageBrokerConnectionError: If messages cannot be sent. + """ + pass + + @abstractmethod + async def _convert(self, data: dict[str, deque]) -> list[_t.Any]: + """Converts field buffer data into broker-specific message format. + + Args: + data: A dictionary mapping field names to deques of field values. + + Returns: + A list of broker-specific message objects ready to send. + """ + pass + + async def _send_with_retry(self, messages: list[_t.Any]) -> None: + """Sends messages with exponential backoff retry and reconnection. + + Args: + messages: The messages to send. + + Raises: + Exception: If all retries are exhausted. + """ + last_exception: _t.Optional[Exception] = None + for attempt in range(self._max_retries + 1): + try: + await self._send(messages) + return + except Exception as e: + last_exception = e + if attempt < self._max_retries: + delay = min( + self._retry_base_delay * (2**attempt), + self._retry_max_delay, + ) + self._logger.warning( + "Transient error sending messages, retrying", + attempt=attempt + 1, + delay=delay, + error=str(e), + ) + await asyncio.sleep(delay) + await self._reconnect() + raise last_exception # type: ignore[misc] + + async def _reconnect(self) -> None: + """Attempts to reconnect to the message broker.""" + self._logger.info("Attempting reconnection to message broker", topic=self._topic) + try: + await self._disconnect() + except Exception: # noqa: S102 + self._logger.warning("Error during disconnect in reconnection", exc_info=True) + await self._connect() + self._logger.info("Reconnected to message broker", topic=self._topic) + + def _bind_inputs(self) -> None: + """Binds input fields to component fields and appends to internal buffer.""" + super()._bind_inputs() + for field in self._field_inputs: + value = getattr(self, field, None) + self._buffer[field].append(value) + + @property + def _completed_rows(self) -> int: + """Calculates how many fully formed rows exist in the buffer.""" + if not self.io.inputs: + return 0 + return min((len(self._buffer[f]) for f in self.io.inputs), default=0) + + @property + def _can_step(self) -> bool: + """We can step if we have at least one fully formed row.""" + return self._completed_rows > 0 + + async def _send_batch(self) -> None: + """Sends completed data rows from the buffer.""" + completed_rows = self._completed_rows + if completed_rows == 0: + return + + if self._task is not None: + await self._task + + # Extract only the completed rows into a new chunk + chunk_data: dict[str, deque] = { + field: deque([self._buffer[field].popleft() for _ in range(completed_rows)]) + for field in self.io.inputs + } + + messages = await self._convert(chunk_data) + self._task = asyncio.create_task(self._send_with_retry(messages)) + + async def init(self) -> None: + """Initialises the `MessageDataWriter`. + + Connects to the message broker. + """ + await self._connect() + self._logger.info("Connected to message broker", topic=self._topic) + + async def step(self) -> None: + """Triggers send when buffer is at target size. + + If `chunk_size` is set and the buffer has reached that size, + sends the buffered data as messages. + """ + if self._chunk_size and self._completed_rows >= self._chunk_size: + await self._send_batch() + + async def run(self) -> None: + """Runs the `MessageDataWriter`. + + Steps until all input is consumed, then flushes any remaining + buffered data. + """ + await super().run() + # Flush any remaining data in the buffer after completion + await self._send_batch() + if self._task is not None: + await self._task + + async def destroy(self) -> None: + """Destroys the `MessageDataWriter` and disconnects from the broker.""" + if self._task is not None: + self._task.cancel() + try: + await self._task + except (asyncio.CancelledError, Exception): + pass + self._task = None + await self._disconnect() + self._logger.info("Disconnected from message broker", topic=self._topic) + await super().destroy() diff --git a/plugboard/utils/settings.py b/plugboard/utils/settings.py index 613f02a2..aa94225f 100644 --- a/plugboard/utils/settings.py +++ b/plugboard/utils/settings.py @@ -57,6 +57,42 @@ class _RedisSettings(BaseSettings): url: _t.Optional[str] = None +class _GCPPubSubSettings(BaseSettings): + """Google Cloud PubSub settings. + + Attributes: + project_id: The GCP project ID for PubSub. + """ + + model_config = SettingsConfigDict(env_prefix="GCP_PUBSUB_") + + project_id: _t.Optional[str] = None + + +class _AWSSettings(BaseSettings): + """AWS settings for SNS/SQS messaging. + + Attributes: + region: The default AWS region. + """ + + model_config = SettingsConfigDict(env_prefix="AWS_") + + region: _t.Optional[str] = None + + +class _KafkaSettings(BaseSettings): + """Apache Kafka settings. + + Attributes: + bootstrap_servers: Kafka broker address(es). + """ + + model_config = SettingsConfigDict(env_prefix="KAFKA_") + + bootstrap_servers: _t.Optional[str] = None + + class Settings(BaseSettings): """Settings for Plugboard. @@ -69,6 +105,9 @@ class Settings(BaseSettings): status checks. rabbitmq: RabbitMQ settings. redis: Redis settings. + gcp_pubsub: Google Cloud PubSub settings. + aws: AWS settings for SNS/SQS messaging. + kafka: Apache Kafka settings. """ model_config = SettingsConfigDict(env_prefix=_ENV_PREFIX) @@ -80,3 +119,6 @@ class Settings(BaseSettings): rabbitmq: _RabbitMQSettings = Field(default_factory=_RabbitMQSettings) redis: _RedisSettings = Field(default_factory=_RedisSettings) + gcp_pubsub: _GCPPubSubSettings = Field(default_factory=_GCPPubSubSettings) + aws: _AWSSettings = Field(default_factory=_AWSSettings) + kafka: _KafkaSettings = Field(default_factory=_KafkaSettings) diff --git a/pyproject.toml b/pyproject.toml index 50bdf483..f81d4c46 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -43,8 +43,11 @@ dependencies = [ [project.optional-dependencies] aws = ["s3fs>=2024.9.0"] +aws-messaging = ["aioboto3>=13.0,<15"] azure = ["adlfs>=2024.7.0"] gcp = ["gcsfs>=2024.9.0"] +gcp-pubsub = ["google-cloud-pubsub>=2.25,<3"] +kafka = ["aiokafka>=0.11,<1"] llm = [ "llama-index-core>=0.12.30,<1", "llama-index-llms-openai>=0.3.33,<1", diff --git a/tests/unit/test_aws_messaging_io.py b/tests/unit/test_aws_messaging_io.py new file mode 100644 index 00000000..ea944ab5 --- /dev/null +++ b/tests/unit/test_aws_messaging_io.py @@ -0,0 +1,331 @@ +"""Unit tests for AWS SQS/SNS message data reader/writer.""" + +from __future__ import annotations + +import importlib.machinery +import json +import sys +import typing as _t +from collections import deque +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from plugboard.exceptions import NoMoreDataException + + +# --------------------------------------------------------------------------- +# Mock the aioboto3 module before importing the implementation +# --------------------------------------------------------------------------- + + +def _make_mock_module(name: str) -> MagicMock: + """Creates a mock module with __spec__ set for find_spec compatibility.""" + mock = MagicMock() + mock.__spec__ = importlib.machinery.ModuleSpec(name, None) + return mock + + +_mock_aioboto3 = _make_mock_module("aioboto3") +_mock_aioboto3_session = MagicMock() +_mock_aioboto3.Session.return_value = _mock_aioboto3_session + +sys.modules.setdefault("aioboto3", _mock_aioboto3) + +from plugboard.library.aws_messaging_io import AWSSQSDataReader, AWSSNSDataWriter # noqa: E402 + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_sqs_message(body: dict[str, _t.Any] | str) -> dict[str, _t.Any]: + """Creates a mock SQS message dict.""" + if isinstance(body, dict): + body_str = json.dumps(body) + else: + body_str = body + return { + "MessageId": f"msg-{id(body)}", + "ReceiptHandle": f"receipt-{id(body)}", + "Body": body_str, + } + + +def _setup_mock_client() -> tuple[AsyncMock, AsyncMock]: + """Sets up a mock boto3 client with async context manager.""" + mock_client = AsyncMock() + mock_client_ctx = AsyncMock() + mock_client_ctx.__aenter__ = AsyncMock(return_value=mock_client) + mock_client_ctx.__aexit__ = AsyncMock(return_value=None) + _mock_aioboto3_session.client.return_value = mock_client_ctx + return mock_client, mock_client_ctx + + +# --------------------------------------------------------------------------- +# Tests: AWSSQSDataReader +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_aws_sqs_reader_connect() -> None: + """Tests that the reader creates an SQS client on connect.""" + mock_client, _ = _setup_mock_client() + + reader = AWSSQSDataReader( + name="test-sqs-reader", + field_names=["x", "y"], + topic="test-queue", + queue_url="https://sqs.us-east-1.amazonaws.com/123456789/test-queue", + region="us-east-1", + ) + await reader._connect() + + _mock_aioboto3_session.client.assert_called_with("sqs", region_name="us-east-1") + assert reader._client is mock_client + + +@pytest.mark.asyncio +async def test_aws_sqs_reader_disconnect() -> None: + """Tests that the reader closes the SQS client on disconnect.""" + mock_client, mock_client_ctx = _setup_mock_client() + + reader = AWSSQSDataReader( + name="test-sqs-reader", + field_names=["x"], + topic="test-queue", + queue_url="https://sqs.us-east-1.amazonaws.com/123456789/test-queue", + region="us-east-1", + ) + await reader._connect() + await reader._disconnect() + + mock_client_ctx.__aexit__.assert_called_once() + assert reader._client is None + + +@pytest.mark.asyncio +async def test_aws_sqs_reader_receive() -> None: + """Tests receiving messages from SQS.""" + mock_client, _ = _setup_mock_client() + + test_data = [{"x": 1, "y": "a"}, {"x": 2, "y": "b"}] + sqs_messages = [_make_sqs_message(d) for d in test_data] + mock_client.receive_message = AsyncMock(return_value={"Messages": sqs_messages}) + + reader = AWSSQSDataReader( + name="test-sqs-reader", + field_names=["x", "y"], + topic="test-queue", + queue_url="https://sqs.us-east-1.amazonaws.com/123456789/test-queue", + region="us-east-1", + chunk_size=10, + ) + await reader._connect() + messages = await reader._receive() + + assert len(messages) == 2 + mock_client.receive_message.assert_called() + + +@pytest.mark.asyncio +async def test_aws_sqs_reader_receive_empty() -> None: + """Tests receiving empty response from SQS.""" + mock_client, _ = _setup_mock_client() + mock_client.receive_message = AsyncMock(return_value={}) + + reader = AWSSQSDataReader( + name="test-sqs-reader", + field_names=["x"], + topic="test-queue", + queue_url="https://sqs.us-east-1.amazonaws.com/123456789/test-queue", + region="us-east-1", + ) + await reader._connect() + messages = await reader._receive() + + assert messages == [] + + +@pytest.mark.asyncio +async def test_aws_sqs_reader_convert_json() -> None: + """Tests converting JSON SQS messages to field buffer.""" + reader = AWSSQSDataReader( + name="test-sqs-reader", + field_names=["x", "y"], + topic="test-queue", + queue_url="https://sqs.us-east-1.amazonaws.com/123456789/test-queue", + region="us-east-1", + parse_json=True, + ) + + sqs_messages = [ + _make_sqs_message({"x": 1, "y": "a"}), + _make_sqs_message({"x": 2, "y": "b"}), + ] + result = await reader._convert(sqs_messages) + assert list(result["x"]) == [1, 2] + assert list(result["y"]) == ["a", "b"] + + +@pytest.mark.asyncio +async def test_aws_sqs_reader_convert_raw() -> None: + """Tests converting raw SQS messages to field buffer.""" + reader = AWSSQSDataReader( + name="test-sqs-reader", + field_names=["data"], + topic="test-queue", + queue_url="https://sqs.us-east-1.amazonaws.com/123456789/test-queue", + region="us-east-1", + parse_json=False, + ) + + sqs_messages = [_make_sqs_message("raw-data-1"), _make_sqs_message("raw-data-2")] + result = await reader._convert(sqs_messages) + assert list(result["data"]) == ["raw-data-1", "raw-data-2"] + + +@pytest.mark.asyncio +async def test_aws_sqs_reader_ack() -> None: + """Tests acknowledging (deleting) SQS messages.""" + mock_client, _ = _setup_mock_client() + mock_client.delete_message = AsyncMock() + + reader = AWSSQSDataReader( + name="test-sqs-reader", + field_names=["x"], + topic="test-queue", + queue_url="https://sqs.us-east-1.amazonaws.com/123456789/test-queue", + region="us-east-1", + ) + await reader._connect() + + sqs_messages = [_make_sqs_message({"x": 1})] + await reader._ack(sqs_messages) + + mock_client.delete_message.assert_called() + + +@pytest.mark.asyncio +async def test_aws_sqs_reader_long_polling() -> None: + """Tests that long polling is configured correctly.""" + mock_client, _ = _setup_mock_client() + mock_client.receive_message = AsyncMock(return_value={}) + + reader = AWSSQSDataReader( + name="test-sqs-reader", + field_names=["x"], + topic="test-queue", + queue_url="https://sqs.us-east-1.amazonaws.com/123456789/test-queue", + region="us-east-1", + wait_time_seconds=15, + ) + await reader._connect() + await reader._receive() + + call_args = mock_client.receive_message.call_args + assert call_args[1]["WaitTimeSeconds"] == 15 + + +# --------------------------------------------------------------------------- +# Tests: AWSSNSDataWriter +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_aws_sns_writer_connect() -> None: + """Tests that the writer creates an SNS client on connect.""" + mock_client, _ = _setup_mock_client() + + writer = AWSSNSDataWriter( + name="test-sns-writer", + field_names=["x"], + topic="test-topic", + topic_arn="arn:aws:sns:us-east-1:123456789:test-topic", + region="us-east-1", + ) + await writer._connect() + + _mock_aioboto3_session.client.assert_called_with("sns", region_name="us-east-1") + assert writer._client is mock_client + + +@pytest.mark.asyncio +async def test_aws_sns_writer_disconnect() -> None: + """Tests that the writer closes the SNS client on disconnect.""" + mock_client, mock_client_ctx = _setup_mock_client() + + writer = AWSSNSDataWriter( + name="test-sns-writer", + field_names=["x"], + topic="test-topic", + topic_arn="arn:aws:sns:us-east-1:123456789:test-topic", + region="us-east-1", + ) + await writer._connect() + await writer._disconnect() + + mock_client_ctx.__aexit__.assert_called() + assert writer._client is None + + +@pytest.mark.asyncio +async def test_aws_sns_writer_send() -> None: + """Tests sending messages to SNS.""" + mock_client, _ = _setup_mock_client() + mock_client.publish = AsyncMock() + + writer = AWSSNSDataWriter( + name="test-sns-writer", + field_names=["x"], + topic="test-topic", + topic_arn="arn:aws:sns:us-east-1:123456789:test-topic", + region="us-east-1", + ) + await writer._connect() + + messages = ['{"x": 1}', '{"x": 2}'] + await writer._send(messages) + + assert mock_client.publish.call_count == 2 + + +@pytest.mark.asyncio +async def test_aws_sns_writer_convert_json() -> None: + """Tests converting field data to JSON messages.""" + writer = AWSSNSDataWriter( + name="test-sns-writer", + field_names=["x", "y"], + topic="test-topic", + topic_arn="arn:aws:sns:us-east-1:123456789:test-topic", + region="us-east-1", + parse_json=True, + ) + + data = {"x": deque([1, 2]), "y": deque(["a", "b"])} + messages = await writer._convert(data) + + assert len(messages) == 2 + assert json.loads(messages[0]) == {"x": 1, "y": "a"} + assert json.loads(messages[1]) == {"x": 2, "y": "b"} + + +@pytest.mark.asyncio +async def test_aws_sns_writer_convert_raw() -> None: + """Tests converting field data to raw string messages.""" + writer = AWSSNSDataWriter( + name="test-sns-writer", + field_names=["data"], + topic="test-topic", + topic_arn="arn:aws:sns:us-east-1:123456789:test-topic", + region="us-east-1", + parse_json=False, + ) + + data = {"data": deque(["raw1", "raw2"])} + messages = await writer._convert(data) + + assert len(messages) == 2 + assert messages[0] == "raw1" + assert messages[1] == "raw2" diff --git a/tests/unit/test_gcp_pubsub_io.py b/tests/unit/test_gcp_pubsub_io.py new file mode 100644 index 00000000..2966d9a8 --- /dev/null +++ b/tests/unit/test_gcp_pubsub_io.py @@ -0,0 +1,385 @@ +"""Unit tests for GCP PubSub message data reader/writer.""" + +from __future__ import annotations + +import importlib.machinery +import json +import sys +import typing as _t +from collections import deque +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from plugboard.exceptions import NoMoreDataException + + +# --------------------------------------------------------------------------- +# Mock the google.cloud.pubsub_v1 module before importing the implementation +# --------------------------------------------------------------------------- + + +def _make_mock_module(name: str) -> MagicMock: + """Creates a mock module with __spec__ set for find_spec compatibility.""" + mock = MagicMock() + mock.__spec__ = importlib.machinery.ModuleSpec(name, None) + return mock + + +_mock_pubsub = _make_mock_module("google.cloud.pubsub_v1") +_mock_pubsub.SubscriberClient = MagicMock() +_mock_pubsub.PublisherClient = MagicMock() + +_mock_google = _make_mock_module("google") +_mock_google_cloud = _make_mock_module("google.cloud") +# Wire up the attribute chain so `from google.cloud import pubsub_v1` works +_mock_google_cloud.pubsub_v1 = _mock_pubsub +_mock_google.cloud = _mock_google_cloud + +_mock_modules = { + "google": _mock_google, + "google.cloud": _mock_google_cloud, + "google.cloud.pubsub_v1": _mock_pubsub, + "google.cloud.pubsub_v1.subscriber": _make_mock_module("google.cloud.pubsub_v1.subscriber"), + "google.cloud.pubsub_v1.subscriber.message": _make_mock_module( + "google.cloud.pubsub_v1.subscriber.message" + ), +} + +# Install mocks before importing the module under test +for _mod_name, _mod in _mock_modules.items(): + sys.modules.setdefault(_mod_name, _mod) + +from plugboard.library.gcp_pubsub_io import GCPPubSubDataReader, GCPPubSubDataWriter # noqa: E402 + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_pubsub_message(data: dict[str, _t.Any] | bytes) -> MagicMock: + """Creates a mock PubSub ReceivedMessage.""" + if isinstance(data, dict): + raw_data = json.dumps(data).encode("utf-8") + else: + raw_data = data + msg = MagicMock() + msg.message.data = raw_data + msg.ack_id = f"ack-{id(msg)}" + return msg + + +def _make_pull_response(messages: list[MagicMock]) -> MagicMock: + """Creates a mock Pull response.""" + response = MagicMock() + response.received_messages = messages + return response + + +# --------------------------------------------------------------------------- +# Tests: GCPPubSubDataReader +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_gcp_pubsub_reader_connect() -> None: + """Tests that the reader creates a subscriber client on connect.""" + mock_subscriber = MagicMock() + _mock_pubsub.SubscriberClient.return_value = mock_subscriber + + reader = GCPPubSubDataReader( + name="test-gcp-reader", + field_names=["x", "y"], + topic="test-topic", + project_id="test-project", + subscription_id="test-sub", + ) + await reader._connect() + + _mock_pubsub.SubscriberClient.assert_called() + assert reader._subscriber is mock_subscriber + + +@pytest.mark.asyncio +async def test_gcp_pubsub_reader_disconnect() -> None: + """Tests that the reader closes the subscriber client on disconnect.""" + mock_subscriber = MagicMock() + _mock_pubsub.SubscriberClient.return_value = mock_subscriber + + reader = GCPPubSubDataReader( + name="test-gcp-reader", + field_names=["x"], + topic="test-topic", + project_id="test-project", + subscription_id="test-sub", + ) + await reader._connect() + await reader._disconnect() + + mock_subscriber.close.assert_called_once() + assert reader._subscriber is None + + +@pytest.mark.asyncio +async def test_gcp_pubsub_reader_receive() -> None: + """Tests receiving messages from PubSub.""" + mock_subscriber = MagicMock() + _mock_pubsub.SubscriberClient.return_value = mock_subscriber + + test_data = [{"x": 1, "y": "a"}, {"x": 2, "y": "b"}] + mock_messages = [_make_pubsub_message(d) for d in test_data] + mock_response = _make_pull_response(mock_messages) + mock_subscriber.pull.return_value = mock_response + + reader = GCPPubSubDataReader( + name="test-gcp-reader", + field_names=["x", "y"], + topic="test-topic", + project_id="test-project", + subscription_id="test-sub", + chunk_size=10, + ) + await reader._connect() + messages = await reader._receive() + + assert len(messages) == 2 + mock_subscriber.pull.assert_called() + + +@pytest.mark.asyncio +async def test_gcp_pubsub_reader_receive_empty() -> None: + """Tests receiving empty response from PubSub.""" + mock_subscriber = MagicMock() + _mock_pubsub.SubscriberClient.return_value = mock_subscriber + + mock_response = _make_pull_response([]) + mock_subscriber.pull.return_value = mock_response + + reader = GCPPubSubDataReader( + name="test-gcp-reader", + field_names=["x"], + topic="test-topic", + project_id="test-project", + subscription_id="test-sub", + ) + await reader._connect() + messages = await reader._receive() + + assert messages == [] + + +@pytest.mark.asyncio +async def test_gcp_pubsub_reader_receive_not_found() -> None: + """Tests that NOT_FOUND error raises NoMoreDataException.""" + mock_subscriber = MagicMock() + _mock_pubsub.SubscriberClient.return_value = mock_subscriber + mock_subscriber.pull.side_effect = Exception("NOT_FOUND: Subscription deleted") + + reader = GCPPubSubDataReader( + name="test-gcp-reader", + field_names=["x"], + topic="test-topic", + project_id="test-project", + subscription_id="test-sub", + ) + await reader._connect() + + with pytest.raises(NoMoreDataException): + await reader._receive() + + +@pytest.mark.asyncio +async def test_gcp_pubsub_reader_convert_json() -> None: + """Tests converting JSON messages to field buffer.""" + reader = GCPPubSubDataReader( + name="test-gcp-reader", + field_names=["x", "y"], + topic="test-topic", + project_id="test-project", + subscription_id="test-sub", + parse_json=True, + ) + + test_data = [{"x": 1, "y": "a"}, {"x": 2, "y": "b"}] + mock_messages = [_make_pubsub_message(d) for d in test_data] + + result = await reader._convert(mock_messages) + assert list(result["x"]) == [1, 2] + assert list(result["y"]) == ["a", "b"] + + +@pytest.mark.asyncio +async def test_gcp_pubsub_reader_convert_raw() -> None: + """Tests converting raw bytes messages to field buffer.""" + reader = GCPPubSubDataReader( + name="test-gcp-reader", + field_names=["data"], + topic="test-topic", + project_id="test-project", + subscription_id="test-sub", + parse_json=False, + ) + + mock_messages = [_make_pubsub_message(b"raw-data-1"), _make_pubsub_message(b"raw-data-2")] + result = await reader._convert(mock_messages) + assert list(result["data"]) == [b"raw-data-1", b"raw-data-2"] + + +@pytest.mark.asyncio +async def test_gcp_pubsub_reader_ack() -> None: + """Tests acknowledging messages.""" + mock_subscriber = MagicMock() + _mock_pubsub.SubscriberClient.return_value = mock_subscriber + + reader = GCPPubSubDataReader( + name="test-gcp-reader", + field_names=["x"], + topic="test-topic", + project_id="test-project", + subscription_id="test-sub", + ) + await reader._connect() + + mock_messages = [_make_pubsub_message({"x": 1})] + mock_messages[0].ack_id = "ack-123" + await reader._ack(mock_messages) + + mock_subscriber.acknowledge.assert_called() + call_args = mock_subscriber.acknowledge.call_args + assert call_args[1]["request"]["ack_ids"] == ["ack-123"] + + +@pytest.mark.asyncio +async def test_gcp_pubsub_reader_subscription_path() -> None: + """Tests that the subscription path is constructed correctly.""" + reader = GCPPubSubDataReader( + name="test-gcp-reader", + field_names=["x"], + topic="test-topic", + project_id="my-project", + subscription_id="my-sub", + ) + assert reader._subscription_path == "projects/my-project/subscriptions/my-sub" + + +# --------------------------------------------------------------------------- +# Tests: GCPPubSubDataWriter +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_gcp_pubsub_writer_connect() -> None: + """Tests that the writer creates a publisher client on connect.""" + mock_publisher = MagicMock() + _mock_pubsub.PublisherClient.return_value = mock_publisher + + writer = GCPPubSubDataWriter( + name="test-gcp-writer", + field_names=["x"], + topic="test-topic", + project_id="test-project", + topic_id="test-topic-id", + ) + await writer._connect() + + _mock_pubsub.PublisherClient.assert_called() + assert writer._publisher is mock_publisher + + +@pytest.mark.asyncio +async def test_gcp_pubsub_writer_disconnect() -> None: + """Tests that the writer closes the publisher client on disconnect.""" + mock_publisher = MagicMock() + _mock_pubsub.PublisherClient.return_value = mock_publisher + + writer = GCPPubSubDataWriter( + name="test-gcp-writer", + field_names=["x"], + topic="test-topic", + project_id="test-project", + topic_id="test-topic-id", + ) + await writer._connect() + await writer._disconnect() + + mock_publisher.close.assert_called_once() + assert writer._publisher is None + + +@pytest.mark.asyncio +async def test_gcp_pubsub_writer_send() -> None: + """Tests sending messages to PubSub.""" + mock_publisher = MagicMock() + _mock_pubsub.PublisherClient.return_value = mock_publisher + + mock_future = MagicMock() + mock_publisher.publish.return_value = mock_future + + writer = GCPPubSubDataWriter( + name="test-gcp-writer", + field_names=["x"], + topic="test-topic", + project_id="test-project", + topic_id="test-topic-id", + ) + await writer._connect() + + messages = [b"msg1", b"msg2"] + await writer._send(messages) + + assert mock_publisher.publish.call_count == 2 + assert mock_future.result.call_count == 2 + + +@pytest.mark.asyncio +async def test_gcp_pubsub_writer_convert_json() -> None: + """Tests converting field data to JSON messages.""" + writer = GCPPubSubDataWriter( + name="test-gcp-writer", + field_names=["x", "y"], + topic="test-topic", + project_id="test-project", + topic_id="test-topic-id", + parse_json=True, + ) + + data = {"x": deque([1, 2]), "y": deque(["a", "b"])} + messages = await writer._convert(data) + + assert len(messages) == 2 + assert json.loads(messages[0]) == {"x": 1, "y": "a"} + assert json.loads(messages[1]) == {"x": 2, "y": "b"} + + +@pytest.mark.asyncio +async def test_gcp_pubsub_writer_convert_raw() -> None: + """Tests converting field data to raw bytes messages.""" + writer = GCPPubSubDataWriter( + name="test-gcp-writer", + field_names=["data"], + topic="test-topic", + project_id="test-project", + topic_id="test-topic-id", + parse_json=False, + ) + + data = {"data": deque([b"raw1", b"raw2"])} + messages = await writer._convert(data) + + assert len(messages) == 2 + assert messages[0] == b"raw1" + assert messages[1] == b"raw2" + + +@pytest.mark.asyncio +async def test_gcp_pubsub_writer_topic_path() -> None: + """Tests that the topic path is constructed correctly.""" + writer = GCPPubSubDataWriter( + name="test-gcp-writer", + field_names=["x"], + topic="test-topic", + project_id="my-project", + topic_id="my-topic", + ) + assert writer._topic_path == "projects/my-project/topics/my-topic" diff --git a/tests/unit/test_kafka_io.py b/tests/unit/test_kafka_io.py new file mode 100644 index 00000000..1eaaa039 --- /dev/null +++ b/tests/unit/test_kafka_io.py @@ -0,0 +1,377 @@ +"""Unit tests for Kafka message data reader/writer.""" + +from __future__ import annotations + +import importlib.machinery +import json +import sys +import typing as _t +from collections import deque +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from plugboard.exceptions import NoMoreDataException + + +# --------------------------------------------------------------------------- +# Mock the aiokafka module before importing the implementation +# --------------------------------------------------------------------------- + + +def _make_mock_module(name: str) -> MagicMock: + """Creates a mock module with __spec__ set for find_spec compatibility.""" + mock = MagicMock() + mock.__spec__ = importlib.machinery.ModuleSpec(name, None) + return mock + + +_mock_aiokafka = _make_mock_module("aiokafka") +_mock_consumer_class = MagicMock() +_mock_producer_class = MagicMock() +_mock_aiokafka.AIOKafkaConsumer = _mock_consumer_class +_mock_aiokafka.AIOKafkaProducer = _mock_producer_class + +sys.modules.setdefault("aiokafka", _mock_aiokafka) + +from plugboard.library.kafka_io import KafkaDataReader, KafkaDataWriter # noqa: E402 + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_kafka_record(value: dict[str, _t.Any] | bytes) -> MagicMock: + """Creates a mock Kafka ConsumerRecord.""" + record = MagicMock() + if isinstance(value, dict): + record.value = json.dumps(value).encode("utf-8") + else: + record.value = value + record.topic = "test-topic" + record.partition = 0 + record.offset = 0 + return record + + +# --------------------------------------------------------------------------- +# Tests: KafkaDataReader +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_kafka_reader_connect() -> None: + """Tests that the reader creates and starts a Kafka consumer.""" + mock_consumer = AsyncMock() + _mock_consumer_class.return_value = mock_consumer + mock_consumer.start = AsyncMock() + + reader = KafkaDataReader( + name="test-kafka-reader", + field_names=["x", "y"], + topic="test-topic", + bootstrap_servers="localhost:9092", + group_id="test-group", + ) + await reader._connect() + + _mock_consumer_class.assert_called() + mock_consumer.start.assert_called() + assert reader._consumer is mock_consumer + + +@pytest.mark.asyncio +async def test_kafka_reader_disconnect() -> None: + """Tests that the reader stops the Kafka consumer.""" + mock_consumer = AsyncMock() + _mock_consumer_class.return_value = mock_consumer + mock_consumer.start = AsyncMock() + mock_consumer.stop = AsyncMock() + + reader = KafkaDataReader( + name="test-kafka-reader", + field_names=["x"], + topic="test-topic", + bootstrap_servers="localhost:9092", + group_id="test-group", + ) + await reader._connect() + await reader._disconnect() + + mock_consumer.stop.assert_called() + assert reader._consumer is None + + +@pytest.mark.asyncio +async def test_kafka_reader_receive() -> None: + """Tests receiving messages from Kafka.""" + mock_consumer = AsyncMock() + _mock_consumer_class.return_value = mock_consumer + mock_consumer.start = AsyncMock() + + test_data = [{"x": 1, "y": "a"}, {"x": 2, "y": "b"}] + mock_records = [_make_kafka_record(d) for d in test_data] + tp = MagicMock() + mock_consumer.getmany = AsyncMock(return_value={tp: mock_records}) + + reader = KafkaDataReader( + name="test-kafka-reader", + field_names=["x", "y"], + topic="test-topic", + bootstrap_servers="localhost:9092", + group_id="test-group", + chunk_size=10, + ) + await reader._connect() + messages = await reader._receive() + + assert len(messages) == 2 + mock_consumer.getmany.assert_called() + + +@pytest.mark.asyncio +async def test_kafka_reader_receive_empty() -> None: + """Tests that empty response raises NoMoreDataException.""" + mock_consumer = AsyncMock() + _mock_consumer_class.return_value = mock_consumer + mock_consumer.start = AsyncMock() + mock_consumer.getmany = AsyncMock(return_value={}) + + reader = KafkaDataReader( + name="test-kafka-reader", + field_names=["x"], + topic="test-topic", + bootstrap_servers="localhost:9092", + group_id="test-group", + ) + await reader._connect() + + with pytest.raises(NoMoreDataException): + await reader._receive() + + +@pytest.mark.asyncio +async def test_kafka_reader_convert_json() -> None: + """Tests converting JSON Kafka messages to field buffer.""" + reader = KafkaDataReader( + name="test-kafka-reader", + field_names=["x", "y"], + topic="test-topic", + bootstrap_servers="localhost:9092", + group_id="test-group", + parse_json=True, + ) + + mock_records = [ + _make_kafka_record({"x": 1, "y": "a"}), + _make_kafka_record({"x": 2, "y": "b"}), + ] + result = await reader._convert(mock_records) + assert list(result["x"]) == [1, 2] + assert list(result["y"]) == ["a", "b"] + + +@pytest.mark.asyncio +async def test_kafka_reader_convert_raw() -> None: + """Tests converting raw Kafka messages to field buffer.""" + reader = KafkaDataReader( + name="test-kafka-reader", + field_names=["data"], + topic="test-topic", + bootstrap_servers="localhost:9092", + group_id="test-group", + parse_json=False, + ) + + mock_records = [_make_kafka_record(b"raw-1"), _make_kafka_record(b"raw-2")] + result = await reader._convert(mock_records) + assert list(result["data"]) == ["raw-1", "raw-2"] + + +@pytest.mark.asyncio +async def test_kafka_reader_ack() -> None: + """Tests committing offsets for Kafka messages.""" + mock_consumer = AsyncMock() + _mock_consumer_class.return_value = mock_consumer + mock_consumer.start = AsyncMock() + mock_consumer.commit = AsyncMock() + + reader = KafkaDataReader( + name="test-kafka-reader", + field_names=["x"], + topic="test-topic", + bootstrap_servers="localhost:9092", + group_id="test-group", + ) + await reader._connect() + + mock_records = [_make_kafka_record({"x": 1})] + await reader._ack(mock_records) + + mock_consumer.commit.assert_called() + + +@pytest.mark.asyncio +async def test_kafka_reader_bootstrap_servers_list() -> None: + """Tests that bootstrap_servers can be a list.""" + mock_consumer = AsyncMock() + _mock_consumer_class.return_value = mock_consumer + mock_consumer.start = AsyncMock() + + reader = KafkaDataReader( + name="test-kafka-reader", + field_names=["x"], + topic="test-topic", + bootstrap_servers=["host1:9092", "host2:9092"], + group_id="test-group", + ) + await reader._connect() + + call_kwargs = _mock_consumer_class.call_args[1] + assert call_kwargs["bootstrap_servers"] == ["host1:9092", "host2:9092"] + + +@pytest.mark.asyncio +async def test_kafka_reader_bootstrap_servers_string() -> None: + """Tests that bootstrap_servers string is converted to list.""" + mock_consumer = AsyncMock() + _mock_consumer_class.return_value = mock_consumer + mock_consumer.start = AsyncMock() + + reader = KafkaDataReader( + name="test-kafka-reader", + field_names=["x"], + topic="test-topic", + bootstrap_servers="localhost:9092", + group_id="test-group", + ) + await reader._connect() + + call_kwargs = _mock_consumer_class.call_args[1] + assert call_kwargs["bootstrap_servers"] == ["localhost:9092"] + + +# --------------------------------------------------------------------------- +# Tests: KafkaDataWriter +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_kafka_writer_connect() -> None: + """Tests that the writer creates and starts a Kafka producer.""" + mock_producer = AsyncMock() + _mock_producer_class.return_value = mock_producer + mock_producer.start = AsyncMock() + + writer = KafkaDataWriter( + name="test-kafka-writer", + field_names=["x"], + topic="test-topic", + bootstrap_servers="localhost:9092", + ) + await writer._connect() + + _mock_producer_class.assert_called() + mock_producer.start.assert_called() + assert writer._producer is mock_producer + + +@pytest.mark.asyncio +async def test_kafka_writer_disconnect() -> None: + """Tests that the writer stops the Kafka producer.""" + mock_producer = AsyncMock() + _mock_producer_class.return_value = mock_producer + mock_producer.start = AsyncMock() + mock_producer.stop = AsyncMock() + + writer = KafkaDataWriter( + name="test-kafka-writer", + field_names=["x"], + topic="test-topic", + bootstrap_servers="localhost:9092", + ) + await writer._connect() + await writer._disconnect() + + mock_producer.stop.assert_called() + assert writer._producer is None + + +@pytest.mark.asyncio +async def test_kafka_writer_send() -> None: + """Tests sending messages to Kafka.""" + mock_producer = AsyncMock() + _mock_producer_class.return_value = mock_producer + mock_producer.start = AsyncMock() + mock_producer.send_and_wait = AsyncMock() + + writer = KafkaDataWriter( + name="test-kafka-writer", + field_names=["x"], + topic="test-topic", + bootstrap_servers="localhost:9092", + ) + await writer._connect() + + messages = [b"msg1", b"msg2"] + await writer._send(messages) + + assert mock_producer.send_and_wait.call_count == 2 + + +@pytest.mark.asyncio +async def test_kafka_writer_convert_json() -> None: + """Tests converting field data to JSON messages.""" + writer = KafkaDataWriter( + name="test-kafka-writer", + field_names=["x", "y"], + topic="test-topic", + bootstrap_servers="localhost:9092", + parse_json=True, + ) + + data = {"x": deque([1, 2]), "y": deque(["a", "b"])} + messages = await writer._convert(data) + + assert len(messages) == 2 + assert json.loads(messages[0]) == {"x": 1, "y": "a"} + assert json.loads(messages[1]) == {"x": 2, "y": "b"} + + +@pytest.mark.asyncio +async def test_kafka_writer_convert_raw() -> None: + """Tests converting field data to raw bytes messages.""" + writer = KafkaDataWriter( + name="test-kafka-writer", + field_names=["data"], + topic="test-topic", + bootstrap_servers="localhost:9092", + parse_json=False, + ) + + data = {"data": deque([b"raw1", b"raw2"])} + messages = await writer._convert(data) + + assert len(messages) == 2 + assert messages[0] == b"raw1" + assert messages[1] == b"raw2" + + +@pytest.mark.asyncio +async def test_kafka_writer_bootstrap_servers_list() -> None: + """Tests that bootstrap_servers can be a list.""" + mock_producer = AsyncMock() + _mock_producer_class.return_value = mock_producer + mock_producer.start = AsyncMock() + + writer = KafkaDataWriter( + name="test-kafka-writer", + field_names=["x"], + topic="test-topic", + bootstrap_servers=["host1:9092", "host2:9092"], + ) + await writer._connect() + + call_kwargs = _mock_producer_class.call_args[1] + assert call_kwargs["bootstrap_servers"] == ["host1:9092", "host2:9092"] diff --git a/tests/unit/test_message_data_reader.py b/tests/unit/test_message_data_reader.py new file mode 100644 index 00000000..6add4ffe --- /dev/null +++ b/tests/unit/test_message_data_reader.py @@ -0,0 +1,366 @@ +"""Unit tests for the `MessageDataReader` base class.""" + +from __future__ import annotations + +from collections import deque +import typing as _t + +import pytest + +from plugboard.exceptions import IOStreamClosedError, NoMoreDataException +from plugboard.library.message_reader import MessageDataReader + + +# --------------------------------------------------------------------------- +# Mock implementation +# --------------------------------------------------------------------------- + + +class MockMessageDataReader(MessageDataReader): + """Mock `MessageDataReader` for testing the base class logic.""" + + def __init__( + self, + *args: _t.Any, + messages: list[dict[str, _t.Any]], + fail_on_connect: bool = False, + fail_on_receive: int | None = None, + **kwargs: _t.Any, + ) -> None: + super().__init__(*args, **kwargs) + self._messages = messages + self._idx = 0 + self._connected = False + self._disconnected = False + self._acknowledged: list[list[dict[str, _t.Any]]] = [] + self._fail_on_connect = fail_on_connect + self._fail_on_receive = fail_on_receive + self._receive_call_count = 0 + self._connect_call_count = 0 + self._disconnect_call_count = 0 + + async def _connect(self) -> None: + self._connect_call_count += 1 + if self._fail_on_connect and self._connect_call_count <= 1: + raise ConnectionError("Simulated connection failure") + self._connected = True + + async def _disconnect(self) -> None: + self._disconnect_call_count += 1 + self._connected = False + self._disconnected = True + + async def _receive(self) -> list[_t.Any]: + self._receive_call_count += 1 + if self._fail_on_receive is not None and self._receive_call_count == self._fail_on_receive: + raise ConnectionError("Simulated receive failure") + if self._chunk_size: + chunk = self._messages[self._idx : self._idx + self._chunk_size] + else: + chunk = self._messages[self._idx :] + self._idx += len(chunk) + if not chunk and self._idx >= len(self._messages): + raise NoMoreDataException + return chunk + + async def _convert(self, messages: list[_t.Any]) -> dict[str, deque]: + converted: dict[str, deque] = {field: deque() for field in self.io.outputs} + for msg in messages: + for field in self.io.outputs: + converted[field].append(msg.get(field)) + return converted + + async def _ack(self, messages: list[_t.Any]) -> None: + self._acknowledged.append(messages) + + +# --------------------------------------------------------------------------- +# Test data +# --------------------------------------------------------------------------- + +TEST_MESSAGES = [ + {"x": 1, "y": "a"}, + {"x": 2, "y": "b"}, + {"x": 3, "y": "c"}, + {"x": 4, "y": "d"}, + {"x": 5, "y": "e"}, +] + + +# --------------------------------------------------------------------------- +# Tests: Basic lifecycle +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_message_data_reader_init() -> None: + """Tests that `init` connects to the broker and pre-fetches data.""" + reader = MockMessageDataReader( + name="test-reader", + field_names=["x", "y"], + topic="test-topic", + messages=TEST_MESSAGES, + ) + await reader.init() + assert reader._connected is True + assert reader._connect_call_count == 1 + # First batch should be pre-fetched + assert reader._receive_call_count == 1 + await reader.destroy() + + +@pytest.mark.asyncio +async def test_message_data_reader_destroy() -> None: + """Tests that `destroy` disconnects from the broker.""" + reader = MockMessageDataReader( + name="test-reader", + field_names=["x", "y"], + topic="test-topic", + messages=TEST_MESSAGES, + ) + await reader.init() + await reader.destroy() + assert reader._disconnected is True + assert reader._disconnect_call_count == 1 + + +@pytest.mark.asyncio +async def test_message_data_reader_step() -> None: + """Tests that `step` reads one record at a time.""" + reader = MockMessageDataReader( + name="test-reader", + field_names=["x", "y"], + topic="test-topic", + messages=TEST_MESSAGES, + ) + await reader.init() + + results: list[dict[str, _t.Any]] = [] + while True: + try: + await reader.step() + results.append({"x": reader.x, "y": reader.y}) # type: ignore[attr-defined] + except IOStreamClosedError: + break + + assert results == TEST_MESSAGES + await reader.destroy() + + +@pytest.mark.asyncio +async def test_message_data_reader_acknowledgment() -> None: + """Tests that messages are acknowledged after processing.""" + reader = MockMessageDataReader( + name="test-reader", + field_names=["x", "y"], + topic="test-topic", + messages=TEST_MESSAGES, + ) + await reader.init() + + # Step through first message + await reader.step() + # First batch should be acknowledged + assert len(reader._acknowledged) >= 1 + + await reader.destroy() + + +# --------------------------------------------------------------------------- +# Tests: Chunked reading +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +@pytest.mark.parametrize("chunk_size", [1, 2, 3, 5, 10]) +async def test_message_data_reader_chunked(chunk_size: int) -> None: + """Tests reading with various chunk sizes.""" + reader = MockMessageDataReader( + name="test-reader", + field_names=["x", "y"], + topic="test-topic", + chunk_size=chunk_size, + messages=TEST_MESSAGES, + ) + await reader.init() + + results: list[dict[str, _t.Any]] = [] + while True: + try: + await reader.step() + results.append({"x": reader.x, "y": reader.y}) # type: ignore[attr-defined] + except IOStreamClosedError: + break + + assert results == TEST_MESSAGES + await reader.destroy() + + +@pytest.mark.asyncio +async def test_message_data_reader_no_chunk_size() -> None: + """Tests reading without chunk size (all messages at once).""" + reader = MockMessageDataReader( + name="test-reader", + field_names=["x", "y"], + topic="test-topic", + chunk_size=None, + messages=TEST_MESSAGES, + ) + await reader.init() + + results: list[dict[str, _t.Any]] = [] + while True: + try: + await reader.step() + results.append({"x": reader.x, "y": reader.y}) # type: ignore[attr-defined] + except IOStreamClosedError: + break + + assert results == TEST_MESSAGES + await reader.destroy() + + +# --------------------------------------------------------------------------- +# Tests: Empty messages +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_message_data_reader_empty_messages() -> None: + """Tests that reader handles empty message source correctly.""" + reader = MockMessageDataReader( + name="test-reader", + field_names=["x", "y"], + topic="test-topic", + messages=[], + ) + await reader.init() + + with pytest.raises(IOStreamClosedError): + await reader.step() + + await reader.destroy() + + +# --------------------------------------------------------------------------- +# Tests: Retry logic +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_message_data_reader_retry_on_receive_failure() -> None: + """Tests that reader retries on transient receive failures.""" + reader = MockMessageDataReader( + name="test-reader", + field_names=["x", "y"], + topic="test-topic", + messages=TEST_MESSAGES[:2], + fail_on_receive=2, # Fail on the second receive call + max_retries=3, + retry_base_delay=0.01, # Fast retries for testing + ) + await reader.init() + + results: list[dict[str, _t.Any]] = [] + while True: + try: + await reader.step() + results.append({"x": reader.x, "y": reader.y}) # type: ignore[attr-defined] + except IOStreamClosedError: + break + + assert results == TEST_MESSAGES[:2] + # Should have attempted reconnection + assert reader._connect_call_count >= 2 + await reader.destroy() + + +@pytest.mark.asyncio +async def test_message_data_reader_retry_exhausted() -> None: + """Tests that reader raises after all retries are exhausted.""" + reader = MockMessageDataReader( + name="test-reader", + field_names=["x", "y"], + topic="test-topic", + messages=TEST_MESSAGES[:1], + fail_on_receive=2, # Always fail on receive + max_retries=2, + retry_base_delay=0.01, + ) + await reader.init() + + # First step should succeed (from pre-fetched data) + await reader.step() + + # Second step should fail after retries exhausted + with pytest.raises((IOStreamClosedError, ConnectionError)): + await reader.step() + + await reader.destroy() + + +# --------------------------------------------------------------------------- +# Tests: Connection failure on init +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_message_data_reader_connection_failure_on_init() -> None: + """Tests that init raises on connection failure.""" + reader = MockMessageDataReader( + name="test-reader", + field_names=["x", "y"], + topic="test-topic", + messages=TEST_MESSAGES, + fail_on_connect=True, + ) + with pytest.raises(ConnectionError): + await reader.init() + + +# --------------------------------------------------------------------------- +# Tests: Single field +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_message_data_reader_single_field() -> None: + """Tests reading with a single output field.""" + messages = [{"value": i} for i in range(3)] + reader = MockMessageDataReader( + name="test-reader", + field_names=["value"], + topic="test-topic", + messages=messages, + ) + await reader.init() + + results: list[_t.Any] = [] + while True: + try: + await reader.step() + results.append(reader.value) # type: ignore[attr-defined] + except IOStreamClosedError: + break + + assert results == [0, 1, 2] + await reader.destroy() + + +# --------------------------------------------------------------------------- +# Tests: Topic attribute +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_message_data_reader_topic() -> None: + """Tests that the topic is stored correctly.""" + reader = MockMessageDataReader( + name="test-reader", + field_names=["x"], + topic="my-topic", + messages=[{"x": 1}], + ) + assert reader._topic == "my-topic" + await reader.init() + await reader.destroy() diff --git a/tests/unit/test_message_data_writer.py b/tests/unit/test_message_data_writer.py new file mode 100644 index 00000000..5412617a --- /dev/null +++ b/tests/unit/test_message_data_writer.py @@ -0,0 +1,364 @@ +"""Unit tests for the `MessageDataWriter` base class.""" + +from __future__ import annotations + +from collections import deque +import typing as _t + +import pytest + +from plugboard.connector import AsyncioConnector +from plugboard.library.message_writer import MessageDataWriter +from plugboard.schemas import ConnectorSpec + + +# --------------------------------------------------------------------------- +# Mock implementation +# --------------------------------------------------------------------------- + + +class MockMessageDataWriter(MessageDataWriter): + """Mock `MessageDataWriter` for testing the base class logic.""" + + def __init__( + self, + *args: _t.Any, + fail_on_connect: bool = False, + fail_on_send: int | None = None, + **kwargs: _t.Any, + ) -> None: + super().__init__(*args, **kwargs) + self._connected = False + self._disconnected = False + self._sent_messages: list[list[_t.Any]] = [] + self._fail_on_connect = fail_on_connect + self._fail_on_send = fail_on_send + self._send_call_count = 0 + self._connect_call_count = 0 + self._disconnect_call_count = 0 + + async def _connect(self) -> None: + self._connect_call_count += 1 + if self._fail_on_connect and self._connect_call_count <= 1: + raise ConnectionError("Simulated connection failure") + self._connected = True + + async def _disconnect(self) -> None: + self._disconnect_call_count += 1 + self._connected = False + self._disconnected = True + + async def _send(self, messages: list[_t.Any]) -> None: + self._send_call_count += 1 + if self._fail_on_send is not None and self._send_call_count == self._fail_on_send: + raise ConnectionError("Simulated send failure") + self._sent_messages.append(messages) + + async def _convert(self, data: dict[str, deque]) -> list[_t.Any]: + completed_rows = min(len(d) for d in data.values()) if data else 0 + messages: list[dict[str, _t.Any]] = [] + for i in range(completed_rows): + record = {field: data[field][i] for field in data} + messages.append(record) + return messages + + +# --------------------------------------------------------------------------- +# Test helpers +# --------------------------------------------------------------------------- + + +async def _setup_writer_with_channels( + writer: MockMessageDataWriter, field_names: list[str] +) -> dict[str, AsyncioConnector]: + """Sets up a writer with connected asyncio channels for sending data.""" + connectors = { + field: AsyncioConnector( + spec=ConnectorSpec(source="none.none", target=f"{writer.name}.{field}"), + ) + for field in field_names + } + await writer.io.connect(list(connectors.values())) + return connectors + + +# --------------------------------------------------------------------------- +# Tests: Basic lifecycle +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_message_data_writer_init() -> None: + """Tests that `init` connects to the broker.""" + writer = MockMessageDataWriter( + name="test-writer", + field_names=["x", "y"], + topic="test-topic", + ) + await writer.init() + assert writer._connected is True + assert writer._connect_call_count == 1 + await writer.destroy() + + +@pytest.mark.asyncio +async def test_message_data_writer_destroy() -> None: + """Tests that `destroy` disconnects from the broker.""" + writer = MockMessageDataWriter( + name="test-writer", + field_names=["x", "y"], + topic="test-topic", + ) + await writer.init() + await writer.destroy() + assert writer._disconnected is True + assert writer._disconnect_call_count == 1 + + +# --------------------------------------------------------------------------- +# Tests: Writing data +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_message_data_writer_step_and_run() -> None: + """Tests that data is written via step and flushed on run.""" + writer = MockMessageDataWriter( + name="test-writer", + field_names=["x", "y"], + topic="test-topic", + chunk_size=2, + ) + connectors = await _setup_writer_with_channels(writer, ["x", "y"]) + await writer.init() + + output_channels = {field: await connectors[field].connect_send() for field in ["x", "y"]} + + # Send data + test_data = [(1, "a"), (2, "b"), (3, "c")] + for x_val, y_val in test_data: + await output_channels["x"].send(x_val) + await output_channels["y"].send(y_val) + await writer.step() + + # Close inputs and run to flush + await writer.io.close() + await writer.run() + + # Verify sent messages + all_sent = [msg for batch in writer._sent_messages for msg in batch] + assert len(all_sent) == 3 + assert all_sent[0] == {"x": 1, "y": "a"} + assert all_sent[1] == {"x": 2, "y": "b"} + assert all_sent[2] == {"x": 3, "y": "c"} + + await writer.destroy() + + +@pytest.mark.asyncio +async def test_message_data_writer_flush_on_run() -> None: + """Tests that remaining buffered data is flushed on `run`.""" + writer = MockMessageDataWriter( + name="test-writer", + field_names=["x"], + topic="test-topic", + chunk_size=10, # Large chunk size so nothing is sent during step + ) + connectors = await _setup_writer_with_channels(writer, ["x"]) + await writer.init() + + output_channels = {"x": await connectors["x"].connect_send()} + + # Send data (less than chunk_size) + for i in range(3): + await output_channels["x"].send(i) + await writer.step() + + # Nothing should be sent yet (buffer < chunk_size) + assert len(writer._sent_messages) == 0 + + # Close and run to flush + await writer.io.close() + await writer.run() + + # Now data should be flushed + all_sent = [msg for batch in writer._sent_messages for msg in batch] + assert len(all_sent) == 3 + assert all_sent == [{"x": 0}, {"x": 1}, {"x": 2}] + + await writer.destroy() + + +# --------------------------------------------------------------------------- +# Tests: Chunked writing +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +@pytest.mark.parametrize("chunk_size", [1, 2, 3, 5]) +async def test_message_data_writer_chunked(chunk_size: int) -> None: + """Tests writing with various chunk sizes.""" + writer = MockMessageDataWriter( + name="test-writer", + field_names=["x", "y"], + topic="test-topic", + chunk_size=chunk_size, + ) + connectors = await _setup_writer_with_channels(writer, ["x", "y"]) + await writer.init() + + output_channels = {field: await connectors[field].connect_send() for field in ["x", "y"]} + + test_data = [(i, f"val_{i}") for i in range(5)] + for x_val, y_val in test_data: + await output_channels["x"].send(x_val) + await output_channels["y"].send(y_val) + await writer.step() + + await writer.io.close() + await writer.run() + + all_sent = [msg for batch in writer._sent_messages for msg in batch] + assert len(all_sent) == 5 + for i, (x_val, y_val) in enumerate(test_data): + assert all_sent[i] == {"x": x_val, "y": y_val} + + await writer.destroy() + + +@pytest.mark.asyncio +async def test_message_data_writer_no_chunk_size() -> None: + """Tests writing without chunk size (flush only on run).""" + writer = MockMessageDataWriter( + name="test-writer", + field_names=["x"], + topic="test-topic", + chunk_size=None, + ) + connectors = await _setup_writer_with_channels(writer, ["x"]) + await writer.init() + + output_channels = {"x": await connectors["x"].connect_send()} + + for i in range(3): + await output_channels["x"].send(i) + await writer.step() + + # Nothing sent yet (no chunk_size trigger) + assert len(writer._sent_messages) == 0 + + await writer.io.close() + await writer.run() + + all_sent = [msg for batch in writer._sent_messages for msg in batch] + assert len(all_sent) == 3 + + await writer.destroy() + + +# --------------------------------------------------------------------------- +# Tests: Retry logic +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_message_data_writer_retry_on_send_failure() -> None: + """Tests that writer retries on transient send failures.""" + writer = MockMessageDataWriter( + name="test-writer", + field_names=["x"], + topic="test-topic", + chunk_size=1, + fail_on_send=1, # Fail on first send + max_retries=3, + retry_base_delay=0.01, + ) + connectors = await _setup_writer_with_channels(writer, ["x"]) + await writer.init() + + output_channels = {"x": await connectors["x"].connect_send()} + + # Send one item and step (triggers send which fails, then retries) + await output_channels["x"].send(0) + await writer.step() + + # Send another item and step (should succeed now) + await output_channels["x"].send(1) + await writer.step() + + await writer.io.close() + await writer.run() + + # Should have retried and eventually succeeded + all_sent = [msg for batch in writer._sent_messages for msg in batch] + assert len(all_sent) == 2 + # Should have reconnected + assert writer._connect_call_count >= 2 + + await writer.destroy() + + +# --------------------------------------------------------------------------- +# Tests: Connection failure on init +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_message_data_writer_connection_failure_on_init() -> None: + """Tests that init raises on connection failure.""" + writer = MockMessageDataWriter( + name="test-writer", + field_names=["x"], + topic="test-topic", + fail_on_connect=True, + ) + with pytest.raises(ConnectionError): + await writer.init() + + +# --------------------------------------------------------------------------- +# Tests: Topic attribute +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_message_data_writer_topic() -> None: + """Tests that the topic is stored correctly.""" + writer = MockMessageDataWriter( + name="test-writer", + field_names=["x"], + topic="my-topic", + ) + assert writer._topic == "my-topic" + + +# --------------------------------------------------------------------------- +# Tests: Single field +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_message_data_writer_single_field() -> None: + """Tests writing with a single input field.""" + writer = MockMessageDataWriter( + name="test-writer", + field_names=["value"], + topic="test-topic", + chunk_size=3, + ) + connectors = await _setup_writer_with_channels(writer, ["value"]) + await writer.init() + + output_channels = {"value": await connectors["value"].connect_send()} + + for i in range(3): + await output_channels["value"].send(i * 10) + await writer.step() + + await writer.io.close() + await writer.run() + + all_sent = [msg for batch in writer._sent_messages for msg in batch] + assert all_sent == [{"value": 0}, {"value": 10}, {"value": 20}] + + await writer.destroy() From e538c242105b73b596bbbf1d0cf5f6be33c7ae9c Mon Sep 17 00:00:00 2001 From: Chris Knight Date: Wed, 9 Sep 2026 16:17:01 +0200 Subject: [PATCH 05/11] fix: resolve lint, format, and mypy errors - Fix ruff lint errors: import sorting, unused imports, S110 noqa comments - Fix ruff format errors in gcp_pubsub_io.py - Fix mypy overlap errors: remove duplicate fields from ArgsDict TypedDicts - Fix mypy multiple values error: use kwargs.setdefault instead of pop - Remove untracked test data files causing lint failures --- plugboard/library/aws_messaging_io.py | 21 ++++++++------------- plugboard/library/gcp_pubsub_io.py | 21 +++++++-------------- plugboard/library/kafka_io.py | 7 ++----- plugboard/library/message_reader.py | 9 ++++----- plugboard/library/message_writer.py | 7 +++---- tests/unit/test_aws_messaging_io.py | 6 ++---- tests/unit/test_gcp_pubsub_io.py | 4 ++-- tests/unit/test_kafka_io.py | 2 +- 8 files changed, 29 insertions(+), 48 deletions(-) diff --git a/plugboard/library/aws_messaging_io.py b/plugboard/library/aws_messaging_io.py index cd30d39d..01852be3 100644 --- a/plugboard/library/aws_messaging_io.py +++ b/plugboard/library/aws_messaging_io.py @@ -28,10 +28,7 @@ class AWSSQSDataReaderArgsDict(MessageDataReaderArgsDict): wait_time_seconds: Long-polling wait time in seconds. """ - queue_url: str - region: str - parse_json: _t.NotRequired[bool] - wait_time_seconds: _t.NotRequired[int] + pass class AWSSNSDataWriterArgsDict(MessageDataWriterArgsDict): @@ -43,9 +40,7 @@ class AWSSNSDataWriterArgsDict(MessageDataWriterArgsDict): parse_json: Whether to encode message data as JSON. """ - topic_arn: str - region: str - parse_json: _t.NotRequired[bool] + pass class AWSSQSDataReader(MessageDataReader): @@ -75,8 +70,8 @@ def __init__( **kwargs: Additional keyword arguments for [`MessageDataReader`][plugboard.library.MessageDataReader]. """ - topic = kwargs.pop("topic", queue_url) - super().__init__(topic=topic, **kwargs) + kwargs.setdefault("topic", queue_url) + super().__init__(**kwargs) self._queue_url = queue_url self._region = region self._parse_json = parse_json @@ -95,7 +90,7 @@ async def _disconnect(self) -> None: if self._client is not None: try: await self._client_ctx.__aexit__(None, None, None) - except Exception: # noqa: S102 + except Exception: # noqa: S110 pass self._client = None self._session = None @@ -184,8 +179,8 @@ def __init__( **kwargs: Additional keyword arguments for [`MessageDataWriter`][plugboard.library.MessageDataWriter]. """ - topic = kwargs.pop("topic", topic_arn) - super().__init__(topic=topic, **kwargs) + kwargs.setdefault("topic", topic_arn) + super().__init__(**kwargs) self._topic_arn = topic_arn self._region = region self._parse_json = parse_json @@ -203,7 +198,7 @@ async def _disconnect(self) -> None: if self._client is not None: try: await self._client_ctx.__aexit__(None, None, None) - except Exception: # noqa: S102 + except Exception: # noqa: S110 pass self._client = None self._session = None diff --git a/plugboard/library/gcp_pubsub_io.py b/plugboard/library/gcp_pubsub_io.py index 57b97efd..a1a1dd1d 100644 --- a/plugboard/library/gcp_pubsub_io.py +++ b/plugboard/library/gcp_pubsub_io.py @@ -14,7 +14,6 @@ try: from google.cloud import pubsub_v1 - from google.cloud.pubsub_v1.subscriber.message import Message as PubSubMessage except ImportError: # pragma: no cover pass @@ -28,9 +27,7 @@ class GCPPubSubDataReaderArgsDict(MessageDataReaderArgsDict): parse_json: Whether to parse message data as JSON. """ - project_id: str - subscription_id: str - parse_json: _t.NotRequired[bool] + pass class GCPPubSubDataWriterArgsDict(MessageDataWriterArgsDict): @@ -42,9 +39,7 @@ class GCPPubSubDataWriterArgsDict(MessageDataWriterArgsDict): parse_json: Whether to encode message data as JSON. """ - project_id: str - topic_id: str - parse_json: _t.NotRequired[bool] + pass class GCPPubSubDataReader(MessageDataReader): @@ -71,13 +66,11 @@ def __init__( **kwargs: Additional keyword arguments for [`MessageDataReader`][plugboard.library.MessageDataReader]. """ - topic = kwargs.pop("topic", f"{project_id}/{subscription_id}") - super().__init__(topic=topic, **kwargs) + kwargs.setdefault("topic", f"{project_id}/{subscription_id}") + super().__init__(**kwargs) self._project_id = project_id self._subscription_id = subscription_id - self._subscription_path = ( - f"projects/{project_id}/subscriptions/{subscription_id}" - ) + self._subscription_path = f"projects/{project_id}/subscriptions/{subscription_id}" self._parse_json = parse_json self._subscriber: _t.Optional[pubsub_v1.SubscriberClient] = None @@ -174,8 +167,8 @@ def __init__( **kwargs: Additional keyword arguments for [`MessageDataWriter`][plugboard.library.MessageDataWriter]. """ - topic = kwargs.pop("topic", f"{project_id}/{topic_id}") - super().__init__(topic=topic, **kwargs) + kwargs.setdefault("topic", f"{project_id}/{topic_id}") + super().__init__(**kwargs) self._project_id = project_id self._topic_id = topic_id self._topic_path = f"projects/{project_id}/topics/{topic_id}" diff --git a/plugboard/library/kafka_io.py b/plugboard/library/kafka_io.py index c8ef5646..ef479c89 100644 --- a/plugboard/library/kafka_io.py +++ b/plugboard/library/kafka_io.py @@ -27,9 +27,7 @@ class KafkaDataReaderArgsDict(MessageDataReaderArgsDict): parse_json: Whether to parse message values as JSON. """ - bootstrap_servers: _t.Union[str, list[str]] - group_id: str - parse_json: _t.NotRequired[bool] + pass class KafkaDataWriterArgsDict(MessageDataWriterArgsDict): @@ -40,8 +38,7 @@ class KafkaDataWriterArgsDict(MessageDataWriterArgsDict): parse_json: Whether to encode message values as JSON. """ - bootstrap_servers: _t.Union[str, list[str]] - parse_json: _t.NotRequired[bool] + pass class KafkaDataReader(MessageDataReader): diff --git a/plugboard/library/message_reader.py b/plugboard/library/message_reader.py index eafd611c..9c222a74 100644 --- a/plugboard/library/message_reader.py +++ b/plugboard/library/message_reader.py @@ -4,8 +4,8 @@ from abc import ABC, abstractmethod import asyncio -from collections import deque from asyncio.tasks import Task +from collections import deque import typing as _t from plugboard.component import Component, IOController @@ -18,7 +18,6 @@ class MessageDataReaderArgsDict(ComponentArgsDict): Attributes: field_names: The names of the fields to read from messages. - topic: The topic/queue to read from. chunk_size: Optional; The number of messages to fetch per batch. max_retries: Maximum number of retry attempts for transient failures. retry_base_delay: Base delay in seconds for exponential backoff. @@ -26,7 +25,7 @@ class MessageDataReaderArgsDict(ComponentArgsDict): """ field_names: list[str] - topic: str + topic: _t.NotRequired[str] chunk_size: _t.NotRequired[int | None] max_retries: _t.NotRequired[int] retry_base_delay: _t.NotRequired[float] @@ -183,7 +182,7 @@ async def _reconnect(self) -> None: self._logger.info("Attempting reconnection to message broker", topic=self._topic) try: await self._disconnect() - except Exception: # noqa: S102 + except Exception: # noqa: S110 self._logger.warning("Error during disconnect in reconnection", exc_info=True) await self._connect() self._logger.info("Reconnected to message broker", topic=self._topic) @@ -258,7 +257,7 @@ async def destroy(self) -> None: self._task.cancel() try: await self._task - except (asyncio.CancelledError, Exception): + except (asyncio.CancelledError, Exception): # noqa: S110 pass self._task = None await self._disconnect() diff --git a/plugboard/library/message_writer.py b/plugboard/library/message_writer.py index 6e1ee7a3..e59f0801 100644 --- a/plugboard/library/message_writer.py +++ b/plugboard/library/message_writer.py @@ -4,8 +4,8 @@ from abc import ABC, abstractmethod import asyncio -from collections import defaultdict, deque from asyncio.tasks import Task +from collections import defaultdict, deque import typing as _t from plugboard.component import Component, IOController @@ -18,7 +18,6 @@ class MessageDataWriterArgsDict(ComponentArgsDict): Attributes: field_names: The names of the fields to include in messages. - topic: The topic/queue to write to. chunk_size: Optional; The number of records to batch into messages. max_retries: Maximum number of retry attempts for transient failures. retry_base_delay: Base delay in seconds for exponential backoff. @@ -26,7 +25,7 @@ class MessageDataWriterArgsDict(ComponentArgsDict): """ field_names: list[str] - topic: str + topic: _t.NotRequired[str] chunk_size: _t.NotRequired[int | None] max_retries: _t.NotRequired[int] retry_base_delay: _t.NotRequired[float] @@ -245,7 +244,7 @@ async def destroy(self) -> None: self._task.cancel() try: await self._task - except (asyncio.CancelledError, Exception): + except (asyncio.CancelledError, Exception): # noqa: S110 pass self._task = None await self._disconnect() diff --git a/tests/unit/test_aws_messaging_io.py b/tests/unit/test_aws_messaging_io.py index ea944ab5..59db971a 100644 --- a/tests/unit/test_aws_messaging_io.py +++ b/tests/unit/test_aws_messaging_io.py @@ -2,17 +2,15 @@ from __future__ import annotations +from collections import deque import importlib.machinery import json import sys import typing as _t -from collections import deque from unittest.mock import AsyncMock, MagicMock import pytest -from plugboard.exceptions import NoMoreDataException - # --------------------------------------------------------------------------- # Mock the aioboto3 module before importing the implementation @@ -32,7 +30,7 @@ def _make_mock_module(name: str) -> MagicMock: sys.modules.setdefault("aioboto3", _mock_aioboto3) -from plugboard.library.aws_messaging_io import AWSSQSDataReader, AWSSNSDataWriter # noqa: E402 +from plugboard.library.aws_messaging_io import AWSSNSDataWriter, AWSSQSDataReader # noqa: E402 # --------------------------------------------------------------------------- diff --git a/tests/unit/test_gcp_pubsub_io.py b/tests/unit/test_gcp_pubsub_io.py index 2966d9a8..db222dbc 100644 --- a/tests/unit/test_gcp_pubsub_io.py +++ b/tests/unit/test_gcp_pubsub_io.py @@ -2,12 +2,12 @@ from __future__ import annotations +from collections import deque import importlib.machinery import json import sys import typing as _t -from collections import deque -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import MagicMock import pytest diff --git a/tests/unit/test_kafka_io.py b/tests/unit/test_kafka_io.py index 1eaaa039..c5738eee 100644 --- a/tests/unit/test_kafka_io.py +++ b/tests/unit/test_kafka_io.py @@ -2,11 +2,11 @@ from __future__ import annotations +from collections import deque import importlib.machinery import json import sys import typing as _t -from collections import deque from unittest.mock import AsyncMock, MagicMock import pytest From e49d472e26badb749975791ad27cc83ebe9c21dd Mon Sep 17 00:00:00 2001 From: Chris Knight Date: Wed, 9 Sep 2026 16:33:32 +0200 Subject: [PATCH 06/11] fix: reduce ZMQ proxy flakiness in Ray integration test Increase connection establishment sleep in _ZMQPipelineConnectorProxy from 0.1s to 0.5s to allow the proxy subprocess's SUB socket subscription to propagate to XPUB before the sender starts publishing (ZMQ slow joiner problem). Also mark the test as flaky with 3 reruns following the existing pattern used elsewhere in the repo. Fixes: test_process_with_components_run[RayProcess-zmq_connector_cls-zmq_pubsub_proxy=True-10-2.0] --- plugboard/connector/zmq_channel.py | 4 +++- tests/integration/test_process_with_components_run.py | 1 + 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/plugboard/connector/zmq_channel.py b/plugboard/connector/zmq_channel.py index d74ad849..5f74e404 100644 --- a/plugboard/connector/zmq_channel.py +++ b/plugboard/connector/zmq_channel.py @@ -353,7 +353,9 @@ async def connect_recv(self) -> ZMQChannel: self._recv_channel = ZMQChannel( recv_socket=recv_socket, topic=self._topic, maxsize=self._maxsize ) - await asyncio.sleep(0.1) # Ensure connections established before first send. Better way? + # Allow extra time for the proxy subprocess's SUB socket subscription to propagate + # to XPUB before the sender starts publishing (ZMQ "slow joiner" problem). + await asyncio.sleep(0.5) return self._recv_channel diff --git a/tests/integration/test_process_with_components_run.py b/tests/integration/test_process_with_components_run.py index ca599d1e..19653596 100644 --- a/tests/integration/test_process_with_components_run.py +++ b/tests/integration/test_process_with_components_run.py @@ -84,6 +84,7 @@ def tempfile_path() -> _t.Generator[Path, None, None]: @pytest.mark.asyncio +@pytest.mark.flaky(reruns=3) # Flaky on Github Actions with Ray + ZMQ proxy (slow joiner) @pytest_cases.parametrize( "process_cls, connector_cls", [ From d315dd0739cd5e6eec6eb7c575fbbbcb447b1ba4 Mon Sep 17 00:00:00 2001 From: Chris Knight Date: Wed, 9 Sep 2026 16:45:57 +0200 Subject: [PATCH 07/11] fix: resolve ty type checker errors - Fix invalid-raise errors in message_reader.py and message_writer.py by initializing last_exception with a non-None default instead of Optional[Exception] - Remove PublisherClient.close() call in gcp_pubsub_io.py (method does not exist on the client); just set reference to None for GC - Update test to match new disconnect behavior --- plugboard/library/gcp_pubsub_io.py | 1 - plugboard/library/message_reader.py | 4 ++-- plugboard/library/message_writer.py | 4 ++-- tests/unit/test_gcp_pubsub_io.py | 1 - 4 files changed, 4 insertions(+), 6 deletions(-) diff --git a/plugboard/library/gcp_pubsub_io.py b/plugboard/library/gcp_pubsub_io.py index a1a1dd1d..90d67094 100644 --- a/plugboard/library/gcp_pubsub_io.py +++ b/plugboard/library/gcp_pubsub_io.py @@ -182,7 +182,6 @@ async def _connect(self) -> None: async def _disconnect(self) -> None: """Closes the PubSub publisher client.""" if self._publisher is not None: - self._publisher.close() # type: ignore[no-untyped-call] self._publisher = None async def _send(self, messages: list[_t.Any]) -> None: diff --git a/plugboard/library/message_reader.py b/plugboard/library/message_reader.py index 9c222a74..94436d8b 100644 --- a/plugboard/library/message_reader.py +++ b/plugboard/library/message_reader.py @@ -154,7 +154,7 @@ async def _receive_with_retry(self) -> list[_t.Any]: NoMoreDataException: If the source is exhausted. MessageBrokerConnectionError: If all retries are exhausted. """ - last_exception: _t.Optional[Exception] = None + last_exception: Exception = RuntimeError("All retries exhausted") for attempt in range(self._max_retries + 1): try: return await self._receive() @@ -175,7 +175,7 @@ async def _receive_with_retry(self) -> list[_t.Any]: ) await asyncio.sleep(delay) await self._reconnect() - raise last_exception # type: ignore[misc] + raise last_exception async def _reconnect(self) -> None: """Attempts to reconnect to the message broker.""" diff --git a/plugboard/library/message_writer.py b/plugboard/library/message_writer.py index e59f0801..cc66674b 100644 --- a/plugboard/library/message_writer.py +++ b/plugboard/library/message_writer.py @@ -140,7 +140,7 @@ async def _send_with_retry(self, messages: list[_t.Any]) -> None: Raises: Exception: If all retries are exhausted. """ - last_exception: _t.Optional[Exception] = None + last_exception: Exception = RuntimeError("All retries exhausted") for attempt in range(self._max_retries + 1): try: await self._send(messages) @@ -160,7 +160,7 @@ async def _send_with_retry(self, messages: list[_t.Any]) -> None: ) await asyncio.sleep(delay) await self._reconnect() - raise last_exception # type: ignore[misc] + raise last_exception async def _reconnect(self) -> None: """Attempts to reconnect to the message broker.""" diff --git a/tests/unit/test_gcp_pubsub_io.py b/tests/unit/test_gcp_pubsub_io.py index db222dbc..c88fa537 100644 --- a/tests/unit/test_gcp_pubsub_io.py +++ b/tests/unit/test_gcp_pubsub_io.py @@ -303,7 +303,6 @@ async def test_gcp_pubsub_writer_disconnect() -> None: await writer._connect() await writer._disconnect() - mock_publisher.close.assert_called_once() assert writer._publisher is None From 0fa21c58f4c5602640fa7a45faeffa32a682c3b8 Mon Sep 17 00:00:00 2001 From: Chris Knight Date: Wed, 9 Sep 2026 16:56:28 +0200 Subject: [PATCH 08/11] chore: remove proposal doc from VCS, keep locally The design proposal is not tracked in the repo. Added to .gitignore to prevent accidental re-commit. --- .gitignore | 1 + docs/message-data-reader-writer-proposal.md | 592 -------------------- 2 files changed, 1 insertion(+), 592 deletions(-) delete mode 100644 docs/message-data-reader-writer-proposal.md diff --git a/.gitignore b/.gitignore index 9107d228..8b7a80bc 100644 --- a/.gitignore +++ b/.gitignore @@ -170,3 +170,4 @@ cython_debug/ # Generated version files */_version.py +docs/message-data-reader-writer-proposal.md diff --git a/docs/message-data-reader-writer-proposal.md b/docs/message-data-reader-writer-proposal.md deleted file mode 100644 index db4180ab..00000000 --- a/docs/message-data-reader-writer-proposal.md +++ /dev/null @@ -1,592 +0,0 @@ -# Proposal: MessageDataReader and MessageDataWriter Base Classes - -## Issue Reference - -[Issue #102: feat: Base component for external communication](https://github.com/plugboard-dev/plugboard/issues/102) - -## Summary - -Develop `MessageDataReader` and `MessageDataWriter` abstract base classes that provide common logic for reading from and writing to pub/sub message broker infrastructure. These are analogous to the existing `DataReader` and `DataWriter` components (which handle chunking/transforming for file access), but focused on message broker communication — including connection management, reconnection, retries, and message acknowledgment. - -Three concrete implementations will be provided: -1. **Google Cloud PubSub** (`GCPPubSubDataReader` / `GCPPubSubDataWriter`) -2. **AWS SNS/SQS** (`AWSSNSQSDataReader` / `AWSSQSDataWriter`) -3. **Apache Kafka** (`KafkaDataReader` / `KafkaDataWriter`) - ---- - -## Design Rationale - -### Why not extend `DataReader`/`DataWriter`? - -The existing `DataReader`/`DataWriter` classes are designed for finite data sources (files, databases) where: -- `_fetch()` raises `NoMoreDataException` when data is exhausted -- Data is read in chunks until the source is depleted -- No connection lifecycle management is needed (connections are per-query) - -Message brokers have fundamentally different semantics: -- Data arrives continuously (no natural "end of data") -- Connections are long-lived and must be managed (connect, reconnect, disconnect) -- Messages require acknowledgment after processing -- Transient failures require retry with exponential backoff - -Therefore, `MessageDataReader`/`MessageDataWriter` will be standalone `Component` subclasses that follow a *similar* pattern to `DataReader`/`DataWriter` (field-based IO, chunking, buffering) but with message-broker-specific lifecycle management. - -### Relationship to existing patterns - -| Pattern | Base Class | Handles | Subclasses implement | -|---------|-----------|---------|---------------------| -| File I/O | `DataReader`/`DataWriter` | Chunking, buffering, field IO | `_fetch()`, `_convert()`, `_save()` | -| WebSocket | `WebsocketBase` | Connection lifecycle, reconnection | `step()` for read/write | -| **Message Broker** | `MessageDataReader`/`MessageDataWriter` | Connection lifecycle, reconnection, retry, chunking, buffering, acknowledgment | `_connect()`, `_disconnect()`, `_receive()`/`_send()`, `_convert()`, `_ack()` | - ---- - -## Interface Design - -### `MessageDataReader` - -```python -class MessageDataReader(Component, ABC): - """Abstract base class for reading data from a pub/sub message broker. - - Provides connection management, reconnection with exponential backoff, - retry logic, message acknowledgment, and chunked/buffered reading - analogous to `DataReader`. - - Subclasses must implement broker-specific methods for connecting, - receiving messages, converting messages to field buffers, and - acknowledging processed messages. - """ - - io = IOController() - - def __init__( - self, - field_names: list[str], - topic: str, - subscription_id: str | None = None, - chunk_size: int | None = None, - max_retries: int = 3, - retry_base_delay: float = 1.0, - retry_max_delay: float = 60.0, - **kwargs: Unpack[ComponentArgsDict], - ) -> None: - """Instantiate the `MessageDataReader`. - - Args: - field_names: The names of the fields to extract from messages. - topic: The topic/queue to read from. - subscription_id: Optional; A subscription ID (required for some brokers like GCP PubSub). - chunk_size: Optional; Number of messages to fetch per batch. - max_retries: Maximum number of retry attempts for transient failures. - retry_base_delay: Base delay in seconds for exponential backoff. - retry_max_delay: Maximum delay in seconds for exponential backoff. - **kwargs: Additional keyword arguments for `Component`. - """ -``` - -#### Abstract methods (implemented by subclasses): - -| Method | Signature | Description | -|--------|-----------|-------------| -| `_connect` | `async def _connect(self) -> None` | Establish connection to the message broker. | -| `_disconnect` | `async def _disconnect(self) -> None` | Close the connection to the message broker. | -| `_receive` | `async def _receive(self) -> list[Any]` | Receive a batch of raw messages from the broker. Should block until at least one message is available or a timeout occurs. Return empty list on timeout. | -| `_convert` | `async def _convert(self, messages: list[Any]) -> dict[str, deque]` | Convert raw messages into a `dict[str, deque]` field buffer. | -| `_ack` | `async def _ack(self, messages: list[Any]) -> None` | Acknowledge successful processing of messages. | - -#### Concrete methods (provided by base class): - -| Method | Description | -|--------|-------------| -| `init()` | Calls `_connect()` with retry logic. Pre-fetches first batch. | -| `step()` | Consumes one record from the buffer. Fetches next batch if buffer empty. Calls `_ack()` on processed messages. | -| `destroy()` | Calls `_disconnect()` to clean up broker connection. | -| `_receive_with_retry()` | Wraps `_receive()` with exponential backoff retry and automatic reconnection. | - -### `MessageDataWriter` - -```python -class MessageDataWriter(Component, ABC): - """Abstract base class for writing data to a pub/sub message broker. - - Provides connection management, reconnection with exponential backoff, - retry logic, and chunked/buffered writing analogous to `DataWriter`. - - Subclasses must implement broker-specific methods for connecting, - sending messages, converting field data to messages, and - broker-specific message formatting. - """ - - io = IOController() - - def __init__( - self, - field_names: list[str], - topic: str, - chunk_size: int | None = None, - max_retries: int = 3, - retry_base_delay: float = 1.0, - retry_max_delay: float = 60.0, - **kwargs: Unpack[ComponentArgsDict], - ) -> None: - """Instantiate the `MessageDataWriter`. - - Args: - field_names: The names of the fields to include in messages. - topic: The topic/queue to write to. - chunk_size: Optional; Number of records to batch into a single message. - max_retries: Maximum number of retry attempts for transient failures. - retry_base_delay: Base delay in seconds for exponential backoff. - retry_max_delay: Maximum delay in seconds for exponential backoff. - **kwargs: Additional keyword arguments for `Component`. - """ -``` - -#### Abstract methods (implemented by subclasses): - -| Method | Signature | Description | -|--------|-----------|-------------| -| `_connect` | `async def _connect(self) -> None` | Establish connection to the message broker. | -| `_disconnect` | `async def _disconnect(self) -> None` | Close the connection to the message broker. | -| `_send` | `async def _send(self, messages: list[Any]) -> None` | Send a batch of messages to the broker. | -| `_convert` | `async def _convert(self, data: dict[str, deque]) -> list[Any]` | Convert field buffer data into broker-specific message format. | - -#### Concrete methods (provided by base class): - -| Method | Description | -|--------|-------------| -| `init()` | Calls `_connect()` with retry logic. | -| `step()` | Buffers input fields. Triggers `_send()` when `chunk_size` reached. | -| `run()` | Runs step loop to completion, then flushes remaining buffered data. | -| `destroy()` | Calls `_disconnect()` to clean up broker connection. | -| `_send_with_retry()` | Wraps `_send()` with exponential backoff retry and automatic reconnection. | - ---- - -## Connection Management & Retry Strategy - -The base classes provide robust connection management: - -### Connection lifecycle - -``` -init() → _connect() [with retry] → ready for step() -step() → _receive_with_retry() / _send_with_retry() → process messages -destroy() → _disconnect() -``` - -### Reconnection with exponential backoff - -```python -async def _receive_with_retry(self) -> list[Any]: - """Receives messages with retry and exponential backoff.""" - last_exception = None - for attempt in range(self._max_retries + 1): - try: - return await self._receive() - except TransientError as e: - last_exception = e - if attempt < self._max_retries: - delay = min( - self._retry_base_delay * (2 ** attempt), - self._retry_max_delay, - ) - self._logger.warning( - "Transient error receiving messages, retrying", - attempt=attempt + 1, - delay=delay, - error=str(e), - ) - await asyncio.sleep(delay) - # Attempt reconnection before retry - await self._reconnect() - raise last_exception # type: ignore[misc] -``` - -### Reconnection strategy - -```python -async def _reconnect(self) -> None: - """Attempts to reconnect to the message broker.""" - self._logger.info("Attempting reconnection to message broker") - try: - await self._disconnect() - except Exception: - pass # Best-effort disconnect - await self._connect() - self._logger.info("Reconnected to message broker") -``` - ---- - -## Concrete Implementations - -### 1. Google Cloud PubSub - -**Dependencies**: `google-cloud-pubsub` (added as optional dependency `gcp-pubsub`) - -#### `GCPPubSubDataReader` - -```python -class GCPPubSubDataReader(MessageDataReader): - """Reads data from Google Cloud PubSub subscription.""" - - def __init__( - self, - project_id: str, - subscription_id: str, - parse_json: bool = True, - **kwargs: Unpack[MessageDataReaderArgsSpec], - ) -> None: - ... - - async def _connect(self) -> None: - # Create AsyncSubscriberClient - # Subscribe to subscription - - async def _disconnect(self) -> None: - # Close subscriber client - - async def _receive(self) -> list[Any]: - # Pull batch of messages (up to chunk_size) - # Return list of PubSubMessage - - async def _convert(self, messages: list[Any]) -> dict[str, deque]: - # Parse message data (JSON or raw bytes) - # Extract fields into dict[str, deque] - - async def _ack(self, messages: list[Any]) -> None: - # Acknowledge messages via subscriber -``` - -#### `GCPPubSubDataWriter` - -```python -class GCPPubSubDataWriter(MessageDataWriter): - """Writes data to Google Cloud PubSub topic.""" - - def __init__( - self, - project_id: str, - topic_id: str, - parse_json: bool = True, - **kwargs: Unpack[MessageDataWriterArgsSpec], - ) -> None: - ... - - async def _connect(self) -> None: - # Create AsyncPublisherClient - - async def _disconnect(self) -> None: - # Close publisher client - - async def _send(self, messages: list[Any]) -> None: - # Publish messages to topic - - async def _convert(self, data: dict[str, deque]) -> list[Any]: - # Convert field data to JSON-encoded bytes -``` - -### 2. AWS SNS/SQS - -**Dependencies**: `aioboto3` or `aws-sdk-pandas` (added as optional dependency `aws-messaging`) - -> **Note**: AWS uses SQS for receiving (queue-based) and SNS for publishing (topic-based). The reader uses SQS; the writer can use either SNS (pub/sub) or SQS (queue). We'll implement both. - -#### `AWSSQSDataReader` - -```python -class AWSSQSDataReader(MessageDataReader): - """Reads data from AWS SQS queue.""" - - def __init__( - self, - queue_url: str, - region: str, - parse_json: bool = True, - wait_time_seconds: int = 20, # Long polling - **kwargs: Unpack[MessageDataReaderArgsSpec], - ) -> None: - ... - - async def _connect(self) -> None: - # Create aioboto3 SQS client - - async def _disconnect(self) -> None: - # Close session - - async def _receive(self) -> list[Any]: - # ReceiveMessage with MaxNumberOfMessages=chunk_size - # Long-polling with WaitTimeSeconds - - async def _convert(self, messages: list[Any]) -> dict[str, deque]: - # Parse message body (JSON) - # Extract fields - - async def _ack(self, messages: list[Any]) -> None: - # DeleteMessage for each processed message -``` - -#### `AWSSNSDataWriter` - -```python -class AWSSNSDataWriter(MessageDataWriter): - """Writes data to AWS SNS topic.""" - - def __init__( - self, - topic_arn: str, - region: str, - parse_json: bool = True, - **kwargs: Unpack[MessageDataWriterArgsSpec], - ) -> None: - ... - - async def _connect(self) -> None: - # Create aioboto3 SNS client - - async def _disconnect(self) -> None: - # Close session - - async def _send(self, messages: list[Any]) -> None: - # Publish each message to SNS topic - - async def _convert(self, data: dict[str, deque]) -> list[Any]: - # Convert field data to JSON strings -``` - -### 3. Apache Kafka - -**Dependencies**: `aiokafka` (added as optional dependency `kafka`) - -#### `KafkaDataReader` - -```python -class KafkaDataReader(MessageDataReader): - """Reads data from Apache Kafka topic.""" - - def __init__( - self, - bootstrap_servers: str | list[str], - topic: str, - group_id: str, - parse_json: bool = True, - **kwargs: Unpack[MessageDataReaderArgsSpec], - ) -> None: - ... - - async def _connect(self) -> None: - # Create AIOKafkaConsumer - # Subscribe to topic - - async def _disconnect(self) -> None: - # Stop consumer - - async def _receive(self) -> list[Any]: - # getmany() with timeout to fetch batch of messages - - async def _convert(self, messages: list[Any]) -> dict[str, deque]: - # Parse message value (JSON or raw bytes) - # Extract fields - - async def _ack(self, messages: list[Any]) -> None: - # Commit offsets for processed messages -``` - -#### `KafkaDataWriter` - -```python -class KafkaDataWriter(MessageDataWriter): - """Writes data to Apache Kafka topic.""" - - def __init__( - self, - bootstrap_servers: str | list[str], - topic: str, - parse_json: bool = True, - **kwargs: Unpack[MessageDataWriterArgsSpec], - ) -> None: - ... - - async def _connect(self) -> None: - # Create AIOKafkaProducer - - async def _disconnect(self) -> None: - # Stop producer - - async def _send(self, messages: list[Any]) -> None: - # send_and_wait for each message - - async def _convert(self, data: dict[str, deque]) -> list[Any]: - # Convert field data to JSON-encoded bytes -``` - ---- - -## Module Structure - -``` -plugboard/library/ -├── __init__.py # Updated exports -├── data_reader.py # Existing DataReader -├── data_writer.py # Existing DataWriter -├── file_io.py # Existing FileReader/FileWriter -├── sql_io.py # Existing SQLReader/SQLWriter -├── websocket_io.py # Existing WebsocketBase/Reader/Writer -├── message_reader.py # NEW: MessageDataReader base class -├── message_writer.py # NEW: MessageDataWriter base class -├── gcp_pubsub_io.py # NEW: GCPPubSubDataReader/Writer -├── aws_messaging_io.py # NEW: AWSSQSDataReader/Writer, AWSSNSDataWriter -└── kafka_io.py # NEW: KafkaDataReader/Writer -``` - ---- - -## Settings & Dependency Injection - -### Settings additions (`utils/settings.py`) - -```python -class _GCPPubSubSettings(BaseSettings): - project_id: str | None = None - model_config = SettingsConfigDict(env_prefix="GCP_PUBSUB_") - -class _AWSSettings(BaseSettings): - region: str | None = None - access_key_id: str | None = None - secret_access_key: str | None = None - model_config = SettingsConfigDict(env_prefix="AWS_") - -class _KafkaSettings(BaseSettings): - bootstrap_servers: str | list[str] | None = None - model_config = SettingsConfigDict(env_prefix="KAFKA_") -``` - -### DI additions (`utils/di.py`) - -No new DI resources are needed initially — each concrete implementation manages its own client lifecycle via `_connect()`/`_disconnect()`. DI resources can be added later when integrating against real infrastructure. - ---- - -## Optional Dependencies (`pyproject.toml`) - -```toml -[project.optional-dependencies] -gcp-pubsub = ["google-cloud-pubsub>=2.25,<3"] -aws-messaging = ["aioboto3>=13.0,<15"] -kafka = ["aiokafka>=0.11,<1"] -``` - ---- - -## Testing Strategy - -### Unit Tests (no cloud infrastructure required) - -For each base class and concrete implementation, we'll create unit tests using mocks: - -1. **`tests/unit/test_message_data_reader.py`**: - - Test `MessageDataReader` base class behavior with a mock implementation - - Test connection lifecycle (init → connect, destroy → disconnect) - - Test retry logic with simulated transient failures - - Test reconnection behavior - - Test chunked reading and buffering - - Test message acknowledgment - - Test field extraction from messages - -2. **`tests/unit/test_message_data_writer.py`**: - - Test `MessageDataWriter` base class behavior with a mock implementation - - Test connection lifecycle - - Test retry logic - - Test chunked writing and buffering - - Test flush on `run()` completion - - Test field data conversion to messages - -3. **`tests/unit/test_gcp_pubsub_io.py`**: - - Test `GCPPubSubDataReader`/`Writer` with mocked `google.cloud.pubsub` clients - - Test connection setup/teardown - - Test message receive/convert/ack - - Test message send/convert - -4. **`tests/unit/test_aws_messaging_io.py`**: - - Test `AWSSQSDataReader`/`AWSSNSDataWriter` with mocked `aioboto3` clients - - Test SQS receive/ack (delete) - - Test SNS publish - - Test long-polling configuration - -5. **`tests/unit/test_kafka_io.py`**: - - Test `KafkaDataReader`/`Writer` with mocked `aiokafka` clients - - Test consumer/producer lifecycle - - Test message receive/convert/commit - - Test message send/convert - -### Integration Tests (require cloud infrastructure — for later) - -Integration tests will be added in `tests/integration/` once cloud infrastructure is set up: -- `tests/integration/test_gcp_pubsub_io.py` -- `tests/integration/test_aws_messaging_io.py` -- `tests/integration/test_kafka_io.py` - -### Test patterns - -Following existing patterns: -- `pytest.mark.asyncio` for async tests -- Mock classes extending the abstract base (like `MockDataReader` in existing tests) -- `pytest.fixture` for test data -- Parametrized tests for chunk_size variations -- `structlog` for test logging - ---- - -## Implementation Order - -1. **Phase 1**: Base classes (`message_reader.py`, `message_writer.py`) + unit tests -2. **Phase 2**: Google Cloud PubSub implementation + unit tests -3. **Phase 3**: AWS SNS/SQS implementation + unit tests -4. **Phase 4**: Kafka implementation + unit tests -5. **Phase 5**: Update `__init__.py` exports, settings, pyproject.toml dependencies -6. **Phase 6**: Integration tests (when cloud infrastructure is available) - ---- - -## Error Handling - -### Custom exceptions - -```python -class MessageBrokerConnectionError(Exception): - """Raised when connection to message broker fails.""" - -class MessageBrokerTransientError(Exception): - """Raised on transient broker errors (eligible for retry).""" - -class MessageBrokerPermanentError(Exception): - """Raised on permanent broker errors (not eligible for retry).""" -``` - -### Error classification - -Each concrete implementation is responsible for classifying broker-specific errors into these categories. The base class handles retry logic based on these classifications. - ---- - -## Serialization - -Messages will be serialized as JSON by default (configurable via `parse_json` flag). This follows the pattern established by `WebsocketReader`/`WebsocketWriter` and ensures interoperability across different broker implementations. - -For the `_convert()` method: -- **Reader**: Parse JSON message data → extract named fields → `dict[str, deque]` -- **Writer**: Take `dict[str, deque]` → combine into records → serialize as JSON - ---- - -## Future Enhancements - -- Dead-letter queue handling -- Message filtering / schema validation -- Metrics collection (message rates, latencies) -- Schema registry integration (Avro, Protobuf) -- DI-managed broker connections (for connection pooling across components) -- Batch acknowledgment optimizations From b569b788908a6730930f7934e8beea774a218d9e Mon Sep 17 00:00:00 2001 From: Chris Knight Date: Sun, 4 Oct 2026 23:19:13 +0200 Subject: [PATCH 09/11] fix: address PR #284 review for message data classes Blocking issues: - An empty broker poll no longer ends the stream. `_receive` returning [] means "nothing yet", so `step()` waits and polls again like `WebsocketReader`; only a subclass raising `NoMoreDataException` closes the IO (B1). - A batch is acknowledged only after its last record is consumed, instead of after the first one. Previously a crash mid-batch lost the unread messages (B2). - `KafkaDataReader._ack` commits the offsets following the processed records, per partition, instead of the consumer position, which could mark unprocessed records as done (B3). - GCP reader and writer dispatch every synchronous PubSub call (client construction, pull, acknowledge, publish, future.result, stop, close) to a worker thread, so they no longer block the shared event loop (B4, B5). Design and consolidation: - `MessageDataReader`/`MessageDataWriter` now extend `DataReader`/`DataWriter`, reusing the buffer, binding and batching logic instead of copying it (W14). - Retry, backoff and reconnect live in one place: `plugboard.utils.retry` with a frozen `RetryPolicy` value object shared by both bases (W13, L1, L2). - A failing reconnect no longer escapes the retry loop; it costs an attempt and the broker error is raised once attempts run out (W12). - Row-to-message encoding is shared as `iter_records`/`encode_records`/ `encode_records_bytes`; the old per-row deque indexing was quadratic in buffer size and duplicated across three writers (W16, W19). - Connection access is serialised with an `asyncio.Lock` across ack and reconnect (W15). - Kafka writer submits all records then flushes once, instead of one round trip per record (W18). SQS deletes ack batches of up to 10 with `delete_message_batch` and reports partial failures rather than pretending success (L3). - The GCP publisher is stopped on disconnect, so a reconnect no longer leaks a channel and its threads (W17). - `MessageBrokerPermanentError`/`TransientError` are now actually raised and honoured: each broker maps its SDK's failures onto the hierarchy, and permanent errors skip the retries that previously burnt against a guaranteed failure (W1). Configuration and packaging: - The GCP/AWS/Kafka settings are read as constructor fallbacks via `resolve_argument`, so `GCP_PUBSUB_PROJECT_ID`, `AWS_REGION` and `KAFKA_BOOTSTRAP_SERVERS` work (W2). - `aws-messaging` now uses `aiobotocore` instead of `aioboto3`, which no longer forces the shared botocore/boto3 stack backwards (W5). - Broker SDKs are declared in the `test` dependency group, which also resolves the three `ty` unresolved-import findings (W10 root cause). - The six concrete broker classes are exported from `plugboard.library` (W3). - Removed the misleading broker-specific ArgsDicts that documented attributes they did not declare (W20). Tests: 126 tests, up from 72, all using teardown-aware monkeypatch fixtures instead of global `sys.modules` mocks, asserting broker calls rather than private attribute identity (W10, W11). New coverage for retry exhaustion and the delay cap, empty-batch handling, error-branch mapping, offset commits, reconnect failure, and the acknowledgment timing that previously lost records (W7, W8, W9). Record encoding is parametrised once rather than copied across three modules (L4). Docs: new Message Data usage page (extras, worked example, delivery semantics), broker settings in the configuration page, and the new extras in the README (W4). Test isolation: the S3 file tests pinned their bucket region so they no longer depend on the developer's ambient AWS profile, which the newer botocore resolves over `AWS_REGION`. --- README.md | 2 +- docs/usage/configuration.md | 20 + docs/usage/message-data.md | 150 +++++++ mkdocs.yaml | 1 + plugboard/library/__init__.py | 13 +- plugboard/library/aws_messaging_io.py | 306 ++++++++----- plugboard/library/gcp_pubsub_io.py | 237 +++++++---- plugboard/library/kafka_io.py | 236 ++++++---- plugboard/library/message_reader.py | 263 ++++++------ plugboard/library/message_writer.py | 274 ++++++------ plugboard/utils/__init__.py | 3 +- plugboard/utils/retry.py | 142 +++++++ plugboard/utils/settings.py | 36 ++ pyproject.toml | 7 +- tests/unit/test_aws_messaging_io.py | 541 ++++++++++++----------- tests/unit/test_file_reader_writer.py | 6 + tests/unit/test_gcp_pubsub_io.py | 567 ++++++++++++------------- tests/unit/test_kafka_io.py | 567 ++++++++++++------------- tests/unit/test_message_data_reader.py | 521 ++++++++++++----------- tests/unit/test_message_data_writer.py | 486 ++++++++++----------- uv.lock | 224 +++++----- 21 files changed, 2612 insertions(+), 1990 deletions(-) create mode 100644 docs/usage/message-data.md create mode 100644 plugboard/utils/retry.py diff --git a/README.md b/README.md index f38593eb..1efbf135 100644 --- a/README.md +++ b/README.md @@ -69,7 +69,7 @@ Optional integrations for different cloud providers can be installed using `plug Support for parallelisation and hyperparameter optimisation can be installed using `plugboard[ray]`. -Additional optional extras: `plugboard[llm]` for LLM components, `plugboard[redis]` for Redis-based connectors, `plugboard[omq]` for the pyomq backend for ZMQ connectors, and `plugboard[websockets]` for WebSocket I/O. +Additional optional extras: `plugboard[llm]` for LLM components, `plugboard[redis]` for Redis-based connectors, `plugboard[omq]` for the pyomq backend for ZMQ connectors, `plugboard[websockets]` for WebSocket I/O, and `plugboard[gcp-pubsub]`, `plugboard[aws-messaging]` or `plugboard[kafka]` for [message data components](https://docs.plugboard.dev/usage/message-data/). ## ⚡ Quickstart with AI diff --git a/docs/usage/configuration.md b/docs/usage/configuration.md index 2f9e70db..5430c2eb 100644 --- a/docs/usage/configuration.md +++ b/docs/usage/configuration.md @@ -24,6 +24,26 @@ Plugboard can make use of a message broker for data exchange between components | `RABBITMQ_URL` | URL for RabbitMQ AMQP message broker (must include credentials if required) | | | `REDIS_URL` | URL for Redis message broker (must include credentials if required) | | +These brokers carry data *between components* within a run. To read or write records +through an external broker as part of the model itself, see +[Message Data](message-data.md). + +### Message data brokers + +The message data components ([`MessageDataReader`][plugboard.library.MessageDataReader] / +[`MessageDataWriter`][plugboard.library.MessageDataWriter] and their broker +implementations) take connection details as constructor arguments, falling back to the +evironment below when they are not supplied: + +| Option Name | Description | Default Value | +|---------------------------|----------------------------------------|---------------| +| `GCP_PUBSUB_PROJECT_ID` | GCP project for PubSub topics and subscriptions | | +| `AWS_REGION` | AWS region for SQS and SNS clients | | +| `KAFKA_BOOTSTRAP_SERVERS` | Kafka broker address(es) | | + +Note that `AWS_REGION` is also read by the AWS SDK itself, so setting it affects both +Plugboard's defaults and the credential/endpoint resolution of the underlying client. + ## Job ID Each plugboard run has a unique job ID associated with it. This is used to: track state for each run; and separate data messages between runs when using a message broker. Typically, a run would be started without explicitly setting the job ID, in which case a unique job ID will be created automatically. However, there are instances when it may be desirable to specify the job ID, such as stopping a run and resuming the same run later with the existing persisted state. In these scenarios the job ID can be set with the below environment variable which will then be used by any `StateBackend`, `Process` and `Component` while the value is set. diff --git a/docs/usage/message-data.md b/docs/usage/message-data.md new file mode 100644 index 00000000..039db62a --- /dev/null +++ b/docs/usage/message-data.md @@ -0,0 +1,150 @@ +# Message Data + +Message data components read and write **records** through a pub/sub message broker, +rather than through a channel wired between components. They are the broker equivalent +of [`FileReader`][plugboard.library.FileReader] / [`FileWriter`][plugboard.library.FileWriter]: +each message carries one record, and the fields named by `field_names` become the +component's outputs (for a reader) or inputs (for a writer). + +| Class | Broker | Reads from | Writes to | +|---|---|---|---| +| [`GCPPubSubDataReader`][plugboard.library.GCPPubSubDataReader] | Google Cloud PubSub | a subscription | | +| [`GCPPubSubDataWriter`][plugboard.library.GCPPubSubDataWriter] | Google Cloud PubSub | | a topic | +| [`AWSSQSDataReader`][plugboard.library.AWSSQSDataReader] | AWS SQS | a queue | | +| [`AWSSNSDataWriter`][plugboard.library.AWSSNSDataWriter] | AWS SNS | | a topic | +| [`KafkaDataReader`][plugboard.library.KafkaDataReader] | Apache Kafka | a topic + consumer group | | +| [`KafkaDataWriter`][plugboard.library.KafkaDataWriter] | Apache Kafka | | a topic | + +All six build on the [`MessageDataReader`][plugboard.library.MessageDataReader] and +[`MessageDataWriter`][plugboard.library.MessageDataWriter] base classes, which handle +connecting, reconnection with backoff, batching and acknowledgment. + +## Installing + +Each broker uses its own client library, supplied as an extra: + +```shell +pip install "plugboard[gcp-pubsub]" # google-cloud-pubsub +pip install "plugboard[aws-messaging]" # aiobotocore +pip install "plugboard[kafka]" # aiokafka +``` + +Constructing a component without its extra raises an `ImportError` naming the extra to +install. + +## Reading and writing records + +Messages are expected to carry a JSON object per record, whose keys match +`field_names`: + +```python +import asyncio + +from plugboard.component import Component, IOController as IO +from plugboard.library import AWSSNSDataWriter, GCPPubSubDataReader +from plugboard.process import LocalProcess +from plugboard.schemas import ConnectorSpec + + +class Bucket(Component): + """Turns a raw measurement into a bucketed label before it leaves the model.""" + + io = IO(inputs=["x", "y"], outputs=["bucket"]) + + async def step(self) -> None: + self.bucket = "high" if self.y > 50 else "low" + + +async def main() -> None: + reader = GCPPubSubDataReader( + name="reader", + subscription_id="measurements-pull", + project_id="plugboard-dev", # Optional; defaults to GCP_PUBSUB_PROJECT_ID + field_names=["x", "y"], + chunk_size=20, + ) + bucket = Bucket(name="bucket") + writer = AWSSNSDataWriter( + name="writer", + topic_arn="arn:aws:sns:eu-west-1:123456789012:measurements", + region="eu-west-1", # Optional; defaults to AWS_REGION + field_names=["bucket"], + chunk_size=20, + ) + + process = LocalProcess( + name="measurements", + components=[reader, bucket, writer], + connectors=[ + ConnectorSpec(source="reader.x", target="bucket.x"), + ConnectorSpec(source="reader.y", target="bucket.y"), + ConnectorSpec(source="bucket.bucket", target="writer.bucket"), + ], + ) + await process.run() + + +asyncio.run(main()) +``` + +The components are ordinary `Component`s, so they also work in a `RayProcess` and with +any connector. + +## Behaviour worth knowing + +**A reader waits; it does not finish.** An empty poll from a broker means "nothing yet", +not "no more data" — unlike a file, which ends. A reader started before its producer +keeps running and picks up messages as they arrive. To end a reader's stream, the broker +must report the source as gone (for example a deleted PubSub subscription or SQS queue), +which raises `IOStreamClosedError` like other data readers do. Stop the process to end a +run that would otherwise wait forever. + +**Messages are acknowledged after they are processed.** A batch fetched from the broker +is acknowledged only once every record in it has been published downstream, so a crash +mid-batch leaves the unread messages to be redelivered (at-least-once). For SQS, +"acknowledge" means deleting the messages; for Kafka it means committing the offsets of +the processed records, per partition. + +**Batching is bounded by the broker.** `chunk_size` sets how many messages are fetched +per poll (and per send for writers). SQS delivers at most 10 messages per call, so a +larger `chunk_size` is capped there and a warning is logged. + +**Failures are retried with backoff.** Transient broker errors trigger a reconnect and a +retry, doubling the delay up to a cap. Permanent errors — such as `AccessDenied` or a +deleted topic — fail immediately instead of burning the retries. Both the classification +and the policy are handled by the base classes: + +```python +from plugboard.library import KafkaDataReader +from plugboard.utils.retry import RetryPolicy + +reader = KafkaDataReader( + name="reader", + topic="measurements", + group_id="plugboard", + field_names=["x", "y"], + bootstrap_servers="localhost:9092", # Optional; defaults to KAFKA_BOOTSTRAP_SERVERS + retry_policy=RetryPolicy(max_retries=5, base_delay=0.5, max_delay=30.0), + idle_poll_delay=0.5, # Pause between empty polls, if the broker returns immediately +) +``` + +**Message encoding.** By default each record is JSON-encoded (`parse_json=True`). With +`parse_json=False` a writer sends only the first field's value, and a reader exposes the +raw payload under a `data` field — useful for single-value or binary messages. + +## Configuration + +Connection details can be passed explicitly or read from the environment, which lets the +same model run against a different account or region without code changes. Explicit +arguments always win. + +| Option Name | Description | Used by | +|---------------------------|------------------------------------------|---------| +| `GCP_PUBSUB_PROJECT_ID` | Default GCP project for PubSub topics | `GCPPubSubDataReader`, `GCPPubSubDataWriter` | +| `AWS_REGION` | Default AWS region for SQS/SNS clients | `AWSSQSDataReader`, `AWSSNSDataWriter` | +| `KAFKA_BOOTSTRAP_SERVERS` | Default Kafka broker address(es) | `KafkaDataReader`, `KafkaDataWriter` | + +Credentials themselves are never configured through Plugboard: the clients use their +normal provider mechanisms (Application Default Credentials for PubSub, the standard AWS +credential chain, SASL settings for Kafka). diff --git a/mkdocs.yaml b/mkdocs.yaml index 1dabbd9c..a78390a9 100644 --- a/mkdocs.yaml +++ b/mkdocs.yaml @@ -129,6 +129,7 @@ nav: - Event-driven models: examples/tutorials/event-driven-models.md - Tuning a process: examples/tutorials/tuning-a-process.md - Configuration: usage/configuration.md + - Message Data: usage/message-data.md - AI-Assisted Development: usage/ai.md - Topics: usage/topics.md - Demos: diff --git a/plugboard/library/__init__.py b/plugboard/library/__init__.py index 51b77ab1..c4bd2ac2 100644 --- a/plugboard/library/__init__.py +++ b/plugboard/library/__init__.py @@ -1,8 +1,11 @@ """Provides implementations of Plugboard objects for use in user models.""" +from .aws_messaging_io import AWSSNSDataWriter, AWSSQSDataReader from .data_reader import DataReader from .data_writer import DataWriter from .file_io import FileReader, FileWriter +from .gcp_pubsub_io import GCPPubSubDataReader, GCPPubSubDataWriter +from .kafka_io import KafkaDataReader, KafkaDataWriter from .llm import LLMChat, LLMImageProcessor from .message_reader import MessageDataReader from .message_writer import MessageDataWriter @@ -11,12 +14,18 @@ __all__ = [ + "AWSSQSDataReader", + "AWSSNSDataWriter", "DataReader", "DataWriter", - "LLMChat", - "LLMImageProcessor", "FileReader", "FileWriter", + "GCPPubSubDataReader", + "GCPPubSubDataWriter", + "KafkaDataReader", + "KafkaDataWriter", + "LLMChat", + "LLMImageProcessor", "MessageDataReader", "MessageDataWriter", "SQLReader", diff --git a/plugboard/library/aws_messaging_io.py b/plugboard/library/aws_messaging_io.py index 01852be3..092104a4 100644 --- a/plugboard/library/aws_messaging_io.py +++ b/plugboard/library/aws_messaging_io.py @@ -3,133 +3,213 @@ from __future__ import annotations from collections import deque +import contextlib import json import typing as _t -from plugboard.exceptions import NoMoreDataException +from plugboard.exceptions import ( + MessageBrokerConnectionError, + MessageBrokerPermanentError, + MessageBrokerTransientError, + NoMoreDataException, +) from plugboard.library.message_reader import MessageDataReader, MessageDataReaderArgsDict -from plugboard.library.message_writer import MessageDataWriter, MessageDataWriterArgsDict -from plugboard.utils import depends_on_optional +from plugboard.library.message_writer import ( + MessageDataWriter, + MessageDataWriterArgsDict, + encode_records, +) +from plugboard.utils import DI, depends_on_optional +from plugboard.utils.settings import resolve_argument try: - import aioboto3 -except ImportError: # pragma: no cover - pass - - -class AWSSQSDataReaderArgsDict(MessageDataReaderArgsDict): - """Specification of the `AWSSQSDataReader` constructor arguments. - - Attributes: - queue_url: The SQS queue URL. - region: The AWS region. - parse_json: Whether to parse message bodies as JSON. - wait_time_seconds: Long-polling wait time in seconds. - """ - - pass - + from aiobotocore.session import get_session + from botocore.exceptions import BotoCoreError, ClientError -class AWSSNSDataWriterArgsDict(MessageDataWriterArgsDict): - """Specification of the `AWSSNSDataWriter` constructor arguments. - - Attributes: - topic_arn: The SNS topic ARN. - region: The AWS region. - parse_json: Whether to encode message data as JSON. + _CLIENT_ERRORS: tuple[type[Exception], ...] = (ClientError,) + # BotoCoreError covers transport failures: timeouts, DNS and connection resets. + _TRANSPORT_ERRORS: tuple[type[Exception], ...] = (BotoCoreError,) +except ImportError: # pragma: no cover + # With the extra not installed, `depends_on_optional` blocks construction and + # classification is moot: `isinstance(error, ())` never matches. + _CLIENT_ERRORS = () + _TRANSPORT_ERRORS = () + + +# SQS returns at most this many messages per `ReceiveMessage` call. +_SQS_MAX_MESSAGES = 10 +# SQS accepts at most this many entries per `DeleteMessageBatch` call. +_SQS_MAX_BATCH = 10 + +_PERMANENT_CODES = frozenset( + { + "AccessDenied", + "AuthorizationError", + "IncompleteSignature", + "InvalidAddress", + "InvalidClientTokenId", + "InvalidMessageContents", + "MalformedXML", + "MessageNotInflight", + "ReceiptHandleIsInvalid", + "SignatureDoesNotMatch", + "UnrecognizedClientException", + } +) + +_TRANSIENT_CODES = frozenset( + { + "InternalError", + "RequestExpired", + "RequestThrottled", + "ServiceUnavailable", + "ThrottledException", + "ThrottlingException", + "TooManyRequestsException", + } +) + +_MISSING_RESOURCE_CODES = frozenset( + { + "AWS.SimpleQueueService.NonExistentQueue", + "NotFound", + "QueueDoesNotExist", + "NoSuchTopic", + } +) + + +def _error_code(error: Exception) -> _t.Optional[str]: + """Extracts the AWS error code from a client error, if it has one.""" + if not isinstance(error, _CLIENT_ERRORS): + return None + response = getattr(error, "response", None) + if not isinstance(response, dict): + return None + return response.get("Error", {}).get("Code") + + +def _wrap_error(error: Exception, *, exhausted: bool = False) -> Exception: + """Maps an AWS SDK error onto the Plugboard message broker exception hierarchy. + + Args: + error: The error raised by the AWS client. + exhausted: Whether a missing queue/topic should be reported as an exhausted + source. Only meaningful for a reader: a writer cannot recover from a + missing topic, so it is a permanent failure there. + + Returns: + A `NoMoreDataException` for a deleted queue (readers only), a + `MessageBrokerPermanentError` for a rejection that will not change on retry, a + `MessageBrokerTransientError` for a failure worth retrying, otherwise the + original error. """ - - pass + code = _error_code(error) + if code in _MISSING_RESOURCE_CODES: + return NoMoreDataException() if exhausted else MessageBrokerPermanentError(str(error)) + if code in _PERMANENT_CODES: + return MessageBrokerPermanentError(str(error)) + if code in _TRANSIENT_CODES or isinstance(error, _TRANSPORT_ERRORS): + return MessageBrokerTransientError(str(error)) + return error class AWSSQSDataReader(MessageDataReader): """Reads data from an AWS SQS queue. - Messages are received from the queue using long-polling and converted - to field values. Messages are deleted from the queue after processing - (acknowledgment). + Messages are received using long polling and converted to field values. + Acknowledging a batch deletes its messages from the queue, which happens once all + of the batch's records have been consumed, so a failure before that point leaves + the remaining messages visible for redelivery. """ - @depends_on_optional("aioboto3", extra="aws-messaging") + @depends_on_optional("aiobotocore", extra="aws-messaging") def __init__( self, queue_url: str, - region: str, + region: _t.Optional[str] = None, parse_json: bool = True, wait_time_seconds: int = 20, - **kwargs: _t.Unpack[AWSSQSDataReaderArgsDict], + **kwargs: _t.Unpack[MessageDataReaderArgsDict], ) -> None: """Instantiates the `AWSSQSDataReader`. Args: queue_url: The SQS queue URL. - region: The AWS region. + region: Optional; The AWS region. Falls back to `AWS_REGION`. parse_json: Whether to parse message bodies as JSON. wait_time_seconds: Long-polling wait time in seconds (max 20). **kwargs: Additional keyword arguments for [`MessageDataReader`][plugboard.library.MessageDataReader]. """ + region = resolve_argument( + region, DI.settings.resolve_sync().aws.region, "region", "AWS_REGION" + ) kwargs.setdefault("topic", queue_url) super().__init__(**kwargs) + if self._chunk_size and self._chunk_size > _SQS_MAX_MESSAGES: + self._logger.warning( + "SQS returns at most 10 messages per receive, so chunk_size is capped", + chunk_size=self._chunk_size, + max_messages=_SQS_MAX_MESSAGES, + ) self._queue_url = queue_url self._region = region self._parse_json = parse_json self._wait_time_seconds = wait_time_seconds - self._session: _t.Any = None self._client: _t.Any = None + self._exit_stack: _t.Optional[contextlib.AsyncExitStack] = None async def _connect(self) -> None: - """Creates an SQS client session.""" - self._session = aioboto3.Session() - self._client_ctx = self._session.client("sqs", region_name=self._region) - self._client = await self._client_ctx.__aenter__() + """Creates an SQS client.""" + self._exit_stack = contextlib.AsyncExitStack() + self._client = await self._exit_stack.enter_async_context( + get_session().create_client("sqs", region_name=self._region) + ) async def _disconnect(self) -> None: - """Closes the SQS client session.""" - if self._client is not None: - try: - await self._client_ctx.__aexit__(None, None, None) - except Exception: # noqa: S110 - pass + """Closes the SQS client.""" + if self._exit_stack is not None: + exit_stack = self._exit_stack + self._exit_stack = None self._client = None - self._session = None + await exit_stack.aclose() async def _receive(self) -> list[_t.Any]: - """Receives a batch of messages from the SQS queue. + """Receives a batch of messages from the SQS queue using long polling. Returns: - A list of SQS message dicts. + A list of SQS message dicts, empty if the long poll timed out. Raises: - NoMoreDataException: If the queue does not exist. + NoMoreDataException: If the queue no longer exists. + MessageBrokerError: If the receive failed for a retryable or permanent reason. """ if self._client is None: - raise RuntimeError("SQS client not initialized") - max_messages = min(self._chunk_size or 10, 10) # SQS max is 10 + raise MessageBrokerConnectionError("SQS client is not connected") + request = { + "QueueUrl": self._queue_url, + "MaxNumberOfMessages": min(self._chunk_size or _SQS_MAX_MESSAGES, _SQS_MAX_MESSAGES), + "WaitTimeSeconds": self._wait_time_seconds, + } try: - response = await self._client.receive_message( - QueueUrl=self._queue_url, - MaxNumberOfMessages=max_messages, - WaitTimeSeconds=self._wait_time_seconds, - ) - except Exception as e: - if "QueueDoesNotExist" in str(type(e).__name__) or "NonExistentQueue" in str(e): - raise NoMoreDataException from e - raise + response = await self._client.receive_message(**request) + except Exception as error: # noqa: BLE001 + raise _wrap_error(error, exhausted=True) from error return response.get("Messages", []) - async def _convert(self, messages: list[_t.Any]) -> dict[str, deque]: + async def _convert(self, data: list[_t.Any]) -> dict[str, deque]: """Converts SQS messages to a field buffer. Args: - messages: A list of SQS message dicts. + data: A list of SQS message dicts. Returns: A dictionary mapping field names to deques of field values. """ converted: dict[str, deque] = {field: deque() for field in self.io.outputs} - for msg in messages: + for msg in data: body = msg.get("Body", "") if self._parse_json: record = json.loads(body) @@ -142,80 +222,107 @@ async def _convert(self, messages: list[_t.Any]) -> dict[str, deque]: async def _ack(self, messages: list[_t.Any]) -> None: """Deletes processed messages from the SQS queue. + Deletes run in batches of up to 10, the SQS batch limit. + Args: messages: The SQS message dicts to delete. + + Raises: + MessageBrokerConnectionError: If the client is not connected. + MessageBrokerTransientError: If any message could not be deleted. """ if self._client is None: - raise RuntimeError("SQS client not initialized") - for msg in messages: - receipt_handle = msg.get("ReceiptHandle") - if receipt_handle: - await self._client.delete_message( - QueueUrl=self._queue_url, ReceiptHandle=receipt_handle + raise MessageBrokerConnectionError("SQS client is not connected") + for start in range(0, len(messages), _SQS_MAX_BATCH): + entries = [ + {"Id": str(index), "ReceiptHandle": msg["ReceiptHandle"]} + for index, msg in enumerate(messages[start : start + _SQS_MAX_BATCH]) + if msg.get("ReceiptHandle") + ] + if not entries: + continue + try: + response = await self._client.delete_message_batch( + QueueUrl=self._queue_url, DeleteMessageRequests=entries + ) + except Exception as error: # noqa: BLE001 + raise _wrap_error(error) from error + if errors := response.get("Errors"): + # A partially failed delete leaves those messages visible for + # redelivery, so report it rather than pretending the batch is done. + raise MessageBrokerTransientError( + f"Failed to delete {len(errors)} SQS message(s): {errors}" ) class AWSSNSDataWriter(MessageDataWriter): """Writes data to an AWS SNS topic. - Field data is converted to JSON-encoded messages and published - to the specified SNS topic. + Field data is converted to JSON-encoded messages and published to the specified + topic. SNS has no batch publish API for standard topics, so records are published + one at a time per send. """ - @depends_on_optional("aioboto3", extra="aws-messaging") + @depends_on_optional("aiobotocore", extra="aws-messaging") def __init__( self, topic_arn: str, - region: str, + region: _t.Optional[str] = None, parse_json: bool = True, - **kwargs: _t.Unpack[AWSSNSDataWriterArgsDict], + **kwargs: _t.Unpack[MessageDataWriterArgsDict], ) -> None: """Instantiates the `AWSSNSDataWriter`. Args: topic_arn: The SNS topic ARN. - region: The AWS region. + region: Optional; The AWS region. Falls back to `AWS_REGION`. parse_json: Whether to encode message data as JSON. **kwargs: Additional keyword arguments for [`MessageDataWriter`][plugboard.library.MessageDataWriter]. """ + region = resolve_argument( + region, DI.settings.resolve_sync().aws.region, "region", "AWS_REGION" + ) kwargs.setdefault("topic", topic_arn) super().__init__(**kwargs) self._topic_arn = topic_arn self._region = region self._parse_json = parse_json - self._session: _t.Any = None self._client: _t.Any = None + self._exit_stack: _t.Optional[contextlib.AsyncExitStack] = None async def _connect(self) -> None: - """Creates an SNS client session.""" - self._session = aioboto3.Session() - self._client_ctx = self._session.client("sns", region_name=self._region) - self._client = await self._client_ctx.__aenter__() + """Creates an SNS client.""" + self._exit_stack = contextlib.AsyncExitStack() + self._client = await self._exit_stack.enter_async_context( + get_session().create_client("sns", region_name=self._region) + ) async def _disconnect(self) -> None: - """Closes the SNS client session.""" - if self._client is not None: - try: - await self._client_ctx.__aexit__(None, None, None) - except Exception: # noqa: S110 - pass + """Closes the SNS client.""" + if self._exit_stack is not None: + exit_stack = self._exit_stack + self._exit_stack = None self._client = None - self._session = None + await exit_stack.aclose() async def _send(self, messages: list[_t.Any]) -> None: """Publishes messages to the SNS topic. Args: messages: A list of message strings to publish. + + Raises: + MessageBrokerConnectionError: If the client is not connected. + MessageBrokerError: If a publish failed. """ if self._client is None: - raise RuntimeError("SNS client not initialized") + raise MessageBrokerConnectionError("SNS client is not connected") for msg_data in messages: - await self._client.publish( - TopicArn=self._topic_arn, - Message=msg_data, - ) + try: + await self._client.publish(TopicArn=self._topic_arn, Message=msg_data) + except Exception as error: # noqa: BLE001 + raise _wrap_error(error) from error async def _convert(self, data: dict[str, deque]) -> list[_t.Any]: """Converts field buffer data to JSON-encoded message strings. @@ -226,13 +333,4 @@ async def _convert(self, data: dict[str, deque]) -> list[_t.Any]: Returns: A list of message strings ready to publish. """ - completed_rows = min(len(d) for d in data.values()) if data else 0 - messages: list[str] = [] - for i in range(completed_rows): - record = {field: data[field][i] for field in data} - if self._parse_json: - messages.append(json.dumps(record)) - else: - first_field = next(iter(record.values())) - messages.append(str(first_field)) - return messages + return encode_records(data, self._parse_json) diff --git a/plugboard/library/gcp_pubsub_io.py b/plugboard/library/gcp_pubsub_io.py index 90d67094..1dcdb795 100644 --- a/plugboard/library/gcp_pubsub_io.py +++ b/plugboard/library/gcp_pubsub_io.py @@ -2,129 +2,178 @@ from __future__ import annotations +import asyncio from collections import deque import json import typing as _t -from plugboard.exceptions import NoMoreDataException +from plugboard.exceptions import ( + MessageBrokerConnectionError, + MessageBrokerPermanentError, + MessageBrokerTransientError, + NoMoreDataException, +) from plugboard.library.message_reader import MessageDataReader, MessageDataReaderArgsDict -from plugboard.library.message_writer import MessageDataWriter, MessageDataWriterArgsDict -from plugboard.utils import depends_on_optional +from plugboard.library.message_writer import ( + MessageDataWriter, + MessageDataWriterArgsDict, + encode_records_bytes, +) +from plugboard.utils import DI, depends_on_optional +from plugboard.utils.settings import resolve_argument try: + from google.api_core import exceptions as gcp_exceptions from google.cloud import pubsub_v1 -except ImportError: # pragma: no cover - pass - - -class GCPPubSubDataReaderArgsDict(MessageDataReaderArgsDict): - """Specification of the `GCPPubSubDataReader` constructor arguments. - Attributes: - project_id: The GCP project ID. - subscription_id: The PubSub subscription ID. - parse_json: Whether to parse message data as JSON. - """ + # PubSub rejections that retrying cannot resolve, so the retry loop should stop + # rather than back off repeatedly against a guaranteed failure. + _PERMANENT_ERRORS: tuple[type[Exception], ...] = ( + gcp_exceptions.PermissionDenied, + gcp_exceptions.Unauthorized, + gcp_exceptions.InvalidArgument, + gcp_exceptions.MethodNotImplemented, + gcp_exceptions.FailedPrecondition, + ) + + # PubSub failures that are expected to clear on their own, so they are worth retrying. + _TRANSIENT_ERRORS: tuple[type[Exception], ...] = ( + gcp_exceptions.ServiceUnavailable, + gcp_exceptions.DeadlineExceeded, + gcp_exceptions.Aborted, + gcp_exceptions.InternalServerError, + gcp_exceptions.BadGateway, + gcp_exceptions.TooManyRequests, + ) + + _NOT_FOUND_ERRORS: tuple[type[Exception], ...] = (gcp_exceptions.NotFound,) +except ImportError: # pragma: no cover + # With the extra not installed, `depends_on_optional` blocks construction and + # classification is moot: `isinstance(error, ())` never matches. + _PERMANENT_ERRORS = () + _TRANSIENT_ERRORS = () + _NOT_FOUND_ERRORS = () - pass +def _wrap_error(error: Exception) -> Exception: + """Maps a PubSub SDK error onto the Plugboard message broker exception hierarchy. -class GCPPubSubDataWriterArgsDict(MessageDataWriterArgsDict): - """Specification of the `GCPPubSubDataWriter` constructor arguments. + Args: + error: The error raised by the PubSub client. - Attributes: - project_id: The GCP project ID. - topic_id: The PubSub topic ID. - parse_json: Whether to encode message data as JSON. + Returns: + A `NoMoreDataException` for a deleted topic/subscription, a + `MessageBrokerPermanentError` for a rejection that will not change on retry, a + `MessageBrokerTransientError` for a failure worth retrying, otherwise the + original error. """ - - pass + if isinstance(error, _NOT_FOUND_ERRORS): + return NoMoreDataException() + if isinstance(error, _PERMANENT_ERRORS): + return MessageBrokerPermanentError(str(error)) + if isinstance(error, _TRANSIENT_ERRORS): + return MessageBrokerTransientError(str(error)) + return error class GCPPubSubDataReader(MessageDataReader): """Reads data from a Google Cloud PubSub subscription. - Messages are pulled from the subscription in batches and converted - to field values. Messages are acknowledged after processing. + Messages are pulled from the subscription in batches and converted to field + values. A batch is acknowledged once all of its records have been consumed, so a + failure before that point leaves the remaining messages for redelivery. + + The PubSub client library is synchronous, so every call is dispatched to a worker + thread to keep the process event loop free for other components. """ @depends_on_optional("google.cloud.pubsub_v1", extra="gcp-pubsub") def __init__( self, - project_id: str, subscription_id: str, + project_id: _t.Optional[str] = None, parse_json: bool = True, - **kwargs: _t.Unpack[GCPPubSubDataReaderArgsDict], + pull_timeout: float = 30.0, + **kwargs: _t.Unpack[MessageDataReaderArgsDict], ) -> None: """Instantiates the `GCPPubSubDataReader`. Args: - project_id: The GCP project ID. subscription_id: The PubSub subscription ID. + project_id: Optional; The GCP project ID. Falls back to + `GCP_PUBSUB_PROJECT_ID`. parse_json: Whether to parse message data as JSON. + pull_timeout: Seconds to wait for a `pull()` call to return. **kwargs: Additional keyword arguments for [`MessageDataReader`][plugboard.library.MessageDataReader]. """ + project_id = resolve_argument( + project_id, + DI.settings.resolve_sync().gcp_pubsub.project_id, + "project_id", + "GCP_PUBSUB_PROJECT_ID", + ) kwargs.setdefault("topic", f"{project_id}/{subscription_id}") super().__init__(**kwargs) self._project_id = project_id self._subscription_id = subscription_id self._subscription_path = f"projects/{project_id}/subscriptions/{subscription_id}" self._parse_json = parse_json + self._pull_timeout = pull_timeout self._subscriber: _t.Optional[pubsub_v1.SubscriberClient] = None async def _connect(self) -> None: """Creates a PubSub subscriber client.""" - self._subscriber = pubsub_v1.SubscriberClient() + self._subscriber = await asyncio.to_thread(pubsub_v1.SubscriberClient) async def _disconnect(self) -> None: """Closes the PubSub subscriber client.""" if self._subscriber is not None: - self._subscriber.close() + subscriber = self._subscriber self._subscriber = None + await asyncio.to_thread(subscriber.close) async def _receive(self) -> list[_t.Any]: """Pulls a batch of messages from the PubSub subscription. Returns: - A list of PubSub `Message` objects. + A list of PubSub `ReceivedMessage` objects, empty if the pull timed out. Raises: - NoMoreDataException: If the subscription is deleted or unreachable. + NoMoreDataException: If the subscription no longer exists. + MessageBrokerError: If the pull failed for a retryable or permanent reason. """ if self._subscriber is None: - raise RuntimeError("Subscriber client not initialized") - max_messages = self._chunk_size or 10 + raise MessageBrokerConnectionError("Subscriber client is not connected") + request = { + "subscription": self._subscription_path, + "max_messages": self._chunk_size or 10, + } try: - response = self._subscriber.pull( - request={"subscription": self._subscription_path, "max_messages": max_messages}, - timeout=30.0, + response = await asyncio.to_thread( + self._subscriber.pull, request=request, timeout=self._pull_timeout ) - except Exception as e: - if "NOT_FOUND" in str(e) or "Subscription not found" in str(e): - raise NoMoreDataException from e - raise - if not response.received_messages: - return [] + except Exception as error: # noqa: BLE001 + raise _wrap_error(error) from error return list(response.received_messages) - async def _convert(self, messages: list[_t.Any]) -> dict[str, deque]: + async def _convert(self, data: list[_t.Any]) -> dict[str, deque]: """Converts PubSub messages to a field buffer. Args: - messages: A list of `ReceivedMessage` objects. + data: A list of `ReceivedMessage` objects. Returns: A dictionary mapping field names to deques of field values. """ converted: dict[str, deque] = {field: deque() for field in self.io.outputs} - for msg_wrapper in messages: - data = msg_wrapper.message.data + for msg_wrapper in data: + payload = msg_wrapper.message.data if self._parse_json: - record = json.loads(data.decode("utf-8")) + record = json.loads(payload.decode("utf-8")) else: - record = {"data": data} + record = {"data": payload} for field in self.io.outputs: converted[field].append(record.get(field)) return converted @@ -134,71 +183,107 @@ async def _ack(self, messages: list[_t.Any]) -> None: Args: messages: The `ReceivedMessage` objects to acknowledge. + + Raises: + MessageBrokerConnectionError: If the client is not connected. + MessageBrokerError: If the acknowledgment failed. """ if self._subscriber is None: - raise RuntimeError("Subscriber client not initialized") + raise MessageBrokerConnectionError("Subscriber client is not connected") ack_ids = [msg_wrapper.ack_id for msg_wrapper in messages] - self._subscriber.acknowledge( - request={"subscription": self._subscription_path, "ack_ids": ack_ids} - ) + request = {"subscription": self._subscription_path, "ack_ids": ack_ids} + try: + await asyncio.to_thread(self._subscriber.acknowledge, request=request) + except Exception as error: # noqa: BLE001 + raise _wrap_error(error) from error class GCPPubSubDataWriter(MessageDataWriter): """Writes data to a Google Cloud PubSub topic. - Field data is converted to JSON-encoded messages and published - to the specified topic. + Field data is converted to JSON-encoded messages and published to the specified + topic. The PubSub client library is synchronous, so publishing and waiting for + broker confirmation run in a worker thread. """ @depends_on_optional("google.cloud.pubsub_v1", extra="gcp-pubsub") def __init__( self, - project_id: str, topic_id: str, + project_id: _t.Optional[str] = None, parse_json: bool = True, - **kwargs: _t.Unpack[GCPPubSubDataWriterArgsDict], + publish_timeout: float = 60.0, + **kwargs: _t.Unpack[MessageDataWriterArgsDict], ) -> None: """Instantiates the `GCPPubSubDataWriter`. Args: - project_id: The GCP project ID. topic_id: The PubSub topic ID. + project_id: Optional; The GCP project ID. Falls back to + `GCP_PUBSUB_PROJECT_ID`. parse_json: Whether to encode message data as JSON. + publish_timeout: Seconds to wait for each publish to be confirmed. **kwargs: Additional keyword arguments for [`MessageDataWriter`][plugboard.library.MessageDataWriter]. """ + project_id = resolve_argument( + project_id, + DI.settings.resolve_sync().gcp_pubsub.project_id, + "project_id", + "GCP_PUBSUB_PROJECT_ID", + ) kwargs.setdefault("topic", f"{project_id}/{topic_id}") super().__init__(**kwargs) self._project_id = project_id self._topic_id = topic_id self._topic_path = f"projects/{project_id}/topics/{topic_id}" self._parse_json = parse_json + self._publish_timeout = publish_timeout self._publisher: _t.Optional[pubsub_v1.PublisherClient] = None async def _connect(self) -> None: """Creates a PubSub publisher client.""" - self._publisher = pubsub_v1.PublisherClient() + self._publisher = await asyncio.to_thread(pubsub_v1.PublisherClient) async def _disconnect(self) -> None: - """Closes the PubSub publisher client.""" + """Stops the PubSub publisher, flushing any queued messages and closing its channel. + + Raises: + MessageBrokerError: If the publisher fails to stop. + """ if self._publisher is not None: + publisher = self._publisher self._publisher = None + try: + await asyncio.to_thread(publisher.stop) + except Exception as error: # noqa: BLE001 + raise _wrap_error(error) from error async def _send(self, messages: list[_t.Any]) -> None: - """Publishes messages to the PubSub topic. + """Publishes messages to the PubSub topic and waits for broker confirmation. Args: messages: A list of bytes objects to publish. + + Raises: + MessageBrokerConnectionError: If the client is not connected. + MessageBrokerError: If a publish failed or was not confirmed in time. """ if self._publisher is None: - raise RuntimeError("Publisher client not initialized") - futures = [] - for msg_data in messages: - future = self._publisher.publish(self._topic_path, data=msg_data) - futures.append(future) - # Wait for all publishes to complete - for future in futures: - future.result(timeout=60.0) + raise MessageBrokerConnectionError("Publisher client is not connected") + publisher = self._publisher + + def _publish() -> None: + futures = [publisher.publish(self._topic_path, data=msg_data) for msg_data in messages] + # The publisher batches internally; each future resolves once the broker + # has accepted that message. + for future in futures: + future.result(timeout=self._publish_timeout) + + try: + await asyncio.to_thread(_publish) + except Exception as error: # noqa: BLE001 + raise _wrap_error(error) from error async def _convert(self, data: dict[str, deque]) -> list[_t.Any]: """Converts field buffer data to JSON-encoded bytes messages. @@ -209,16 +294,4 @@ async def _convert(self, data: dict[str, deque]) -> list[_t.Any]: Returns: A list of bytes objects ready to publish. """ - completed_rows = min(len(d) for d in data.values()) if data else 0 - messages: list[bytes] = [] - for i in range(completed_rows): - record = {field: data[field][i] for field in data} - if self._parse_json: - messages.append(json.dumps(record).encode("utf-8")) - else: - # Send raw data from the first field - first_field = next(iter(record.values())) - messages.append( - first_field if isinstance(first_field, bytes) else str(first_field).encode() - ) - return messages + return encode_records_bytes(data, self._parse_json) diff --git a/plugboard/library/kafka_io.py b/plugboard/library/kafka_io.py index ef479c89..133a34b3 100644 --- a/plugboard/library/kafka_io.py +++ b/plugboard/library/kafka_io.py @@ -2,80 +2,128 @@ from __future__ import annotations +import asyncio from collections import deque import json import typing as _t -from plugboard.exceptions import NoMoreDataException +from plugboard.exceptions import ( + MessageBrokerConnectionError, + MessageBrokerPermanentError, + MessageBrokerTransientError, +) from plugboard.library.message_reader import MessageDataReader, MessageDataReaderArgsDict -from plugboard.library.message_writer import MessageDataWriter, MessageDataWriterArgsDict -from plugboard.utils import depends_on_optional +from plugboard.library.message_writer import ( + MessageDataWriter, + MessageDataWriterArgsDict, + encode_records_bytes, +) +from plugboard.utils import DI, depends_on_optional +from plugboard.utils.settings import resolve_argument try: from aiokafka import AIOKafkaConsumer, AIOKafkaProducer + from aiokafka.errors import ( + CoordinatorNotAvailableError, + KafkaConnectionError, + NodeNotReadyError, + NotLeaderForPartitionError, + RecordTooLargeError, + RequestTimedOutError, + TopicAuthorizationFailedError, + UnknownTopicOrPartitionError, + ) + from aiokafka.structs import OffsetAndMetadata, TopicPartition + + # Kafka failures that retrying cannot resolve, such as missing authorization or a + # record the broker will always reject. + _PERMANENT_ERRORS: tuple[type[Exception], ...] = ( + TopicAuthorizationFailedError, + RecordTooLargeError, + ) + + # Kafka failures raised while the cluster is rebalancing or a broker is unreachable, + # which are expected to clear. + _TRANSIENT_ERRORS: tuple[type[Exception], ...] = ( + KafkaConnectionError, + NodeNotReadyError, + CoordinatorNotAvailableError, + NotLeaderForPartitionError, + RequestTimedOutError, + UnknownTopicOrPartitionError, + ) except ImportError: # pragma: no cover - pass + # With the extra not installed, `depends_on_optional` blocks construction and + # classification is moot: `isinstance(error, ())` never matches. + _PERMANENT_ERRORS = () + _TRANSIENT_ERRORS = () -class KafkaDataReaderArgsDict(MessageDataReaderArgsDict): - """Specification of the `KafkaDataReader` constructor arguments. - - Attributes: - bootstrap_servers: Kafka broker address(es). - group_id: Consumer group ID. - parse_json: Whether to parse message values as JSON. - """ - - pass +def _wrap_error(error: Exception) -> Exception: + """Maps an aiokafka error onto the Plugboard message broker exception hierarchy. + Args: + error: The error raised by the Kafka client. -class KafkaDataWriterArgsDict(MessageDataWriterArgsDict): - """Specification of the `KafkaDataWriter` constructor arguments. - - Attributes: - bootstrap_servers: Kafka broker address(es). - parse_json: Whether to encode message values as JSON. + Returns: + A `MessageBrokerPermanentError` for a rejection that will not change on retry, + a `MessageBrokerTransientError` for a failure worth retrying, otherwise the + original error. """ - - pass + if isinstance(error, _PERMANENT_ERRORS): + return MessageBrokerPermanentError(str(error)) + if isinstance(error, _TRANSIENT_ERRORS): + return MessageBrokerTransientError(str(error)) + return error class KafkaDataReader(MessageDataReader): """Reads data from an Apache Kafka topic. - Messages are consumed from the topic using a consumer group and converted - to field values. Offsets are committed after processing (acknowledgment). + Messages are consumed from the topic using a consumer group. Acknowledging a batch + commits the offsets following the records that were consumed, so a failure before + that point leaves the remaining records for redelivery. """ @depends_on_optional("aiokafka", extra="kafka") def __init__( self, - bootstrap_servers: str | list[str], group_id: str, + bootstrap_servers: _t.Optional[str | list[str]] = None, parse_json: bool = True, - **kwargs: _t.Unpack[KafkaDataReaderArgsDict], + poll_timeout_ms: int = 30_000, + **kwargs: _t.Unpack[MessageDataReaderArgsDict], ) -> None: """Instantiates the `KafkaDataReader`. Args: - bootstrap_servers: Kafka broker address(es) (e.g. `"localhost:9092"`). group_id: Consumer group ID. + bootstrap_servers: Optional; Kafka broker address(es), for example + `"localhost:9092"`. Falls back to `KAFKA_BOOTSTRAP_SERVERS`. parse_json: Whether to parse message values as JSON. + poll_timeout_ms: Milliseconds to wait for a batch of records. **kwargs: Additional keyword arguments for [`MessageDataReader`][plugboard.library.MessageDataReader]. """ + bootstrap_servers = resolve_argument( + bootstrap_servers, + DI.settings.resolve_sync().kafka.bootstrap_servers, + "bootstrap_servers", + "KAFKA_BOOTSTRAP_SERVERS", + ) super().__init__(**kwargs) - if isinstance(bootstrap_servers, str): - bootstrap_servers = [bootstrap_servers] - self._bootstrap_servers = bootstrap_servers + self._bootstrap_servers = ( + [bootstrap_servers] if isinstance(bootstrap_servers, str) else bootstrap_servers + ) self._group_id = group_id self._parse_json = parse_json + self._poll_timeout_ms = poll_timeout_ms self._consumer: _t.Optional[AIOKafkaConsumer] = None async def _connect(self) -> None: """Creates and starts a Kafka consumer.""" - self._consumer = AIOKafkaConsumer( + consumer = AIOKafkaConsumer( self._topic, bootstrap_servers=self._bootstrap_servers, group_id=self._group_id, @@ -83,46 +131,47 @@ async def _connect(self) -> None: enable_auto_commit=False, max_poll_records=self._chunk_size or 10, ) - await self._consumer.start() + await consumer.start() + self._consumer = consumer async def _disconnect(self) -> None: """Stops and closes the Kafka consumer.""" if self._consumer is not None: - await self._consumer.stop() + consumer = self._consumer self._consumer = None + await consumer.stop() async def _receive(self) -> list[_t.Any]: """Receives a batch of messages from the Kafka topic. Returns: - A list of Kafka `ConsumerRecord` objects. + A list of Kafka `ConsumerRecord` objects, empty if the poll timed out. Raises: - NoMoreDataException: If the consumer has been closed. + MessageBrokerError: If the poll failed for a retryable or permanent reason. """ if self._consumer is None: - raise RuntimeError("Kafka consumer not initialized") + raise MessageBrokerConnectionError("Kafka consumer is not connected") max_messages = self._chunk_size or 10 - # Use getmany to fetch a batch with timeout - data = await self._consumer.getmany(timeout_ms=30000, max_records=max_messages) - messages: list[_t.Any] = [] - for _tp, records in data.items(): - messages.extend(records) - if not messages: - raise NoMoreDataException - return messages[:max_messages] - - async def _convert(self, messages: list[_t.Any]) -> dict[str, deque]: + try: + data = await self._consumer.getmany( + timeout_ms=self._poll_timeout_ms, max_records=max_messages + ) + except Exception as error: # noqa: BLE001 + raise _wrap_error(error) from error + return [record for records in data.values() for record in records] + + async def _convert(self, data: list[_t.Any]) -> dict[str, deque]: """Converts Kafka messages to a field buffer. Args: - messages: A list of `ConsumerRecord` objects. + data: A list of `ConsumerRecord` objects. Returns: A dictionary mapping field names to deques of field values. """ converted: dict[str, deque] = {field: deque() for field in self.io.outputs} - for record in messages: + for record in data: value = record.value if isinstance(value, bytes): value = value.decode("utf-8") @@ -135,68 +184,106 @@ async def _convert(self, messages: list[_t.Any]) -> dict[str, deque]: return converted async def _ack(self, messages: list[_t.Any]) -> None: - """Commits offsets for processed Kafka messages. + """Commits the offsets that follow the processed Kafka messages. + + Only the offsets of the records that were actually consumed are committed, so + records still waiting in the buffer stay eligible for redelivery. Args: - messages: The `ConsumerRecord` objects to acknowledge. + messages: The `ConsumerRecord` objects that have been processed. + + Raises: + MessageBrokerConnectionError: If the consumer is not connected. + MessageBrokerError: If the commit failed. """ if self._consumer is None: - raise RuntimeError("Kafka consumer not initialized") - await self._consumer.commit() + raise MessageBrokerConnectionError("Kafka consumer is not connected") + offsets: dict[TopicPartition, OffsetAndMetadata] = {} + for record in messages: + partition = TopicPartition(record.topic, record.partition) + next_offset = OffsetAndMetadata(record.offset + 1, "") + current = offsets.get(partition) + if current is None or current.offset < next_offset.offset: + offsets[partition] = next_offset + if not offsets: + return + try: + await self._consumer.commit(offsets=offsets) + except Exception as error: # noqa: BLE001 + raise _wrap_error(error) from error class KafkaDataWriter(MessageDataWriter): """Writes data to an Apache Kafka topic. - Field data is converted to JSON-encoded messages and produced - to the specified Kafka topic. + Field data is converted to JSON-encoded messages and produced to the specified + topic. Records are submitted together and flushed once, so the producer's internal + batching is used instead of one round trip per record. """ @depends_on_optional("aiokafka", extra="kafka") def __init__( self, - bootstrap_servers: str | list[str], + bootstrap_servers: _t.Optional[str | list[str]] = None, parse_json: bool = True, - **kwargs: _t.Unpack[KafkaDataWriterArgsDict], + **kwargs: _t.Unpack[MessageDataWriterArgsDict], ) -> None: """Instantiates the `KafkaDataWriter`. Args: - bootstrap_servers: Kafka broker address(es) (e.g. `"localhost:9092"`). + bootstrap_servers: Optional; Kafka broker address(es), for example + `"localhost:9092"`. Falls back to `KAFKA_BOOTSTRAP_SERVERS`. parse_json: Whether to encode message values as JSON. **kwargs: Additional keyword arguments for [`MessageDataWriter`][plugboard.library.MessageDataWriter]. """ + bootstrap_servers = resolve_argument( + bootstrap_servers, + DI.settings.resolve_sync().kafka.bootstrap_servers, + "bootstrap_servers", + "KAFKA_BOOTSTRAP_SERVERS", + ) super().__init__(**kwargs) - if isinstance(bootstrap_servers, str): - bootstrap_servers = [bootstrap_servers] - self._bootstrap_servers = bootstrap_servers + self._bootstrap_servers = ( + [bootstrap_servers] if isinstance(bootstrap_servers, str) else bootstrap_servers + ) self._parse_json = parse_json self._producer: _t.Optional[AIOKafkaProducer] = None async def _connect(self) -> None: """Creates and starts a Kafka producer.""" - self._producer = AIOKafkaProducer( - bootstrap_servers=self._bootstrap_servers, - ) - await self._producer.start() + producer = AIOKafkaProducer(bootstrap_servers=self._bootstrap_servers) + await producer.start() + self._producer = producer async def _disconnect(self) -> None: """Stops and closes the Kafka producer.""" if self._producer is not None: - await self._producer.stop() + producer = self._producer self._producer = None + await producer.stop() async def _send(self, messages: list[_t.Any]) -> None: """Sends messages to the Kafka topic. Args: messages: A list of bytes objects to send. + + Raises: + MessageBrokerConnectionError: If the producer is not connected. + MessageBrokerError: If a send failed or the buffer could not be flushed. """ if self._producer is None: - raise RuntimeError("Kafka producer not initialized") - for msg_data in messages: - await self._producer.send_and_wait(self._topic, value=msg_data) + raise MessageBrokerConnectionError("Kafka producer is not connected") + try: + # `send` returns a future that resolves once the broker acknowledges the + # record; awaiting them together keeps the producer's batching intact. + await asyncio.gather( + *(self._producer.send(self._topic, value=msg_data) for msg_data in messages) + ) + await self._producer.flush() + except Exception as error: # noqa: BLE001 + raise _wrap_error(error) from error async def _convert(self, data: dict[str, deque]) -> list[_t.Any]: """Converts field buffer data to JSON-encoded bytes messages. @@ -207,15 +294,4 @@ async def _convert(self, data: dict[str, deque]) -> list[_t.Any]: Returns: A list of bytes objects ready to send. """ - completed_rows = min(len(d) for d in data.values()) if data else 0 - messages: list[bytes] = [] - for i in range(completed_rows): - record = {field: data[field][i] for field in data} - if self._parse_json: - messages.append(json.dumps(record).encode("utf-8")) - else: - first_field = next(iter(record.values())) - messages.append( - first_field if isinstance(first_field, bytes) else str(first_field).encode() - ) - return messages + return encode_records_bytes(data, self._parse_json) diff --git a/plugboard/library/message_reader.py b/plugboard/library/message_reader.py index 94436d8b..1b22ec06 100644 --- a/plugboard/library/message_reader.py +++ b/plugboard/library/message_reader.py @@ -2,15 +2,17 @@ from __future__ import annotations -from abc import ABC, abstractmethod +from abc import abstractmethod import asyncio -from asyncio.tasks import Task from collections import deque +import contextlib import typing as _t -from plugboard.component import Component, IOController -from plugboard.exceptions import IOSetupError, IOStreamClosedError, NoMoreDataException +from plugboard.component import IOController +from plugboard.exceptions import IOStreamClosedError, NoMoreDataException +from plugboard.library.data_reader import DataReader from plugboard.schemas import ComponentArgsDict +from plugboard.utils.retry import RetryPolicy, attempt_reconnect, with_retry class MessageDataReaderArgsDict(ComponentArgsDict): @@ -18,30 +20,36 @@ class MessageDataReaderArgsDict(ComponentArgsDict): Attributes: field_names: The names of the fields to read from messages. + topic: The topic/queue/subscription to read from. chunk_size: Optional; The number of messages to fetch per batch. - max_retries: Maximum number of retry attempts for transient failures. - retry_base_delay: Base delay in seconds for exponential backoff. - retry_max_delay: Maximum delay in seconds for exponential backoff. + retry_policy: Optional; The backoff policy for retrying failed broker calls. + idle_poll_delay: Optional; Delay in seconds between polls that return no + messages. """ field_names: list[str] - topic: _t.NotRequired[str] + topic: str chunk_size: _t.NotRequired[int | None] - max_retries: _t.NotRequired[int] - retry_base_delay: _t.NotRequired[float] - retry_max_delay: _t.NotRequired[float] + retry_policy: _t.NotRequired[RetryPolicy] + idle_poll_delay: _t.NotRequired[float] -class MessageDataReader(Component, ABC): +class MessageDataReader(DataReader): """Abstract base class for reading data from a pub/sub message broker. - Provides connection management, reconnection with exponential backoff, - retry logic, message acknowledgment, and chunked/buffered reading - analogous to [`DataReader`][plugboard.library.DataReader]. + Extends [`DataReader`][plugboard.library.DataReader] with the parts a message + broker needs and a finite source does not: a long-lived connection, reconnection + with exponential backoff, and acknowledgment of the messages that were consumed. - Subclasses must implement broker-specific methods for connecting, - receiving messages, converting messages to field buffers, and - acknowledging processed messages. + Unlike a file or table reader, an empty poll does not mean the stream is over - + it means the broker had nothing to hand out yet. `step()` therefore waits for the + next message instead of closing the stream, so a reader started before its + producer keeps running. A subclass signals a genuinely exhausted source by + raising [`NoMoreDataException`][plugboard.exceptions.NoMoreDataException] from + `_receive`, which closes the IO stream like `DataReader` does. + + Subclasses implement the broker-specific methods for connecting, receiving raw + messages, converting them to field buffers, and acknowledging processed messages. """ io = IOController() @@ -51,47 +59,27 @@ def __init__( field_names: list[str], topic: str, chunk_size: _t.Optional[int] = None, - max_retries: int = 3, - retry_base_delay: float = 1.0, - retry_max_delay: float = 60.0, + retry_policy: RetryPolicy = RetryPolicy(), + idle_poll_delay: float = 0.1, **kwargs: _t.Unpack[ComponentArgsDict], ) -> None: """Instantiates the `MessageDataReader`. Args: field_names: The names of the fields to extract from messages. - topic: The topic/queue to read from. + topic: The topic/queue/subscription to read from. chunk_size: Optional; The number of messages to fetch per batch. - max_retries: Maximum number of retry attempts for transient failures. - retry_base_delay: Base delay in seconds for exponential backoff. - retry_max_delay: Maximum delay in seconds for exponential backoff. + retry_policy: The backoff policy for retrying failed broker calls. + idle_poll_delay: Delay in seconds before polling again after a poll that + returned no messages. **kwargs: Additional keyword arguments for [`Component`][plugboard.component.Component]. """ - super().__init__(**kwargs) + super().__init__(field_names=field_names, chunk_size=chunk_size, **kwargs) self._topic = topic - self._buffer: dict[str, deque] = dict() - self._chunk_size = chunk_size - self._max_retries = max_retries - self._retry_base_delay = retry_base_delay - self._retry_max_delay = retry_max_delay + self._retry_policy = retry_policy + self._idle_poll_delay = idle_poll_delay self._pending_ack: list[_t.Any] = [] - self._task: _t.Optional[Task] = None - self.io = IOController( - inputs=None, - outputs=field_names, - input_events=self.__class__.io.input_events, - output_events=self.__class__.io.output_events, - namespace=self.name, - component=self, - ) - - def __init_subclass__(cls, *args: _t.Any, **kwargs: _t.Any) -> None: - try: - return super().__init_subclass__(*args, **kwargs) - except IOSetupError: - # Concrete subclasses of the abstract data io classes represent a special case for io - # setup. They receive io args at run time, not declaration time, so skip error. - pass + self._connection_lock = asyncio.Lock() @abstractmethod async def _connect(self) -> None: @@ -111,154 +99,155 @@ async def _disconnect(self) -> None: async def _receive(self) -> list[_t.Any]: """Receives a batch of raw messages from the broker. - Should block until at least one message is available or a timeout occurs. - Returns an empty list on timeout. + Should block until at least one message is available or a timeout occurs, and + return an empty list on timeout. Returning an empty list is *not* how a + subclass reports an exhausted source - raising `NoMoreDataException` is. Returns: A list of raw broker-specific message objects. Raises: - NoMoreDataException: If the subscription/source is exhausted and no - more messages will arrive. + NoMoreDataException: If the source is exhausted and no more messages will + arrive. + MessageBrokerPermanentError: If the failure cannot be recovered by retrying. """ pass @abstractmethod - async def _convert(self, messages: list[_t.Any]) -> dict[str, deque]: - """Converts raw messages into a `dict[str, deque]` field buffer. + async def _ack(self, messages: list[_t.Any]) -> None: + """Acknowledges successful processing of messages. Args: - messages: Raw broker-specific message objects. - - Returns: - A dictionary mapping field names to deques of field values. + messages: The raw messages to acknowledge. """ pass @abstractmethod - async def _ack(self, messages: list[_t.Any]) -> None: - """Acknowledges successful processing of messages. + async def _convert(self, data: list[_t.Any]) -> dict[str, deque]: + """Converts raw messages into a `dict[str, deque]` field buffer. Args: - messages: The raw messages to acknowledge. + data: Raw broker-specific message objects, as returned by `_receive`. + + Returns: + A dictionary mapping field names to deques of field values. """ pass - async def _receive_with_retry(self) -> list[_t.Any]: - """Receives messages with exponential backoff retry and reconnection. + async def _fetch(self) -> list[_t.Any]: + """Receives a batch of messages, retrying transient failures. Returns: - A list of raw broker-specific message objects. + A list of raw broker-specific message objects, possibly empty. Raises: NoMoreDataException: If the source is exhausted. MessageBrokerConnectionError: If all retries are exhausted. """ - last_exception: Exception = RuntimeError("All retries exhausted") - for attempt in range(self._max_retries + 1): - try: - return await self._receive() - except NoMoreDataException: - raise - except Exception as e: - last_exception = e - if attempt < self._max_retries: - delay = min( - self._retry_base_delay * (2**attempt), - self._retry_max_delay, - ) - self._logger.warning( - "Transient error receiving messages, retrying", - attempt=attempt + 1, - delay=delay, - error=str(e), - ) - await asyncio.sleep(delay) - await self._reconnect() - raise last_exception + return await with_retry( + self._receive, + self._reconnect, + policy=self._retry_policy, + logger=self._logger, + description="receiving messages", + ) async def _reconnect(self) -> None: - """Attempts to reconnect to the message broker.""" - self._logger.info("Attempting reconnection to message broker", topic=self._topic) - try: - await self._disconnect() - except Exception: # noqa: S110 - self._logger.warning("Error during disconnect in reconnection", exc_info=True) - await self._connect() - self._logger.info("Reconnected to message broker", topic=self._topic) + """Attempts to reconnect to the message broker. - async def _fetch_batch(self) -> None: - """Fetches a batch of messages and updates the internal buffer.""" + Held under the connection lock so an in-flight acknowledgment cannot use a + client that is being replaced. + """ + async with self._connection_lock: + await attempt_reconnect( + self._connect, + self._disconnect, + logger=self._logger, + topic=self._topic, + ) + + async def _fetch_chunk(self) -> None: + """Fetches a batch of messages and updates the internal buffer. + + An empty batch leaves the buffer untouched, so the caller can poll again: a + broker with nothing to deliver has not ended the stream. The next batch is + requested concurrently so that waiting on the broker overlaps with processing. + """ if self._task is None: - self._task = asyncio.create_task(self._receive_with_retry()) - messages = await self._task - # Start fetching next batch concurrently - self._task = asyncio.create_task(self._receive_with_retry()) - if len(messages) == 0: - raise NoMoreDataException + self._task = asyncio.create_task(self._fetch()) + task = self._task + # Clear the reference before awaiting so a failed fetch leaves no dangling task. + self._task = None + messages = await task + if not messages: + await asyncio.sleep(self._idle_poll_delay) + return new_buffer = await self._convert(messages) self._buffer = {field_name: new_buffer[field_name] for field_name in self.io.outputs} self._pending_ack = messages + # Prefetch the next batch while the current one is being consumed. + self._task = asyncio.create_task(self._fetch()) - def _consume_record(self) -> None: - """Consumes one record from the buffer and sets field attributes.""" - for field in self.io.outputs: - setattr(self, field, self._buffer[field].popleft()) + @property + def _records_available(self) -> int: + """Calculates how many consumed-but-unpublished records the buffer holds.""" + if not self._buffer: + return 0 + return min(len(values) for values in self._buffer.values()) async def _ack_pending(self) -> None: - """Acknowledges all pending messages.""" - if self._pending_ack: + """Acknowledges the batch whose records have all been consumed. + + Acknowledgment waits for the connection lock, so it never runs against a client + that a concurrent reconnect is replacing. + """ + if not self._pending_ack: + return + async with self._connection_lock: await self._ack(self._pending_ack) - self._pending_ack = [] + self._pending_ack = [] async def init(self) -> None: """Initialises the `MessageDataReader`. - Connects to the message broker and pre-fetches the first batch of messages. - If no messages are available, the reader will raise `IOStreamClosedError` - on the first `step()` call. + Connects to the message broker and starts pre-fetching the first batch so that + the first `step()` does not wait on a cold poll. """ await self._connect() self._logger.info("Connected to message broker", topic=self._topic) - try: - await self._fetch_batch() - except NoMoreDataException: - # No messages available at init time; step() will raise IOStreamClosedError - pass + self._task = asyncio.create_task(self._fetch()) async def step(self) -> None: - """Reads data from the message broker and updates outputs. + """Reads the next message from the broker and updates outputs. - Consumes one record from the buffer. If the buffer is empty, - fetches the next batch. Acknowledges processed messages. + Waits until a message is available, then consumes one record. Once the batch's + last record has been consumed the whole batch is acknowledged, so a failure + before that point leaves the remaining messages for redelivery. Raises: - IOStreamClosedError: If there is no more data to read. + IOStreamClosedError: If the source is exhausted. """ - if not self._buffer: - # Buffer was never populated (e.g. empty source at init) - await self.io.close() - raise IOStreamClosedError("No more messages from broker") - try: - self._consume_record() - await self._ack_pending() - except IndexError: + while not self._records_available: try: - await self._fetch_batch() - self._consume_record() - await self._ack_pending() + await self._fetch_chunk() except NoMoreDataException: + await self._ack_pending() await self.io.close() - raise IOStreamClosedError("No more messages from broker") + raise IOStreamClosedError("No more messages from broker") from None + self._consume_record() + if not self._records_available: + await self._ack_pending() async def destroy(self) -> None: - """Destroys the `MessageDataReader` and disconnects from the broker.""" + """Destroys the `MessageDataReader` and disconnects from the broker. + + Messages that were pre-fetched but never consumed are left unacknowledged so + the broker redelivers them to whoever reads the queue next. + """ if self._task is not None: self._task.cancel() - try: + with contextlib.suppress(asyncio.CancelledError, Exception): await self._task - except (asyncio.CancelledError, Exception): # noqa: S110 - pass self._task = None await self._disconnect() self._logger.info("Disconnected from message broker", topic=self._topic) diff --git a/plugboard/library/message_writer.py b/plugboard/library/message_writer.py index cc66674b..6309b497 100644 --- a/plugboard/library/message_writer.py +++ b/plugboard/library/message_writer.py @@ -2,15 +2,16 @@ from __future__ import annotations -from abc import ABC, abstractmethod +from abc import abstractmethod import asyncio -from asyncio.tasks import Task -from collections import defaultdict, deque +from collections import deque +import json import typing as _t -from plugboard.component import Component, IOController -from plugboard.exceptions import IOSetupError +from plugboard.component import IOController +from plugboard.library.data_writer import DataWriter from plugboard.schemas import ComponentArgsDict +from plugboard.utils.retry import RetryPolicy, attempt_reconnect, with_retry class MessageDataWriterArgsDict(ComponentArgsDict): @@ -18,30 +19,96 @@ class MessageDataWriterArgsDict(ComponentArgsDict): Attributes: field_names: The names of the fields to include in messages. - chunk_size: Optional; The number of records to batch into messages. - max_retries: Maximum number of retry attempts for transient failures. - retry_base_delay: Base delay in seconds for exponential backoff. - retry_max_delay: Maximum delay in seconds for exponential backoff. + topic: The topic/queue to write to. + chunk_size: Optional; The number of records to batch into each send. + retry_policy: Optional; The backoff policy for retrying failed broker calls. """ field_names: list[str] - topic: _t.NotRequired[str] + topic: str chunk_size: _t.NotRequired[int | None] - max_retries: _t.NotRequired[int] - retry_base_delay: _t.NotRequired[float] - retry_max_delay: _t.NotRequired[float] + retry_policy: _t.NotRequired[RetryPolicy] -class MessageDataWriter(Component, ABC): +def iter_records(data: dict[str, deque]) -> _t.Iterator[dict[str, _t.Any]]: + """Iterates over the complete records held in a field buffer. + + Zipping the deques walks each field once, which keeps record building linear in + the number of buffered values - indexing the deques per row would be quadratic, + and with the default `chunk_size` the buffer holds the whole run. + + Args: + data: A dictionary mapping field names to deques of field values. + + Yields: + One dictionary per complete record, keyed by field name. + """ + fields = tuple(data) + if not fields: + return + for row in zip(*data.values()): + yield dict(zip(fields, row)) + + +def _raw_value(record: dict[str, _t.Any]) -> str: + """Returns the first field of a record as text, for non-JSON payloads.""" + value = next(iter(record.values())) + return value.decode("utf-8") if isinstance(value, bytes) else str(value) + + +def _raw_bytes(record: dict[str, _t.Any]) -> bytes: + """Returns the first field of a record as bytes, for non-JSON payloads. + + Byte values are passed through unchanged so binary payloads are not corrupted by + a text round trip. + """ + value = next(iter(record.values())) + return bytes(value) if isinstance(value, (bytes, bytearray)) else str(value).encode("utf-8") + + +def encode_records(data: dict[str, deque], parse_json: bool) -> list[str]: + """Encodes a field buffer as one text message per record. + + Args: + data: A dictionary mapping field names to deques of field values. + parse_json: Whether to encode each record as a JSON object. When false, the + first field's value is used as the message body. + + Returns: + A list of message strings ready to send. + """ + return [ + json.dumps(record) if parse_json else _raw_value(record) for record in iter_records(data) + ] + + +def encode_records_bytes(data: dict[str, deque], parse_json: bool) -> list[bytes]: + """Encodes a field buffer as one binary message per record. + + Args: + data: A dictionary mapping field names to deques of field values. + parse_json: Whether to encode each record as a JSON object. When false, the + first field's value is used as the message body. + + Returns: + A list of bytes objects ready to send. + """ + return [ + json.dumps(record).encode("utf-8") if parse_json else _raw_bytes(record) + for record in iter_records(data) + ] + + +class MessageDataWriter(DataWriter): """Abstract base class for writing data to a pub/sub message broker. - Provides connection management, reconnection with exponential backoff, - retry logic, and chunked/buffered writing analogous to - [`DataWriter`][plugboard.library.DataWriter]. + Extends [`DataWriter`][plugboard.library.DataWriter] with the parts a message + broker needs and a file or table does not: a long-lived connection, and + reconnection with exponential backoff around each send. Buffered records are + batched into messages by the broker-specific `_convert`. - Subclasses must implement broker-specific methods for connecting, - sending messages, and converting field data to broker-specific - message format. + Subclasses implement the broker-specific methods for connecting, sending encoded + messages, and converting field data to the broker's message format. """ io = IOController() @@ -51,9 +118,7 @@ def __init__( field_names: list[str], topic: str, chunk_size: _t.Optional[int] = None, - max_retries: int = 3, - retry_base_delay: float = 1.0, - retry_max_delay: float = 60.0, + retry_policy: RetryPolicy = RetryPolicy(), **kwargs: _t.Unpack[ComponentArgsDict], ) -> None: """Instantiates the `MessageDataWriter`. @@ -61,37 +126,13 @@ def __init__( Args: field_names: The names of the fields to include in messages. topic: The topic/queue to write to. - chunk_size: Optional; The number of records to batch into a single send operation. - max_retries: Maximum number of retry attempts for transient failures. - retry_base_delay: Base delay in seconds for exponential backoff. - retry_max_delay: Maximum delay in seconds for exponential backoff. + chunk_size: Optional; The number of records to batch into a single send. + retry_policy: The backoff policy for retrying failed broker calls. **kwargs: Additional keyword arguments for [`Component`][plugboard.component.Component]. """ - super().__init__(**kwargs) + super().__init__(field_names=field_names, chunk_size=chunk_size, **kwargs) self._topic = topic - self._buffer: dict[str, deque] = defaultdict(deque) - self._chunk_size = chunk_size - self._max_retries = max_retries - self._retry_base_delay = retry_base_delay - self._retry_max_delay = retry_max_delay - self._task: _t.Optional[Task] = None - self.io = IOController( - inputs=field_names, - outputs=None, - input_events=self.__class__.io.input_events, - output_events=self.__class__.io.output_events, - event_field_coverage=self.__class__.io.event_field_coverage, - namespace=self.name, - component=self, - ) - - def __init_subclass__(cls, *args: _t.Any, **kwargs: _t.Any) -> None: - try: - return super().__init_subclass__(*args, **kwargs) - except IOSetupError: - # Concrete subclasses of the abstract data io classes represent a special case for io - # setup. They receive io args at run time, not declaration time, so skip error. - pass + self._retry_policy = retry_policy @abstractmethod async def _connect(self) -> None: @@ -115,7 +156,9 @@ async def _send(self, messages: list[_t.Any]) -> None: messages: A list of broker-specific message objects to send. Raises: - MessageBrokerConnectionError: If messages cannot be sent. + MessageBrokerConnectionError: If the connection to the broker fails. + MessageBrokerPermanentError: If the messages cannot be delivered and + retrying cannot help. """ pass @@ -131,121 +174,54 @@ async def _convert(self, data: dict[str, deque]) -> list[_t.Any]: """ pass - async def _send_with_retry(self, messages: list[_t.Any]) -> None: - """Sends messages with exponential backoff retry and reconnection. + async def _save(self, data: list[_t.Any]) -> None: + """Sends encoded messages to the broker, retrying transient failures. Args: - messages: The messages to send. + data: The converted messages to send, as produced by `_convert`. Raises: - Exception: If all retries are exhausted. + MessageBrokerConnectionError: If all retries are exhausted. """ - last_exception: Exception = RuntimeError("All retries exhausted") - for attempt in range(self._max_retries + 1): - try: - await self._send(messages) - return - except Exception as e: - last_exception = e - if attempt < self._max_retries: - delay = min( - self._retry_base_delay * (2**attempt), - self._retry_max_delay, - ) - self._logger.warning( - "Transient error sending messages, retrying", - attempt=attempt + 1, - delay=delay, - error=str(e), - ) - await asyncio.sleep(delay) - await self._reconnect() - raise last_exception + await with_retry( + lambda: self._send(data), + self._reconnect, + policy=self._retry_policy, + logger=self._logger, + description="sending messages", + ) async def _reconnect(self) -> None: - """Attempts to reconnect to the message broker.""" - self._logger.info("Attempting reconnection to message broker", topic=self._topic) - try: - await self._disconnect() - except Exception: # noqa: S102 - self._logger.warning("Error during disconnect in reconnection", exc_info=True) - await self._connect() - self._logger.info("Reconnected to message broker", topic=self._topic) - - def _bind_inputs(self) -> None: - """Binds input fields to component fields and appends to internal buffer.""" - super()._bind_inputs() - for field in self._field_inputs: - value = getattr(self, field, None) - self._buffer[field].append(value) - - @property - def _completed_rows(self) -> int: - """Calculates how many fully formed rows exist in the buffer.""" - if not self.io.inputs: - return 0 - return min((len(self._buffer[f]) for f in self.io.inputs), default=0) - - @property - def _can_step(self) -> bool: - """We can step if we have at least one fully formed row.""" - return self._completed_rows > 0 - - async def _send_batch(self) -> None: - """Sends completed data rows from the buffer.""" - completed_rows = self._completed_rows - if completed_rows == 0: - return - - if self._task is not None: - await self._task + """Attempts to reconnect to the message broker. - # Extract only the completed rows into a new chunk - chunk_data: dict[str, deque] = { - field: deque([self._buffer[field].popleft() for _ in range(completed_rows)]) - for field in self.io.inputs - } - - messages = await self._convert(chunk_data) - self._task = asyncio.create_task(self._send_with_retry(messages)) + No lock is needed here: unlike a reader, a writer never acknowledges + concurrently, and sends are serialised through `self._task`. + """ + await attempt_reconnect( + self._connect, + self._disconnect, + logger=self._logger, + topic=self._topic, + ) async def init(self) -> None: - """Initialises the `MessageDataWriter`. - - Connects to the message broker. - """ + """Initialises the `MessageDataWriter` by connecting to the message broker.""" await self._connect() self._logger.info("Connected to message broker", topic=self._topic) - async def step(self) -> None: - """Triggers send when buffer is at target size. - - If `chunk_size` is set and the buffer has reached that size, - sends the buffered data as messages. - """ - if self._chunk_size and self._completed_rows >= self._chunk_size: - await self._send_batch() - - async def run(self) -> None: - """Runs the `MessageDataWriter`. + async def destroy(self) -> None: + """Destroys the `MessageDataWriter` and disconnects from the broker. - Steps until all input is consumed, then flushes any remaining - buffered data. + Any send still in flight is awaited rather than cancelled, so buffered data is + not lost on teardown. """ - await super().run() - # Flush any remaining data in the buffer after completion - await self._send_batch() - if self._task is not None: - await self._task - - async def destroy(self) -> None: - """Destroys the `MessageDataWriter` and disconnects from the broker.""" if self._task is not None: - self._task.cancel() try: await self._task - except (asyncio.CancelledError, Exception): # noqa: S110 - pass + except asyncio.CancelledError: # pragma: no cover + raise + except Exception as error: # noqa: BLE001 + self._logger.warning("Pending send failed during destroy", error=str(error)) self._task = None await self._disconnect() self._logger.info("Disconnected from message broker", topic=self._topic) diff --git a/plugboard/utils/__init__.py b/plugboard/utils/__init__.py index 3c0d792e..e8d398aa 100644 --- a/plugboard/utils/__init__.py +++ b/plugboard/utils/__init__.py @@ -9,7 +9,7 @@ from plugboard.utils.random import gen_rand_str from plugboard.utils.ray import build_actor_wrapper, is_on_ray_worker from plugboard.utils.registry import ClassRegistry -from plugboard.utils.settings import Settings +from plugboard.utils.settings import Settings, resolve_argument __all__ = [ @@ -19,6 +19,7 @@ "gather_except", "gen_rand_str", "is_on_ray_worker", + "resolve_argument", "run_coro_sync", "ClassRegistry", "DI", diff --git a/plugboard/utils/retry.py b/plugboard/utils/retry.py new file mode 100644 index 00000000..112d9515 --- /dev/null +++ b/plugboard/utils/retry.py @@ -0,0 +1,142 @@ +"""Provides retry and backoff helpers for long-lived connections.""" + +import asyncio +from collections.abc import Awaitable, Callable +import typing as _t + +import msgspec +import structlog + +from plugboard.exceptions import MessageBrokerPermanentError, NoMoreDataException + + +_T = _t.TypeVar("_T") + + +class RetryPolicy(msgspec.Struct, frozen=True): + """Exponential backoff parameters for a retryable operation. + + The three values always travel together and drive a single backoff formula, so + they are grouped into one immutable value object. + + Attributes: + max_retries: Number of times to retry after a failure, so the operation runs + at most `max_retries + 1` times. + base_delay: Base delay in seconds, doubled on each consecutive attempt. + max_delay: Upper bound for the delay between attempts. + """ + + max_retries: int = 3 + base_delay: float = 1.0 + max_delay: float = 60.0 + + def delay_for(self, attempt: int) -> float: + """Calculates the delay to wait after the given zero-based attempt failed. + + Args: + attempt: The zero-based index of the attempt that failed. + + Returns: + The backoff delay in seconds, capped by `max_delay`. + """ + return min(self.base_delay * (2**attempt), self.max_delay) + + +async def attempt_reconnect( + connect: Callable[[], Awaitable[None]], + disconnect: Callable[[], Awaitable[None]], + *, + logger: structlog.BoundLogger, + topic: str, +) -> None: + """Re-establishes a message broker connection by disconnecting then connecting. + + A failure to disconnect is logged and ignored: the connection is assumed to be + broken already, and the new connection attempt is what matters. A failure to + connect propagates so the caller can count it as another failed attempt. + + Args: + connect: The coroutine function that establishes the connection. + disconnect: The coroutine function that closes the connection. + logger: Logger to report reconnection progress to. + topic: The topic/queue being connected to, used for log context. + """ + logger.info("Attempting reconnection to message broker", topic=topic) + try: + await disconnect() + except Exception as error: # noqa: BLE001 + logger.warning("Error during disconnect in reconnection", error=str(error)) + await connect() + logger.info("Reconnected to message broker", topic=topic) + + +async def _backoff_and_reconnect( + policy: RetryPolicy, + logger: structlog.BoundLogger, + description: str, + attempt: int, + error: Exception, + reconnect: Callable[[], Awaitable[None]], +) -> None: + """Waits out the backoff delay for a failed attempt and then reconnects. + + A failing reconnect is recorded and swallowed: the next attempt may succeed, and + the caller raises the broker error once the attempts are used up. Without this the + reconnect error would escape the retry loop on the first transient failure. + """ + delay = policy.delay_for(attempt) + logger.warning( + f"Transient error {description}, retrying", + attempt=attempt + 1, + delay=delay, + error=str(error), + ) + await asyncio.sleep(delay) + try: + await reconnect() + except Exception as reconnect_error: # noqa: BLE001 + logger.warning( + f"Reconnection during {description} failed", + attempt=attempt + 1, + error=str(reconnect_error), + ) + + +async def with_retry( + operation: Callable[[], Awaitable[_T]], + reconnect: Callable[[], Awaitable[None]], + *, + policy: RetryPolicy, + logger: structlog.BoundLogger, + description: str, +) -> _T: + """Runs an operation, retrying transient failures with backoff and reconnection. + + Only transient failures are retried: `NoMoreDataException` means the source is + exhausted and `MessageBrokerPermanentError` means retrying cannot help, so both + propagate immediately. + + Args: + operation: The coroutine function to run. + reconnect: The coroutine function to call between attempts. + policy: Backoff parameters, including the retry bound. + logger: Logger to report transient failures to. + description: Human-readable name of the operation, used in log messages. + + Returns: + Whatever `operation` returns. + + Raises: + NoMoreDataException: If the source is exhausted. + MessageBrokerPermanentError: If the failure is not retryable. + Exception: The error raised by the final attempt, if every attempt failed. + """ + for attempt in range(policy.max_retries): + try: + return await operation() + except (NoMoreDataException, MessageBrokerPermanentError): + raise + except Exception as error: # noqa: BLE001 + await _backoff_and_reconnect(policy, logger, description, attempt, error, reconnect) + # Final attempt: no backoff, so its error propagates unchanged. + return await operation() diff --git a/plugboard/utils/settings.py b/plugboard/utils/settings.py index aa94225f..93c6c683 100644 --- a/plugboard/utils/settings.py +++ b/plugboard/utils/settings.py @@ -122,3 +122,39 @@ class Settings(BaseSettings): gcp_pubsub: _GCPPubSubSettings = Field(default_factory=_GCPPubSubSettings) aws: _AWSSettings = Field(default_factory=_AWSSettings) kafka: _KafkaSettings = Field(default_factory=_KafkaSettings) + + +_ArgumentT = _t.TypeVar("_ArgumentT") + + +def resolve_argument( + value: _t.Optional[_ArgumentT], + fallback: _t.Optional[_ArgumentT], + arg_name: str, + env_var: str, +) -> _ArgumentT: + """Resolves a component argument against its configured fallback. + + Components accept connection details explicitly, but falling back to settings lets + the same model run in a different environment without code changes. + + Args: + value: The value passed to the component, if any. + fallback: The value loaded from `Settings`, if any. + arg_name: The name of the argument, used in the error message. + env_var: The environment variable that backs the fallback, used in the error + message. + + Returns: + The explicit value, or the configured value when none was passed. + + Raises: + ValueError: If neither an explicit value nor a configured value is available. + """ + resolved = value if value else fallback + if resolved is None: + raise ValueError( + f"'{arg_name}' is required: pass it explicitly or set the '{env_var}' " + f"environment variable." + ) + return resolved diff --git a/pyproject.toml b/pyproject.toml index 5ca74e49..b0ad7b72 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -43,7 +43,7 @@ dependencies = [ [project.optional-dependencies] aws = ["s3fs>=2024.9.0"] -aws-messaging = ["aioboto3>=13.0,<15"] +aws-messaging = ["aiobotocore>=3.4,<4"] azure = ["adlfs>=2024.7.0"] gcp = ["gcsfs>=2024.9.0"] gcp-pubsub = ["google-cloud-pubsub>=2.25,<3"] @@ -78,7 +78,12 @@ dev = [ test = [ "aiofile>=3.9,<4", "aiosqlite>=0.20,<1", + # Mirrors the optional broker extras so tests can import the real client classes and + # inject fakes, instead of leaving them unresolved to the type checker. + "aiobotocore>=3.4,<4", + "aiokafka>=0.11,<1", "coverage>=7.7.1,<8", + "google-cloud-pubsub>=2.25,<3", "llama-index-core>=0.12.30,<1", "llama-index-llms-openai>=0.3.33,<1", "moto[server]>=5.0,<6", diff --git a/tests/unit/test_aws_messaging_io.py b/tests/unit/test_aws_messaging_io.py index 59db971a..a145a28d 100644 --- a/tests/unit/test_aws_messaging_io.py +++ b/tests/unit/test_aws_messaging_io.py @@ -1,329 +1,388 @@ -"""Unit tests for AWS SQS/SNS message data reader/writer.""" +"""Unit tests for the AWS SQS/SNS message data implementations.""" from __future__ import annotations from collections import deque -import importlib.machinery -import json -import sys +import contextlib import typing as _t -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import AsyncMock, MagicMock, call +from botocore.exceptions import BotoCoreError, ClientError import pytest +from plugboard.exceptions import ( + MessageBrokerConnectionError, + MessageBrokerPermanentError, + MessageBrokerTransientError, + NoMoreDataException, +) +from plugboard.library import aws_messaging_io +from plugboard.library.aws_messaging_io import AWSSNSDataWriter, AWSSQSDataReader +from plugboard.utils.settings import Settings +from tests import conftest + + +def client_error(code: str, operation: str = "ReceiveMessage") -> ClientError: + """Builds a botocore `ClientError` carrying the given AWS error code.""" + return ClientError({"Error": {"Code": code, "Message": code}}, operation) + + +@contextlib.asynccontextmanager +async def _managed(client: MagicMock, closed: list[str]) -> _t.AsyncIterator[MagicMock]: + """Yields `client` and records that its context was exited.""" + try: + yield client + finally: + closed.append(getattr(client, "service_name", "client")) + + +@pytest.fixture +def aws_client(monkeypatch: pytest.MonkeyPatch) -> _t.Iterator[MagicMock]: + """Replaces the aiobotocore session with one handing out a single client double. + + The double's operations are `AsyncMock`s, so the call arguments are the observable + behaviour of the implementation. + """ + client = MagicMock(name="aws_client") + client.receive_message = AsyncMock(return_value={}) + client.delete_message_batch = AsyncMock(return_value={}) + client.publish = AsyncMock(return_value={"MessageId": "m1"}) + closed: list[str] = [] + created: list[tuple[str, dict[str, _t.Any]]] = [] + + def _create_client(service_name: str, **kwargs: _t.Any) -> _t.Any: + created.append((service_name, kwargs)) + client.service_name = service_name + return _managed(client, closed) + + session = MagicMock(name="session") + session.create_client.side_effect = _create_client + monkeypatch.setattr(aws_messaging_io, "get_session", lambda: session) + client.created = created + client.closed = closed + yield client + + +def make_reader(**kwargs: _t.Any) -> AWSSQSDataReader: + """Builds a reader for the fixed test queue.""" + kwargs.setdefault("field_names", ["x", "y"]) + return AWSSQSDataReader( + name="sqs-reader", + queue_url="https://sqs.test/1/messages", + region="eu-west-1", + **kwargs, + ) + + +def make_writer(**kwargs: _t.Any) -> AWSSNSDataWriter: + """Builds a writer for the fixed test topic.""" + kwargs.setdefault("field_names", ["x", "y"]) + return AWSSNSDataWriter( + name="sns-writer", + topic_arn="arn:aws:sns:eu-west-1:1:messages", + region="eu-west-1", + **kwargs, + ) + + +def sqs_message(body: str, handle: str) -> dict[str, str]: + """Builds an SQS message dict as returned by `ReceiveMessage`.""" + return {"Body": body, "ReceiptHandle": handle, "MessageId": handle} + # --------------------------------------------------------------------------- -# Mock the aioboto3 module before importing the implementation +# Reader: connection and receiving # --------------------------------------------------------------------------- -def _make_mock_module(name: str) -> MagicMock: - """Creates a mock module with __spec__ set for find_spec compatibility.""" - mock = MagicMock() - mock.__spec__ = importlib.machinery.ModuleSpec(name, None) - return mock +async def test_reader_connects_sqs_client_in_region(aws_client: MagicMock) -> None: + """Tests that the client is created for SQS in the configured region.""" + reader = make_reader() + await reader._connect() + assert aws_client.created == [("sqs", {"region_name": "eu-west-1"})] -_mock_aioboto3 = _make_mock_module("aioboto3") -_mock_aioboto3_session = MagicMock() -_mock_aioboto3.Session.return_value = _mock_aioboto3_session +async def test_reader_disconnect_exits_client_context(aws_client: MagicMock) -> None: + """Tests that teardown closes the client rather than dropping the reference.""" + reader = make_reader() + await reader._connect() + await reader._disconnect() + assert aws_client.closed == ["sqs"] -sys.modules.setdefault("aioboto3", _mock_aioboto3) -from plugboard.library.aws_messaging_io import AWSSNSDataWriter, AWSSQSDataReader # noqa: E402 +async def test_reader_receives_with_long_polling(aws_client: MagicMock) -> None: + """Tests the receive request carries the queue, batch size and wait time.""" + aws_client.receive_message.return_value = {"Messages": [sqs_message('{"x": 1}', "h1")]} + reader = make_reader(chunk_size=5, wait_time_seconds=7) + await reader._connect() + messages = await reader._receive() -# --------------------------------------------------------------------------- -# Helpers -# --------------------------------------------------------------------------- + aws_client.receive_message.assert_awaited_once_with( + QueueUrl="https://sqs.test/1/messages", MaxNumberOfMessages=5, WaitTimeSeconds=7 + ) + assert [msg["ReceiptHandle"] for msg in messages] == ["h1"] -def _make_sqs_message(body: dict[str, _t.Any] | str) -> dict[str, _t.Any]: - """Creates a mock SQS message dict.""" - if isinstance(body, dict): - body_str = json.dumps(body) - else: - body_str = body - return { - "MessageId": f"msg-{id(body)}", - "ReceiptHandle": f"receipt-{id(body)}", - "Body": body_str, - } +async def test_reader_caps_batch_at_the_sqs_maximum(aws_client: MagicMock) -> None: + """Tests that a larger `chunk_size` is capped, since SQS delivers at most 10.""" + reader = make_reader(chunk_size=25) + await reader._connect() + await reader._receive() + assert aws_client.receive_message.await_args.kwargs["MaxNumberOfMessages"] == 10 + +def test_reader_warns_when_capping_chunk_size(capsys: pytest.CaptureFixture[str]) -> None: + """Tests that capping `chunk_size` is reported rather than applied silently.""" + make_reader(chunk_size=25) + assert "chunk_size is capped" in capsys.readouterr().out -def _setup_mock_client() -> tuple[AsyncMock, AsyncMock]: - """Sets up a mock boto3 client with async context manager.""" - mock_client = AsyncMock() - mock_client_ctx = AsyncMock() - mock_client_ctx.__aenter__ = AsyncMock(return_value=mock_client) - mock_client_ctx.__aexit__ = AsyncMock(return_value=None) - _mock_aioboto3_session.client.return_value = mock_client_ctx - return mock_client, mock_client_ctx + +async def test_reader_empty_poll_returns_no_messages(aws_client: MagicMock) -> None: + """Tests that an empty long poll reads as 'nothing yet', not as an exhausted queue.""" + aws_client.receive_message.return_value = {} + reader = make_reader() + await reader._connect() + assert await reader._receive() == [] # --------------------------------------------------------------------------- -# Tests: AWSSQSDataReader +# Reader: conversion and acknowledgment # --------------------------------------------------------------------------- -@pytest.mark.asyncio -async def test_aws_sqs_reader_connect() -> None: - """Tests that the reader creates an SQS client on connect.""" - mock_client, _ = _setup_mock_client() - - reader = AWSSQSDataReader( - name="test-sqs-reader", - field_names=["x", "y"], - topic="test-queue", - queue_url="https://sqs.us-east-1.amazonaws.com/123456789/test-queue", - region="us-east-1", - ) +async def test_reader_converts_json_bodies() -> None: + """Tests that JSON message bodies become field buffers.""" + reader = make_reader() await reader._connect() + batch = [sqs_message('{"x": 1, "y": "a"}', "h1"), sqs_message('{"x": 2, "y": "b"}', "h2")] + assert await reader._convert(batch) == {"x": deque([1, 2]), "y": deque(["a", "b"])} - _mock_aioboto3_session.client.assert_called_with("sqs", region_name="us-east-1") - assert reader._client is mock_client +async def test_reader_non_json_body_uses_data_field() -> None: + """Tests that a raw body is exposed through the `data` field.""" + reader = make_reader(field_names=["data"], parse_json=False) + await reader._connect() + converted = await reader._convert([sqs_message("hello", "h1")]) + assert converted == {"data": deque(["hello"])} -@pytest.mark.asyncio -async def test_aws_sqs_reader_disconnect() -> None: - """Tests that the reader closes the SQS client on disconnect.""" - mock_client, mock_client_ctx = _setup_mock_client() - reader = AWSSQSDataReader( - name="test-sqs-reader", - field_names=["x"], - topic="test-queue", - queue_url="https://sqs.us-east-1.amazonaws.com/123456789/test-queue", - region="us-east-1", - ) +async def test_reader_deletes_processed_messages(aws_client: MagicMock) -> None: + """Tests that acking deletes by receipt handle, in one batch.""" + reader = make_reader() await reader._connect() - await reader._disconnect() - - mock_client_ctx.__aexit__.assert_called_once() - assert reader._client is None + messages = [sqs_message('{"x": 1}', "h1"), sqs_message('{"x": 2}', "h2")] + await reader._ack(messages) -@pytest.mark.asyncio -async def test_aws_sqs_reader_receive() -> None: - """Tests receiving messages from SQS.""" - mock_client, _ = _setup_mock_client() + aws_client.delete_message_batch.assert_awaited_once_with( + QueueUrl="https://sqs.test/1/messages", + DeleteMessageRequests=[ + {"Id": "0", "ReceiptHandle": "h1"}, + {"Id": "1", "ReceiptHandle": "h2"}, + ], + ) - test_data = [{"x": 1, "y": "a"}, {"x": 2, "y": "b"}] - sqs_messages = [_make_sqs_message(d) for d in test_data] - mock_client.receive_message = AsyncMock(return_value={"Messages": sqs_messages}) - reader = AWSSQSDataReader( - name="test-sqs-reader", - field_names=["x", "y"], - topic="test-queue", - queue_url="https://sqs.us-east-1.amazonaws.com/123456789/test-queue", - region="us-east-1", - chunk_size=10, - ) +async def test_reader_splits_deletes_into_batch_size_chunks(aws_client: MagicMock) -> None: + """Tests that more than 10 messages are deleted over several batch calls.""" + reader = make_reader() await reader._connect() - messages = await reader._receive() + messages = [sqs_message('{"x": 1}', f"h{index}") for index in range(12)] - assert len(messages) == 2 - mock_client.receive_message.assert_called() + await reader._ack(messages) + calls = aws_client.delete_message_batch.await_args_list + assert [len(call.kwargs["DeleteMessageRequests"]) for call in calls] == [10, 2] + assert calls[1].kwargs["DeleteMessageRequests"][0]["ReceiptHandle"] == "h10" -@pytest.mark.asyncio -async def test_aws_sqs_reader_receive_empty() -> None: - """Tests receiving empty response from SQS.""" - mock_client, _ = _setup_mock_client() - mock_client.receive_message = AsyncMock(return_value={}) - reader = AWSSQSDataReader( - name="test-sqs-reader", - field_names=["x"], - topic="test-queue", - queue_url="https://sqs.us-east-1.amazonaws.com/123456789/test-queue", - region="us-east-1", - ) +async def test_reader_skips_messages_without_receipt_handle(aws_client: MagicMock) -> None: + """Tests that a message with no receipt handle is not sent to the delete API.""" + reader = make_reader() await reader._connect() - messages = await reader._receive() + await reader._ack([{"Body": '{"x": 1}'}]) + aws_client.delete_message_batch.assert_not_awaited() - assert messages == [] +async def test_reader_reports_failed_deletes_as_transient(aws_client: MagicMock) -> None: + """Tests that a partially failed delete retries instead of claiming success.""" + aws_client.delete_message_batch.return_value = { + "Errors": [{"Id": "0", "Code": "ReceiptHandleIsInvalid"}] + } + reader = make_reader() + await reader._connect() -@pytest.mark.asyncio -async def test_aws_sqs_reader_convert_json() -> None: - """Tests converting JSON SQS messages to field buffer.""" - reader = AWSSQSDataReader( - name="test-sqs-reader", - field_names=["x", "y"], - topic="test-queue", - queue_url="https://sqs.us-east-1.amazonaws.com/123456789/test-queue", - region="us-east-1", - parse_json=True, - ) + with pytest.raises(MessageBrokerTransientError, match="Failed to delete 1"): + await reader._ack([sqs_message('{"x": 1}', "h1")]) - sqs_messages = [ - _make_sqs_message({"x": 1, "y": "a"}), - _make_sqs_message({"x": 2, "y": "b"}), - ] - result = await reader._convert(sqs_messages) - assert list(result["x"]) == [1, 2] - assert list(result["y"]) == ["a", "b"] - - -@pytest.mark.asyncio -async def test_aws_sqs_reader_convert_raw() -> None: - """Tests converting raw SQS messages to field buffer.""" - reader = AWSSQSDataReader( - name="test-sqs-reader", - field_names=["data"], - topic="test-queue", - queue_url="https://sqs.us-east-1.amazonaws.com/123456789/test-queue", - region="us-east-1", - parse_json=False, - ) - sqs_messages = [_make_sqs_message("raw-data-1"), _make_sqs_message("raw-data-2")] - result = await reader._convert(sqs_messages) - assert list(result["data"]) == ["raw-data-1", "raw-data-2"] +async def test_reader_without_connection_raises_connection_error(aws_client: MagicMock) -> None: + """Tests that receiving before connecting fails loudly.""" + reader = make_reader() + with pytest.raises(MessageBrokerConnectionError): + await reader._receive() -@pytest.mark.asyncio -async def test_aws_sqs_reader_ack() -> None: - """Tests acknowledging (deleting) SQS messages.""" - mock_client, _ = _setup_mock_client() - mock_client.delete_message = AsyncMock() +# --------------------------------------------------------------------------- +# Reader: error mapping +# --------------------------------------------------------------------------- - reader = AWSSQSDataReader( - name="test-sqs-reader", - field_names=["x"], - topic="test-queue", - queue_url="https://sqs.us-east-1.amazonaws.com/123456789/test-queue", - region="us-east-1", - ) + +@pytest.mark.parametrize( + ("code", "expected"), + [ + ("AWS.SimpleQueueService.NonExistentQueue", NoMoreDataException), + ("AccessDenied", MessageBrokerPermanentError), + ("InvalidAddress", MessageBrokerPermanentError), + ("ThrottlingException", MessageBrokerTransientError), + ("InternalError", MessageBrokerTransientError), + ], +) +async def test_reader_maps_aws_error_codes( + aws_client: MagicMock, code: str, expected: type[Exception] +) -> None: + """Tests that AWS error codes land on the right broker exception type.""" + aws_client.receive_message.side_effect = client_error(code) + reader = make_reader() await reader._connect() - sqs_messages = [_make_sqs_message({"x": 1})] - await reader._ack(sqs_messages) + with pytest.raises(expected): + await reader._receive() - mock_client.delete_message.assert_called() +async def test_reader_transport_error_is_transient(aws_client: MagicMock) -> None: + """Tests that a connection-level failure is retryable.""" + aws_client.receive_message.side_effect = BotoCoreError() + reader = make_reader() + await reader._connect() -@pytest.mark.asyncio -async def test_aws_sqs_reader_long_polling() -> None: - """Tests that long polling is configured correctly.""" - mock_client, _ = _setup_mock_client() - mock_client.receive_message = AsyncMock(return_value={}) + with pytest.raises(MessageBrokerTransientError): + await reader._receive() - reader = AWSSQSDataReader( - name="test-sqs-reader", - field_names=["x"], - topic="test-queue", - queue_url="https://sqs.us-east-1.amazonaws.com/123456789/test-queue", - region="us-east-1", - wait_time_seconds=15, - ) + +async def test_reader_unclassified_error_propagates(aws_client: MagicMock) -> None: + """Tests that an unknown error is not rewritten into a broker exception.""" + aws_client.receive_message.side_effect = ValueError("unexpected") + reader = make_reader() await reader._connect() - await reader._receive() - call_args = mock_client.receive_message.call_args - assert call_args[1]["WaitTimeSeconds"] == 15 + with pytest.raises(ValueError, match="unexpected"): + await reader._receive() # --------------------------------------------------------------------------- -# Tests: AWSSNSDataWriter +# Writer # --------------------------------------------------------------------------- -@pytest.mark.asyncio -async def test_aws_sns_writer_connect() -> None: - """Tests that the writer creates an SNS client on connect.""" - mock_client, _ = _setup_mock_client() +async def test_writer_connects_sns_client(aws_client: MagicMock) -> None: + """Tests that the writer creates a client for SNS in the configured region.""" + writer = make_writer() + await writer._connect() + assert aws_client.created == [("sns", {"region_name": "eu-west-1"})] - writer = AWSSNSDataWriter( - name="test-sns-writer", - field_names=["x"], - topic="test-topic", - topic_arn="arn:aws:sns:us-east-1:123456789:test-topic", - region="us-east-1", - ) + +async def test_writer_publishes_each_message(aws_client: MagicMock) -> None: + """Tests that every encoded record is published to the topic.""" + writer = make_writer() await writer._connect() - _mock_aioboto3_session.client.assert_called_with("sns", region_name="us-east-1") - assert writer._client is mock_client + await writer._send(['{"x": 1}', '{"x": 2}']) + assert aws_client.publish.await_args_list == [ + call(TopicArn="arn:aws:sns:eu-west-1:1:messages", Message='{"x": 1}'), + call(TopicArn="arn:aws:sns:eu-west-1:1:messages", Message='{"x": 2}'), + ] -@pytest.mark.asyncio -async def test_aws_sns_writer_disconnect() -> None: - """Tests that the writer closes the SNS client on disconnect.""" - mock_client, mock_client_ctx = _setup_mock_client() - writer = AWSSNSDataWriter( - name="test-sns-writer", - field_names=["x"], - topic="test-topic", - topic_arn="arn:aws:sns:us-east-1:123456789:test-topic", - region="us-east-1", - ) +async def test_writer_converts_fields_to_json_strings() -> None: + """Tests that buffered fields become one JSON message per record.""" + writer = make_writer() await writer._connect() - await writer._disconnect() + converted = await writer._convert({"x": deque([1, 2]), "y": deque(["a", "b"])}) + assert converted == ['{"x": 1, "y": "a"}', '{"x": 2, "y": "b"}'] - mock_client_ctx.__aexit__.assert_called() - assert writer._client is None +async def test_writer_non_json_sends_first_field() -> None: + """Tests that a raw writer publishes the first field's value.""" + writer = make_writer(field_names=["x"], parse_json=False) + await writer._connect() + assert await writer._convert({"x": deque([7])}) == ["7"] -@pytest.mark.asyncio -async def test_aws_sns_writer_send() -> None: - """Tests sending messages to SNS.""" - mock_client, _ = _setup_mock_client() - mock_client.publish = AsyncMock() - writer = AWSSNSDataWriter( - name="test-sns-writer", - field_names=["x"], - topic="test-topic", - topic_arn="arn:aws:sns:us-east-1:123456789:test-topic", - region="us-east-1", - ) +async def test_writer_missing_topic_is_permanent(aws_client: MagicMock) -> None: + """Tests that a deleted topic is not treated as an exhausted source.""" + aws_client.publish.side_effect = client_error("NotFound", "Publish") + writer = make_writer() await writer._connect() - messages = ['{"x": 1}', '{"x": 2}'] - await writer._send(messages) + with pytest.raises(MessageBrokerPermanentError): + await writer._send(['{"x": 1}']) - assert mock_client.publish.call_count == 2 +async def test_writer_disconnect_exits_client_context(aws_client: MagicMock) -> None: + """Tests that teardown closes the client.""" + writer = make_writer() + await writer._connect() + await writer._disconnect() + assert aws_client.closed == ["sns"] -@pytest.mark.asyncio -async def test_aws_sns_writer_convert_json() -> None: - """Tests converting field data to JSON messages.""" - writer = AWSSNSDataWriter( - name="test-sns-writer", - field_names=["x", "y"], - topic="test-topic", - topic_arn="arn:aws:sns:us-east-1:123456789:test-topic", - region="us-east-1", - parse_json=True, - ) - - data = {"x": deque([1, 2]), "y": deque(["a", "b"])} - messages = await writer._convert(data) - assert len(messages) == 2 - assert json.loads(messages[0]) == {"x": 1, "y": "a"} - assert json.loads(messages[1]) == {"x": 2, "y": "b"} +async def test_writer_without_connection_raises_connection_error(aws_client: MagicMock) -> None: + """Tests that sending before connecting fails loudly.""" + writer = make_writer() + with pytest.raises(MessageBrokerConnectionError): + await writer._send(["one"]) -@pytest.mark.asyncio -async def test_aws_sns_writer_convert_raw() -> None: - """Tests converting field data to raw string messages.""" - writer = AWSSNSDataWriter( - name="test-sns-writer", - field_names=["data"], - topic="test-topic", - topic_arn="arn:aws:sns:us-east-1:123456789:test-topic", - region="us-east-1", - parse_json=False, - ) +# --------------------------------------------------------------------------- +# Settings resolution +# --------------------------------------------------------------------------- - data = {"data": deque(["raw1", "raw2"])} - messages = await writer._convert(data) - assert len(messages) == 2 - assert messages[0] == "raw1" - assert messages[1] == "raw2" +async def test_region_falls_back_to_settings() -> None: + """Tests that `AWS_REGION` supplies the region when none is passed.""" + settings = Settings.model_validate({"aws": {"region": "ap-south-1"}}) + with conftest.override_settings(settings): + reader = AWSSQSDataReader( + name="sqs-reader", queue_url="https://sqs.test/1/messages", field_names=["x"] + ) + assert reader._region == "ap-south-1" + + +def test_explicit_region_wins_over_settings() -> None: + """Tests that an explicit region is not overridden by configuration.""" + settings = Settings.model_validate({"aws": {"region": "ap-south-1"}}) + with conftest.override_settings(settings): + reader = AWSSQSDataReader( + name="sqs-reader", + queue_url="https://sqs.test/1/messages", + field_names=["x"], + region="us-west-2", + ) + assert reader._region == "us-west-2" + + +def test_missing_region_names_the_environment_variable() -> None: + """Tests that an unresolvable region reports what to configure.""" + with conftest.override_settings(Settings.model_validate({})): + with pytest.raises(ValueError, match="AWS_REGION"): + AWSSQSDataReader( + name="sqs-reader", + queue_url="https://sqs.test/1/messages", + field_names=["x"], + region=None, + ) + + +def test_reader_topic_defaults_to_queue_url() -> None: + """Tests that the queue URL identifies the reader's topic.""" + reader = make_reader() + assert reader._topic == "https://sqs.test/1/messages" + + +def test_writer_topic_defaults_to_topic_arn() -> None: + """Tests that the topic ARN identifies the writer's topic.""" + writer = make_writer() + assert writer._topic == "arn:aws:sns:eu-west-1:1:messages" diff --git a/tests/unit/test_file_reader_writer.py b/tests/unit/test_file_reader_writer.py index 859a1fcd..60617ec6 100644 --- a/tests/unit/test_file_reader_writer.py +++ b/tests/unit/test_file_reader_writer.py @@ -30,6 +30,12 @@ def mock_s3_server() -> _t.Generator[None, None, None]: """ server = ThreadedMotoServer(ip_address=S3_IP_ADDRESS, port=S3_PORT) server.start() + if "AWS_CONFIG_FILE" not in os.environ: + # These tests only ever talk to the local moto server, so they must not pick up + # a developer's real AWS profile: its region changes how S3 validates buckets. + empty_config = tempfile.NamedTemporaryFile(prefix="plugboard-empty-aws-", delete=False) + os.environ["AWS_CONFIG_FILE"] = empty_config.name + os.environ["AWS_SHARED_CREDENTIALS_FILE"] = empty_config.name if "AWS_SECRET_ACCESS_KEY" not in os.environ: os.environ["AWS_SECRET_ACCESS_KEY"] = "test-access-key" # noqa: S105 if "AWS_ACCESS_KEY_ID" not in os.environ: diff --git a/tests/unit/test_gcp_pubsub_io.py b/tests/unit/test_gcp_pubsub_io.py index c88fa537..b037c954 100644 --- a/tests/unit/test_gcp_pubsub_io.py +++ b/tests/unit/test_gcp_pubsub_io.py @@ -1,384 +1,361 @@ -"""Unit tests for GCP PubSub message data reader/writer.""" +"""Unit tests for the Google Cloud PubSub message data implementations.""" from __future__ import annotations from collections import deque -import importlib.machinery -import json -import sys +import threading +from types import SimpleNamespace import typing as _t -from unittest.mock import MagicMock +from unittest.mock import MagicMock, call +from google.api_core import exceptions as gcp_exceptions import pytest -from plugboard.exceptions import NoMoreDataException - - -# --------------------------------------------------------------------------- -# Mock the google.cloud.pubsub_v1 module before importing the implementation -# --------------------------------------------------------------------------- - - -def _make_mock_module(name: str) -> MagicMock: - """Creates a mock module with __spec__ set for find_spec compatibility.""" - mock = MagicMock() - mock.__spec__ = importlib.machinery.ModuleSpec(name, None) - return mock - - -_mock_pubsub = _make_mock_module("google.cloud.pubsub_v1") -_mock_pubsub.SubscriberClient = MagicMock() -_mock_pubsub.PublisherClient = MagicMock() +from plugboard.exceptions import ( + MessageBrokerConnectionError, + MessageBrokerPermanentError, + MessageBrokerTransientError, + NoMoreDataException, +) +from plugboard.library import gcp_pubsub_io +from plugboard.library.gcp_pubsub_io import GCPPubSubDataReader, GCPPubSubDataWriter +from plugboard.utils.settings import Settings +from tests import conftest + + +def make_message(data: bytes) -> MagicMock: + """Builds a PubSub `ReceivedMessage` double carrying `data`.""" + wrapper = MagicMock() + wrapper.message.data = data + wrapper.ack_id = f"ack-{data!r}" + return wrapper + + +@pytest.fixture +def subscriber(monkeypatch: pytest.MonkeyPatch) -> MagicMock: + """Replaces the PubSub subscriber client class with a per-test double. + + The double is synchronous, so any call it records also proves the implementation + left the event loop thread. + """ + client = MagicMock(name="subscriber_client") + client.pull.return_value = MagicMock(received_messages=[]) + factory = MagicMock(name="SubscriberClient", return_value=client) + monkeypatch.setattr( + gcp_pubsub_io, + "pubsub_v1", + SimpleNamespace(SubscriberClient=factory, PublisherClient=MagicMock()), + ) + return client + + +@pytest.fixture +def publisher(monkeypatch: pytest.MonkeyPatch) -> MagicMock: + """Replaces the PubSub publisher client class with a per-test double.""" + client = MagicMock(name="publisher_client") + factory = MagicMock(name="PublisherClient", return_value=client) + monkeypatch.setattr( + gcp_pubsub_io, + "pubsub_v1", + SimpleNamespace(SubscriberClient=MagicMock(), PublisherClient=factory), + ) + return client -_mock_google = _make_mock_module("google") -_mock_google_cloud = _make_mock_module("google.cloud") -# Wire up the attribute chain so `from google.cloud import pubsub_v1` works -_mock_google_cloud.pubsub_v1 = _mock_pubsub -_mock_google.cloud = _mock_google_cloud -_mock_modules = { - "google": _mock_google, - "google.cloud": _mock_google_cloud, - "google.cloud.pubsub_v1": _mock_pubsub, - "google.cloud.pubsub_v1.subscriber": _make_mock_module("google.cloud.pubsub_v1.subscriber"), - "google.cloud.pubsub_v1.subscriber.message": _make_mock_module( - "google.cloud.pubsub_v1.subscriber.message" - ), -} +def make_reader(**kwargs: _t.Any) -> GCPPubSubDataReader: + """Builds a reader for the fixed test project/subscription.""" + kwargs.setdefault("field_names", ["x", "y"]) + return GCPPubSubDataReader( + name="gcp-reader", + project_id="test-project", + subscription_id="test-subscription", + **kwargs, + ) -# Install mocks before importing the module under test -for _mod_name, _mod in _mock_modules.items(): - sys.modules.setdefault(_mod_name, _mod) -from plugboard.library.gcp_pubsub_io import GCPPubSubDataReader, GCPPubSubDataWriter # noqa: E402 +def make_writer(**kwargs: _t.Any) -> GCPPubSubDataWriter: + """Builds a writer for the fixed test project/topic.""" + kwargs.setdefault("field_names", ["x", "y"]) + return GCPPubSubDataWriter( + name="gcp-writer", + project_id="test-project", + topic_id="test-topic", + **kwargs, + ) # --------------------------------------------------------------------------- -# Helpers +# Reader # --------------------------------------------------------------------------- -def _make_pubsub_message(data: dict[str, _t.Any] | bytes) -> MagicMock: - """Creates a mock PubSub ReceivedMessage.""" - if isinstance(data, dict): - raw_data = json.dumps(data).encode("utf-8") - else: - raw_data = data - msg = MagicMock() - msg.message.data = raw_data - msg.ack_id = f"ack-{id(msg)}" - return msg - +async def test_reader_connect_builds_client_off_the_event_loop( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Tests that constructing the subscriber client does not run on the event loop. -def _make_pull_response(messages: list[MagicMock]) -> MagicMock: - """Creates a mock Pull response.""" - response = MagicMock() - response.received_messages = messages - return response + Building the client opens a gRPC channel and starts its threads, which is blocking. + """ + build_threads: list[str] = [] + client = MagicMock(name="subscriber_client") + def _build(*args: _t.Any, **kwargs: _t.Any) -> MagicMock: + build_threads.append(threading.current_thread().name) + return client -# --------------------------------------------------------------------------- -# Tests: GCPPubSubDataReader -# --------------------------------------------------------------------------- + monkeypatch.setattr( + gcp_pubsub_io, + "pubsub_v1", + SimpleNamespace(SubscriberClient=_build, PublisherClient=MagicMock()), + ) + reader = make_reader() + await reader._connect() + assert build_threads == [build_threads[0]] + assert build_threads[0] != threading.current_thread().name -@pytest.mark.asyncio -async def test_gcp_pubsub_reader_connect() -> None: - """Tests that the reader creates a subscriber client on connect.""" - mock_subscriber = MagicMock() - _mock_pubsub.SubscriberClient.return_value = mock_subscriber - reader = GCPPubSubDataReader( - name="test-gcp-reader", - field_names=["x", "y"], - topic="test-topic", - project_id="test-project", - subscription_id="test-sub", - ) +async def test_reader_pulls_from_subscription_path(subscriber: MagicMock) -> None: + """Tests that a pull asks for the configured subscription and batch size.""" + subscriber.pull.return_value = MagicMock(received_messages=[make_message(b'{"x": 1}')]) + reader = make_reader(chunk_size=5) await reader._connect() - _mock_pubsub.SubscriberClient.assert_called() - assert reader._subscriber is mock_subscriber + messages = await reader._receive() + subscriber.pull.assert_called_once_with( + request={ + "subscription": "projects/test-project/subscriptions/test-subscription", + "max_messages": 5, + }, + timeout=30.0, + ) + assert len(messages) == 1 -@pytest.mark.asyncio -async def test_gcp_pubsub_reader_disconnect() -> None: - """Tests that the reader closes the subscriber client on disconnect.""" - mock_subscriber = MagicMock() - _mock_pubsub.SubscriberClient.return_value = mock_subscriber - reader = GCPPubSubDataReader( - name="test-gcp-reader", - field_names=["x"], - topic="test-topic", - project_id="test-project", - subscription_id="test-sub", - ) +async def test_reader_empty_pull_returns_no_messages(subscriber: MagicMock) -> None: + """Tests that an idle pull reads as 'nothing yet', not as an exhausted source.""" + subscriber.pull.return_value = MagicMock(received_messages=[]) + reader = make_reader() await reader._connect() - await reader._disconnect() + assert await reader._receive() == [] - mock_subscriber.close.assert_called_once() - assert reader._subscriber is None +async def test_reader_pulls_off_the_event_loop(subscriber: MagicMock) -> None: + """Tests that the blocking PubSub pull runs in a worker thread. -@pytest.mark.asyncio -async def test_gcp_pubsub_reader_receive() -> None: - """Tests receiving messages from PubSub.""" - mock_subscriber = MagicMock() - _mock_pubsub.SubscriberClient.return_value = mock_subscriber + `pull` is synchronous gRPC in the real client; if it ran on the event loop it would + stall every other component in the process. + """ + caller_threads: list[str] = [] - test_data = [{"x": 1, "y": "a"}, {"x": 2, "y": "b"}] - mock_messages = [_make_pubsub_message(d) for d in test_data] - mock_response = _make_pull_response(mock_messages) - mock_subscriber.pull.return_value = mock_response + def _pull(request: dict[str, _t.Any], timeout: float) -> MagicMock: + caller_threads.append(threading.current_thread().name) + return MagicMock(received_messages=[]) - reader = GCPPubSubDataReader( - name="test-gcp-reader", - field_names=["x", "y"], - topic="test-topic", - project_id="test-project", - subscription_id="test-sub", - chunk_size=10, - ) + subscriber.pull.side_effect = _pull + reader = make_reader() await reader._connect() - messages = await reader._receive() - - assert len(messages) == 2 - mock_subscriber.pull.assert_called() - + await reader._receive() -@pytest.mark.asyncio -async def test_gcp_pubsub_reader_receive_empty() -> None: - """Tests receiving empty response from PubSub.""" - mock_subscriber = MagicMock() - _mock_pubsub.SubscriberClient.return_value = mock_subscriber + assert len(caller_threads) == 1 + assert caller_threads[0] != threading.current_thread().name - mock_response = _make_pull_response([]) - mock_subscriber.pull.return_value = mock_response - reader = GCPPubSubDataReader( - name="test-gcp-reader", - field_names=["x"], - topic="test-topic", - project_id="test-project", - subscription_id="test-sub", - ) +async def test_reader_converts_json_messages(subscriber: MagicMock) -> None: + """Tests that JSON payloads become field buffers keyed by the declared fields.""" + reader = make_reader() await reader._connect() - messages = await reader._receive() + batch = [make_message(b'{"x": 1, "y": "a"}'), make_message(b'{"x": 2, "y": "b"}')] + assert await reader._convert(batch) == {"x": deque([1, 2]), "y": deque(["a", "b"])} - assert messages == [] +async def test_reader_non_json_messages_fill_single_field() -> None: + """Tests that a raw payload is exposed through the `data` field.""" + reader = make_reader(field_names=["data"], parse_json=False) + await reader._connect() + converted = await reader._convert([make_message(b"hello")]) + assert converted == {"data": deque([b"hello"])} -@pytest.mark.asyncio -async def test_gcp_pubsub_reader_receive_not_found() -> None: - """Tests that NOT_FOUND error raises NoMoreDataException.""" - mock_subscriber = MagicMock() - _mock_pubsub.SubscriberClient.return_value = mock_subscriber - mock_subscriber.pull.side_effect = Exception("NOT_FOUND: Subscription deleted") - reader = GCPPubSubDataReader( - name="test-gcp-reader", - field_names=["x"], - topic="test-topic", - project_id="test-project", - subscription_id="test-sub", - ) +async def test_reader_ack_sends_ack_ids(subscriber: MagicMock) -> None: + """Tests that acknowledging a batch sends exactly its ack ids.""" + reader = make_reader() await reader._connect() + messages = [make_message(b'{"x": 1}'), make_message(b'{"x": 2}')] - with pytest.raises(NoMoreDataException): - await reader._receive() - + await reader._ack(messages) -@pytest.mark.asyncio -async def test_gcp_pubsub_reader_convert_json() -> None: - """Tests converting JSON messages to field buffer.""" - reader = GCPPubSubDataReader( - name="test-gcp-reader", - field_names=["x", "y"], - topic="test-topic", - project_id="test-project", - subscription_id="test-sub", - parse_json=True, + subscriber.acknowledge.assert_called_once_with( + request={ + "subscription": "projects/test-project/subscriptions/test-subscription", + "ack_ids": [msg.ack_id for msg in messages], + } ) - test_data = [{"x": 1, "y": "a"}, {"x": 2, "y": "b"}] - mock_messages = [_make_pubsub_message(d) for d in test_data] - result = await reader._convert(mock_messages) - assert list(result["x"]) == [1, 2] - assert list(result["y"]) == ["a", "b"] +async def test_reader_disconnect_closes_subscriber(subscriber: MagicMock) -> None: + """Tests that teardown closes the client, releasing its channel.""" + reader = make_reader() + await reader._connect() + await reader._disconnect() + subscriber.close.assert_called_once_with() + + +@pytest.mark.parametrize( + ("error", "expected"), + [ + (gcp_exceptions.NotFound("gone"), NoMoreDataException), + (gcp_exceptions.PermissionDenied("nope"), MessageBrokerPermanentError), + (gcp_exceptions.InvalidArgument("bad"), MessageBrokerPermanentError), + (gcp_exceptions.ServiceUnavailable("busy"), MessageBrokerTransientError), + (gcp_exceptions.DeadlineExceeded("late"), MessageBrokerTransientError), + ], +) +async def test_reader_maps_broker_errors( + subscriber: MagicMock, error: Exception, expected: type[Exception] +) -> None: + """Tests that PubSub failures land on the right broker exception type.""" + subscriber.pull.side_effect = error + reader = make_reader() + await reader._connect() + with pytest.raises(expected): + await reader._receive() -@pytest.mark.asyncio -async def test_gcp_pubsub_reader_convert_raw() -> None: - """Tests converting raw bytes messages to field buffer.""" - reader = GCPPubSubDataReader( - name="test-gcp-reader", - field_names=["data"], - topic="test-topic", - project_id="test-project", - subscription_id="test-sub", - parse_json=False, - ) - mock_messages = [_make_pubsub_message(b"raw-data-1"), _make_pubsub_message(b"raw-data-2")] - result = await reader._convert(mock_messages) - assert list(result["data"]) == [b"raw-data-1", b"raw-data-2"] +async def test_reader_unclassified_error_propagates(subscriber: MagicMock) -> None: + """Tests that an unknown error is not rewritten into a broker exception.""" + subscriber.pull.side_effect = ValueError("unexpected") + reader = make_reader() + await reader._connect() + with pytest.raises(ValueError, match="unexpected"): + await reader._receive() -@pytest.mark.asyncio -async def test_gcp_pubsub_reader_ack() -> None: - """Tests acknowledging messages.""" - mock_subscriber = MagicMock() - _mock_pubsub.SubscriberClient.return_value = mock_subscriber - reader = GCPPubSubDataReader( - name="test-gcp-reader", - field_names=["x"], - topic="test-topic", - project_id="test-project", - subscription_id="test-sub", - ) - await reader._connect() +async def test_reader_without_connection_raises_connection_error(subscriber: MagicMock) -> None: + """Tests that receiving before connecting fails loudly rather than silently.""" + reader = make_reader() + with pytest.raises(MessageBrokerConnectionError): + await reader._receive() - mock_messages = [_make_pubsub_message({"x": 1})] - mock_messages[0].ack_id = "ack-123" - await reader._ack(mock_messages) - mock_subscriber.acknowledge.assert_called() - call_args = mock_subscriber.acknowledge.call_args - assert call_args[1]["request"]["ack_ids"] == ["ack-123"] +async def test_reader_missing_project_id(monkeypatch: pytest.MonkeyPatch) -> None: + """Tests that a missing project id reports the environment variable to set.""" + with conftest.override_settings(Settings.model_validate({})): + with pytest.raises(ValueError, match="GCP_PUBSUB_PROJECT_ID"): + GCPPubSubDataReader( + name="gcp-reader", + subscription_id="test-subscription", + field_names=["x"], + project_id=None, + ) -@pytest.mark.asyncio -async def test_gcp_pubsub_reader_subscription_path() -> None: - """Tests that the subscription path is constructed correctly.""" - reader = GCPPubSubDataReader( - name="test-gcp-reader", - field_names=["x"], - topic="test-topic", - project_id="my-project", - subscription_id="my-sub", - ) - assert reader._subscription_path == "projects/my-project/subscriptions/my-sub" +async def test_reader_project_id_from_settings() -> None: + """Tests that `GCP_PUBSUB_PROJECT_ID` supplies the project when not passed.""" + settings = Settings.model_validate({"gcp_pubsub": {"project_id": "settings-project"}}) + with conftest.override_settings(settings): + reader = GCPPubSubDataReader( + name="gcp-reader", + subscription_id="test-subscription", + field_names=["x"], + ) + assert reader._subscription_path == "projects/settings-project/subscriptions/test-subscription" # --------------------------------------------------------------------------- -# Tests: GCPPubSubDataWriter +# Writer # --------------------------------------------------------------------------- -@pytest.mark.asyncio -async def test_gcp_pubsub_writer_connect() -> None: - """Tests that the writer creates a publisher client on connect.""" - mock_publisher = MagicMock() - _mock_pubsub.PublisherClient.return_value = mock_publisher - - writer = GCPPubSubDataWriter( - name="test-gcp-writer", - field_names=["x"], - topic="test-topic", - project_id="test-project", - topic_id="test-topic-id", - ) +async def test_writer_publishes_encoded_records(publisher: MagicMock) -> None: + """Tests that each record is published to the topic path as JSON bytes.""" + futures = [MagicMock(), MagicMock()] + publisher.publish.side_effect = futures + writer = make_writer() await writer._connect() - _mock_pubsub.PublisherClient.assert_called() - assert writer._publisher is mock_publisher + await writer._send([b'{"x": 1}', b'{"x": 2}']) + assert [published.kwargs["data"] for published in publisher.publish.call_args_list] == [ + b'{"x": 1}', + b'{"x": 2}', + ] + assert [published.args for published in publisher.publish.call_args_list] == [ + ("projects/test-project/topics/test-topic",) + ] * 2 + assert [future.result.call_args for future in futures] == [ + call(timeout=60.0), + call(timeout=60.0), + ] -@pytest.mark.asyncio -async def test_gcp_pubsub_writer_disconnect() -> None: - """Tests that the writer closes the publisher client on disconnect.""" - mock_publisher = MagicMock() - _mock_pubsub.PublisherClient.return_value = mock_publisher - writer = GCPPubSubDataWriter( - name="test-gcp-writer", - field_names=["x"], - topic="test-topic", - project_id="test-project", - topic_id="test-topic-id", - ) - await writer._connect() - await writer._disconnect() +async def test_writer_waits_off_the_event_loop(publisher: MagicMock) -> None: + """Tests that waiting for broker confirmation does not block the event loop.""" + caller_threads: list[str] = [] - assert writer._publisher is None + class _Future: + def result(self, timeout: float) -> None: + caller_threads.append(threading.current_thread().name) + publisher.publish.return_value = _Future() + writer = make_writer() + await writer._connect() + await writer._send([b"one", b"two"]) -@pytest.mark.asyncio -async def test_gcp_pubsub_writer_send() -> None: - """Tests sending messages to PubSub.""" - mock_publisher = MagicMock() - _mock_pubsub.PublisherClient.return_value = mock_publisher + assert len(caller_threads) == 2 + assert all(name != threading.current_thread().name for name in caller_threads) - mock_future = MagicMock() - mock_publisher.publish.return_value = mock_future - writer = GCPPubSubDataWriter( - name="test-gcp-writer", - field_names=["x"], - topic="test-topic", - project_id="test-project", - topic_id="test-topic-id", - ) +async def test_writer_converts_fields_to_json_bytes(publisher: MagicMock) -> None: + """Tests that buffered fields become one JSON message per record.""" + writer = make_writer() await writer._connect() + converted = await writer._convert({"x": deque([1, 2]), "y": deque(["a", "b"])}) + assert converted == [b'{"x": 1, "y": "a"}', b'{"x": 2, "y": "b"}'] - messages = [b"msg1", b"msg2"] - await writer._send(messages) - assert mock_publisher.publish.call_count == 2 - assert mock_future.result.call_count == 2 +async def test_writer_non_json_sends_first_field(publisher: MagicMock) -> None: + """Tests that a raw writer publishes the first field's value.""" + writer = make_writer(field_names=["x"], parse_json=False) + await writer._connect() + assert await writer._convert({"x": deque(["hello"])}) == [b"hello"] -@pytest.mark.asyncio -async def test_gcp_pubsub_writer_convert_json() -> None: - """Tests converting field data to JSON messages.""" - writer = GCPPubSubDataWriter( - name="test-gcp-writer", - field_names=["x", "y"], - topic="test-topic", - project_id="test-project", - topic_id="test-topic-id", - parse_json=True, - ) +async def test_writer_disconnect_stops_publisher(publisher: MagicMock) -> None: + """Tests that teardown stops the publisher so its channel and threads are released.""" + writer = make_writer() + await writer._connect() + await writer._disconnect() + publisher.stop.assert_called_once_with() - data = {"x": deque([1, 2]), "y": deque(["a", "b"])} - messages = await writer._convert(data) - assert len(messages) == 2 - assert json.loads(messages[0]) == {"x": 1, "y": "a"} - assert json.loads(messages[1]) == {"x": 2, "y": "b"} +async def test_writer_publish_failure_maps_error(publisher: MagicMock) -> None: + """Tests that a publish failure surfaces as a broker exception.""" + publisher.publish.side_effect = gcp_exceptions.ServiceUnavailable("busy") + writer = make_writer() + await writer._connect() + with pytest.raises(MessageBrokerTransientError): + await writer._send([b"one"]) -@pytest.mark.asyncio -async def test_gcp_pubsub_writer_convert_raw() -> None: - """Tests converting field data to raw bytes messages.""" - writer = GCPPubSubDataWriter( - name="test-gcp-writer", - field_names=["data"], - topic="test-topic", - project_id="test-project", - topic_id="test-topic-id", - parse_json=False, - ) - data = {"data": deque([b"raw1", b"raw2"])} - messages = await writer._convert(data) +async def test_writer_without_connection_raises_connection_error(publisher: MagicMock) -> None: + """Tests that sending before connecting fails loudly.""" + writer = make_writer() + with pytest.raises(MessageBrokerConnectionError): + await writer._send([b"one"]) - assert len(messages) == 2 - assert messages[0] == b"raw1" - assert messages[1] == b"raw2" +async def test_writer_maps_unconfirmed_publish(publisher: MagicMock) -> None: + """Tests that a future that fails on `result` is reported as a broker error.""" + future = MagicMock() + future.result.side_effect = gcp_exceptions.FailedPrecondition("rejected") + publisher.publish.return_value = future + writer = make_writer() + await writer._connect() -@pytest.mark.asyncio -async def test_gcp_pubsub_writer_topic_path() -> None: - """Tests that the topic path is constructed correctly.""" - writer = GCPPubSubDataWriter( - name="test-gcp-writer", - field_names=["x"], - topic="test-topic", - project_id="my-project", - topic_id="my-topic", - ) - assert writer._topic_path == "projects/my-project/topics/my-topic" + with pytest.raises(MessageBrokerPermanentError): + await writer._send([b"one"]) diff --git a/tests/unit/test_kafka_io.py b/tests/unit/test_kafka_io.py index c5738eee..425ab396 100644 --- a/tests/unit/test_kafka_io.py +++ b/tests/unit/test_kafka_io.py @@ -1,377 +1,372 @@ -"""Unit tests for Kafka message data reader/writer.""" +"""Unit tests for the Apache Kafka message data implementations.""" from __future__ import annotations from collections import deque -import importlib.machinery -import json -import sys import typing as _t -from unittest.mock import AsyncMock, MagicMock - +from unittest.mock import AsyncMock, MagicMock, call + +from aiokafka.errors import ( + KafkaConnectionError, + RecordTooLargeError, + TopicAuthorizationFailedError, +) +from aiokafka.structs import OffsetAndMetadata, TopicPartition import pytest -from plugboard.exceptions import NoMoreDataException - - -# --------------------------------------------------------------------------- -# Mock the aiokafka module before importing the implementation -# --------------------------------------------------------------------------- - - -def _make_mock_module(name: str) -> MagicMock: - """Creates a mock module with __spec__ set for find_spec compatibility.""" - mock = MagicMock() - mock.__spec__ = importlib.machinery.ModuleSpec(name, None) - return mock - - -_mock_aiokafka = _make_mock_module("aiokafka") -_mock_consumer_class = MagicMock() -_mock_producer_class = MagicMock() -_mock_aiokafka.AIOKafkaConsumer = _mock_consumer_class -_mock_aiokafka.AIOKafkaProducer = _mock_producer_class - -sys.modules.setdefault("aiokafka", _mock_aiokafka) - -from plugboard.library.kafka_io import KafkaDataReader, KafkaDataWriter # noqa: E402 +from plugboard.exceptions import ( + MessageBrokerConnectionError, + MessageBrokerPermanentError, + MessageBrokerTransientError, +) +from plugboard.library import kafka_io +from plugboard.library.kafka_io import KafkaDataReader, KafkaDataWriter +from plugboard.utils.settings import Settings +from tests import conftest + + +def consumer_record(topic: str, partition: int, offset: int, value: bytes) -> MagicMock: + """Builds a Kafka `ConsumerRecord` double.""" + record = MagicMock(name="consumer_record") + record.topic = topic + record.partition = partition + record.offset = offset + record.value = value + return record -# --------------------------------------------------------------------------- -# Helpers -# --------------------------------------------------------------------------- +@pytest.fixture +def consumer(monkeypatch: pytest.MonkeyPatch) -> MagicMock: + """Replaces `AIOKafkaConsumer` with a double that records its construction.""" + client = MagicMock(name="consumer") + client.start = AsyncMock() + client.stop = AsyncMock() + client.commit = AsyncMock() + client.getmany = AsyncMock(return_value={}) + factory = MagicMock(name="AIOKafkaConsumer", return_value=client) + monkeypatch.setattr(kafka_io, "AIOKafkaConsumer", factory) + client.factory = factory + return client + + +@pytest.fixture +def producer(monkeypatch: pytest.MonkeyPatch) -> MagicMock: + """Replaces `AIOKafkaProducer` with a double that records its construction.""" + client = MagicMock(name="producer") + client.start = AsyncMock() + client.stop = AsyncMock() + client.send = AsyncMock() + client.flush = AsyncMock() + client.send_and_wait = AsyncMock() + factory = MagicMock(name="AIOKafkaProducer", return_value=client) + monkeypatch.setattr(kafka_io, "AIOKafkaProducer", factory) + client.factory = factory + return client + + +def make_reader(**kwargs: _t.Any) -> KafkaDataReader: + """Builds a reader for the fixed test topic/group.""" + kwargs.setdefault("field_names", ["x", "y"]) + return KafkaDataReader( + name="kafka-reader", + topic="test-topic", + group_id="test-group", + bootstrap_servers="localhost:9092", + **kwargs, + ) -def _make_kafka_record(value: dict[str, _t.Any] | bytes) -> MagicMock: - """Creates a mock Kafka ConsumerRecord.""" - record = MagicMock() - if isinstance(value, dict): - record.value = json.dumps(value).encode("utf-8") - else: - record.value = value - record.topic = "test-topic" - record.partition = 0 - record.offset = 0 - return record +def make_writer(**kwargs: _t.Any) -> KafkaDataWriter: + """Builds a writer for the fixed test topic.""" + kwargs.setdefault("field_names", ["x", "y"]) + return KafkaDataWriter( + name="kafka-writer", + topic="test-topic", + bootstrap_servers="localhost:9092", + **kwargs, + ) # --------------------------------------------------------------------------- -# Tests: KafkaDataReader +# Reader: connection and consuming # --------------------------------------------------------------------------- -@pytest.mark.asyncio -async def test_kafka_reader_connect() -> None: - """Tests that the reader creates and starts a Kafka consumer.""" - mock_consumer = AsyncMock() - _mock_consumer_class.return_value = mock_consumer - mock_consumer.start = AsyncMock() +async def test_reader_starts_consumer_on_topic_group(consumer: MagicMock) -> None: + """Tests the consumer is built for the topic, group and batch size.""" + reader = make_reader(chunk_size=5) + await reader._connect() - reader = KafkaDataReader( - name="test-kafka-reader", - field_names=["x", "y"], - topic="test-topic", - bootstrap_servers="localhost:9092", + consumer.factory.assert_called_once_with( + "test-topic", + bootstrap_servers=["localhost:9092"], group_id="test-group", + auto_offset_reset="earliest", + enable_auto_commit=False, + max_poll_records=5, ) - await reader._connect() - - _mock_consumer_class.assert_called() - mock_consumer.start.assert_called() - assert reader._consumer is mock_consumer + consumer.start.assert_awaited_once_with() -@pytest.mark.asyncio -async def test_kafka_reader_disconnect() -> None: - """Tests that the reader stops the Kafka consumer.""" - mock_consumer = AsyncMock() - _mock_consumer_class.return_value = mock_consumer - mock_consumer.start = AsyncMock() - mock_consumer.stop = AsyncMock() - - reader = KafkaDataReader( - name="test-kafka-reader", - field_names=["x"], - topic="test-topic", - bootstrap_servers="localhost:9092", - group_id="test-group", - ) +async def test_reader_stops_consumer_on_disconnect(consumer: MagicMock) -> None: + """Tests that teardown stops the consumer, releasing its group membership.""" + reader = make_reader() await reader._connect() await reader._disconnect() + consumer.stop.assert_awaited_once_with() + + +async def test_reader_flattens_records_from_partitions(consumer: MagicMock) -> None: + """Tests that records from every assigned partition are returned.""" + consumer.getmany.return_value = { + TopicPartition("test-topic", 0): [consumer_record("test-topic", 0, 1, b"{}")], + TopicPartition("test-topic", 1): [ + consumer_record("test-topic", 1, 7, b"{}"), + consumer_record("test-topic", 1, 8, b"{}"), + ], + } + reader = make_reader() + await reader._connect() - mock_consumer.stop.assert_called() - assert reader._consumer is None + messages = await reader._receive() + consumer.getmany.assert_awaited_once_with(timeout_ms=30_000, max_records=10) + assert [(r.partition, r.offset) for r in messages] == [(0, 1), (1, 7), (1, 8)] -@pytest.mark.asyncio -async def test_kafka_reader_receive() -> None: - """Tests receiving messages from Kafka.""" - mock_consumer = AsyncMock() - _mock_consumer_class.return_value = mock_consumer - mock_consumer.start = AsyncMock() - test_data = [{"x": 1, "y": "a"}, {"x": 2, "y": "b"}] - mock_records = [_make_kafka_record(d) for d in test_data] - tp = MagicMock() - mock_consumer.getmany = AsyncMock(return_value={tp: mock_records}) +async def test_reader_empty_poll_returns_no_messages(consumer: MagicMock) -> None: + """Tests that a timed-out poll reads as 'nothing yet', not as an exhausted topic. - reader = KafkaDataReader( - name="test-kafka-reader", - field_names=["x", "y"], - topic="test-topic", - bootstrap_servers="localhost:9092", - group_id="test-group", - chunk_size=10, - ) + A consumer catching up between messages must keep running, so an empty poll must + not raise `NoMoreDataException`. + """ + consumer.getmany.return_value = {} + reader = make_reader() await reader._connect() - messages = await reader._receive() + assert await reader._receive() == [] - assert len(messages) == 2 - mock_consumer.getmany.assert_called() +async def test_reader_converts_json_values(consumer: MagicMock) -> None: + """Tests that JSON record values become field buffers.""" + reader = make_reader() + await reader._connect() + batch = [ + consumer_record("test-topic", 0, 1, b'{"x": 1, "y": "a"}'), + consumer_record("test-topic", 0, 2, b'{"x": 2, "y": "b"}'), + ] + assert await reader._convert(batch) == {"x": deque([1, 2]), "y": deque(["a", "b"])} -@pytest.mark.asyncio -async def test_kafka_reader_receive_empty() -> None: - """Tests that empty response raises NoMoreDataException.""" - mock_consumer = AsyncMock() - _mock_consumer_class.return_value = mock_consumer - mock_consumer.start = AsyncMock() - mock_consumer.getmany = AsyncMock(return_value={}) - reader = KafkaDataReader( - name="test-kafka-reader", - field_names=["x"], - topic="test-topic", - bootstrap_servers="localhost:9092", - group_id="test-group", - ) +async def test_reader_non_json_value_uses_data_field(consumer: MagicMock) -> None: + """Tests that a raw value is exposed through the `data` field.""" + reader = make_reader(field_names=["data"], parse_json=False) await reader._connect() + converted = await reader._convert([consumer_record("test-topic", 0, 1, b"hello")]) + assert converted == {"data": deque(["hello"])} - with pytest.raises(NoMoreDataException): - await reader._receive() +# --------------------------------------------------------------------------- +# Reader: offset commits +# --------------------------------------------------------------------------- -@pytest.mark.asyncio -async def test_kafka_reader_convert_json() -> None: - """Tests converting JSON Kafka messages to field buffer.""" - reader = KafkaDataReader( - name="test-kafka-reader", - field_names=["x", "y"], - topic="test-topic", - bootstrap_servers="localhost:9092", - group_id="test-group", - parse_json=True, - ) - mock_records = [ - _make_kafka_record({"x": 1, "y": "a"}), - _make_kafka_record({"x": 2, "y": "b"}), +async def test_reader_commits_offsets_of_processed_records(consumer: MagicMock) -> None: + """Tests that an ack commits the offset *after* each processed record. + + Committing the consumer position instead of these offsets would mark records that + were never processed as done, and they would never be redelivered. + """ + reader = make_reader() + await reader._connect() + records = [ + consumer_record("test-topic", 0, 3, b"{}"), + consumer_record("test-topic", 0, 4, b"{}"), ] - result = await reader._convert(mock_records) - assert list(result["x"]) == [1, 2] - assert list(result["y"]) == ["a", "b"] + await reader._ack(records) -@pytest.mark.asyncio -async def test_kafka_reader_convert_raw() -> None: - """Tests converting raw Kafka messages to field buffer.""" - reader = KafkaDataReader( - name="test-kafka-reader", - field_names=["data"], - topic="test-topic", - bootstrap_servers="localhost:9092", - group_id="test-group", - parse_json=False, + consumer.commit.assert_awaited_once_with( + offsets={TopicPartition("test-topic", 0): OffsetAndMetadata(5, "")} ) - mock_records = [_make_kafka_record(b"raw-1"), _make_kafka_record(b"raw-2")] - result = await reader._convert(mock_records) - assert list(result["data"]) == ["raw-1", "raw-2"] +async def test_reader_commits_highest_offset_per_partition(consumer: MagicMock) -> None: + """Tests that out-of-order records commit the furthest offset for each partition.""" + reader = make_reader() + await reader._connect() + records = [ + consumer_record("test-topic", 1, 2, b"{}"), + consumer_record("test-topic", 0, 10, b"{}"), + consumer_record("test-topic", 1, 5, b"{}"), + consumer_record("test-topic", 0, 8, b"{}"), + ] -@pytest.mark.asyncio -async def test_kafka_reader_ack() -> None: - """Tests committing offsets for Kafka messages.""" - mock_consumer = AsyncMock() - _mock_consumer_class.return_value = mock_consumer - mock_consumer.start = AsyncMock() - mock_consumer.commit = AsyncMock() + await reader._ack(records) - reader = KafkaDataReader( - name="test-kafka-reader", - field_names=["x"], - topic="test-topic", - bootstrap_servers="localhost:9092", - group_id="test-group", - ) + committed = consumer.commit.await_args.kwargs["offsets"] + assert committed == { + TopicPartition("test-topic", 0): OffsetAndMetadata(11, ""), + TopicPartition("test-topic", 1): OffsetAndMetadata(6, ""), + } + + +async def test_reader_commits_nothing_without_records(consumer: MagicMock) -> None: + """Tests that an empty batch does not issue a commit.""" + reader = make_reader() await reader._connect() + await reader._ack([]) + consumer.commit.assert_not_awaited() - mock_records = [_make_kafka_record({"x": 1})] - await reader._ack(mock_records) - mock_consumer.commit.assert_called() +async def test_reader_without_connection_raises_connection_error(consumer: MagicMock) -> None: + """Tests that consuming before connecting fails loudly.""" + reader = make_reader() + with pytest.raises(MessageBrokerConnectionError): + await reader._receive() + with pytest.raises(MessageBrokerConnectionError): + await reader._ack([]) -@pytest.mark.asyncio -async def test_kafka_reader_bootstrap_servers_list() -> None: - """Tests that bootstrap_servers can be a list.""" - mock_consumer = AsyncMock() - _mock_consumer_class.return_value = mock_consumer - mock_consumer.start = AsyncMock() +# --------------------------------------------------------------------------- +# Reader: error mapping +# --------------------------------------------------------------------------- - reader = KafkaDataReader( - name="test-kafka-reader", - field_names=["x"], - topic="test-topic", - bootstrap_servers=["host1:9092", "host2:9092"], - group_id="test-group", - ) - await reader._connect() - call_kwargs = _mock_consumer_class.call_args[1] - assert call_kwargs["bootstrap_servers"] == ["host1:9092", "host2:9092"] +@pytest.mark.parametrize( + ("error", "expected"), + [ + (TopicAuthorizationFailedError("denied"), MessageBrokerPermanentError), + (RecordTooLargeError("too big"), MessageBrokerPermanentError), + (KafkaConnectionError("down"), MessageBrokerTransientError), + ], +) +async def test_reader_maps_kafka_errors( + consumer: MagicMock, error: Exception, expected: type[Exception] +) -> None: + """Tests that Kafka failures land on the right broker exception type.""" + consumer.getmany.side_effect = error + reader = make_reader() + await reader._connect() + with pytest.raises(expected): + await reader._receive() -@pytest.mark.asyncio -async def test_kafka_reader_bootstrap_servers_string() -> None: - """Tests that bootstrap_servers string is converted to list.""" - mock_consumer = AsyncMock() - _mock_consumer_class.return_value = mock_consumer - mock_consumer.start = AsyncMock() - reader = KafkaDataReader( - name="test-kafka-reader", - field_names=["x"], - topic="test-topic", - bootstrap_servers="localhost:9092", - group_id="test-group", - ) +async def test_reader_unclassified_error_propagates(consumer: MagicMock) -> None: + """Tests that an unknown error is not rewritten into a broker exception.""" + consumer.getmany.side_effect = ValueError("unexpected") + reader = make_reader() await reader._connect() - call_kwargs = _mock_consumer_class.call_args[1] - assert call_kwargs["bootstrap_servers"] == ["localhost:9092"] + with pytest.raises(ValueError, match="unexpected"): + await reader._receive() # --------------------------------------------------------------------------- -# Tests: KafkaDataWriter +# Writer # --------------------------------------------------------------------------- -@pytest.mark.asyncio -async def test_kafka_writer_connect() -> None: - """Tests that the writer creates and starts a Kafka producer.""" - mock_producer = AsyncMock() - _mock_producer_class.return_value = mock_producer - mock_producer.start = AsyncMock() - - writer = KafkaDataWriter( - name="test-kafka-writer", - field_names=["x"], - topic="test-topic", - bootstrap_servers="localhost:9092", - ) +async def test_writer_starts_producer(producer: MagicMock) -> None: + """Tests that the producer is created and started.""" + writer = make_writer() await writer._connect() + producer.factory.assert_called_once_with(bootstrap_servers=["localhost:9092"]) + producer.start.assert_awaited_once_with() - _mock_producer_class.assert_called() - mock_producer.start.assert_called() - assert writer._producer is mock_producer +async def test_writer_stops_producer_on_disconnect(producer: MagicMock) -> None: + """Tests that teardown stops the producer.""" + writer = make_writer() + await writer._connect() + await writer._disconnect() + producer.stop.assert_awaited_once_with() -@pytest.mark.asyncio -async def test_kafka_writer_disconnect() -> None: - """Tests that the writer stops the Kafka producer.""" - mock_producer = AsyncMock() - _mock_producer_class.return_value = mock_producer - mock_producer.start = AsyncMock() - mock_producer.stop = AsyncMock() - writer = KafkaDataWriter( - name="test-kafka-writer", - field_names=["x"], - topic="test-topic", - bootstrap_servers="localhost:9092", - ) +async def test_writer_sends_batch_then_flushes_once(producer: MagicMock) -> None: + """Tests that records are submitted together and the buffer flushed once. + + Awaiting each send individually (as `send_and_wait` does) costs one round trip per + record and defeats the producer's batching. + """ + writer = make_writer() await writer._connect() - await writer._disconnect() - mock_producer.stop.assert_called() - assert writer._producer is None + await writer._send([b"one", b"two", b"three"]) + assert producer.send.await_args_list == [ + call("test-topic", value=b"one"), + call("test-topic", value=b"two"), + call("test-topic", value=b"three"), + ] + producer.send_and_wait.assert_not_awaited() + producer.flush.assert_awaited_once_with() -@pytest.mark.asyncio -async def test_kafka_writer_send() -> None: - """Tests sending messages to Kafka.""" - mock_producer = AsyncMock() - _mock_producer_class.return_value = mock_producer - mock_producer.start = AsyncMock() - mock_producer.send_and_wait = AsyncMock() - writer = KafkaDataWriter( - name="test-kafka-writer", - field_names=["x"], - topic="test-topic", - bootstrap_servers="localhost:9092", - ) +async def test_writer_converts_fields_to_json_bytes(producer: MagicMock) -> None: + """Tests that buffered fields become one JSON message per record.""" + writer = make_writer() await writer._connect() + converted = await writer._convert({"x": deque([1, 2]), "y": deque(["a", "b"])}) + assert converted == [b'{"x": 1, "y": "a"}', b'{"x": 2, "y": "b"}'] - messages = [b"msg1", b"msg2"] - await writer._send(messages) - assert mock_producer.send_and_wait.call_count == 2 +async def test_writer_non_json_sends_first_field(producer: MagicMock) -> None: + """Tests that a raw writer publishes the first field's value.""" + writer = make_writer(field_names=["x"], parse_json=False) + await writer._connect() + assert await writer._convert({"x": deque(["hello"])}) == [b"hello"] -@pytest.mark.asyncio -async def test_kafka_writer_convert_json() -> None: - """Tests converting field data to JSON messages.""" - writer = KafkaDataWriter( - name="test-kafka-writer", - field_names=["x", "y"], - topic="test-topic", - bootstrap_servers="localhost:9092", - parse_json=True, - ) +async def test_writer_send_failure_maps_error(producer: MagicMock) -> None: + """Tests that a failed send surfaces as a broker exception.""" + producer.send.side_effect = KafkaConnectionError("down") + writer = make_writer() + await writer._connect() - data = {"x": deque([1, 2]), "y": deque(["a", "b"])} - messages = await writer._convert(data) + with pytest.raises(MessageBrokerTransientError): + await writer._send([b"one"]) - assert len(messages) == 2 - assert json.loads(messages[0]) == {"x": 1, "y": "a"} - assert json.loads(messages[1]) == {"x": 2, "y": "b"} +async def test_writer_without_connection_raises_connection_error(producer: MagicMock) -> None: + """Tests that sending before connecting fails loudly.""" + writer = make_writer() + with pytest.raises(MessageBrokerConnectionError): + await writer._send([b"one"]) -@pytest.mark.asyncio -async def test_kafka_writer_convert_raw() -> None: - """Tests converting field data to raw bytes messages.""" - writer = KafkaDataWriter( - name="test-kafka-writer", - field_names=["data"], - topic="test-topic", - bootstrap_servers="localhost:9092", - parse_json=False, - ) - data = {"data": deque([b"raw1", b"raw2"])} - messages = await writer._convert(data) +# --------------------------------------------------------------------------- +# Settings resolution +# --------------------------------------------------------------------------- - assert len(messages) == 2 - assert messages[0] == b"raw1" - assert messages[1] == b"raw2" +async def test_bootstrap_servers_fall_back_to_settings() -> None: + """Tests that `KAFKA_BOOTSTRAP_SERVERS` supplies the brokers when not passed.""" + settings = Settings.model_validate({"kafka": {"bootstrap_servers": "kafka:9092"}}) + with conftest.override_settings(settings): + writer = KafkaDataWriter(name="kafka-writer", topic="test-topic", field_names=["x"]) + assert writer._bootstrap_servers == ["kafka:9092"] -@pytest.mark.asyncio -async def test_kafka_writer_bootstrap_servers_list() -> None: - """Tests that bootstrap_servers can be a list.""" - mock_producer = AsyncMock() - _mock_producer_class.return_value = mock_producer - mock_producer.start = AsyncMock() - writer = KafkaDataWriter( - name="test-kafka-writer", - field_names=["x"], +def test_bootstrap_servers_accepts_a_list() -> None: + """Tests that an explicit list of brokers is kept as given.""" + reader = KafkaDataReader( + name="kafka-reader", topic="test-topic", - bootstrap_servers=["host1:9092", "host2:9092"], + group_id="test-group", + field_names=["x"], + bootstrap_servers=["a:9092", "b:9092"], ) - await writer._connect() - - call_kwargs = _mock_producer_class.call_args[1] - assert call_kwargs["bootstrap_servers"] == ["host1:9092", "host2:9092"] + assert reader._bootstrap_servers == ["a:9092", "b:9092"] + + +def test_missing_bootstrap_servers_names_the_environment_variable() -> None: + """Tests that unresolvable brokers report what to configure.""" + with conftest.override_settings(Settings.model_validate({})): + with pytest.raises(ValueError, match="KAFKA_BOOTSTRAP_SERVERS"): + KafkaDataWriter( + name="kafka-writer", + topic="test-topic", + field_names=["x"], + bootstrap_servers=None, + ) diff --git a/tests/unit/test_message_data_reader.py b/tests/unit/test_message_data_reader.py index 6add4ffe..3ea34ce0 100644 --- a/tests/unit/test_message_data_reader.py +++ b/tests/unit/test_message_data_reader.py @@ -3,82 +3,77 @@ from __future__ import annotations from collections import deque +from types import SimpleNamespace import typing as _t import pytest -from plugboard.exceptions import IOStreamClosedError, NoMoreDataException +from plugboard.exceptions import ( + IOStreamClosedError, + MessageBrokerPermanentError, + MessageBrokerTransientError, + NoMoreDataException, +) from plugboard.library.message_reader import MessageDataReader +from plugboard.utils import retry as retry_utils +from plugboard.utils.retry import RetryPolicy -# --------------------------------------------------------------------------- -# Mock implementation -# --------------------------------------------------------------------------- +class ScriptedReader(MessageDataReader): + """A `MessageDataReader` driven by a script of poll outcomes. + Each `_receive` call consumes the next entry of `script`: a list of records is + delivered, and an exception is raised. Once the script runs out, polls return no + messages - which is how a real idle broker behaves, and must not end the stream. + """ -class MockMessageDataReader(MessageDataReader): - """Mock `MessageDataReader` for testing the base class logic.""" + def __init__(self, *args: _t.Any, script: list[_t.Any], **kwargs: _t.Any) -> None: + """Instantiates the reader. - def __init__( - self, - *args: _t.Any, - messages: list[dict[str, _t.Any]], - fail_on_connect: bool = False, - fail_on_receive: int | None = None, - **kwargs: _t.Any, - ) -> None: + Args: + *args: Positional arguments for + [`MessageDataReader`][plugboard.library.MessageDataReader]. + script: The per-poll outcomes to play back. + **kwargs: Keyword arguments for + [`MessageDataReader`][plugboard.library.MessageDataReader]. + """ super().__init__(*args, **kwargs) - self._messages = messages - self._idx = 0 - self._connected = False - self._disconnected = False - self._acknowledged: list[list[dict[str, _t.Any]]] = [] - self._fail_on_connect = fail_on_connect - self._fail_on_receive = fail_on_receive - self._receive_call_count = 0 - self._connect_call_count = 0 - self._disconnect_call_count = 0 + self.script = script + self.polls = 0 + self.connects = 0 + self.disconnects = 0 + self.acked: list[list[_t.Any]] = [] async def _connect(self) -> None: - self._connect_call_count += 1 - if self._fail_on_connect and self._connect_call_count <= 1: - raise ConnectionError("Simulated connection failure") - self._connected = True + self.connects += 1 async def _disconnect(self) -> None: - self._disconnect_call_count += 1 - self._connected = False - self._disconnected = True + self.disconnects += 1 async def _receive(self) -> list[_t.Any]: - self._receive_call_count += 1 - if self._fail_on_receive is not None and self._receive_call_count == self._fail_on_receive: - raise ConnectionError("Simulated receive failure") - if self._chunk_size: - chunk = self._messages[self._idx : self._idx + self._chunk_size] - else: - chunk = self._messages[self._idx :] - self._idx += len(chunk) - if not chunk and self._idx >= len(self._messages): - raise NoMoreDataException - return chunk - - async def _convert(self, messages: list[_t.Any]) -> dict[str, deque]: - converted: dict[str, deque] = {field: deque() for field in self.io.outputs} - for msg in messages: - for field in self.io.outputs: - converted[field].append(msg.get(field)) - return converted + outcome = self.script[self.polls] if self.polls < len(self.script) else [] + self.polls += 1 + if isinstance(outcome, Exception): + raise outcome + return _t.cast(list[_t.Any], outcome) + + async def _convert(self, data: list[_t.Any]) -> dict[str, deque]: + return {field: deque(record.get(field) for record in data) for field in self.io.outputs} async def _ack(self, messages: list[_t.Any]) -> None: - self._acknowledged.append(messages) + self.acked.append(list(messages)) -# --------------------------------------------------------------------------- -# Test data -# --------------------------------------------------------------------------- +class UnreconnectableReader(ScriptedReader): + """A reader whose first connect succeeds and every reconnect afterwards fails.""" -TEST_MESSAGES = [ + async def _connect(self) -> None: + self.connects += 1 + if self.connects > 1: + raise OSError("broker unreachable") + + +MESSAGES = [ {"x": 1, "y": "a"}, {"x": 2, "y": "b"}, {"x": 3, "y": "c"}, @@ -87,280 +82,316 @@ async def _ack(self, messages: list[_t.Any]) -> None: ] +def make_reader(script: list[_t.Any], *, max_retries: int = 3, **kwargs: _t.Any) -> ScriptedReader: + """Builds a `ScriptedReader` over the `x`/`y` test fields.""" + kwargs.setdefault("field_names", ["x", "y"]) + kwargs.setdefault("topic", "test-topic") + kwargs.setdefault("idle_poll_delay", 0.0) + kwargs.setdefault("retry_policy", RetryPolicy(max_retries=max_retries)) + return ScriptedReader(name="test-reader", script=script, **kwargs) + + +async def drain(reader: ScriptedReader, limit: int = 100) -> list[dict[str, _t.Any]]: + """Steps a reader until its stream closes, collecting the records it published.""" + records: list[dict[str, _t.Any]] = [] + for _ in range(limit): + try: + await reader.step() + except IOStreamClosedError: + return records + records.append({"x": reader.x, "y": reader.y}) # type: ignore[attr-defined] + raise AssertionError(f"Reader did not finish within {limit} steps") + + +@pytest.fixture +def backoff_delays(monkeypatch: pytest.MonkeyPatch) -> _t.Iterator[list[float]]: + """Captures the non-zero backoff delays applied by the retry helper. + + Only the retry module's `asyncio` reference is replaced: patching `asyncio.sleep` + globally would also break the event loop teardown that waits for pending tasks. + """ + delays: list[float] = [] + + async def _fake_sleep(delay: float, *args: _t.Any) -> None: + if delay: # Idle poll pacing is not a backoff delay. + delays.append(delay) + + monkeypatch.setattr(retry_utils, "asyncio", _t.cast(_t.Any, SimpleNamespace(sleep=_fake_sleep))) + yield delays + + # --------------------------------------------------------------------------- -# Tests: Basic lifecycle +# Lifecycle # --------------------------------------------------------------------------- -@pytest.mark.asyncio -async def test_message_data_reader_init() -> None: - """Tests that `init` connects to the broker and pre-fetches data.""" - reader = MockMessageDataReader( - name="test-reader", - field_names=["x", "y"], - topic="test-topic", - messages=TEST_MESSAGES, - ) +async def test_init_connects_and_starts_prefetch() -> None: + """Tests that `init` connects and pre-fetches without waiting for data.""" + reader = make_reader([MESSAGES[:2]]) await reader.init() - assert reader._connected is True - assert reader._connect_call_count == 1 - # First batch should be pre-fetched - assert reader._receive_call_count == 1 + assert reader.connects == 1 + prefetch = reader._task + assert prefetch is not None + # The pre-fetch runs as a task, so `init` never blocks on the broker. + assert await prefetch == MESSAGES[:2] + assert reader.polls == 1 await reader.destroy() -@pytest.mark.asyncio -async def test_message_data_reader_destroy() -> None: - """Tests that `destroy` disconnects from the broker.""" - reader = MockMessageDataReader( - name="test-reader", - field_names=["x", "y"], - topic="test-topic", - messages=TEST_MESSAGES, - ) +async def test_destroy_disconnects() -> None: + """Tests that `destroy` stops the pre-fetch and disconnects.""" + reader = make_reader([MESSAGES]) await reader.init() + await reader.step() await reader.destroy() - assert reader._disconnected is True - assert reader._disconnect_call_count == 1 + assert reader.disconnects == 1 + assert reader._task is None -@pytest.mark.asyncio -async def test_message_data_reader_step() -> None: - """Tests that `step` reads one record at a time.""" - reader = MockMessageDataReader( - name="test-reader", - field_names=["x", "y"], - topic="test-topic", - messages=TEST_MESSAGES, - ) +async def test_topic_is_retained() -> None: + """Tests that the topic is kept for broker identification and logging.""" + reader = make_reader([], topic="my-topic") + assert reader._topic == "my-topic" + await reader.destroy() + + +# --------------------------------------------------------------------------- +# Reading records +# --------------------------------------------------------------------------- + + +async def test_step_publishes_records_in_order() -> None: + """Tests that records are published one per step, in broker order.""" + reader = make_reader([MESSAGES, NoMoreDataException()]) await reader.init() + assert await drain(reader) == MESSAGES + await reader.destroy() - results: list[dict[str, _t.Any]] = [] - while True: - try: - await reader.step() - results.append({"x": reader.x, "y": reader.y}) # type: ignore[attr-defined] - except IOStreamClosedError: - break - assert results == TEST_MESSAGES +@pytest.mark.parametrize("chunk_size", [1, 2, 3, 5]) +async def test_reads_across_batch_boundaries(chunk_size: int) -> None: + """Tests that splitting messages over several polls loses nothing.""" + batches = [MESSAGES[i : i + chunk_size] for i in range(0, len(MESSAGES), chunk_size)] + reader = make_reader([*batches, NoMoreDataException()], chunk_size=chunk_size) + await reader.init() + assert await drain(reader) == MESSAGES await reader.destroy() -@pytest.mark.asyncio -async def test_message_data_reader_acknowledgment() -> None: - """Tests that messages are acknowledged after processing.""" - reader = MockMessageDataReader( +async def test_missing_fields_read_as_none() -> None: + """Tests that a field absent from a message reads as `None`.""" + reader = make_reader([[{"x": 1}], NoMoreDataException()]) + await reader.init() + await reader.step() + assert reader.x == 1 # type: ignore[attr-defined] + assert reader.y is None # type: ignore[attr-defined] + await reader.destroy() + + +async def test_single_field_reader() -> None: + """Tests a reader with a single output field.""" + reader = ScriptedReader( name="test-reader", - field_names=["x", "y"], + field_names=["value"], topic="test-topic", - messages=TEST_MESSAGES, + idle_poll_delay=0.0, + script=[[{"value": 0}, {"value": 1}], NoMoreDataException()], ) await reader.init() - - # Step through first message - await reader.step() - # First batch should be acknowledged - assert len(reader._acknowledged) >= 1 - + values = [] + for _ in range(2): + await reader.step() + values.append(reader.value) # type: ignore[attr-defined] + assert values == [0, 1] + with pytest.raises(IOStreamClosedError): + await reader.step() await reader.destroy() # --------------------------------------------------------------------------- -# Tests: Chunked reading +# Empty polls are not the end of the stream # --------------------------------------------------------------------------- -@pytest.mark.asyncio -@pytest.mark.parametrize("chunk_size", [1, 2, 3, 5, 10]) -async def test_message_data_reader_chunked(chunk_size: int) -> None: - """Tests reading with various chunk sizes.""" - reader = MockMessageDataReader( - name="test-reader", - field_names=["x", "y"], - topic="test-topic", - chunk_size=chunk_size, - messages=TEST_MESSAGES, - ) +async def test_empty_poll_waits_for_late_messages() -> None: + """Tests that idle polls do not end the stream, so late messages are still read. + + A reader that starts before its producer must keep running: two empty polls are + followed by a batch, and `step()` returns that record rather than raising. + """ + reader = make_reader([[], [], [{"x": 9, "y": "z"}], NoMoreDataException()]) await reader.init() - results: list[dict[str, _t.Any]] = [] - while True: - try: - await reader.step() - results.append({"x": reader.x, "y": reader.y}) # type: ignore[attr-defined] - except IOStreamClosedError: - break + await reader.step() - assert results == TEST_MESSAGES + assert (reader.x, reader.y) == (9, "z") # type: ignore[attr-defined] + assert reader.polls >= 3 + assert not reader.io.is_closed await reader.destroy() -@pytest.mark.asyncio -async def test_message_data_reader_no_chunk_size() -> None: - """Tests reading without chunk size (all messages at once).""" - reader = MockMessageDataReader( - name="test-reader", - field_names=["x", "y"], - topic="test-topic", - chunk_size=None, - messages=TEST_MESSAGES, - ) +async def test_idle_polls_do_not_publish_or_ack() -> None: + """Tests that empty polls neither emit records nor acknowledge anything.""" + reader = make_reader([[], [], MESSAGES[:1], NoMoreDataException()]) await reader.init() + records = await drain(reader) + assert records == MESSAGES[:1] + assert reader.polls == 4 + assert len(reader.acked) == 1 + await reader.destroy() - results: list[dict[str, _t.Any]] = [] - while True: - try: - await reader.step() - results.append({"x": reader.x, "y": reader.y}) # type: ignore[attr-defined] - except IOStreamClosedError: - break - assert results == TEST_MESSAGES +async def test_no_dangling_task_when_source_exhausted() -> None: + """Tests that exhaustion closes the stream and leaves no task behind.""" + reader = make_reader([NoMoreDataException()]) + await reader.init() + with pytest.raises(IOStreamClosedError): + await reader.step() + assert reader._task is None await reader.destroy() # --------------------------------------------------------------------------- -# Tests: Empty messages +# Acknowledgment timing # --------------------------------------------------------------------------- -@pytest.mark.asyncio -async def test_message_data_reader_empty_messages() -> None: - """Tests that reader handles empty message source correctly.""" - reader = MockMessageDataReader( - name="test-reader", - field_names=["x", "y"], - topic="test-topic", - messages=[], - ) +async def test_batch_acked_only_after_its_last_record_is_consumed() -> None: + """Tests that a batch is acknowledged once, after every record was published. + + Acknowledging earlier would lose the remaining messages on a crash, so this is the + at-least-once guarantee for the batch. + """ + reader = make_reader([MESSAGES[:3], NoMoreDataException()], chunk_size=3) await reader.init() - with pytest.raises(IOStreamClosedError): - await reader.step() + await reader.step() + await reader.step() + assert reader.acked == [] # Two of three records consumed: nothing acked yet. + + await reader.step() # The third record drains the batch. + assert [[msg["x"] for msg in batch] for batch in reader.acked] == [[1, 2, 3]] + await reader.destroy() + +async def test_every_consumed_record_is_acknowledged_once() -> None: + """Tests that the acked messages are exactly the messages that were published.""" + batches = [MESSAGES[:2], MESSAGES[2:]] + reader = make_reader([*batches, NoMoreDataException()]) + await reader.init() + await drain(reader) + + acked = [msg["x"] for batch in reader.acked for msg in batch] + assert acked == [msg["x"] for msg in MESSAGES] + await reader.destroy() + + +async def test_unconsumed_messages_are_not_acknowledged() -> None: + """Tests that destroying mid-batch leaves the unread messages for redelivery.""" + reader = make_reader([MESSAGES]) + await reader.init() + await reader.step() await reader.destroy() + assert reader.acked == [] # --------------------------------------------------------------------------- -# Tests: Retry logic +# Retry and reconnection # --------------------------------------------------------------------------- -@pytest.mark.asyncio -async def test_message_data_reader_retry_on_receive_failure() -> None: - """Tests that reader retries on transient receive failures.""" - reader = MockMessageDataReader( - name="test-reader", - field_names=["x", "y"], - topic="test-topic", - messages=TEST_MESSAGES[:2], - fail_on_receive=2, # Fail on the second receive call +async def test_transient_receive_failure_reconnects_and_recovers( + backoff_delays: list[float], +) -> None: + """Tests that a transient failure is retried after reconnecting.""" + reader = make_reader( + [MessageBrokerTransientError("blip"), MESSAGES[:2], NoMoreDataException()], max_retries=3, - retry_base_delay=0.01, # Fast retries for testing ) await reader.init() - - results: list[dict[str, _t.Any]] = [] - while True: - try: - await reader.step() - results.append({"x": reader.x, "y": reader.y}) # type: ignore[attr-defined] - except IOStreamClosedError: - break - - assert results == TEST_MESSAGES[:2] - # Should have attempted reconnection - assert reader._connect_call_count >= 2 + assert await drain(reader) == MESSAGES[:2] + assert reader.connects == 2 # Initial connect plus one reconnect. + assert reader.disconnects == 1 + assert backoff_delays == [1.0] # Base delay on the first retry. await reader.destroy() -@pytest.mark.asyncio -async def test_message_data_reader_retry_exhausted() -> None: - """Tests that reader raises after all retries are exhausted.""" - reader = MockMessageDataReader( - name="test-reader", - field_names=["x", "y"], - topic="test-topic", - messages=TEST_MESSAGES[:1], - fail_on_receive=2, # Always fail on receive - max_retries=2, - retry_base_delay=0.01, - ) +async def test_retry_exhaustion_raises_last_broker_error(backoff_delays: list[float]) -> None: + """Tests that retries are bounded and the broker's own error is raised.""" + failures = [ConnectionError(f"failure {index}") for index in range(5)] + reader = make_reader(failures, max_retries=2) await reader.init() - # First step should succeed (from pre-fetched data) - await reader.step() - - # Second step should fail after retries exhausted - with pytest.raises((IOStreamClosedError, ConnectionError)): + with pytest.raises(ConnectionError, match="failure 2"): await reader.step() + assert reader.polls == 3 # Initial attempt plus two retries. + assert reader.connects == 3 # Initial connect plus one reconnect per retry. + assert reader.disconnects == 2 + assert backoff_delays == [1.0, 2.0] await reader.destroy() -# --------------------------------------------------------------------------- -# Tests: Connection failure on init -# --------------------------------------------------------------------------- - +async def test_backoff_delay_is_capped(backoff_delays: list[float]) -> None: + """Tests that exponential backoff stops growing at `max_delay`.""" + policy = RetryPolicy(max_retries=5, base_delay=1.0, max_delay=3.0) + reader = make_reader([ConnectionError("down")] * 6, retry_policy=policy) + await reader.init() -@pytest.mark.asyncio -async def test_message_data_reader_connection_failure_on_init() -> None: - """Tests that init raises on connection failure.""" - reader = MockMessageDataReader( - name="test-reader", - field_names=["x", "y"], - topic="test-topic", - messages=TEST_MESSAGES, - fail_on_connect=True, - ) with pytest.raises(ConnectionError): - await reader.init() + await reader.step() + assert backoff_delays == [1.0, 2.0, 3.0, 3.0, 3.0] + await reader.destroy() -# --------------------------------------------------------------------------- -# Tests: Single field -# --------------------------------------------------------------------------- +@pytest.mark.parametrize( + ("attempt", "expected"), + [(0, 1.0), (1, 2.0), (2, 4.0), (10, 60.0)], +) +def test_retry_policy_delay_for(attempt: int, expected: float) -> None: + """Tests that the backoff formula doubles per attempt and respects the cap.""" + assert RetryPolicy().delay_for(attempt) == expected -@pytest.mark.asyncio -async def test_message_data_reader_single_field() -> None: - """Tests reading with a single output field.""" - messages = [{"value": i} for i in range(3)] - reader = MockMessageDataReader( - name="test-reader", - field_names=["value"], - topic="test-topic", - messages=messages, + +async def test_permanent_error_is_not_retried() -> None: + """Tests that a permanent broker error propagates after a single attempt.""" + reader = make_reader( + [MessageBrokerPermanentError("access denied"), MESSAGES], + max_retries=3, ) await reader.init() - results: list[_t.Any] = [] - while True: - try: - await reader.step() - results.append(reader.value) # type: ignore[attr-defined] - except IOStreamClosedError: - break + with pytest.raises(MessageBrokerPermanentError, match="access denied"): + await reader.step() - assert results == [0, 1, 2] + assert reader.polls == 1 + assert reader.connects == 1 # No reconnect attempted. await reader.destroy() -# --------------------------------------------------------------------------- -# Tests: Topic attribute -# --------------------------------------------------------------------------- - +async def test_reconnect_failure_does_not_abort_retries( + backoff_delays: list[float], +) -> None: + """Tests that a failing reconnect costs an attempt instead of ending the retrying. -@pytest.mark.asyncio -async def test_message_data_reader_topic() -> None: - """Tests that the topic is stored correctly.""" - reader = MockMessageDataReader( + During a network partition the read and the reconnect can both fail; the reader must + still use up its retries and then report the broker error. + """ + reader = UnreconnectableReader( name="test-reader", - field_names=["x"], - topic="my-topic", - messages=[{"x": 1}], + field_names=["x", "y"], + topic="test-topic", + idle_poll_delay=0.0, + script=[ConnectionError("read failed")] * 3, + retry_policy=RetryPolicy(max_retries=2), ) - assert reader._topic == "my-topic" await reader.init() + + with pytest.raises(ConnectionError, match="read failed"): + await reader.step() + + assert reader.polls == 3 # All three attempts ran, despite the reconnect failures. + assert reader.connects == 3 # Initial connect plus one per retry. + assert backoff_delays == [1.0, 2.0] await reader.destroy() diff --git a/tests/unit/test_message_data_writer.py b/tests/unit/test_message_data_writer.py index 5412617a..d072ab38 100644 --- a/tests/unit/test_message_data_writer.py +++ b/tests/unit/test_message_data_writer.py @@ -1,4 +1,4 @@ -"""Unit tests for the `MessageDataWriter` base class.""" +"""Unit tests for the `MessageDataWriter` base class and its shared encoding helpers.""" from __future__ import annotations @@ -8,357 +8,335 @@ import pytest from plugboard.connector import AsyncioConnector -from plugboard.library.message_writer import MessageDataWriter +from plugboard.exceptions import ( + MessageBrokerPermanentError, + MessageBrokerTransientError, +) +from plugboard.library.message_writer import ( + MessageDataWriter, + encode_records, + encode_records_bytes, + iter_records, +) from plugboard.schemas import ConnectorSpec +from plugboard.utils.retry import RetryPolicy -# --------------------------------------------------------------------------- -# Mock implementation -# --------------------------------------------------------------------------- +class RecordingWriter(MessageDataWriter): + """A `MessageDataWriter` that records what it sends and can script send failures. + Each `_send` call consumes the next entry of `script`: an exception is raised, and + anything else is recorded as a delivered batch. When the script runs out, sends + succeed and are recorded. + """ -class MockMessageDataWriter(MessageDataWriter): - """Mock `MessageDataWriter` for testing the base class logic.""" + def __init__(self, *args: _t.Any, script: list[_t.Any] | None = None, **kwargs: _t.Any) -> None: + """Instantiates the writer. - def __init__( - self, - *args: _t.Any, - fail_on_connect: bool = False, - fail_on_send: int | None = None, - **kwargs: _t.Any, - ) -> None: + Args: + *args: Positional arguments for + [`MessageDataWriter`][plugboard.library.MessageDataWriter]. + script: The per-send outcomes to play back. + **kwargs: Keyword arguments for + [`MessageDataWriter`][plugboard.library.MessageDataWriter]. + """ super().__init__(*args, **kwargs) - self._connected = False - self._disconnected = False - self._sent_messages: list[list[_t.Any]] = [] - self._fail_on_connect = fail_on_connect - self._fail_on_send = fail_on_send - self._send_call_count = 0 - self._connect_call_count = 0 - self._disconnect_call_count = 0 + self.script = script or [] + self.sends: list[list[_t.Any]] = [] + self.attempts = 0 + self.connects = 0 + self.disconnects = 0 async def _connect(self) -> None: - self._connect_call_count += 1 - if self._fail_on_connect and self._connect_call_count <= 1: - raise ConnectionError("Simulated connection failure") - self._connected = True + self.connects += 1 async def _disconnect(self) -> None: - self._disconnect_call_count += 1 - self._connected = False - self._disconnected = True + self.disconnects += 1 + + async def _convert(self, data: dict[str, deque]) -> list[_t.Any]: + return list(iter_records(data)) async def _send(self, messages: list[_t.Any]) -> None: - self._send_call_count += 1 - if self._fail_on_send is not None and self._send_call_count == self._fail_on_send: - raise ConnectionError("Simulated send failure") - self._sent_messages.append(messages) + outcome = self.script[self.attempts] if self.attempts < len(self.script) else None + self.attempts += 1 + if isinstance(outcome, Exception): + raise outcome + self.sends.append(list(messages)) - async def _convert(self, data: dict[str, deque]) -> list[_t.Any]: - completed_rows = min(len(d) for d in data.values()) if data else 0 - messages: list[dict[str, _t.Any]] = [] - for i in range(completed_rows): - record = {field: data[field][i] for field in data} - messages.append(record) - return messages +class UnreconnectableWriter(RecordingWriter): + """A writer whose first connect succeeds and every reconnect afterwards fails.""" -# --------------------------------------------------------------------------- -# Test helpers -# --------------------------------------------------------------------------- + async def _connect(self) -> None: + self.connects += 1 + if self.connects > 1: + raise OSError("broker unreachable") + + +def make_writer( + field_names: list[str], + *, + script: list[_t.Any] | None = None, + max_retries: int = 3, + writer_cls: type[RecordingWriter] = RecordingWriter, + **kwargs: _t.Any, +) -> RecordingWriter: + """Builds a writer over the given fields.""" + return writer_cls( + name="test-writer", + field_names=field_names, + topic="test-topic", + script=script, + retry_policy=RetryPolicy(max_retries=max_retries), + **kwargs, + ) -async def _setup_writer_with_channels( - writer: MockMessageDataWriter, field_names: list[str] -) -> dict[str, AsyncioConnector]: - """Sets up a writer with connected asyncio channels for sending data.""" +async def connect_writer(writer: RecordingWriter) -> dict[str, AsyncioConnector]: + """Connects input channels to a writer and returns them by field name.""" connectors = { field: AsyncioConnector( - spec=ConnectorSpec(source="none.none", target=f"{writer.name}.{field}"), + spec=ConnectorSpec(source="none.none", target=f"{writer.name}.{field}") ) - for field in field_names + for field in writer.io.inputs } await writer.io.connect(list(connectors.values())) return connectors +async def write_records( + writer: RecordingWriter, + channels: dict[str, AsyncioConnector], + records: list[dict[str, _t.Any]], +) -> None: + """Pushes records into the writer one at a time, stepping after each.""" + senders = {field: await channel.connect_send() for field, channel in channels.items()} + for record in records: + for field, value in record.items(): + await senders[field].send(value) + await writer.step() + + +def delivered(writer: RecordingWriter) -> list[_t.Any]: + """Flattens everything the writer sent.""" + return [record for batch in writer.sends for record in batch] + + +RECORDS = [{"x": 1, "y": "a"}, {"x": 2, "y": "b"}, {"x": 3, "y": "c"}] +SINGLE = [{"x": 1}] # For the single-field writers used by the retry tests. + + # --------------------------------------------------------------------------- -# Tests: Basic lifecycle +# Lifecycle # --------------------------------------------------------------------------- -@pytest.mark.asyncio -async def test_message_data_writer_init() -> None: - """Tests that `init` connects to the broker.""" - writer = MockMessageDataWriter( - name="test-writer", - field_names=["x", "y"], - topic="test-topic", - ) +async def test_init_connects_to_broker() -> None: + """Tests that `init` opens the broker connection.""" + writer = make_writer(["x"]) await writer.init() - assert writer._connected is True - assert writer._connect_call_count == 1 + assert writer.connects == 1 await writer.destroy() -@pytest.mark.asyncio -async def test_message_data_writer_destroy() -> None: - """Tests that `destroy` disconnects from the broker.""" - writer = MockMessageDataWriter( - name="test-writer", - field_names=["x", "y"], - topic="test-topic", - ) +async def test_destroy_disconnects() -> None: + """Tests that `destroy` closes the broker connection.""" + writer = make_writer(["x"]) + await writer.init() + await writer.destroy() + assert writer.disconnects == 1 + + +async def test_destroy_waits_for_in_flight_send() -> None: + """Tests that a send still running at teardown completes instead of being dropped.""" + writer = make_writer(["x"]) + channels = await connect_writer(writer) await writer.init() + await write_records(writer, channels, [{"x": 1}]) + await writer._save_chunk() # Buffer one batch without waiting for it to land. await writer.destroy() - assert writer._disconnected is True - assert writer._disconnect_call_count == 1 + assert delivered(writer) == [{"x": 1}] +@pytest.mark.parametrize("chunk_size", [1, 2, 3]) # --------------------------------------------------------------------------- -# Tests: Writing data +# Batching # --------------------------------------------------------------------------- -@pytest.mark.asyncio -async def test_message_data_writer_step_and_run() -> None: - """Tests that data is written via step and flushed on run.""" - writer = MockMessageDataWriter( - name="test-writer", - field_names=["x", "y"], - topic="test-topic", - chunk_size=2, - ) - connectors = await _setup_writer_with_channels(writer, ["x", "y"]) +async def test_records_are_sent_in_batches(chunk_size: int) -> None: + """Tests that every record is delivered exactly once, in order.""" + writer = make_writer(["x", "y"], chunk_size=chunk_size) + channels = await connect_writer(writer) await writer.init() - output_channels = {field: await connectors[field].connect_send() for field in ["x", "y"]} - - # Send data - test_data = [(1, "a"), (2, "b"), (3, "c")] - for x_val, y_val in test_data: - await output_channels["x"].send(x_val) - await output_channels["y"].send(y_val) - await writer.step() - - # Close inputs and run to flush + await write_records(writer, channels, RECORDS) await writer.io.close() await writer.run() - # Verify sent messages - all_sent = [msg for batch in writer._sent_messages for msg in batch] - assert len(all_sent) == 3 - assert all_sent[0] == {"x": 1, "y": "a"} - assert all_sent[1] == {"x": 2, "y": "b"} - assert all_sent[2] == {"x": 3, "y": "c"} + assert delivered(writer) == RECORDS + assert all(len(batch) <= chunk_size for batch in writer.sends) - await writer.destroy() - -@pytest.mark.asyncio -async def test_message_data_writer_flush_on_run() -> None: - """Tests that remaining buffered data is flushed on `run`.""" - writer = MockMessageDataWriter( - name="test-writer", - field_names=["x"], - topic="test-topic", - chunk_size=10, # Large chunk size so nothing is sent during step - ) - connectors = await _setup_writer_with_channels(writer, ["x"]) +async def test_run_flushes_remaining_buffer() -> None: + """Tests that data left under the chunk size is still sent when the run ends.""" + writer = make_writer(["x", "y"], chunk_size=10) + channels = await connect_writer(writer) await writer.init() - output_channels = {"x": await connectors["x"].connect_send()} - - # Send data (less than chunk_size) - for i in range(3): - await output_channels["x"].send(i) - await writer.step() + await write_records(writer, channels, RECORDS) + assert writer.sends == [] # Below the chunk size: still buffered. - # Nothing should be sent yet (buffer < chunk_size) - assert len(writer._sent_messages) == 0 - - # Close and run to flush await writer.io.close() await writer.run() - - # Now data should be flushed - all_sent = [msg for batch in writer._sent_messages for msg in batch] - assert len(all_sent) == 3 - assert all_sent == [{"x": 0}, {"x": 1}, {"x": 2}] - - await writer.destroy() + assert delivered(writer) == RECORDS # --------------------------------------------------------------------------- -# Tests: Chunked writing +# Retry and reconnection # --------------------------------------------------------------------------- -@pytest.mark.asyncio -@pytest.mark.parametrize("chunk_size", [1, 2, 3, 5]) -async def test_message_data_writer_chunked(chunk_size: int) -> None: - """Tests writing with various chunk sizes.""" - writer = MockMessageDataWriter( - name="test-writer", - field_names=["x", "y"], - topic="test-topic", - chunk_size=chunk_size, - ) - connectors = await _setup_writer_with_channels(writer, ["x", "y"]) +async def test_transient_send_failure_reconnects_and_recovers() -> None: + """Tests that a transient failure is retried after reconnecting.""" + writer = make_writer(["x"], script=[MessageBrokerTransientError("blip")], max_retries=3) + channels = await connect_writer(writer) await writer.init() - output_channels = {field: await connectors[field].connect_send() for field in ["x", "y"]} - - test_data = [(i, f"val_{i}") for i in range(5)] - for x_val, y_val in test_data: - await output_channels["x"].send(x_val) - await output_channels["y"].send(y_val) - await writer.step() - + await write_records(writer, channels, SINGLE) await writer.io.close() await writer.run() - all_sent = [msg for batch in writer._sent_messages for msg in batch] - assert len(all_sent) == 5 - for i, (x_val, y_val) in enumerate(test_data): - assert all_sent[i] == {"x": x_val, "y": y_val} - - await writer.destroy() + assert delivered(writer) == [{"x": 1}] + assert writer.attempts == 2 # Failed once, then succeeded. + assert writer.connects == 2 # Initial connect plus one reconnect. + assert writer.disconnects == 1 -@pytest.mark.asyncio -async def test_message_data_writer_no_chunk_size() -> None: - """Tests writing without chunk size (flush only on run).""" - writer = MockMessageDataWriter( - name="test-writer", - field_names=["x"], - topic="test-topic", - chunk_size=None, +async def test_permanent_send_failure_is_not_retried() -> None: + """Tests that a permanent broker error surfaces immediately.""" + writer = make_writer( + ["x"], + script=[MessageBrokerPermanentError("topic deleted"), None], + max_retries=3, ) - connectors = await _setup_writer_with_channels(writer, ["x"]) + channels = await connect_writer(writer) await writer.init() + await write_records(writer, channels, SINGLE) - output_channels = {"x": await connectors["x"].connect_send()} + await writer._save_chunk() + send = writer._task + assert send is not None + with pytest.raises(MessageBrokerPermanentError, match="topic deleted"): + await send - for i in range(3): - await output_channels["x"].send(i) - await writer.step() + assert writer.attempts == 1 + assert writer.connects == 1 # No reconnect attempted. + writer._task = None # The failed send is already inspected here. + await writer.destroy() - # Nothing sent yet (no chunk_size trigger) - assert len(writer._sent_messages) == 0 - await writer.io.close() - await writer.run() +async def test_retry_exhaustion_raises_last_send_error() -> None: + """Tests that sends are bounded and the broker's own error is raised.""" + failures = [ConnectionError(f"failure {index}") for index in range(5)] + writer = make_writer(["x"], script=failures, max_retries=2) + channels = await connect_writer(writer) + await writer.init() + await write_records(writer, channels, SINGLE) - all_sent = [msg for batch in writer._sent_messages for msg in batch] - assert len(all_sent) == 3 + await writer._save_chunk() + send = writer._task + assert send is not None + with pytest.raises(ConnectionError, match="failure 2"): + await send + assert writer.attempts == 3 # Initial attempt plus two retries. + assert writer.connects == 3 # Initial connect plus one reconnect per retry. + writer._task = None await writer.destroy() -# --------------------------------------------------------------------------- -# Tests: Retry logic -# --------------------------------------------------------------------------- - - -@pytest.mark.asyncio -async def test_message_data_writer_retry_on_send_failure() -> None: - """Tests that writer retries on transient send failures.""" - writer = MockMessageDataWriter( - name="test-writer", - field_names=["x"], - topic="test-topic", - chunk_size=1, - fail_on_send=1, # Fail on first send - max_retries=3, - retry_base_delay=0.01, +async def test_reconnect_failure_does_not_abort_retries() -> None: + """Tests that a failing reconnect costs an attempt instead of ending the retrying.""" + writer = make_writer( + ["x"], + script=[ConnectionError("send failed")] * 3, + max_retries=2, + writer_cls=UnreconnectableWriter, ) - connectors = await _setup_writer_with_channels(writer, ["x"]) + channels = await connect_writer(writer) await writer.init() + await write_records(writer, channels, SINGLE) - output_channels = {"x": await connectors["x"].connect_send()} - - # Send one item and step (triggers send which fails, then retries) - await output_channels["x"].send(0) - await writer.step() - - # Send another item and step (should succeed now) - await output_channels["x"].send(1) - await writer.step() - - await writer.io.close() - await writer.run() - - # Should have retried and eventually succeeded - all_sent = [msg for batch in writer._sent_messages for msg in batch] - assert len(all_sent) == 2 - # Should have reconnected - assert writer._connect_call_count >= 2 + await writer._save_chunk() + send = writer._task + assert send is not None + with pytest.raises(ConnectionError, match="send failed"): + await send + assert writer.attempts == 3 # Every attempt ran despite the reconnect failures. + assert writer.connects == 3 + writer._task = None await writer.destroy() # --------------------------------------------------------------------------- -# Tests: Connection failure on init +# Shared record encoding helpers # --------------------------------------------------------------------------- -@pytest.mark.asyncio -async def test_message_data_writer_connection_failure_on_init() -> None: - """Tests that init raises on connection failure.""" - writer = MockMessageDataWriter( - name="test-writer", - field_names=["x"], - topic="test-topic", - fail_on_connect=True, - ) - with pytest.raises(ConnectionError): - await writer.init() +def test_iter_records_yields_one_dict_per_row() -> None: + """Tests that a field buffer becomes one record per row, in order.""" + buffer = {"x": deque([1, 2]), "y": deque(["a", "b"])} + assert list(iter_records(buffer)) == [{"x": 1, "y": "a"}, {"x": 2, "y": "b"}] -# --------------------------------------------------------------------------- -# Tests: Topic attribute -# --------------------------------------------------------------------------- +def test_iter_records_stops_at_the_shortest_field() -> None: + """Tests that an incomplete trailing row is not emitted.""" + buffer = {"x": deque([1, 2, 3]), "y": deque(["a"])} + assert list(iter_records(buffer)) == [{"x": 1, "y": "a"}] -@pytest.mark.asyncio -async def test_message_data_writer_topic() -> None: - """Tests that the topic is stored correctly.""" - writer = MockMessageDataWriter( - name="test-writer", - field_names=["x"], - topic="my-topic", - ) - assert writer._topic == "my-topic" +def test_iter_records_handles_empty_buffers() -> None: + """Tests that empty buffers produce no records.""" + assert list(iter_records({})) == [] + assert list(iter_records({"x": deque()})) == [] -# --------------------------------------------------------------------------- -# Tests: Single field -# --------------------------------------------------------------------------- +def test_iter_records_is_linear_in_row_count() -> None: + """Tests that many buffered rows convert in one pass, not one index per row.""" + rows = 20_000 + buffer = {"x": deque(range(rows)), "y": deque(range(rows))} + records = list(iter_records(buffer)) + assert len(records) == rows + assert records[-1] == {"x": rows - 1, "y": rows - 1} -@pytest.mark.asyncio -async def test_message_data_writer_single_field() -> None: - """Tests writing with a single input field.""" - writer = MockMessageDataWriter( - name="test-writer", - field_names=["value"], - topic="test-topic", - chunk_size=3, - ) - connectors = await _setup_writer_with_channels(writer, ["value"]) - await writer.init() +@pytest.mark.parametrize( + ("parse_json", "expected"), + [ + (True, ['{"x": 1, "y": "a"}', '{"x": 2, "y": "b"}']), + (False, ["1", "2"]), + ], +) +def test_encode_records_as_text(parse_json: bool, expected: list[str]) -> None: + """Tests text payloads, JSON-encoded or the raw first field.""" + assert encode_records({"x": deque([1, 2]), "y": deque(["a", "b"])}, parse_json) == expected - output_channels = {"value": await connectors["value"].connect_send()} - for i in range(3): - await output_channels["value"].send(i * 10) - await writer.step() +@pytest.mark.parametrize( + ("parse_json", "expected"), + [ + (True, [b'{"x": 1}', b'{"x": 2}']), + (False, [b"1", b"2"]), + ], +) +def test_encode_records_as_bytes(parse_json: bool, expected: list[bytes]) -> None: + """Tests binary payloads, JSON-encoded or the raw first field.""" + assert encode_records_bytes({"x": deque([1, 2])}, parse_json) == expected - await writer.io.close() - await writer.run() - all_sent = [msg for batch in writer._sent_messages for msg in batch] - assert all_sent == [{"value": 0}, {"value": 10}, {"value": 20}] - - await writer.destroy() +def test_encode_records_bytes_preserves_binary_payloads() -> None: + """Tests that non-UTF-8 bytes survive a raw (non-JSON) payload untouched.""" + payload = bytes([0xFF, 0xFE, 0x00]) + assert encode_records_bytes({"x": deque([payload])}, parse_json=False) == [payload] diff --git a/uv.lock b/uv.lock index 102d9900..93c2b18c 100644 --- a/uv.lock +++ b/uv.lock @@ -50,22 +50,9 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/65/3f/329d0e52f994349ff7449c714c242ad65f14586b0e205ca632ac817fda72/aio_pika-10.0.1-py3-none-any.whl", hash = "sha256:12120a3cf8022d2a8bc5dc89e716512a38bf742c24c5562f54764af27eec7edd", size = 56332, upload-time = "2026-07-09T13:31:33.634Z" }, ] -[[package]] -name = "aioboto3" -version = "14.3.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "aiobotocore", extra = ["boto3"] }, - { name = "aiofiles" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/79/b7/2f0d45cf31f77f8432102d7225d189e6e65cc7a16a32a8ac929eabd719a7/aioboto3-14.3.0.tar.gz", hash = "sha256:1d18f88bb56835c607b62bb6cb907754d717bedde3ddfff6935727cb48a80135", size = 322658, upload-time = "2025-05-07T15:23:59.262Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/0e/b0/f2415f03af890693ba8cb669c67f30b9ffa8b2065ecf91cc92e6782b5aa2/aioboto3-14.3.0-py3-none-any.whl", hash = "sha256:aec5de94e9edc1ffbdd58eead38a37f00ddac59a519db749a910c20b7b81bca7", size = 35697, upload-time = "2025-05-07T15:23:57.539Z" }, -] - [[package]] name = "aiobotocore" -version = "2.22.0" +version = "3.9.2" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "aiohttp" }, @@ -76,14 +63,9 @@ dependencies = [ { name = "python-dateutil" }, { name = "wrapt" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/9c/4c/113c4f5611103bba8e5252805fbee7944f5d9541addba9a96b091c0c4308/aiobotocore-2.22.0.tar.gz", hash = "sha256:11091477266b75c2b5d28421c1f2bc9a87d175d0b8619cb830805e7a113a170b", size = 110322, upload-time = "2025-05-01T16:45:45.484Z" } +sdist = { url = "https://files.pythonhosted.org/packages/a5/95/4f7e0453d5e81bdb1b40ecf09f135bb88e12ce27fdb3efa125022dcca5d7/aiobotocore-3.9.2.tar.gz", hash = "sha256:8e32238c5bd77717ab1ac90077ff51ab141d025528115ccc7cc9b50b2d176e5f", size = 524343, upload-time = "2026-10-01T02:01:59.522Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/00/8e/ffa5840cb7de19ada85bda1fae1ae22671a18992e9373f2e2df9db5389b5/aiobotocore-2.22.0-py3-none-any.whl", hash = "sha256:b4e6306f79df9d81daff1f9d63189a2dbee4b77ce3ab937304834e35eaaeeccf", size = 78930, upload-time = "2025-05-01T16:45:43.508Z" }, -] - -[package.optional-dependencies] -boto3 = [ - { name = "boto3" }, + { url = "https://files.pythonhosted.org/packages/09/ab/c8f982c4dfaab49200e31820a4ffe6275356f7dcaaa1024735d8f8b4e048/aiobotocore-3.9.2-py3-none-any.whl", hash = "sha256:363b4892423b272eb84afbb58000da0be6ccd0cf8a68b98438c957d05b841efc", size = 102571, upload-time = "2026-10-01T02:01:57.883Z" }, ] [[package]] @@ -98,15 +80,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/4e/79/6e45e778c4c3cab39e0937b007b720c15f76c50c6453d153282d0fcc3588/aiofile-3.12.3-py3-none-any.whl", hash = "sha256:5c1bcc9e929c50834608e8cc1a4cc1d7503eb60c15a535b779fd39e2f372c017", size = 22122, upload-time = "2026-08-04T22:59:25.838Z" }, ] -[[package]] -name = "aiofiles" -version = "25.1.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/41/c3/534eac40372d8ee36ef40df62ec129bee4fdb5ad9706e58a29be53b2c970/aiofiles-25.1.0.tar.gz", hash = "sha256:a8d728f0a29de45dc521f18f07297428d56992a742f0cd2701ba86e44d23d5b2", size = 46354, upload-time = "2025-10-09T20:51:04.358Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/bc/8a/340a1555ae33d7354dbca4faa54948d76d89a27ceef032c8c3bc661d003e/aiofiles-25.1.0-py3-none-any.whl", hash = "sha256:abe311e527c862958650f9438e859c1fa7568a141b22abcd015e120e86a85695", size = 14668, upload-time = "2025-10-09T20:51:03.174Z" }, -] - [[package]] name = "aiohappyeyeballs" version = "2.6.1" @@ -118,7 +91,7 @@ wheels = [ [[package]] name = "aiohttp" -version = "3.13.5" +version = "3.14.3" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "aiohappyeyeballs" }, @@ -127,78 +100,93 @@ dependencies = [ { name = "frozenlist" }, { name = "multidict" }, { name = "propcache" }, + { name = "typing-extensions", marker = "python_full_version < '3.13'" }, { name = "yarl" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/77/9a/152096d4808df8e4268befa55fba462f440f14beab85e8ad9bf990516918/aiohttp-3.13.5.tar.gz", hash = "sha256:9d98cc980ecc96be6eb4c1994ce35d28d8b1f5e5208a23b421187d1209dbb7d1", size = 7858271, upload-time = "2026-03-31T22:01:03.343Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/be/6f/353954c29e7dcce7cf00280a02c75f30e133c00793c7a2ed3776d7b2f426/aiohttp-3.13.5-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:023ecba036ddd840b0b19bf195bfae970083fd7024ce1ac22e9bba90464620e9", size = 748876, upload-time = "2026-03-31T21:57:36.319Z" }, - { url = "https://files.pythonhosted.org/packages/f5/1b/428a7c64687b3b2e9cd293186695affc0e1e54a445d0361743b231f11066/aiohttp-3.13.5-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:15c933ad7920b7d9a20de151efcd05a6e38302cbf0e10c9b2acb9a42210a2416", size = 499557, upload-time = "2026-03-31T21:57:38.236Z" }, - { url = "https://files.pythonhosted.org/packages/29/47/7be41556bfbb6917069d6a6634bb7dd5e163ba445b783a90d40f5ac7e3a7/aiohttp-3.13.5-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:ab2899f9fa2f9f741896ebb6fa07c4c883bfa5c7f2ddd8cf2aafa86fa981b2d2", size = 500258, upload-time = "2026-03-31T21:57:39.923Z" }, - { url = "https://files.pythonhosted.org/packages/67/84/c9ecc5828cb0b3695856c07c0a6817a99d51e2473400f705275a2b3d9239/aiohttp-3.13.5-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:a60eaa2d440cd4707696b52e40ed3e2b0f73f65be07fd0ef23b6b539c9c0b0b4", size = 1749199, upload-time = "2026-03-31T21:57:41.938Z" }, - { url = "https://files.pythonhosted.org/packages/f0/d3/3c6d610e66b495657622edb6ae7c7fd31b2e9086b4ec50b47897ad6042a9/aiohttp-3.13.5-cp312-cp312-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:55b3bdd3292283295774ab585160c4004f4f2f203946997f49aac032c84649e9", size = 1721013, upload-time = "2026-03-31T21:57:43.904Z" }, - { url = "https://files.pythonhosted.org/packages/49/a0/24409c12217456df0bae7babe3b014e460b0b38a8e60753d6cb339f6556d/aiohttp-3.13.5-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:c2b2355dc094e5f7d45a7bb262fe7207aa0460b37a0d87027dcf21b5d890e7d5", size = 1781501, upload-time = "2026-03-31T21:57:46.285Z" }, - { url = "https://files.pythonhosted.org/packages/98/9d/b65ec649adc5bccc008b0957a9a9c691070aeac4e41cea18559fef49958b/aiohttp-3.13.5-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:b38765950832f7d728297689ad78f5f2cf79ff82487131c4d26fe6ceecdc5f8e", size = 1878981, upload-time = "2026-03-31T21:57:48.734Z" }, - { url = "https://files.pythonhosted.org/packages/57/d8/8d44036d7eb7b6a8ec4c5494ea0c8c8b94fbc0ed3991c1a7adf230df03bf/aiohttp-3.13.5-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b18f31b80d5a33661e08c89e202edabf1986e9b49c42b4504371daeaa11b47c1", size = 1767934, upload-time = "2026-03-31T21:57:51.171Z" }, - { url = "https://files.pythonhosted.org/packages/31/04/d3f8211f273356f158e3464e9e45484d3fb8c4ce5eb2f6fe9405c3273983/aiohttp-3.13.5-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:33add2463dde55c4f2d9635c6ab33ce154e5ecf322bd26d09af95c5f81cfa286", size = 1566671, upload-time = "2026-03-31T21:57:53.326Z" }, - { url = "https://files.pythonhosted.org/packages/41/db/073e4ebe00b78e2dfcacff734291651729a62953b48933d765dc513bf798/aiohttp-3.13.5-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:327cc432fdf1356fb4fbc6fe833ad4e9f6aacb71a8acaa5f1855e4b25910e4a9", size = 1705219, upload-time = "2026-03-31T21:57:55.385Z" }, - { url = "https://files.pythonhosted.org/packages/48/45/7dfba71a2f9fd97b15c95c06819de7eb38113d2cdb6319669195a7d64270/aiohttp-3.13.5-cp312-cp312-musllinux_1_2_armv7l.whl", hash = "sha256:7c35b0bf0b48a70b4cb4fc5d7bed9b932532728e124874355de1a0af8ec4bc88", size = 1743049, upload-time = "2026-03-31T21:57:57.341Z" }, - { url = "https://files.pythonhosted.org/packages/18/71/901db0061e0f717d226386a7f471bb59b19566f2cae5f0d93874b017271f/aiohttp-3.13.5-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:df23d57718f24badef8656c49743e11a89fd6f5358fa8a7b96e728fda2abf7d3", size = 1749557, upload-time = "2026-03-31T21:57:59.626Z" }, - { url = "https://files.pythonhosted.org/packages/08/d5/41eebd16066e59cd43728fe74bce953d7402f2b4ddfdfef2c0e9f17ca274/aiohttp-3.13.5-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:02e048037a6501a5ec1f6fc9736135aec6eb8a004ce48838cb951c515f32c80b", size = 1558931, upload-time = "2026-03-31T21:58:01.972Z" }, - { url = "https://files.pythonhosted.org/packages/30/e6/4a799798bf05740e66c3a1161079bda7a3dd8e22ca392481d7a7f9af82a6/aiohttp-3.13.5-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:31cebae8b26f8a615d2b546fee45d5ffb76852ae6450e2a03f42c9102260d6fe", size = 1774125, upload-time = "2026-03-31T21:58:04.007Z" }, - { url = "https://files.pythonhosted.org/packages/84/63/7749337c90f92bc2cb18f9560d67aa6258c7060d1397d21529b8004fcf6f/aiohttp-3.13.5-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:888e78eb5ca55a615d285c3c09a7a91b42e9dd6fc699b166ebd5dee87c9ccf14", size = 1732427, upload-time = "2026-03-31T21:58:06.337Z" }, - { url = "https://files.pythonhosted.org/packages/98/de/cf2f44ff98d307e72fb97d5f5bbae3bfcb442f0ea9790c0bf5c5c2331404/aiohttp-3.13.5-cp312-cp312-win32.whl", hash = "sha256:8bd3ec6376e68a41f9f95f5ed170e2fcf22d4eb27a1f8cb361d0508f6e0557f3", size = 433534, upload-time = "2026-03-31T21:58:08.712Z" }, - { url = "https://files.pythonhosted.org/packages/aa/ca/eadf6f9c8fa5e31d40993e3db153fb5ed0b11008ad5d9de98a95045bed84/aiohttp-3.13.5-cp312-cp312-win_amd64.whl", hash = "sha256:110e448e02c729bcebb18c60b9214a87ba33bac4a9fa5e9a5f139938b56c6cb1", size = 460446, upload-time = "2026-03-31T21:58:10.945Z" }, - { url = "https://files.pythonhosted.org/packages/78/e9/d76bf503005709e390122d34e15256b88f7008e246c4bdbe915cd4f1adce/aiohttp-3.13.5-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:a5029cc80718bbd545123cd8fe5d15025eccaaaace5d0eeec6bd556ad6163d61", size = 742930, upload-time = "2026-03-31T21:58:13.155Z" }, - { url = "https://files.pythonhosted.org/packages/57/00/4b7b70223deaebd9bb85984d01a764b0d7bd6526fcdc73cca83bcbe7243e/aiohttp-3.13.5-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:4bb6bf5811620003614076bdc807ef3b5e38244f9d25ca5fe888eaccea2a9832", size = 496927, upload-time = "2026-03-31T21:58:15.073Z" }, - { url = "https://files.pythonhosted.org/packages/9c/f5/0fb20fb49f8efdcdce6cd8127604ad2c503e754a8f139f5e02b01626523f/aiohttp-3.13.5-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:a84792f8631bf5a94e52d9cc881c0b824ab42717165a5579c760b830d9392ac9", size = 497141, upload-time = "2026-03-31T21:58:17.009Z" }, - { url = "https://files.pythonhosted.org/packages/3b/86/b7c870053e36a94e8951b803cb5b909bfbc9b90ca941527f5fcafbf6b0fa/aiohttp-3.13.5-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:57653eac22c6a4c13eb22ecf4d673d64a12f266e72785ab1c8b8e5940d0e8090", size = 1732476, upload-time = "2026-03-31T21:58:18.925Z" }, - { url = "https://files.pythonhosted.org/packages/b5/e5/4e161f84f98d80c03a238671b4136e6530453d65262867d989bbe78244d0/aiohttp-3.13.5-cp313-cp313-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:e5e5f7debc7a57af53fdf5c5009f9391d9f4c12867049d509bf7bb164a6e295b", size = 1706507, upload-time = "2026-03-31T21:58:21.094Z" }, - { url = "https://files.pythonhosted.org/packages/d4/56/ea11a9f01518bd5a2a2fcee869d248c4b8a0cfa0bb13401574fa31adf4d4/aiohttp-3.13.5-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:c719f65bebcdf6716f10e9eff80d27567f7892d8988c06de12bbbd39307c6e3a", size = 1773465, upload-time = "2026-03-31T21:58:23.159Z" }, - { url = "https://files.pythonhosted.org/packages/eb/40/333ca27fb74b0383f17c90570c748f7582501507307350a79d9f9f3c6eb1/aiohttp-3.13.5-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:d97f93fdae594d886c5a866636397e2bcab146fd7a132fd6bb9ce182224452f8", size = 1873523, upload-time = "2026-03-31T21:58:25.59Z" }, - { url = "https://files.pythonhosted.org/packages/f0/d2/e2f77eef1acb7111405433c707dc735e63f67a56e176e72e9e7a2cd3f493/aiohttp-3.13.5-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:3df334e39d4c2f899a914f1dba283c1aadc311790733f705182998c6f7cae665", size = 1754113, upload-time = "2026-03-31T21:58:27.624Z" }, - { url = "https://files.pythonhosted.org/packages/fb/56/3f653d7f53c89669301ec9e42c95233e2a0c0a6dd051269e6e678db4fdb0/aiohttp-3.13.5-cp313-cp313-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:fe6970addfea9e5e081401bcbadf865d2b6da045472f58af08427e108d618540", size = 1562351, upload-time = "2026-03-31T21:58:29.918Z" }, - { url = "https://files.pythonhosted.org/packages/ec/a6/9b3e91eb8ae791cce4ee736da02211c85c6f835f1bdfac0594a8a3b7018c/aiohttp-3.13.5-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:7becdf835feff2f4f335d7477f121af787e3504b48b449ff737afb35869ba7bb", size = 1693205, upload-time = "2026-03-31T21:58:32.214Z" }, - { url = "https://files.pythonhosted.org/packages/98/fc/bfb437a99a2fcebd6b6eaec609571954de2ed424f01c352f4b5504371dd3/aiohttp-3.13.5-cp313-cp313-musllinux_1_2_armv7l.whl", hash = "sha256:676e5651705ad5d8a70aeb8eb6936c436d8ebbd56e63436cb7dd9bb36d2a9a46", size = 1730618, upload-time = "2026-03-31T21:58:34.728Z" }, - { url = "https://files.pythonhosted.org/packages/e4/b6/c8534862126191a034f68153194c389addc285a0f1347d85096d349bbc15/aiohttp-3.13.5-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:9b16c653d38eb1a611cc898c41e76859ca27f119d25b53c12875fd0474ae31a8", size = 1745185, upload-time = "2026-03-31T21:58:36.909Z" }, - { url = "https://files.pythonhosted.org/packages/0b/93/4ca8ee2ef5236e2707e0fd5fecb10ce214aee1ff4ab307af9c558bda3b37/aiohttp-3.13.5-cp313-cp313-musllinux_1_2_riscv64.whl", hash = "sha256:999802d5fa0389f58decd24b537c54aa63c01c3219ce17d1214cbda3c2b22d2d", size = 1557311, upload-time = "2026-03-31T21:58:39.38Z" }, - { url = "https://files.pythonhosted.org/packages/57/ae/76177b15f18c5f5d094f19901d284025db28eccc5ae374d1d254181d33f4/aiohttp-3.13.5-cp313-cp313-musllinux_1_2_s390x.whl", hash = "sha256:ec707059ee75732b1ba130ed5f9580fe10ff75180c812bc267ded039db5128c6", size = 1773147, upload-time = "2026-03-31T21:58:41.476Z" }, - { url = "https://files.pythonhosted.org/packages/01/a4/62f05a0a98d88af59d93b7fcac564e5f18f513cb7471696ac286db970d6a/aiohttp-3.13.5-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:2d6d44a5b48132053c2f6cd5c8cb14bc67e99a63594e336b0f2af81e94d5530c", size = 1730356, upload-time = "2026-03-31T21:58:44.049Z" }, - { url = "https://files.pythonhosted.org/packages/e4/85/fc8601f59dfa8c9523808281f2da571f8b4699685f9809a228adcc90838d/aiohttp-3.13.5-cp313-cp313-win32.whl", hash = "sha256:329f292ed14d38a6c4c435e465f48bebb47479fd676a0411936cc371643225cc", size = 432637, upload-time = "2026-03-31T21:58:46.167Z" }, - { url = "https://files.pythonhosted.org/packages/c0/1b/ac685a8882896acf0f6b31d689e3792199cfe7aba37969fa91da63a7fa27/aiohttp-3.13.5-cp313-cp313-win_amd64.whl", hash = "sha256:69f571de7500e0557801c0b51f4780482c0ec5fe2ac851af5a92cfce1af1cb83", size = 458896, upload-time = "2026-03-31T21:58:48.119Z" }, - { url = "https://files.pythonhosted.org/packages/5d/ce/46572759afc859e867a5bc8ec3487315869013f59281ce61764f76d879de/aiohttp-3.13.5-cp314-cp314-macosx_10_13_universal2.whl", hash = "sha256:eb4639f32fd4a9904ab8fb45bf3383ba71137f3d9d4ba25b3b3f3109977c5b8c", size = 745721, upload-time = "2026-03-31T21:58:50.229Z" }, - { url = "https://files.pythonhosted.org/packages/13/fe/8a2efd7626dbe6049b2ef8ace18ffda8a4dfcbe1bcff3ac30c0c7575c20b/aiohttp-3.13.5-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:7e5dc4311bd5ac493886c63cbf76ab579dbe4641268e7c74e48e774c74b6f2be", size = 497663, upload-time = "2026-03-31T21:58:52.232Z" }, - { url = "https://files.pythonhosted.org/packages/9b/91/cc8cc78a111826c54743d88651e1687008133c37e5ee615fee9b57990fac/aiohttp-3.13.5-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:756c3c304d394977519824449600adaf2be0ccee76d206ee339c5e76b70ded25", size = 499094, upload-time = "2026-03-31T21:58:54.566Z" }, - { url = "https://files.pythonhosted.org/packages/0a/33/a8362cb15cf16a3af7e86ed11962d5cd7d59b449202dc576cdc731310bde/aiohttp-3.13.5-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ecc26751323224cf8186efcf7fbcbc30f4e1d8c7970659daf25ad995e4032a56", size = 1726701, upload-time = "2026-03-31T21:58:56.864Z" }, - { url = "https://files.pythonhosted.org/packages/45/0c/c091ac5c3a17114bd76cbf85d674650969ddf93387876cf67f754204bd77/aiohttp-3.13.5-cp314-cp314-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:10a75acfcf794edf9d8db50e5a7ec5fc818b2a8d3f591ce93bc7b1210df016d2", size = 1683360, upload-time = "2026-03-31T21:58:59.072Z" }, - { url = "https://files.pythonhosted.org/packages/23/73/bcee1c2b79bc275e964d1446c55c54441a461938e70267c86afaae6fba27/aiohttp-3.13.5-cp314-cp314-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:0f7a18f258d124cd678c5fe072fe4432a4d5232b0657fca7c1847f599233c83a", size = 1773023, upload-time = "2026-03-31T21:59:01.776Z" }, - { url = "https://files.pythonhosted.org/packages/c7/ef/720e639df03004fee2d869f771799d8c23046dec47d5b81e396c7cda583a/aiohttp-3.13.5-cp314-cp314-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:df6104c009713d3a89621096f3e3e88cc323fd269dbd7c20afe18535094320be", size = 1853795, upload-time = "2026-03-31T21:59:04.568Z" }, - { url = "https://files.pythonhosted.org/packages/bd/c9/989f4034fb46841208de7aeeac2c6d8300745ab4f28c42f629ba77c2d916/aiohttp-3.13.5-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:241a94f7de7c0c3b616627aaad530fe2cb620084a8b144d3be7b6ecfe95bae3b", size = 1730405, upload-time = "2026-03-31T21:59:07.221Z" }, - { url = "https://files.pythonhosted.org/packages/ce/75/ee1fd286ca7dc599d824b5651dad7b3be7ff8d9a7e7b3fe9820d9180f7db/aiohttp-3.13.5-cp314-cp314-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:c974fb66180e58709b6fc402846f13791240d180b74de81d23913abe48e96d94", size = 1558082, upload-time = "2026-03-31T21:59:09.484Z" }, - { url = "https://files.pythonhosted.org/packages/c3/20/1e9e6650dfc436340116b7aa89ff8cb2bbdf0abc11dfaceaad8f74273a10/aiohttp-3.13.5-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:6e27ea05d184afac78aabbac667450c75e54e35f62238d44463131bd3f96753d", size = 1692346, upload-time = "2026-03-31T21:59:12.068Z" }, - { url = "https://files.pythonhosted.org/packages/d8/40/8ebc6658d48ea630ac7903912fe0dd4e262f0e16825aa4c833c56c9f1f56/aiohttp-3.13.5-cp314-cp314-musllinux_1_2_armv7l.whl", hash = "sha256:a79a6d399cef33a11b6f004c67bb07741d91f2be01b8d712d52c75711b1e07c7", size = 1698891, upload-time = "2026-03-31T21:59:14.552Z" }, - { url = "https://files.pythonhosted.org/packages/d8/78/ea0ae5ec8ba7a5c10bdd6e318f1ba5e76fcde17db8275188772afc7917a4/aiohttp-3.13.5-cp314-cp314-musllinux_1_2_ppc64le.whl", hash = "sha256:c632ce9c0b534fbe25b52c974515ed674937c5b99f549a92127c85f771a78772", size = 1742113, upload-time = "2026-03-31T21:59:17.068Z" }, - { url = "https://files.pythonhosted.org/packages/8a/66/9d308ed71e3f2491be1acb8769d96c6f0c47d92099f3bc9119cada27b357/aiohttp-3.13.5-cp314-cp314-musllinux_1_2_riscv64.whl", hash = "sha256:fceedde51fbd67ee2bcc8c0b33d0126cc8b51ef3bbde2f86662bd6d5a6f10ec5", size = 1553088, upload-time = "2026-03-31T21:59:19.541Z" }, - { url = "https://files.pythonhosted.org/packages/da/a6/6cc25ed8dfc6e00c90f5c6d126a98e2cf28957ad06fa1036bd34b6f24a2c/aiohttp-3.13.5-cp314-cp314-musllinux_1_2_s390x.whl", hash = "sha256:f92995dfec9420bb69ae629abf422e516923ba79ba4403bc750d94fb4a6c68c1", size = 1757976, upload-time = "2026-03-31T21:59:22.311Z" }, - { url = "https://files.pythonhosted.org/packages/c1/2b/cce5b0ffe0de99c83e5e36d8f828e4161e415660a9f3e58339d07cce3006/aiohttp-3.13.5-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:20ae0ff08b1f2c8788d6fb85afcb798654ae6ba0b747575f8562de738078457b", size = 1712444, upload-time = "2026-03-31T21:59:24.635Z" }, - { url = "https://files.pythonhosted.org/packages/6c/cf/9e1795b4160c58d29421eafd1a69c6ce351e2f7c8d3c6b7e4ca44aea1a5b/aiohttp-3.13.5-cp314-cp314-win32.whl", hash = "sha256:b20df693de16f42b2472a9c485e1c948ee55524786a0a34345511afdd22246f3", size = 438128, upload-time = "2026-03-31T21:59:27.291Z" }, - { url = "https://files.pythonhosted.org/packages/22/4d/eaedff67fc805aeba4ba746aec891b4b24cebb1a7d078084b6300f79d063/aiohttp-3.13.5-cp314-cp314-win_amd64.whl", hash = "sha256:f85c6f327bf0b8c29da7d93b1cabb6363fb5e4e160a32fa241ed2dce21b73162", size = 464029, upload-time = "2026-03-31T21:59:29.429Z" }, - { url = "https://files.pythonhosted.org/packages/79/11/c27d9332ee20d68dd164dc12a6ecdef2e2e35ecc97ed6cf0d2442844624b/aiohttp-3.13.5-cp314-cp314t-macosx_10_13_universal2.whl", hash = "sha256:1efb06900858bb618ff5cee184ae2de5828896c448403d51fb633f09e109be0a", size = 778758, upload-time = "2026-03-31T21:59:31.547Z" }, - { url = "https://files.pythonhosted.org/packages/04/fb/377aead2e0a3ba5f09b7624f702a964bdf4f08b5b6728a9799830c80041e/aiohttp-3.13.5-cp314-cp314t-macosx_10_13_x86_64.whl", hash = "sha256:fee86b7c4bd29bdaf0d53d14739b08a106fdda809ca5fe032a15f52fae5fe254", size = 512883, upload-time = "2026-03-31T21:59:34.098Z" }, - { url = "https://files.pythonhosted.org/packages/bb/a6/aa109a33671f7a5d3bd78b46da9d852797c5e665bfda7d6b373f56bff2ec/aiohttp-3.13.5-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:20058e23909b9e65f9da62b396b77dfa95965cbe840f8def6e572538b1d32e36", size = 516668, upload-time = "2026-03-31T21:59:36.497Z" }, - { url = "https://files.pythonhosted.org/packages/79/b3/ca078f9f2fa9563c36fb8ef89053ea2bb146d6f792c5104574d49d8acb63/aiohttp-3.13.5-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:8cf20a8d6868cb15a73cab329ffc07291ba8c22b1b88176026106ae39aa6df0f", size = 1883461, upload-time = "2026-03-31T21:59:38.723Z" }, - { url = "https://files.pythonhosted.org/packages/b7/e3/a7ad633ca1ca497b852233a3cce6906a56c3225fb6d9217b5e5e60b7419d/aiohttp-3.13.5-cp314-cp314t-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:330f5da04c987f1d5bdb8ae189137c77139f36bd1cb23779ca1a354a4b027800", size = 1747661, upload-time = "2026-03-31T21:59:41.187Z" }, - { url = "https://files.pythonhosted.org/packages/33/b9/cd6fe579bed34a906d3d783fe60f2fa297ef55b27bb4538438ee49d4dc41/aiohttp-3.13.5-cp314-cp314t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:6f1cbf0c7926d315c3c26c2da41fd2b5d2fe01ac0e157b78caefc51a782196cf", size = 1863800, upload-time = "2026-03-31T21:59:43.84Z" }, - { url = "https://files.pythonhosted.org/packages/c0/3f/2c1e2f5144cefa889c8afd5cf431994c32f3b29da9961698ff4e3811b79a/aiohttp-3.13.5-cp314-cp314t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:53fc049ed6390d05423ba33103ded7281fe897cf97878f369a527070bd95795b", size = 1958382, upload-time = "2026-03-31T21:59:46.187Z" }, - { url = "https://files.pythonhosted.org/packages/66/1d/f31ec3f1013723b3babe3609e7f119c2c2fb6ef33da90061a705ef3e1bc8/aiohttp-3.13.5-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:898703aa2667e3c5ca4c54ca36cd73f58b7a38ef87a5606414799ebce4d3fd3a", size = 1803724, upload-time = "2026-03-31T21:59:48.656Z" }, - { url = "https://files.pythonhosted.org/packages/0e/b4/57712dfc6f1542f067daa81eb61da282fab3e6f1966fca25db06c4fc62d5/aiohttp-3.13.5-cp314-cp314t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:0494a01ca9584eea1e5fbd6d748e61ecff218c51b576ee1999c23db7066417d8", size = 1640027, upload-time = "2026-03-31T21:59:51.284Z" }, - { url = "https://files.pythonhosted.org/packages/25/3c/734c878fb43ec083d8e31bf029daae1beafeae582d1b35da234739e82ee7/aiohttp-3.13.5-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:6cf81fe010b8c17b09495cbd15c1d35afbc8fb405c0c9cf4738e5ae3af1d65be", size = 1806644, upload-time = "2026-03-31T21:59:53.753Z" }, - { url = "https://files.pythonhosted.org/packages/20/a5/f671e5cbec1c21d044ff3078223f949748f3a7f86b14e34a365d74a5d21f/aiohttp-3.13.5-cp314-cp314t-musllinux_1_2_armv7l.whl", hash = "sha256:c564dd5f09ddc9d8f2c2d0a301cd30a79a2cc1b46dd1a73bef8f0038863d016b", size = 1791630, upload-time = "2026-03-31T21:59:56.239Z" }, - { url = "https://files.pythonhosted.org/packages/0b/63/fb8d0ad63a0b8a99be97deac8c04dacf0785721c158bdf23d679a87aa99e/aiohttp-3.13.5-cp314-cp314t-musllinux_1_2_ppc64le.whl", hash = "sha256:2994be9f6e51046c4f864598fd9abeb4fba6e88f0b2152422c9666dcd4aea9c6", size = 1809403, upload-time = "2026-03-31T21:59:59.103Z" }, - { url = "https://files.pythonhosted.org/packages/59/0c/bfed7f30662fcf12206481c2aac57dedee43fe1c49275e85b3a1e1742294/aiohttp-3.13.5-cp314-cp314t-musllinux_1_2_riscv64.whl", hash = "sha256:157826e2fa245d2ef46c83ea8a5faf77ca19355d278d425c29fda0beb3318037", size = 1634924, upload-time = "2026-03-31T22:00:02.116Z" }, - { url = "https://files.pythonhosted.org/packages/17/d6/fd518d668a09fd5a3319ae5e984d4d80b9a4b3df4e21c52f02251ef5a32e/aiohttp-3.13.5-cp314-cp314t-musllinux_1_2_s390x.whl", hash = "sha256:a8aca50daa9493e9e13c0f566201a9006f080e7c50e5e90d0b06f53146a54500", size = 1836119, upload-time = "2026-03-31T22:00:04.756Z" }, - { url = "https://files.pythonhosted.org/packages/78/b7/15fb7a9d52e112a25b621c67b69c167805cb1f2ab8f1708a5c490d1b52fe/aiohttp-3.13.5-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:3b13560160d07e047a93f23aaa30718606493036253d5430887514715b67c9d9", size = 1772072, upload-time = "2026-03-31T22:00:07.494Z" }, - { url = "https://files.pythonhosted.org/packages/7e/df/57ba7f0c4a553fc2bd8b6321df236870ec6fd64a2a473a8a13d4f733214e/aiohttp-3.13.5-cp314-cp314t-win32.whl", hash = "sha256:9a0f4474b6ea6818b41f82172d799e4b3d29e22c2c520ce4357856fced9af2f8", size = 471819, upload-time = "2026-03-31T22:00:10.277Z" }, - { url = "https://files.pythonhosted.org/packages/62/29/2f8418269e46454a26171bfdd6a055d74febf32234e474930f2f60a17145/aiohttp-3.13.5-cp314-cp314t-win_amd64.whl", hash = "sha256:18a2f6c1182c51baa1d28d68fea51513cb2a76612f038853c0ad3c145423d3d9", size = 505441, upload-time = "2026-03-31T22:00:12.791Z" }, +sdist = { url = "https://files.pythonhosted.org/packages/58/d9/22ce5786ac0c1653ae8b6c23bded02c1686d11f0dbb45b31ce128e0df985/aiohttp-3.14.3.tar.gz", hash = "sha256:9491196535a88924a60afd5b5f434b5b203b6cc616250878dbdb223a8f7844bc", size = 7971213, upload-time = "2026-07-23T01:57:27.037Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/18/d4/eb96299230e20acf2efae207cb8d69051f1f68e357e5ea5e479bf6fb097a/aiohttp-3.14.3-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:39aded8c7f3b935b54aab1d8d73c70ec0ee2d3ec3b943e0e86611bc150ba47f5", size = 754690, upload-time = "2026-07-23T01:53:47.332Z" }, + { url = "https://files.pythonhosted.org/packages/88/11/e7a70a209eb9a067c0d3212b518a0134e3484f5178c7533878b6b514d469/aiohttp-3.14.3-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:5bcb6ff3fdab1258a192679ff1a05d44f59626430aa05cd1a9d2447423599228", size = 509484, upload-time = "2026-07-23T01:53:51.159Z" }, + { url = "https://files.pythonhosted.org/packages/30/07/4bbc222cc8dbe31d4c3e8a5baad2286e4d42026ac0c570027b89afce6344/aiohttp-3.14.3-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:617105e2c3018ee38d0c8ce5ee3c84f621a6d8b9f723202aacaff28449ca91ee", size = 511949, upload-time = "2026-07-23T01:53:55.083Z" }, + { url = "https://files.pythonhosted.org/packages/54/b9/42e74c46b7b7c794b995bbc1f573fb48950c38b19d8600c62a6804ee2d67/aiohttp-3.14.3-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f631fe87a6f30df5fbe6d79640b25e4cffb38c31c7fb6f10871517b84b0f8c1a", size = 1765282, upload-time = "2026-07-23T01:53:59.662Z" }, + { url = "https://files.pythonhosted.org/packages/6b/ed/62bc4d74363ad346d518e0720363a949f63e2e23439a79eb5813d4d29bb3/aiohttp-3.14.3-cp312-cp312-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:a94dbaae5ae27bd849c93570669bff91e0510f33a80805738e3de72a7be0447b", size = 1741511, upload-time = "2026-07-23T01:54:04.063Z" }, + { url = "https://files.pythonhosted.org/packages/d0/9f/181e8a8bc79e47d13c7fc4540bd7a3b729d9505609c61f392a8dd2fbfe55/aiohttp-3.14.3-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:8f2f1c4c032c7cedd7d8da6f54c97b70266c6570c3108d3fdffee7188bb70529", size = 1810680, upload-time = "2026-07-23T01:54:09.882Z" }, + { url = "https://files.pythonhosted.org/packages/5c/9a/dec94d6ad694552fe3424e3f1928d7a606a5d9d9433a04e7ecdd9d38ae7f/aiohttp-3.14.3-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:ea05e1f97ceea523942d9b2a7d7c0359d781d683d6b043f5943a602b14da4787", size = 1905646, upload-time = "2026-07-23T01:54:13.475Z" }, + { url = "https://files.pythonhosted.org/packages/52/b7/7cd31f29d6055bd711ae6e669367fba6f5ae9de463910a793e30556a8db7/aiohttp-3.14.3-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:543906c127fb1d929b95076db19b83fa2d46751006ff1e23b093aa5ac4d8db42", size = 1792122, upload-time = "2026-07-23T01:54:15.752Z" }, + { url = "https://files.pythonhosted.org/packages/66/73/10b1ef93afa61f4963c746257b70ced619cf31a4798671de5fdb2608501d/aiohttp-3.14.3-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:0a5ff2dfbb9ce645fa5b8ef3e02c6c0b9cc3f6030ff863d0c51fffc50cb5541b", size = 1591127, upload-time = "2026-07-23T01:54:19.489Z" }, + { url = "https://files.pythonhosted.org/packages/49/ed/3b203fa6de1b338c14acdc06bf6ca9b043b7944f005966958c2ced932cde/aiohttp-3.14.3-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:041badb8f84396357c4d3ad26de6afd7a32b112f43d3c63045c0c8278cfd2043", size = 1725210, upload-time = "2026-07-23T01:54:24.129Z" }, + { url = "https://files.pythonhosted.org/packages/28/b7/1c2aab8c706436dcc28598452488ac9cd7c409da815237c28c27d58993e6/aiohttp-3.14.3-cp312-cp312-musllinux_1_2_armv7l.whl", hash = "sha256:530125ee1163c4219af35dc3aa1206e541e7b31b6efc1a3f93b70a136f65d427", size = 1764848, upload-time = "2026-07-23T01:54:27.973Z" }, + { url = "https://files.pythonhosted.org/packages/54/50/94c28f08b131c4bf10984ea2c7a536c9920608bb2d6e7f95642c30cc87b7/aiohttp-3.14.3-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:c8653fd547c93a61aadc612007790f5555cdd18946fa48cf45e26d8ea4ea473d", size = 1777102, upload-time = "2026-07-23T01:54:31.775Z" }, + { url = "https://files.pythonhosted.org/packages/13/d4/e7d09ba7d345fb2d74440fd2fa033c5e079fac05552927705986f41a364f/aiohttp-3.14.3-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:89176250f686cb9853c0fb7ead90e639e915b84a6f43eedc2a4e7ec21f1037f0", size = 1580205, upload-time = "2026-07-23T01:54:34.518Z" }, + { url = "https://files.pythonhosted.org/packages/a3/84/072a91d68e1e1eb587985b54baab94221277f877e8ef274fc213a0ceae28/aiohttp-3.14.3-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:3a26434dafe408229ff3403458ca58de24fb51936504decac49ce6755f77e59d", size = 1797219, upload-time = "2026-07-23T01:54:36.995Z" }, + { url = "https://files.pythonhosted.org/packages/e0/eb/aad34e897e668424d6e995da5dff8a4a09af93363d3392488772957a63aa/aiohttp-3.14.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:d1558173930a5a8d3069cee5c92fc91c87c4dbcb099debbb3622053717145a19", size = 1768629, upload-time = "2026-07-23T01:54:40.103Z" }, + { url = "https://files.pythonhosted.org/packages/b6/2b/6bb88ddba0fecd9122aa3ebcad25996cf6c083a4a7040dbb3a4f97972af6/aiohttp-3.14.3-cp312-cp312-win32.whl", hash = "sha256:16100ad3ab8d649fdfbee87602d9d2dcdca9df0b9eda8a1b5fdc0d41f96da559", size = 451481, upload-time = "2026-07-23T01:54:42.547Z" }, + { url = "https://files.pythonhosted.org/packages/76/9b/f2f8f108da17ecef2cc3efc424e8b7ad3782b1a8360f7b8eae8ced84f6ea/aiohttp-3.14.3-cp312-cp312-win_amd64.whl", hash = "sha256:33a2d7c28d33797a2e99923dffa63f83d908a19b6bf26cfe80fa790aa5e1a75a", size = 476845, upload-time = "2026-07-23T01:54:44.853Z" }, + { url = "https://files.pythonhosted.org/packages/3e/44/28dac80a8941b604f4da10ce21097614ca1bf905ce93dca28d8d7de9c1e7/aiohttp-3.14.3-cp312-cp312-win_arm64.whl", hash = "sha256:362a3fd481769cac1a824514bcd86fda51c65e8fe6e051099e008fddde6db17c", size = 448050, upload-time = "2026-07-23T01:54:47.087Z" }, + { url = "https://files.pythonhosted.org/packages/57/be/5afd201cc0ab139029aadb75392efe85a293403d9dd3a3226161c21ce00c/aiohttp-3.14.3-cp313-cp313-android_21_arm64_v8a.whl", hash = "sha256:2e9878ae68e4a5f1c0abe4dd497dbc3d51946f5837b56759e2a02e78fa90ef86", size = 506269, upload-time = "2026-07-23T01:54:49.075Z" }, + { url = "https://files.pythonhosted.org/packages/22/09/dec8189d62b45ade009f6792a2264b942a90cb88aeaf181239933cd72c3c/aiohttp-3.14.3-cp313-cp313-android_21_x86_64.whl", hash = "sha256:f3d2669fe7dec7fc359ecdb5984b29b50d85d5d00f8c1cb61de4f4a24ee42627", size = 515166, upload-time = "2026-07-23T01:54:51.894Z" }, + { url = "https://files.pythonhosted.org/packages/28/24/2854869d29ed8a8b19d74f9ec6629515f7e04d02dd329d9d179201e58e47/aiohttp-3.14.3-cp313-cp313-ios_13_0_arm64_iphoneos.whl", hash = "sha256:cc7cb243a68167172f48c1fd43cee91ec4b1d40cefd190edd43369d1a6bc9c82", size = 486263, upload-time = "2026-07-23T01:54:54.223Z" }, + { url = "https://files.pythonhosted.org/packages/d4/dd/57187c8be2a35aea65eaee3bd2c3dcbbcf0204f5106c89637e3610380cd1/aiohttp-3.14.3-cp313-cp313-ios_13_0_arm64_iphonesimulator.whl", hash = "sha256:78253b573e6ffab5028924fc98bc281aae05445969982a10864bc360dea2016c", size = 492299, upload-time = "2026-07-23T01:54:56.236Z" }, + { url = "https://files.pythonhosted.org/packages/b9/11/06ae6ed8f0d414edf4068861e233d8fe23ee699bfd4b3ceb8663db948a62/aiohttp-3.14.3-cp313-cp313-ios_13_0_x86_64_iphonesimulator.whl", hash = "sha256:7041d52c3a7fa20c9e8c182b534704abb19502c8bdcbde7ab23bfda6f642394f", size = 502235, upload-time = "2026-07-23T01:54:58.377Z" }, + { url = "https://files.pythonhosted.org/packages/7e/a3/559639c34a345d2cf7c52dff6838119f2eaf29eb508227b5b83f573af813/aiohttp-3.14.3-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:ac74facc01463f138b0da5580329cfcc82818dea5656e83ddcd11268fc12ff80", size = 750883, upload-time = "2026-07-23T01:55:00.65Z" }, + { url = "https://files.pythonhosted.org/packages/91/cd/41e131f13afd1e7b0172a9d9eda085ef90eb8439f41f0d279db81ed3ae60/aiohttp-3.14.3-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:d6218d92e450824e9b4881f44e8c09f1853b490f9a64130801024a4793b1b3b0", size = 508473, upload-time = "2026-07-23T01:55:02.945Z" }, + { url = "https://files.pythonhosted.org/packages/bc/6b/e7f13410d391c6e55b4c007a8de024355389d7d459e3d64c42b2d33617e5/aiohttp-3.14.3-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:11fb37ef075669eee52ab1928fbf6e1741fada40409fa309ebde9607a962aebf", size = 509190, upload-time = "2026-07-23T01:55:05.173Z" }, + { url = "https://files.pythonhosted.org/packages/97/21/6464573e53d69672cc1eada3e5c5cb2d2efa82701e8305a0f2047a576967/aiohttp-3.14.3-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:55bdcc472aafe2de4a253045cc128007a64f1e0264fb675791e132ea5edaa3bd", size = 1761478, upload-time = "2026-07-23T01:55:07.383Z" }, + { url = "https://files.pythonhosted.org/packages/1a/81/d217043a4c17fbce360905e3b2bdd20139ebc9a2de836d035d179c4da006/aiohttp-3.14.3-cp313-cp313-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:c39846c3aad97a8530c89d7a3869a8f8e9e3762c6ac0504481e5c80948f7e807", size = 1735092, upload-time = "2026-07-23T01:55:09.803Z" }, + { url = "https://files.pythonhosted.org/packages/a1/66/e13a02d0eeb1a9a502402a977abb4e4abff9fe4051c26f80558c57a7c975/aiohttp-3.14.3-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:5895ef58c4620afe02fa16044f023dc4dafec08158f9d08874a46a7dbc0341b8", size = 1800546, upload-time = "2026-07-23T01:55:12.012Z" }, + { url = "https://files.pythonhosted.org/packages/26/5e/57d42fca1d18cb5acc1cad945d017fabc5d6ae71d8a08ad66be8dc3ee544/aiohttp-3.14.3-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:fa9467a8113aa69d3d7c55a70ef0b7c636010a40993f3df9d9d0d73b3eb7ef24", size = 1895250, upload-time = "2026-07-23T01:55:14.357Z" }, + { url = "https://files.pythonhosted.org/packages/ca/1c/7da8d08e74d56f00070822f9638ff3f1c563f8ad87d1efa996c87bfc8644/aiohttp-3.14.3-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d7d2deec16eeedf55f2c7cf75b521ea3856a5177e123844f8fd0f114ce252cb5", size = 1789289, upload-time = "2026-07-23T01:55:16.668Z" }, + { url = "https://files.pythonhosted.org/packages/cd/0f/cf16bcf56896981c1a0319f5d5db9337994b5165730c48a8fa07e9b34be6/aiohttp-3.14.3-cp313-cp313-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:dd54d0e8717de95939766febac482ac0474d8ac3b048115f9f2b1d23a16e7db4", size = 1586706, upload-time = "2026-07-23T01:55:18.913Z" }, + { url = "https://files.pythonhosted.org/packages/fe/6f/76eac12a7f2480e1e304f842efdb07db33256b0d9165b866b6ef0806c202/aiohttp-3.14.3-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:df82f3787c940c94986b34222d59c9e38843fba85139f36e85255a82ad5355a9", size = 1724652, upload-time = "2026-07-23T01:55:21.296Z" }, + { url = "https://files.pythonhosted.org/packages/39/b6/19c8c592baeeb94b75f966547d40c02ac7590902306ec5863d5c027cf506/aiohttp-3.14.3-cp313-cp313-musllinux_1_2_armv7l.whl", hash = "sha256:42a67efc36300d052fb4508a53e8b6901b9284b599ae63945c377569c5fcc1e1", size = 1756239, upload-time = "2026-07-23T01:55:23.705Z" }, + { url = "https://files.pythonhosted.org/packages/dc/c9/4e9383150296f97f873b680c4de8fb2cd88608fb9f48c79edcb111611abc/aiohttp-3.14.3-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:7a75aa63cbf9b21cfaf60dc2657e19df2c2867d91707d653fee171ffeedd1371", size = 1769161, upload-time = "2026-07-23T01:55:26.082Z" }, + { url = "https://files.pythonhosted.org/packages/aa/1e/147bdc6cc5de5f3ab011be8bf5d6e786633249f22c20bae06f85e45f5387/aiohttp-3.14.3-cp313-cp313-musllinux_1_2_riscv64.whl", hash = "sha256:e92eb8acc45eb6a9f4935071a77edf5b85cc6f8dfad5cd99e97653c26593cdde", size = 1578759, upload-time = "2026-07-23T01:55:28.846Z" }, + { url = "https://files.pythonhosted.org/packages/fd/31/78388a9d6040ece2e11df62ea229a822cf5e52d238374b220ae9975b2623/aiohttp-3.14.3-cp313-cp313-musllinux_1_2_s390x.whl", hash = "sha256:b014a6ed7cf912e787149fdc529166d3ceabac23f26efeea3158c9aba2354e7e", size = 1792025, upload-time = "2026-07-23T01:55:31.457Z" }, + { url = "https://files.pythonhosted.org/packages/03/51/a3d29fdf2c25d796746af8ad6fe56a45d6256c38b0a8a2ed752e1160b3a2/aiohttp-3.14.3-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:3d4f72af88ac2474bb5bca640030320e3d38a0163a1d7533500e87be458eef71", size = 1768477, upload-time = "2026-07-23T01:55:33.87Z" }, + { url = "https://files.pythonhosted.org/packages/29/a6/442e18b5afeade534d877a2dc3c3e392aff8d49787890b0cf84790410267/aiohttp-3.14.3-cp313-cp313-win32.whl", hash = "sha256:5f08ec777f35ee70720233b8b9811d3bb5d728137f30ac91b7457709c3261ac0", size = 451069, upload-time = "2026-07-23T01:55:36.121Z" }, + { url = "https://files.pythonhosted.org/packages/9d/69/3d876ac02659f271cf7f6769f14a8e3de5b6e888ed8b5a7e998086a4cec8/aiohttp-3.14.3-cp313-cp313-win_amd64.whl", hash = "sha256:dff9461ec275f22135650d5ba4b4931a11f3958df7dfbb8db630000d4dee0883", size = 476518, upload-time = "2026-07-23T01:55:38.303Z" }, + { url = "https://files.pythonhosted.org/packages/b2/0e/50d6e6471cd31edce8b282bdec59375a3a69124d8a989a0b1313355cae52/aiohttp-3.14.3-cp313-cp313-win_arm64.whl", hash = "sha256:ddcac3c6b382e81f1dd0499199d4136b877beb4cb5ef770bbbfba56c4b8f55d2", size = 447676, upload-time = "2026-07-23T01:55:40.451Z" }, + { url = "https://files.pythonhosted.org/packages/c8/20/887fdcf832326571b370ffc347b3e70abe101096f3720126aac161b1d872/aiohttp-3.14.3-cp314-cp314-android_24_arm64_v8a.whl", hash = "sha256:49f7325beb0f85ef4aef5f48f490269575f83e6e2acad00a1d80b807eb027062", size = 509067, upload-time = "2026-07-23T01:55:42.618Z" }, + { url = "https://files.pythonhosted.org/packages/ad/a3/92cec936f78cc4bf0fa5554ebe593b73459d94e3c62303e1902a4cccb6f7/aiohttp-3.14.3-cp314-cp314-android_24_x86_64.whl", hash = "sha256:e3be98a7c30b8c25d573dafba7171d66dfb05ee6a9070fc46535464ff97700a6", size = 514774, upload-time = "2026-07-23T01:55:44.937Z" }, + { url = "https://files.pythonhosted.org/packages/29/ba/2a0c38df3fc557620b6a5acd98364af050053b6285b4dc7ee74100c63c18/aiohttp-3.14.3-cp314-cp314-ios_13_0_arm64_iphoneos.whl", hash = "sha256:614c61d478b83953e261d02bb2df750f17227cd33ef8002945bf5aebbde21919", size = 488134, upload-time = "2026-07-23T01:55:47.135Z" }, + { url = "https://files.pythonhosted.org/packages/48/d6/d51b7d4bf309af3693940d8ffd2b9ed0b682434ef85959b7c9c137f60cf8/aiohttp-3.14.3-cp314-cp314-ios_13_0_arm64_iphonesimulator.whl", hash = "sha256:1caa7b0d05f3e3a36f87788c59e970a7ee1cefcfcbb924a9f138c4a6551c9cb7", size = 494201, upload-time = "2026-07-23T01:55:49.451Z" }, + { url = "https://files.pythonhosted.org/packages/3f/5a/8f624384e5f1efabb5229b94157eb966b021e97bdb188c62860c2ae243c2/aiohttp-3.14.3-cp314-cp314-ios_13_0_x86_64_iphonesimulator.whl", hash = "sha256:dfa68deb2a443bdaa3ea5297b0699c1464f08aef3812b486d1348eee61b07dc0", size = 502766, upload-time = "2026-07-23T01:55:51.656Z" }, + { url = "https://files.pythonhosted.org/packages/a6/26/4ff0164370deec18fb19254ee4ab10b7a73304ac0c860b13f5f84663759b/aiohttp-3.14.3-cp314-cp314-macosx_10_15_universal2.whl", hash = "sha256:e72ee89e28d907a18f46959b4eb0bb06701cc7f8cf4366e00029e2ccfaaf5924", size = 756557, upload-time = "2026-07-23T01:55:53.964Z" }, + { url = "https://files.pythonhosted.org/packages/97/a3/7056b86dc0d9ec709ea9777eae3b0161428f943372f8b98c01c11593b682/aiohttp-3.14.3-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:ad4c8b7488d745d2ca4838ebd8ae5ba9b56341d30b1da43640e4ce87f9f49646", size = 510168, upload-time = "2026-07-23T01:55:56.22Z" }, + { url = "https://files.pythonhosted.org/packages/85/ed/0357a015892fd68058bf2d39d3fd1958e459b997a7db30aaa6aaa434ae96/aiohttp-3.14.3-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:db332af25642007330fca8be5c4d194caf2bea7a7fc84415aff3497af5dfee6b", size = 512957, upload-time = "2026-07-23T01:55:58.437Z" }, + { url = "https://files.pythonhosted.org/packages/47/d1/8aba53f15ccb2238405f5e9d30e2a8ca44f93878c26e7165ade00d374b1c/aiohttp-3.14.3-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:25bd2708db6bdf6a6630dd37bdcdfcb47c4434d22ac69c64665b802910140b30", size = 1750149, upload-time = "2026-07-23T01:56:00.856Z" }, + { url = "https://files.pythonhosted.org/packages/49/bd/40c3fee327529284375c6701cbb0fa4600cc2e8432af1378f897e2ef7d3a/aiohttp-3.14.3-cp314-cp314-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:cef89a58e628c4efcac3275c2d68083f82426dcdc89c1492a6f654f9f7ea6ab9", size = 1707685, upload-time = "2026-07-23T01:56:03.371Z" }, + { url = "https://files.pythonhosted.org/packages/2a/a3/ca0cc6724cca8114b05694abd916060758c79894c3aa5b012cdadc1bc28e/aiohttp-3.14.3-cp314-cp314-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:c23ec8ee9d5ab2f5421f9c7fffce208435607af27fd46d4a44e031954352838f", size = 1803911, upload-time = "2026-07-23T01:56:05.817Z" }, + { url = "https://files.pythonhosted.org/packages/95/b5/85b099c299c3ffd38ad9b3e43694c8a346934e4a30c88c4fd5a841234f77/aiohttp-3.14.3-cp314-cp314-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:e2667f0bbe7eb6c74eae5e9691441ad186e5845ca3cff63230fc09c4e7514f5d", size = 1876929, upload-time = "2026-07-23T01:56:08.413Z" }, + { url = "https://files.pythonhosted.org/packages/d5/b7/1da684a04175473fa4cddbf9a2f572e79514c3fd27a74597f43057d4f3da/aiohttp-3.14.3-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:18cb43369747b2ae007bd2655fb8e63a099c2ff1d207962943636dac989b3147", size = 1761112, upload-time = "2026-07-23T01:56:10.918Z" }, + { url = "https://files.pythonhosted.org/packages/d1/16/bc4b55e3e5cb175fd69c53c90d60d2f47797cb343da5106e23863dc4dba4/aiohttp-3.14.3-cp314-cp314-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:d77640cc618c1d99fc4f8589c0f24a730adfa54eb1e57ef7bf0c8dfb78da898c", size = 1583500, upload-time = "2026-07-23T01:56:13.613Z" }, + { url = "https://files.pythonhosted.org/packages/2a/e8/13a9d957a1ee40837f46aa30f0f4c657e673ad86a2e6362a9f9be20d26d9/aiohttp-3.14.3-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:53e5179d8abb5710f8e83ba207c41c8d1261fcffd4616500e15ca2b7a33be10a", size = 1713940, upload-time = "2026-07-23T01:56:15.969Z" }, + { url = "https://files.pythonhosted.org/packages/38/05/d33c680c1bcf1c7e130f9cbfc1fc02fe8bb0c4af2a94a53dd5fb56131e5c/aiohttp-3.14.3-cp314-cp314-musllinux_1_2_armv7l.whl", hash = "sha256:cd817772b2fcf2b8c0905795318485f9ec16eae60b29feb7f4c77085311637f0", size = 1724413, upload-time = "2026-07-23T01:56:18.591Z" }, + { url = "https://files.pythonhosted.org/packages/85/1d/af798d306f7a74b6a632dbcabcf62a4c91391b7582d2a8c6d7712e2cc54e/aiohttp-3.14.3-cp314-cp314-musllinux_1_2_ppc64le.whl", hash = "sha256:4e3ac92d90e92773b2362d506068e9a948192bd553e743c5b2429e28527c8661", size = 1770748, upload-time = "2026-07-23T01:56:21.074Z" }, + { url = "https://files.pythonhosted.org/packages/a8/92/ad720d472556a995049206867765e9410969684f86ee09423ff9969044c1/aiohttp-3.14.3-cp314-cp314-musllinux_1_2_riscv64.whl", hash = "sha256:3f42e9b78301f11c8f861746175d8b9c1ccef713fcad9eab396e2f6db8ed4a22", size = 1577564, upload-time = "2026-07-23T01:56:23.475Z" }, + { url = "https://files.pythonhosted.org/packages/60/ad/0ed7586cbef7a884e23a752fa2bb987a122e6a5dd50dab109258d0a95193/aiohttp-3.14.3-cp314-cp314-musllinux_1_2_s390x.whl", hash = "sha256:9d9edccfe496b476db5f398d97b865e9a6752bcf8aec4eef8390ce20fb64bb41", size = 1782080, upload-time = "2026-07-23T01:56:25.994Z" }, + { url = "https://files.pythonhosted.org/packages/97/ea/dbaed0d73e8a69aad653b045dab451c67c2454bb731a37b45a86593e9422/aiohttp-3.14.3-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:1c5ec8fb1bcc31a8466f74aaf26c345d5c386fa4bd08a3f0eb9c7a4a3fe8b5bf", size = 1745813, upload-time = "2026-07-23T01:56:28.604Z" }, + { url = "https://files.pythonhosted.org/packages/81/1b/6893d4bc57e434fc93a6c9217c637d967a0b651d989f6e3265179375754a/aiohttp-3.14.3-cp314-cp314-win32.whl", hash = "sha256:38901a84da3ce22249f6e860bf8f90d141bcab7da090cc398f8bb58c0e44b7da", size = 455872, upload-time = "2026-07-23T01:56:31.031Z" }, + { url = "https://files.pythonhosted.org/packages/f5/8b/c7baa1ba1eda4db6989baefe5de6d99834921b84ebd7918624febcb9f290/aiohttp-3.14.3-cp314-cp314-win_amd64.whl", hash = "sha256:8b3b60de05f3dcb6f6a00f818bb2ec781cee4de0645f59ccaf99b1d1823b6100", size = 481030, upload-time = "2026-07-23T01:56:33.365Z" }, + { url = "https://files.pythonhosted.org/packages/22/8c/c29d067df825a2df88ca432db848aa2fe8199598359cc06c12b09320cac9/aiohttp-3.14.3-cp314-cp314-win_arm64.whl", hash = "sha256:1576145bdceeb92382d899751e12743a3a5b8e460a841e3e50543859e54864dc", size = 453669, upload-time = "2026-07-23T01:56:35.731Z" }, + { url = "https://files.pythonhosted.org/packages/6a/a4/9c033beb355d39b6147980597ec9645e4729243f686ee4dc73945de72030/aiohttp-3.14.3-cp314-cp314t-macosx_10_15_universal2.whl", hash = "sha256:8800c996b01c2772a783e3e46f3e1abd5823029adca0df54231960de9bfefa5b", size = 791403, upload-time = "2026-07-23T01:56:37.972Z" }, + { url = "https://files.pythonhosted.org/packages/80/ca/87c32a0a7704583cfc49660bd817889bae5b830bf53b5dcb4e92145ac2da/aiohttp-3.14.3-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:ebe8e504f058fe91223351cecd2d9d6946c9d241bb0250d898ffbdf584cc72b0", size = 526413, upload-time = "2026-07-23T01:56:40.523Z" }, + { url = "https://files.pythonhosted.org/packages/9e/d8/8ec0e471248c500acdce2be3f46db8fb62b5eb60efef072529cc85ee1d26/aiohttp-3.14.3-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:30402d03a7c0ff52bce290b57e564e9079fd9d0cb545c8aba73f86a103162d2e", size = 532135, upload-time = "2026-07-23T01:56:42.876Z" }, + { url = "https://files.pythonhosted.org/packages/fe/45/f8919fd936e8b79fcd9bda7b6d8e62613462a713f4f17987fd7c34399142/aiohttp-3.14.3-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:9fc7b5bfec6573f3ae844f457fdde5adeb713f8b8e4a81ad64fc207b49383716", size = 1922742, upload-time = "2026-07-23T01:56:45.528Z" }, + { url = "https://files.pythonhosted.org/packages/f6/ec/9ca76b28a27525b0cc53e20842e0228b022f301ce1f436b7d814b4aaf2df/aiohttp-3.14.3-cp314-cp314t-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:8a5fd34f7f7410d1730d5c2ba873cacb2eed3fede366feb268a70ba22581ed8f", size = 1787371, upload-time = "2026-07-23T01:56:48.045Z" }, + { url = "https://files.pythonhosted.org/packages/b1/04/6acdbf17315f7b55f1937e3387acb89a3cddeb4995689553d064af8e92ab/aiohttp-3.14.3-cp314-cp314t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:270d3dace9ca2f10f0da5d8ebe519b7a310fc6112ed916e32df5866df0888553", size = 1912623, upload-time = "2026-07-23T01:56:50.605Z" }, + { url = "https://files.pythonhosted.org/packages/86/e6/438b0c79ca6f45eb9fd9817dd4c01a91919a38c0de5ee9e05e2b4dc0ece7/aiohttp-3.14.3-cp314-cp314t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:3ae5b3a59436d089b5395d910121a390feed4d00578eb95a0fd1a329fe963100", size = 2005515, upload-time = "2026-07-23T01:56:53.153Z" }, + { url = "https://files.pythonhosted.org/packages/bb/6b/62cbd6577758699525f5c712d1ddef57d9875fbab0ae8d5f5a202fd598f8/aiohttp-3.14.3-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:2498f0fe69ead802f9675beca44a7c21c62fdaa4ec5145ea1c3ad6edbee29f85", size = 1879906, upload-time = "2026-07-23T01:56:55.818Z" }, + { url = "https://files.pythonhosted.org/packages/00/95/18bcbf830a21dc3aae24d8f6b6feaf3db1d2090242d00a7868db2ffb0b67/aiohttp-3.14.3-cp314-cp314t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:a0dc483c00da8b673abbb367eb6f8d8f4bcec30eb58529ea13cb42e7fd2dfa33", size = 1675849, upload-time = "2026-07-23T01:56:58.861Z" }, + { url = "https://files.pythonhosted.org/packages/a9/19/47f4968659c5e23606c3790c80fc624e691c153d036148449ee84d31b287/aiohttp-3.14.3-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:c7d3a97c678d34fc5b59da671ee9cd630096ddc643e7b5a30d54a2a6f3574d3f", size = 1843496, upload-time = "2026-07-23T01:57:01.591Z" }, + { url = "https://files.pythonhosted.org/packages/64/af/38c33c4dd82fddcb4e56c4653b6f1072a8edbc6b7fa15809f14932c41e2d/aiohttp-3.14.3-cp314-cp314t-musllinux_1_2_armv7l.whl", hash = "sha256:f8fb78a83c9e5f741ca3a68cfb455c1f5bb83b4e7249a3848b3cd78d0a8563b0", size = 1827746, upload-time = "2026-07-23T01:57:05.131Z" }, + { url = "https://files.pythonhosted.org/packages/a1/9d/0537cda4885ac8f5b7053d164dd06312f4c483a4edcb8ee5b8aaf2a989bf/aiohttp-3.14.3-cp314-cp314t-musllinux_1_2_ppc64le.whl", hash = "sha256:74ab5b6a9fb13e873e5a90946588baecaf488745e1db1a4a5c433f971f035098", size = 1853810, upload-time = "2026-07-23T01:57:08.043Z" }, + { url = "https://files.pythonhosted.org/packages/19/fe/26f9c5e6458385aa86497836b0dea6fb2f027827d63f37c7856cce9286ee/aiohttp-3.14.3-cp314-cp314t-musllinux_1_2_riscv64.whl", hash = "sha256:bd52f811e65f6fb634b1047159657c98f52b407f8efec907bcfc09da9a4c0a25", size = 1668895, upload-time = "2026-07-23T01:57:10.837Z" }, + { url = "https://files.pythonhosted.org/packages/ec/4c/618b1db9b9ba079b8875d2cdf78e7c4a3bf72903bd5850fee7dd9544600a/aiohttp-3.14.3-cp314-cp314t-musllinux_1_2_s390x.whl", hash = "sha256:f0f177d1b195b9e06376cfd7d308d8a1b920909a609d03ac82a8c73bbb16d3b9", size = 1883833, upload-time = "2026-07-23T01:57:13.672Z" }, + { url = "https://files.pythonhosted.org/packages/94/c6/bd959bd1e4771f9fd944e9e436224c48c77b018b73b519b5aad346335bcc/aiohttp-3.14.3-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:498c6c623134f8e09a3c4e60bcd607a0b4590dd7dbf08dd40851b27cbb520ccb", size = 1844251, upload-time = "2026-07-23T01:57:16.593Z" }, + { url = "https://files.pythonhosted.org/packages/5e/19/08d41839658bdd44a0ed2480f3891705ecb487ce28c0dde62c9040c997e0/aiohttp-3.14.3-cp314-cp314t-win32.whl", hash = "sha256:b304db572b4368edd8dda8a2274f73156fe15558fca4a917cb8a09fc47af5963", size = 474180, upload-time = "2026-07-23T01:57:19.306Z" }, + { url = "https://files.pythonhosted.org/packages/99/5d/3cd6ef0a2b2851f7ab913b5b079334781bd50ff56a323e4454063377a080/aiohttp-3.14.3-cp314-cp314t-win_amd64.whl", hash = "sha256:b20032766aedf6261c7a566585a40867d092ac03a0d81592d5370ef9b054f99b", size = 500528, upload-time = "2026-07-23T01:57:21.762Z" }, + { url = "https://files.pythonhosted.org/packages/a4/37/cfd1ed540a4d318da025590d96b728e63713c09e9377950fc655dadeb856/aiohttp-3.14.3-cp314-cp314t-win_arm64.whl", hash = "sha256:2e1161602f45a54de2ce0905243a95f58cb42dcd378402f3697f5e0b21e9d2e7", size = 469280, upload-time = "2026-07-23T01:57:24.241Z" }, ] [[package]] @@ -714,30 +702,30 @@ wheels = [ [[package]] name = "boto3" -version = "1.37.3" +version = "1.43.106" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "botocore" }, { name = "jmespath" }, { name = "s3transfer" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/7e/3f/135ec0771e6d0e1af2ad7023a15df6677d96112072838d948c9b5075efe1/boto3-1.37.3.tar.gz", hash = "sha256:21f3ce0ef111297e63a6eb998a25197b8c10982970c320d4c6e8db08be2157be", size = 111160, upload-time = "2025-02-27T20:28:15.588Z" } +sdist = { url = "https://files.pythonhosted.org/packages/49/01/97aaee4d3e94467983a0c1b986ed4f4da48960d7ebc948e7d739c818cb59/boto3-1.43.106.tar.gz", hash = "sha256:c11ad4c429a983493ba10014c7af9831a455c2c0eea91c1cefff74530e480277", size = 112730, upload-time = "2026-09-30T19:36:41.684Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/62/8c/213511a505af2239a673de4de145d013379275c569185187922f93dbdf14/boto3-1.37.3-py3-none-any.whl", hash = "sha256:2063b40af99fd02f6228ff52397b552ff3353831edaf8d25cc04801827ab9794", size = 139344, upload-time = "2025-02-27T20:28:13.085Z" }, + { url = "https://files.pythonhosted.org/packages/72/d8/a01764ffd32fbcd66c5575eb85764ae9ea7b7d03dc591592e132c43af1f5/boto3-1.43.106-py3-none-any.whl", hash = "sha256:fdf2e304f9e8864b8613591b18b36ee995560c9666033622eda1cb9057866880", size = 140045, upload-time = "2026-09-30T19:36:40.228Z" }, ] [[package]] name = "botocore" -version = "1.37.3" +version = "1.43.106" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "jmespath" }, { name = "python-dateutil" }, { name = "urllib3" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/74/fb/b243ab806d2e1e6b8a475b731cc59a1f1e4709eded4884b988a27bbc996b/botocore-1.37.3.tar.gz", hash = "sha256:fe8403eb55a88faf9b0f9da6615e5bee7be056d75e17af66c3c8f0a3b0648da4", size = 13574648, upload-time = "2025-02-27T20:27:59.559Z" } +sdist = { url = "https://files.pythonhosted.org/packages/11/b9/10ca68d0092895d5ea60f485a61a9840d5aff9d732c66ed60da53a20b1d4/botocore-1.43.106.tar.gz", hash = "sha256:006870b3b4e40547232ad12c3bb4faec91bbbe0659aafaa3b7fa48a112c4ee97", size = 16271399, upload-time = "2026-09-30T19:36:35.836Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/88/54/772118f15b5990173aa5264946cc8c9ff70c8f02d72ee6d63167a985188c/botocore-1.37.3-py3-none-any.whl", hash = "sha256:d01bd3bf4c80e61fa88d636ad9f5c9f60a551d71549b481386c6b4efe0bb2b2e", size = 13342066, upload-time = "2025-02-27T20:27:53.137Z" }, + { url = "https://files.pythonhosted.org/packages/a8/eb/c51d3b54dfaa9267a687359d946dc2a1ca6a27fb4d381773e59b7bd72bf4/botocore-1.43.106-py3-none-any.whl", hash = "sha256:c1fb8818f9957cb5037f73db34ac1a12ba43562e31d6e7f9cb7e5fc64e1dba78", size = 15965506, upload-time = "2026-09-30T19:36:32.625Z" }, ] [[package]] @@ -3921,7 +3909,7 @@ aws = [ { name = "s3fs" }, ] aws-messaging = [ - { name = "aioboto3" }, + { name = "aiobotocore" }, ] azure = [ { name = "adlfs" }, @@ -3956,9 +3944,12 @@ websockets = [ [package.dev-dependencies] all = [ + { name = "aiobotocore" }, { name = "aiofile" }, + { name = "aiokafka" }, { name = "aiosqlite" }, { name = "coverage" }, + { name = "google-cloud-pubsub" }, { name = "hatch" }, { name = "hatch-vcs" }, { name = "ipython" }, @@ -4015,9 +4006,12 @@ docs = [ { name = "mkdocstrings", extra = ["python"] }, ] test = [ + { name = "aiobotocore" }, { name = "aiofile" }, + { name = "aiokafka" }, { name = "aiosqlite" }, { name = "coverage" }, + { name = "google-cloud-pubsub" }, { name = "llama-index-core" }, { name = "llama-index-llms-openai" }, { name = "moto", extra = ["server"] }, @@ -4040,7 +4034,7 @@ test = [ requires-dist = [ { name = "adlfs", marker = "extra == 'azure'", specifier = ">=2024.7.0" }, { name = "aio-pika", specifier = ">=9.5.5,<11" }, - { name = "aioboto3", marker = "extra == 'aws-messaging'", specifier = ">=13.0,<15" }, + { name = "aiobotocore", marker = "extra == 'aws-messaging'", specifier = ">=3.4,<4" }, { name = "aiokafka", marker = "extra == 'kafka'", specifier = ">=0.11,<1" }, { name = "aiosqlite", specifier = ">=0.20,<1" }, { name = "annotated-types", specifier = ">=0.7,<1" }, @@ -4077,9 +4071,12 @@ provides-extras = ["aws", "aws-messaging", "azure", "gcp", "gcp-pubsub", "kafka" [package.metadata.requires-dev] all = [ + { name = "aiobotocore", specifier = ">=3.4,<4" }, { name = "aiofile", specifier = ">=3.9,<4" }, + { name = "aiokafka", specifier = ">=0.11,<1" }, { name = "aiosqlite", specifier = ">=0.20,<1" }, { name = "coverage", specifier = ">=7.7.1,<8" }, + { name = "google-cloud-pubsub", specifier = ">=2.25,<3" }, { name = "hatch", specifier = ">=1.14.1" }, { name = "hatch-vcs", specifier = ">=0.5.0" }, { name = "ipython", specifier = ">=8.26,<10" }, @@ -4136,9 +4133,12 @@ docs = [ { name = "mkdocstrings", extras = ["python"], specifier = ">=0.25,<2" }, ] test = [ + { name = "aiobotocore", specifier = ">=3.4,<4" }, { name = "aiofile", specifier = ">=3.9,<4" }, + { name = "aiokafka", specifier = ">=0.11,<1" }, { name = "aiosqlite", specifier = ">=0.20,<1" }, { name = "coverage", specifier = ">=7.7.1,<8" }, + { name = "google-cloud-pubsub", specifier = ">=2.25,<3" }, { name = "llama-index-core", specifier = ">=0.12.30,<1" }, { name = "llama-index-llms-openai", specifier = ">=0.3.33,<1" }, { name = "moto", extras = ["server"], specifier = ">=5.0,<6" }, @@ -5336,14 +5336,14 @@ wheels = [ [[package]] name = "s3transfer" -version = "0.11.3" +version = "0.19.2" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "botocore" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/39/24/1390172471d569e281fcfd29b92f2f73774e95972c965d14b6c802ff2352/s3transfer-0.11.3.tar.gz", hash = "sha256:edae4977e3a122445660c7c114bba949f9d191bae3b34a096f18a1c8c354527a", size = 148042, upload-time = "2025-02-26T20:44:57.459Z" } +sdist = { url = "https://files.pythonhosted.org/packages/76/43/35e4d8aa320bffe8287fe8f65f578fa2d2db0a64212f0e710dce58267854/s3transfer-0.19.2.tar.gz", hash = "sha256:ba0309fd86be3c27dbf78cdd813c13c5e1df16e5874b99d2535ebbdfb9892993", size = 165592, upload-time = "2026-07-22T19:30:44.432Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/e4/81/48c41b554a54d75d4407740abb60e3a102ae416284df04d1dbdcbe3dbf24/s3transfer-0.11.3-py3-none-any.whl", hash = "sha256:ca855bdeb885174b5ffa95b9913622459d4ad8e331fc98eb01e6d5eb6a30655d", size = 84246, upload-time = "2025-02-26T20:44:55.509Z" }, + { url = "https://files.pythonhosted.org/packages/bc/e7/5c595c75e9f41a44f30e526eda465ea0b4eec93470e074e4a111b253f13a/s3transfer-0.19.2-py3-none-any.whl", hash = "sha256:d8168eccca828cbb2cd573675333f3bddd254313a9c42494b84c76b539e8ba25", size = 90216, upload-time = "2026-07-22T19:30:43.251Z" }, ] [[package]] From 6ac647392292f71d1cacb40395682c5ee9bcabea Mon Sep 17 00:00:00 2001 From: Chris Knight Date: Sun, 4 Oct 2026 23:19:27 +0200 Subject: [PATCH 10/11] chore: keep ZMQ slow-joiner workaround out of this PR The connection delay and the flaky marker on the Ray + ZMQ integration test fix an unrelated ZMQ connector issue; they now move with it to their own change. --- plugboard/connector/zmq_channel.py | 4 +--- tests/integration/test_process_with_components_run.py | 1 - 2 files changed, 1 insertion(+), 4 deletions(-) diff --git a/plugboard/connector/zmq_channel.py b/plugboard/connector/zmq_channel.py index f20dae0d..4e773214 100644 --- a/plugboard/connector/zmq_channel.py +++ b/plugboard/connector/zmq_channel.py @@ -357,9 +357,7 @@ async def connect_recv(self) -> ZMQChannel: self._recv_channel = ZMQChannel( recv_socket=recv_socket, topic=self._topic, maxsize=self._maxsize ) - # Allow extra time for the proxy subprocess's SUB socket subscription to propagate - # to XPUB before the sender starts publishing (ZMQ "slow joiner" problem). - await asyncio.sleep(0.5) + await asyncio.sleep(0.1) # Ensure connections established before first send. Better way? return self._recv_channel diff --git a/tests/integration/test_process_with_components_run.py b/tests/integration/test_process_with_components_run.py index 3f7737e1..2bb5e0de 100644 --- a/tests/integration/test_process_with_components_run.py +++ b/tests/integration/test_process_with_components_run.py @@ -84,7 +84,6 @@ def tempfile_path() -> _t.Generator[Path, None, None]: @pytest.mark.asyncio -@pytest.mark.flaky(reruns=3) # Flaky on Github Actions with Ray + ZMQ proxy (slow joiner) @pytest.mark.parametrize( "process_cls, connector_cls", [ From 3af63fee490cc8d3fe44ed88e257523ea97eab98 Mon Sep 17 00:00:00 2001 From: Chris Knight Date: Sun, 4 Oct 2026 23:35:55 +0200 Subject: [PATCH 11/11] test: cover disconnect failure during reconnect and pending-send failure at teardown Closes the two remaining uncovered branches from the error-handling paths the review asked about: a reconnect whose disconnect raises must still connect and let the retry proceed, and a send still in flight when destroy runs must be reported rather than swallowed. plugboard/utils/retry.py is now at 100% branch coverage. --- tests/unit/test_message_data_reader.py | 33 ++++++++++++++++++++++++++ tests/unit/test_message_data_writer.py | 16 +++++++++++++ 2 files changed, 49 insertions(+) diff --git a/tests/unit/test_message_data_reader.py b/tests/unit/test_message_data_reader.py index 3ea34ce0..27befb9a 100644 --- a/tests/unit/test_message_data_reader.py +++ b/tests/unit/test_message_data_reader.py @@ -73,6 +73,14 @@ async def _connect(self) -> None: raise OSError("broker unreachable") +class UndisconnectableReader(ScriptedReader): + """A reader whose disconnect always fails, as it does when the link is already gone.""" + + async def _disconnect(self) -> None: + self.disconnects += 1 + raise OSError("connection already lost") + + MESSAGES = [ {"x": 1, "y": "a"}, {"x": 2, "y": "b"}, @@ -370,6 +378,31 @@ async def test_permanent_error_is_not_retried() -> None: await reader.destroy() +async def test_failing_disconnect_still_reconnects() -> None: + """Tests that a broken connection is not mistaken for a reason to give up reconnecting. + + The disconnect during a reconnect is expected to fail when the link already dropped, + so it must not prevent the new connection from being made. + """ + reader = UndisconnectableReader( + name="test-reader", + field_names=["x", "y"], + topic="test-topic", + idle_poll_delay=0.0, + script=[ConnectionError("blip"), MESSAGES[:1], NoMoreDataException()], + retry_policy=RetryPolicy(max_retries=3), + ) + await reader.init() + + assert await drain(reader) == MESSAGES[:1] + assert reader.connects == 2 # Initial connect plus one reconnect. + assert reader.disconnects == 1 + + # A disconnect that fails during teardown is surfaced rather than hidden. + with pytest.raises(OSError, match="connection already lost"): + await reader.destroy() + + async def test_reconnect_failure_does_not_abort_retries( backoff_delays: list[float], ) -> None: diff --git a/tests/unit/test_message_data_writer.py b/tests/unit/test_message_data_writer.py index d072ab38..a1926e10 100644 --- a/tests/unit/test_message_data_writer.py +++ b/tests/unit/test_message_data_writer.py @@ -147,6 +147,22 @@ async def test_destroy_disconnects() -> None: assert writer.disconnects == 1 +async def test_destroy_reports_failed_pending_send( + capsys: pytest.CaptureFixture[str], +) -> None: + """Tests that a send which fails at teardown is reported, not swallowed.""" + writer = make_writer(["x"], script=[ConnectionError("lost")], max_retries=0) + channels = await connect_writer(writer) + await writer.init() + await write_records(writer, channels, SINGLE) + await writer._save_chunk() + + await writer.destroy() # Must not raise: teardown still has to disconnect. + + assert "Pending send failed during destroy" in capsys.readouterr().out + assert writer.disconnects == 1 + + async def test_destroy_waits_for_in_flight_send() -> None: """Tests that a send still running at teardown completes instead of being dropped.""" writer = make_writer(["x"])