From d74d30a47ae2bea298b133c0a7b3cc567d6f0c5e Mon Sep 17 00:00:00 2001 From: Chris Knight Date: Sat, 13 Jun 2026 21:35:28 +0200 Subject: [PATCH 01/21] 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/21] 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/21] 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/21] 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/21] 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/21] 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/21] 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/21] 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/21] 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/21] 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/21] 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"]) From 1be92fad17a079efcbea0b2d468a21ed4e610a86 Mon Sep 17 00:00:00 2001 From: Chris Knight Date: Sun, 18 Oct 2026 17:05:29 +0200 Subject: [PATCH 12/21] fix: split the Kafka implementation out into its own change Kafka is not part of the current testing round: there is no broker available to integration-test against, and whether to run one (managed MSK/Confluent versus a container) is still an open decision. Everything Kafka-specific moves to `feat/kafka-message-data`, which stacks on this branch: - `plugboard/library/kafka_io.py` and `tests/unit/test_kafka_io.py` - the `kafka` extra and `aiokafka` in the test dependency group (dropped from the lockfile, so `async-timeout` goes too) - `_KafkaSettings` and the `Settings.kafka` field - the Kafka exports, doc rows, install extra, settings table row and example The message data base classes, the PubSub and SQS/SNS implementations, the shared retry/backoff helper and the docs structure all stay: none of them depended on Kafka. --- README.md | 2 +- docs/usage/configuration.md | 1 - docs/usage/message-data.md | 22 +- plugboard/library/__init__.py | 3 - plugboard/library/kafka_io.py | 297 --------------------------- plugboard/utils/settings.py | 14 -- pyproject.toml | 6 +- tests/unit/test_kafka_io.py | 372 ---------------------------------- uv.lock | 50 +---- 9 files changed, 12 insertions(+), 755 deletions(-) delete mode 100644 plugboard/library/kafka_io.py delete mode 100644 tests/unit/test_kafka_io.py diff --git a/README.md b/README.md index 1efbf135..dc82f7ba 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, `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/). +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]` or `plugboard[aws-messaging]` 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 5430c2eb..a8e18ed0 100644 --- a/docs/usage/configuration.md +++ b/docs/usage/configuration.md @@ -39,7 +39,6 @@ evironment below when they are not supplied: |---------------------------|----------------------------------------|---------------| | `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. diff --git a/docs/usage/message-data.md b/docs/usage/message-data.md index 039db62a..da7a2171 100644 --- a/docs/usage/message-data.md +++ b/docs/usage/message-data.md @@ -12,10 +12,8 @@ component's outputs (for a reader) or inputs (for a writer). | [`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 +All four build on the [`MessageDataReader`][plugboard.library.MessageDataReader] and [`MessageDataWriter`][plugboard.library.MessageDataWriter] base classes, which handle connecting, reconnection with backoff, batching and acknowledgment. @@ -26,7 +24,6 @@ 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 @@ -102,8 +99,7 @@ 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. +"acknowledge" means deleting the messages. **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 @@ -115,15 +111,14 @@ deleted topic — fail immediately instead of burning the retries. Both the clas and the policy are handled by the base classes: ```python -from plugboard.library import KafkaDataReader +from plugboard.library import GCPPubSubDataReader from plugboard.utils.retry import RetryPolicy -reader = KafkaDataReader( +reader = GCPPubSubDataReader( name="reader", - topic="measurements", - group_id="plugboard", + subscription_id="measurements-pull", + project_id="plugboard-dev", 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 ) @@ -143,8 +138,7 @@ arguments always win. |---------------------------|------------------------------------------|---------| | `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). +normal provider mechanisms (Application Default Credentials for PubSub, and the standard +AWS credential chain for SQS/SNS). diff --git a/plugboard/library/__init__.py b/plugboard/library/__init__.py index c4bd2ac2..22999240 100644 --- a/plugboard/library/__init__.py +++ b/plugboard/library/__init__.py @@ -5,7 +5,6 @@ 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 @@ -22,8 +21,6 @@ "FileWriter", "GCPPubSubDataReader", "GCPPubSubDataWriter", - "KafkaDataReader", - "KafkaDataWriter", "LLMChat", "LLMImageProcessor", "MessageDataReader", diff --git a/plugboard/library/kafka_io.py b/plugboard/library/kafka_io.py deleted file mode 100644 index 133a34b3..00000000 --- a/plugboard/library/kafka_io.py +++ /dev/null @@ -1,297 +0,0 @@ -"""Provides `KafkaDataReader` and `KafkaDataWriter` for Apache Kafka messaging.""" - -from __future__ import annotations - -import asyncio -from collections import deque -import json -import typing as _t - -from plugboard.exceptions import ( - MessageBrokerConnectionError, - MessageBrokerPermanentError, - MessageBrokerTransientError, -) -from plugboard.library.message_reader import MessageDataReader, MessageDataReaderArgsDict -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 - # With the extra not installed, `depends_on_optional` blocks construction and - # classification is moot: `isinstance(error, ())` never matches. - _PERMANENT_ERRORS = () - _TRANSIENT_ERRORS = () - - -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. - - Returns: - A `MessageBrokerPermanentError` for a rejection that will not change on retry, - a `MessageBrokerTransientError` for a failure worth retrying, otherwise the - original error. - """ - 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. 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, - group_id: str, - bootstrap_servers: _t.Optional[str | list[str]] = None, - parse_json: bool = True, - poll_timeout_ms: int = 30_000, - **kwargs: _t.Unpack[MessageDataReaderArgsDict], - ) -> None: - """Instantiates the `KafkaDataReader`. - - Args: - 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) - 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.""" - 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 consumer.start() - self._consumer = consumer - - async def _disconnect(self) -> None: - """Stops and closes the Kafka consumer.""" - if self._consumer is not None: - 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, empty if the poll timed out. - - Raises: - MessageBrokerError: If the poll failed for a retryable or permanent reason. - """ - if self._consumer is None: - raise MessageBrokerConnectionError("Kafka consumer is not connected") - max_messages = self._chunk_size or 10 - 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: - 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 data: - 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 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 that have been processed. - - Raises: - MessageBrokerConnectionError: If the consumer is not connected. - MessageBrokerError: If the commit failed. - """ - if self._consumer is None: - 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 - 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: _t.Optional[str | list[str]] = None, - parse_json: bool = True, - **kwargs: _t.Unpack[MessageDataWriterArgsDict], - ) -> None: - """Instantiates the `KafkaDataWriter`. - - Args: - 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) - 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.""" - 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: - 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 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. - - Args: - data: A dictionary mapping field names to deques of field values. - - Returns: - A list of bytes objects ready to send. - """ - return encode_records_bytes(data, self._parse_json) diff --git a/plugboard/utils/settings.py b/plugboard/utils/settings.py index 93c6c683..42d0cb49 100644 --- a/plugboard/utils/settings.py +++ b/plugboard/utils/settings.py @@ -81,18 +81,6 @@ class _AWSSettings(BaseSettings): 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. @@ -107,7 +95,6 @@ class Settings(BaseSettings): 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) @@ -121,7 +108,6 @@ class Settings(BaseSettings): 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) _ArgumentT = _t.TypeVar("_ArgumentT") diff --git a/pyproject.toml b/pyproject.toml index b0ad7b72..bff3ec42 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -47,7 +47,6 @@ aws-messaging = ["aiobotocore>=3.4,<4"] 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", @@ -78,10 +77,9 @@ 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. + # 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", diff --git a/tests/unit/test_kafka_io.py b/tests/unit/test_kafka_io.py deleted file mode 100644 index 425ab396..00000000 --- a/tests/unit/test_kafka_io.py +++ /dev/null @@ -1,372 +0,0 @@ -"""Unit tests for the Apache Kafka message data implementations.""" - -from __future__ import annotations - -from collections import deque -import typing as _t -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 ( - 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 - - -@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_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, - ) - - -# --------------------------------------------------------------------------- -# Reader: connection and consuming -# --------------------------------------------------------------------------- - - -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() - - 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, - ) - consumer.start.assert_awaited_once_with() - - -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() - - 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)] - - -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. - - 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() - assert await reader._receive() == [] - - -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"])} - - -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"])} - - -# --------------------------------------------------------------------------- -# Reader: offset commits -# --------------------------------------------------------------------------- - - -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"{}"), - ] - - await reader._ack(records) - - consumer.commit.assert_awaited_once_with( - offsets={TopicPartition("test-topic", 0): OffsetAndMetadata(5, "")} - ) - - -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"{}"), - ] - - await reader._ack(records) - - 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() - - -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([]) - - -# --------------------------------------------------------------------------- -# Reader: error mapping -# --------------------------------------------------------------------------- - - -@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() - - -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() - - with pytest.raises(ValueError, match="unexpected"): - await reader._receive() - - -# --------------------------------------------------------------------------- -# Writer -# --------------------------------------------------------------------------- - - -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() - - -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() - - -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._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() - - -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"}'] - - -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"] - - -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() - - with pytest.raises(MessageBrokerTransientError): - await writer._send([b"one"]) - - -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"]) - - -# --------------------------------------------------------------------------- -# Settings resolution -# --------------------------------------------------------------------------- - - -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"] - - -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", - group_id="test-group", - field_names=["x"], - bootstrap_servers=["a:9092", "b: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/uv.lock b/uv.lock index 93c2b18c..0c1a25de 100644 --- a/uv.lock +++ b/uv.lock @@ -210,37 +210,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/10/a1/510b0a7fadc6f43a6ce50152e69dbd86415240835868bb0bd9b5b88b1e06/aioitertools-0.13.0-py3-none-any.whl", hash = "sha256:0be0292b856f08dfac90e31f4739432f4cb6d7520ab9eb73e143f4f2fa5259be", size = 24182, upload-time = "2025-11-06T22:17:06.502Z" }, ] -[[package]] -name = "aiokafka" -version = "0.14.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "async-timeout" }, - { name = "packaging" }, - { name = "typing-extensions" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/89/5f/dfc1180fd22d1acdc91949ec36e97199c43742dacb057cb8efed3679ed04/aiokafka-0.14.0.tar.gz", hash = "sha256:8ffdc945798ba4d3d132b705d4244d0a1f493925efb57c637a2ca88ee82794e1", size = 601374, upload-time = "2026-04-29T10:43:03.574Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/f9/9d/3441db94829f9feb802a2f4052df61c0d1a01272accd174c351d7e9e1f6a/aiokafka-0.14.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:284a90d617584d7e42688a181aaa8c2a909d9c658ab9b69c6cf92f4df5c4b320", size = 348458, upload-time = "2026-04-29T10:42:37.243Z" }, - { url = "https://files.pythonhosted.org/packages/a4/10/7297589aac95654596af13301b31da2c9502c80e7e308530ee7a9bd5b9f1/aiokafka-0.14.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:b4f211d9e03a1fc83871a37eefcf307bc0943ee99adae25aa39bd1722e70747b", size = 351057, upload-time = "2026-04-29T10:42:38.69Z" }, - { url = "https://files.pythonhosted.org/packages/26/4e/5c0aa8db717fff0ffb8f3e16deece8f98ded6ca17c6a543b6b20cc9a7f84/aiokafka-0.14.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:be517b9b9513eba43ba19961dd770a6e26d08325743093feb47182770d235dd9", size = 1142238, upload-time = "2026-04-29T10:42:39.96Z" }, - { url = "https://files.pythonhosted.org/packages/88/78/322f797b9593a4cc8afd647342fa66b9ad732ee55098e5e084188c6202aa/aiokafka-0.14.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:219d2dc66b97b1aaea100697c928024b6a0348b7baa370b824900054bf86916e", size = 1131567, upload-time = "2026-04-29T10:42:41.542Z" }, - { url = "https://files.pythonhosted.org/packages/8d/0a/a45320778385142299a7fc3ae402152ec1f383537130b8aa8e8587742fad/aiokafka-0.14.0-cp312-cp312-win32.whl", hash = "sha256:1086b470f6c452471603a2d9c8d6933739230c75758d777d8d113ff8112bad68", size = 312160, upload-time = "2026-04-29T10:42:42.811Z" }, - { url = "https://files.pythonhosted.org/packages/a3/fb/7802a0ed69200e3e8e8791df06bd6daf9b00523839d045662de4ff061b18/aiokafka-0.14.0-cp312-cp312-win_amd64.whl", hash = "sha256:bcf3a8f6592d73f45965ca0750bfdfccf2555c8625358175c92f75f2cce1261a", size = 331897, upload-time = "2026-04-29T10:42:43.984Z" }, - { url = "https://files.pythonhosted.org/packages/30/b0/c9384541b2e4cc52a16402fc53fb9d44af0d78d37954cf8c7271c376ad47/aiokafka-0.14.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:db16e43fac4c1c5006131046c1bf370c580d6ac4495a10ac7778245710943179", size = 345859, upload-time = "2026-04-29T10:42:45.449Z" }, - { url = "https://files.pythonhosted.org/packages/7f/d1/fc266d9f4ffba4f197356c6ffdfbb0fe32e7cb874e240f299935d058ac06/aiokafka-0.14.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:32a8e91d88cf3ccf0778927715610d6579888c5f4748db4c2022cda25d628a48", size = 348284, upload-time = "2026-04-29T10:42:47.104Z" }, - { url = "https://files.pythonhosted.org/packages/b2/8e/0c4c270786dac79f3fca74c6166c3a25b61b0d26132be0d69f0d7f206f0a/aiokafka-0.14.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:aad4a575a506e7784e25e430f27026fe2f4378560b21b7f4e8c9a54f0d06eaee", size = 1117867, upload-time = "2026-04-29T10:42:48.394Z" }, - { url = "https://files.pythonhosted.org/packages/9d/7f/3b89fbd0a3be9edfd5b51e20bb5cd695c851219b63c501c051cf84367fa9/aiokafka-0.14.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:75e4a003502c9c3b5c705fa7c00d634ba146bf38fa5d525b80bb6ff6e3e779fe", size = 1108860, upload-time = "2026-04-29T10:42:50.249Z" }, - { url = "https://files.pythonhosted.org/packages/b3/59/849aba75cff93277bf6bf8b630de79e902949ff7ec48e4b12a64e6e32cae/aiokafka-0.14.0-cp313-cp313-win32.whl", hash = "sha256:a128e213cbc2bce0ea3db65a68920e52cebeeb8209bf001ac7aa022a8bd54d7d", size = 310889, upload-time = "2026-04-29T10:42:52.038Z" }, - { url = "https://files.pythonhosted.org/packages/c4/e5/52eab8f8515d23da7b5d90e2c5ba10eab9494a0314f749e3f73e003f4a50/aiokafka-0.14.0-cp313-cp313-win_amd64.whl", hash = "sha256:d6fa16bef3544be87bd1a7a8317b9d85e3da59f3202326d9ff22735ed052746e", size = 329470, upload-time = "2026-04-29T10:42:53.536Z" }, - { url = "https://files.pythonhosted.org/packages/50/9d/984803315fe2b883ea6e08b1d9c8a752bd5c16e966d8714bacc67c72c417/aiokafka-0.14.0-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:5d70615d1530ad19d0c4da8d87abaec0a12b9fdaabffdcd4e400efa0c50ef80c", size = 346672, upload-time = "2026-04-29T10:42:55.267Z" }, - { url = "https://files.pythonhosted.org/packages/49/df/da314966b7f3c3117bd78b082563cb03dbe3007848cb8f4b0932faf390a0/aiokafka-0.14.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:7e2392360c370b1ba6564c57d2889e154ecdb43157a8f7b7d7afe5e3c02fcc1a", size = 349594, upload-time = "2026-04-29T10:42:56.565Z" }, - { url = "https://files.pythonhosted.org/packages/57/7a/160516944ea0e0f68ea78e38f944c52f5248c7c7df26cba22a40b9f25709/aiokafka-0.14.0-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:201e38ecc595f9f65a945f1ef9085157ddf28f25cd2e482fd9efa1fcf4638213", size = 1114112, upload-time = "2026-04-29T10:42:57.869Z" }, - { url = "https://files.pythonhosted.org/packages/68/c4/9841118a2157e913e8ebfbc0a2b58f7b60f1f7202040c3e1df8925ed1184/aiokafka-0.14.0-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:1cd651e1f56571baae306fdd0b5509047ab9625797a24cd75902e139c5a20318", size = 1098571, upload-time = "2026-04-29T10:42:59.356Z" }, - { url = "https://files.pythonhosted.org/packages/d6/a1/0af8a37849a4108ae227f46c4c62f6beab31863cf66ba318fb73b0be5b26/aiokafka-0.14.0-cp314-cp314-win32.whl", hash = "sha256:128127eb96dab98150b636bb5f480c80e15f02f82a118eec206a521c8cf7cf7c", size = 314107, upload-time = "2026-04-29T10:43:01.111Z" }, - { url = "https://files.pythonhosted.org/packages/fa/18/fb46c65f758900c71d0f1c73b7802720f99cabcb1f4a11676573f9bc1b8f/aiokafka-0.14.0-cp314-cp314-win_amd64.whl", hash = "sha256:aa385039aa9b235359319bbdcf48c9c86a75d81c9c547d645056d00361238903", size = 333320, upload-time = "2026-04-29T10:43:02.424Z" }, -] - [[package]] name = "aiormq" version = "7.0.0" @@ -413,15 +382,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/e5/e2/c2e3abf398f80732e58b03be77bde9022550d221dd8781bf586bd4d97cc1/async_lru-2.3.0-py3-none-any.whl", hash = "sha256:eea27b01841909316f2cc739807acea1c623df2be8c5cfad7583286397bb8315", size = 8403, upload-time = "2026-03-19T01:04:30.883Z" }, ] -[[package]] -name = "async-timeout" -version = "5.0.1" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/a5/ae/136395dfbfe00dfc94da3f3e136d0b13f394cba8f4841120e34226265780/async_timeout-5.0.1.tar.gz", hash = "sha256:d9321a7a3d5a6a5e187e824d2fa0793ce379a202935782d555d6e9d2735677d3", size = 9274, upload-time = "2024-11-06T16:41:39.6Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/fe/ba/e2081de779ca30d473f21f5b30e0e737c438205440784c7dfc81efc2b029/async_timeout-5.0.1-py3-none-any.whl", hash = "sha256:39e3809566ff85354557ec2398b55e096c8364bacac9405a7a1fa429e77fe76c", size = 6233, upload-time = "2024-11-06T16:41:37.9Z" }, -] - [[package]] name = "asyncpg" version = "0.31.0" @@ -3920,9 +3880,6 @@ gcp = [ gcp-pubsub = [ { name = "google-cloud-pubsub" }, ] -kafka = [ - { name = "aiokafka" }, -] llm = [ { name = "llama-index-core" }, { name = "llama-index-llms-openai" }, @@ -3946,7 +3903,6 @@ websockets = [ all = [ { name = "aiobotocore" }, { name = "aiofile" }, - { name = "aiokafka" }, { name = "aiosqlite" }, { name = "coverage" }, { name = "google-cloud-pubsub" }, @@ -4008,7 +3964,6 @@ docs = [ test = [ { name = "aiobotocore" }, { name = "aiofile" }, - { name = "aiokafka" }, { name = "aiosqlite" }, { name = "coverage" }, { name = "google-cloud-pubsub" }, @@ -4035,7 +3990,6 @@ requires-dist = [ { name = "adlfs", marker = "extra == 'azure'", specifier = ">=2024.7.0" }, { name = "aio-pika", specifier = ">=9.5.5,<11" }, { 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" }, { name = "asyncpg", specifier = ">=0.29.0,<1" }, @@ -4067,13 +4021,12 @@ requires-dist = [ { name = "uvloop", marker = "sys_platform != 'win32'", specifier = ">=0.21.0,<1" }, { name = "websockets", marker = "extra == 'websockets'", specifier = ">=14.2,<18" }, ] -provides-extras = ["aws", "aws-messaging", "azure", "gcp", "gcp-pubsub", "kafka", "llm", "omq", "ray", "redis", "websockets"] +provides-extras = ["aws", "aws-messaging", "azure", "gcp", "gcp-pubsub", "llm", "omq", "ray", "redis", "websockets"] [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" }, @@ -4135,7 +4088,6 @@ docs = [ 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" }, From cc5bd585f41ca22e153d7dcadbbeedc1312b0037 Mon Sep 17 00:00:00 2001 From: Chris Knight Date: Sun, 18 Oct 2026 19:46:51 +0200 Subject: [PATCH 13/21] fix: SQS ack parameter, SNS envelope, and add broker integration tests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fix two bugs found by running against real infrastructure: - AWSSQSDataReader._ack used the wrong parameter name (DeleteMessageRequests instead of Entries) for delete_message_batch, causing every ack to fail with a validation error. - AWSSQSDataReader._convert did not unwrap the SNS notification envelope that SQS delivers when a queue is subscribed to an SNS topic. Add _decode_body to transparently extract the nested record. Add gated integration tests (tests/integration/test_message_data_brokers.py) that verify real SNS→SQS and PubSub topic→subscription round-trips. Tests are skipped unless PLUGBOARD_MESSAGE_BROKER_TESTS is set, and the CI workflow enables them in the test-integration job with the Terraform-provisioned credentials. --- .github/workflows/lint-test.yaml | 22 +++ docs/usage/message-data.md | 6 + plugboard/library/aws_messaging_io.py | 41 ++++- .../integration/test_message_data_brokers.py | 168 ++++++++++++++++++ tests/unit/test_aws_messaging_io.py | 23 ++- 5 files changed, 251 insertions(+), 9 deletions(-) create mode 100644 tests/integration/test_message_data_brokers.py diff --git a/.github/workflows/lint-test.yaml b/.github/workflows/lint-test.yaml index 257bac09..4f2dacf7 100644 --- a/.github/workflows/lint-test.yaml +++ b/.github/workflows/lint-test.yaml @@ -141,6 +141,13 @@ jobs: docker compose up -d sleep 10 # Wait for services to start + - name: Configure GCP Pub/Sub credentials + if: ${{ vars.PLUGBOARD_MESSAGE_BROKER_TESTS == 'true' }} + env: + GCP_SA_KEY_JSON: ${{ secrets.GCP_SA_KEY_JSON }} + run: | + echo "$GCP_SA_KEY_JSON" > /tmp/gcp-sa-key.json + - name: Run integration tests run: COVERAGE_FILE=.coverage.py${{ matrix.python_version }}.integration uv run coverage run -m pytest ./tests/integration/ -m "not tuner" env: @@ -148,6 +155,21 @@ jobs: REDIS_URL: redis://default:password@localhost:6379/ RAY_ENABLE_UV_RUN_RUNTIME_ENV: 0 PLUGBOARD_IO_READ_TIMEOUT: 5.0 + # Run the message broker integration tests (AWS SQS/SNS + GCP Pub/Sub). + PLUGBOARD_MESSAGE_BROKER_TESTS: ${{ vars.PLUGBOARD_MESSAGE_BROKER_TESTS }} + # AWS SQS/SNS identifiers. + AWS_REGION: ${{ vars.AWS_REGION }} + AWS_SQS_QUEUE_URL: ${{ vars.AWS_SQS_QUEUE_URL }} + AWS_SNS_TOPIC_ARN: ${{ vars.AWS_SNS_TOPIC_ARN }} + # GCP Pub/Sub identifiers. + GCP_PROJECT_ID: ${{ vars.GCP_PROJECT_ID }} + GCP_PUBSUB_PROJECT_ID: ${{ vars.GCP_PUBSUB_PROJECT_ID }} + GCP_TOPIC_ID: ${{ vars.GCP_TOPIC_ID }} + GCP_SUBSCRIPTION_ID: ${{ vars.GCP_SUBSCRIPTION_ID }} + # Broker credentials (AWS credential chain + GCP Application Default Credentials). + AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }} + AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }} + GOOGLE_APPLICATION_CREDENTIALS: /tmp/gcp-sa-key.json - name: Upload integration test coverage uses: actions/upload-artifact@v4 diff --git a/docs/usage/message-data.md b/docs/usage/message-data.md index da7a2171..fe24a8e0 100644 --- a/docs/usage/message-data.md +++ b/docs/usage/message-data.md @@ -105,6 +105,12 @@ mid-batch leaves the unread messages to be redelivered (at-least-once). For SQS, 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. +**SNS fan-out is transparent to the SQS reader.** When you subscribe an SQS queue to an +SNS topic (the usual `AWSSNSDataWriter` -> `AWSSQSDataReader` topology), SQS delivers the +SNS notification envelope rather than the raw payload. `AWSSQSDataReader` unwraps the +envelope automatically, so the reader still yields the published fields. A body that is +not a JSON object (a raw string, or a JSON scalar) is exposed under a `data` field. + **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 diff --git a/plugboard/library/aws_messaging_io.py b/plugboard/library/aws_messaging_io.py index 092104a4..22d6c448 100644 --- a/plugboard/library/aws_messaging_io.py +++ b/plugboard/library/aws_messaging_io.py @@ -80,6 +80,39 @@ ) +def _decode_body(body: str, *, parse_json: bool) -> dict[str, _t.Any]: + """Decodes an SQS message body into a record dict. + + Handles SNS fan-out: when an SNS topic is subscribed to the queue, SQS delivers the + SNS notification envelope rather than the raw payload, so the actual record is nested + under the envelope's ``Message`` field. A record that is not an object (a raw string + body, or a JSON scalar) is exposed under the ``data`` field. + + Args: + body: The raw SQS message body. + parse_json: Whether the record is JSON. + + Returns: + The decoded record as a dict. + """ + if not parse_json: + return {"data": body} + try: + record = json.loads(body) + except (json.JSONDecodeError, TypeError): + return {"data": body} + if isinstance(record, dict) and record.get("Type") == "Notification" and "Message" in record: + record = record["Message"] + if isinstance(record, str): + try: + record = json.loads(record) + except (json.JSONDecodeError, TypeError): + return {"data": record} + if not isinstance(record, dict): + return {"data": record} + return record + + 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): @@ -210,11 +243,7 @@ async def _convert(self, data: list[_t.Any]) -> dict[str, deque]: """ converted: dict[str, deque] = {field: deque() for field in self.io.outputs} for msg in data: - body = msg.get("Body", "") - if self._parse_json: - record = json.loads(body) - else: - record = {"data": body} + record = _decode_body(msg.get("Body", ""), parse_json=self._parse_json) for field in self.io.outputs: converted[field].append(record.get(field)) return converted @@ -243,7 +272,7 @@ async def _ack(self, messages: list[_t.Any]) -> None: continue try: response = await self._client.delete_message_batch( - QueueUrl=self._queue_url, DeleteMessageRequests=entries + QueueUrl=self._queue_url, Entries=entries ) except Exception as error: # noqa: BLE001 raise _wrap_error(error) from error diff --git a/tests/integration/test_message_data_brokers.py b/tests/integration/test_message_data_brokers.py new file mode 100644 index 00000000..451c8e59 --- /dev/null +++ b/tests/integration/test_message_data_brokers.py @@ -0,0 +1,168 @@ +"""Integration tests for the message broker `MessageDataReader`/`MessageDataWriter` implementations. + +These tests exercise the AWS SQS/SNS and GCP PubSub readers and writers against +real broker infrastructure. They are gated behind the ``PLUGBOARD_MESSAGE_BROKER_TESTS`` +environment variable (default: not run) so the normal CI suite never touches the brokers. +The flag is enabled in the ``test-integration`` job of ``.github/workflows/lint-test.yaml``, +which also injects the ``AWS_*``/``GCP_*`` credentials and resource identifiers provisioned +by ``plugboard-infra``. + +AWS uses the credential chain (``AWS_ACCESS_KEY_ID``/``AWS_SECRET_ACCESS_KEY``), and the GCP +PubSub client uses Application Default Credentials (``GOOGLE_APPLICATION_CREDENTIALS``). + +The write side is driven through each writer's own ``_convert``/``_send`` pair - the exact +broker write path - so these tests validate the broker wiring and field<->message encoding +round trip, while the generic buffering/batching logic is already covered by the unit tests. +""" + +from __future__ import annotations + +import asyncio +from collections import deque +import os +import typing as _t + +import pytest + +from plugboard.library import ( + AWSSNSDataWriter, + AWSSQSDataReader, + GCPPubSubDataReader, + GCPPubSubDataWriter, +) + + +# Gate the whole module behind an opt-in env var (default: skip). The message broker +# integration tests only run when CI has provisioned the real broker infra. +requires_broker = pytest.mark.skipif( + os.environ.get("PLUGBOARD_MESSAGE_BROKER_TESTS", "").lower() not in ("1", "true", "yes"), + reason="set PLUGBOARD_MESSAGE_BROKER_TESTS=1 to run message broker integration tests", +) + +FIELD_NAMES = ["x", "y"] +# Enough to verify batch delivery, but small so a broken broker fails fast. +RECORDS = [{"x": i, "y": f"value-{i}"} for i in range(1, 6)] + + +def _broker_env_keys() -> list[str]: + """The env vars each broker needs, so a missing one is a clear skip, not a crash.""" + aws = ("AWS_SQS_QUEUE_URL", "AWS_SNS_TOPIC_ARN", "AWS_REGION") + gcp = ("GCP_TOPIC_ID", "GCP_SUBSCRIPTION_ID", "GCP_PUBSUB_PROJECT_ID") + return list(aws) + list(gcp) + + +@pytest.fixture +def broker_env() -> _t.Iterator[None]: + """Skips if any broker identifier is missing, so a partial setup is a clear skip.""" + missing = [key for key in _broker_env_keys() if not os.environ.get(key)] + if missing: + pytest.skip(f"missing broker env vars: {', '.join(missing)}") + yield + + +@requires_broker +async def test_aws_sns_sqs_round_trip(broker_env: None) -> None: + """Writes records to the SNS topic and reads them back from the linked SQS queue.""" + queue_url = os.environ["AWS_SQS_QUEUE_URL"] + topic_arn = os.environ["AWS_SNS_TOPIC_ARN"] + region = os.environ["AWS_REGION"] + + writer = AWSSNSDataWriter( + name="aws-sns-writer", + field_names=FIELD_NAMES, + topic_arn=topic_arn, + region=region, + parse_json=True, + ) + reader = AWSSQSDataReader( + name="aws-sqs-reader", + field_names=FIELD_NAMES, + queue_url=queue_url, + region=region, + parse_json=True, + wait_time_seconds=2, + chunk_size=10, + ) + + await reader.init() + await writer.init() + try: + await write_records(writer, RECORDS) + finally: + await writer.destroy() + await reader.destroy() + + received = await read_records(reader, len(RECORDS)) + assert sorted(map(tuple, received)) == sorted(map(tuple, RECORDS)) + + +@requires_broker +async def test_gcp_pubsub_round_trip(broker_env: None) -> None: + """Writes records to the PubSub topic and reads them back from the linked subscription.""" + project_id = os.environ["GCP_PUBSUB_PROJECT_ID"] + topic_id = os.environ["GCP_TOPIC_ID"] + subscription_id = os.environ["GCP_SUBSCRIPTION_ID"] + + writer = GCPPubSubDataWriter( + name="gcp-pubsub-writer", + field_names=FIELD_NAMES, + topic_id=topic_id, + project_id=project_id, + parse_json=True, + ) + reader = GCPPubSubDataReader( + name="gcp-pubsub-reader", + field_names=FIELD_NAMES, + subscription_id=subscription_id, + project_id=project_id, + parse_json=True, + chunk_size=10, + pull_timeout=3, + ) + + await reader.init() + await writer.init() + try: + await write_records(writer, RECORDS) + finally: + await writer.destroy() + await reader.destroy() + + received = await read_records(reader, len(RECORDS)) + assert sorted(map(tuple, received)) == sorted(map(tuple, RECORDS)) + + +async def write_records( + writer: AWSSNSDataWriter | GCPPubSubDataWriter, + records: list[dict[str, _t.Any]], +) -> None: + """Publishes records through the writer's own convert + send path. + + Building the field buffer as one ``deque`` per field and calling ``_convert`` then + ``_send`` is exactly how the writer reaches the broker, so this validates the real + message encoding without depending on the Component input-IO machinery. + """ + buffer = {field: deque(record[field] for record in records) for field in writer.io.inputs} + messages = await writer._convert(buffer) + await writer._send(messages) + + +async def read_records( + reader: AWSSQSDataReader | GCPPubSubDataReader, + expected: int, + timeout: float = 60.0, +) -> list[dict[str, _t.Any]]: + """Steps the reader until ``expected`` records are collected or ``timeout`` elapses. + + The reader waits for the next message rather than closing on an empty poll, so an + absent message surfaces as a timeout here instead of a silent pass. + """ + collected: list[dict[str, _t.Any]] = [] + loop = asyncio.get_running_loop() + deadline = loop.time() + timeout + while len(collected) < expected: + if loop.time() > deadline: + pytest.fail(f"read {len(collected)}/{expected} records before timeout") + await reader.step() + collected.append({field: getattr(reader, field) for field in FIELD_NAMES}) + return collected diff --git a/tests/unit/test_aws_messaging_io.py b/tests/unit/test_aws_messaging_io.py index a145a28d..d1329f87 100644 --- a/tests/unit/test_aws_messaging_io.py +++ b/tests/unit/test_aws_messaging_io.py @@ -4,6 +4,7 @@ from collections import deque import contextlib +import json import typing as _t from unittest.mock import AsyncMock, MagicMock, call @@ -167,6 +168,22 @@ async def test_reader_non_json_body_uses_data_field() -> None: assert converted == {"data": deque(["hello"])} +async def test_reader_unwraps_sns_envelope() -> None: + """Tests that an SNS notification envelope delivers the nested record.""" + reader = make_reader() + await reader._connect() + envelope = json.dumps( + { + "Type": "Notification", + "MessageId": "abc-123", + "Message": json.dumps({"x": 1, "y": "a"}), + "TopicArn": "arn:aws:sns:us-east-1:123:topic", + } + ) + batch = [sqs_message(envelope, "h1")] + assert await reader._convert(batch) == {"x": deque([1]), "y": deque(["a"])} + + async def test_reader_deletes_processed_messages(aws_client: MagicMock) -> None: """Tests that acking deletes by receipt handle, in one batch.""" reader = make_reader() @@ -177,7 +194,7 @@ async def test_reader_deletes_processed_messages(aws_client: MagicMock) -> None: aws_client.delete_message_batch.assert_awaited_once_with( QueueUrl="https://sqs.test/1/messages", - DeleteMessageRequests=[ + Entries=[ {"Id": "0", "ReceiptHandle": "h1"}, {"Id": "1", "ReceiptHandle": "h2"}, ], @@ -193,8 +210,8 @@ async def test_reader_splits_deletes_into_batch_size_chunks(aws_client: MagicMoc 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" + assert [len(call.kwargs["Entries"]) for call in calls] == [10, 2] + assert calls[1].kwargs["Entries"][0]["ReceiptHandle"] == "h10" async def test_reader_skips_messages_without_receipt_handle(aws_client: MagicMock) -> None: From 19a55895a74ea14ae5c6ce7a0f1396bd1b6d1e2f Mon Sep 17 00:00:00 2001 From: Chris Knight Date: Sun, 18 Oct 2026 20:49:51 +0200 Subject: [PATCH 14/21] test: drain message broker infra before and after each test Before each test the queue/subscription is checked and, if non-empty, a warning is emitted and it is drained, then the clean state is asserted. After each test it is drained again so a failed run leaves no messages to interfere with the next one. - AWS: check via ApproximateNumberOfMessages, drain via purge + receive-and-delete (SQS is the sink for the SNS fan-out, so anything left on it is stray). - GCP: a single timed pull-and-ack drain both checks and empties the subscription (only pubsub.subscriptions.pull/acknowledge are needed). --- .../integration/test_message_data_brokers.py | 197 +++++++++++++++++- 1 file changed, 190 insertions(+), 7 deletions(-) diff --git a/tests/integration/test_message_data_brokers.py b/tests/integration/test_message_data_brokers.py index 451c8e59..238fbf82 100644 --- a/tests/integration/test_message_data_brokers.py +++ b/tests/integration/test_message_data_brokers.py @@ -13,15 +13,24 @@ The write side is driven through each writer's own ``_convert``/``_send`` pair - the exact broker write path - so these tests validate the broker wiring and field<->message encoding round trip, while the generic buffering/batching logic is already covered by the unit tests. + +Hygiene: before each test the queue/subscription is checked, and if it is not empty a +warning is emitted and it is drained; after each test it is drained again so no messages +from a failed run interfere with the next one. """ from __future__ import annotations import asyncio from collections import deque +import contextlib +import logging import os import typing as _t +import warnings +from aiobotocore.session import get_session as get_boto_session +from google.cloud import pubsub_v1 import pytest from plugboard.library import ( @@ -42,6 +51,19 @@ FIELD_NAMES = ["x", "y"] # Enough to verify batch delivery, but small so a broken broker fails fast. RECORDS = [{"x": i, "y": f"value-{i}"} for i in range(1, 6)] +EXPECTED = sorted((record["x"], record["y"]) for record in RECORDS) +_LOG = logging.getLogger(__name__) + +# The infra sets the SQS visibility timeout to 30s, so a message a concurrent test received +# but has not deleted is redelivered within ~30s and picked up by a later drain pass. The +# drain declares the queue empty only after several consecutive long-poll timeouts, which +# also absorbs staggered SNS->SQS fan-out delivery. +_DRAIN_TIMEOUT = 90.0 +_DRAIN_EMPTY_POLLS = 3 # consecutive empty polls before declaring the queue drained +_DRAIN_POLL_WAIT = 2 # seconds per SQS long poll (also a natural settle between polls) +_DEPTH_CHECK_RETRIES = 3 +_DEPTH_CHECK_INTERVAL = 2.0 +_READ_TIMEOUT = 60.0 def _broker_env_keys() -> list[str]: @@ -60,6 +82,151 @@ def broker_env() -> _t.Iterator[None]: yield +# --------------------------------------------------------------------------- +# AWS SQS helpers +# --------------------------------------------------------------------------- + + +async def _aws_queue_depth(queue_url: str, region: str) -> int: + """Reads the number of messages in the queue (visible + in-flight).""" + async with get_boto_session().create_client("sqs", region_name=region) as client: + resp = await client.get_queue_attributes( + QueueUrl=queue_url, + AttributeNames=["ApproximateNumberOfMessages", "ApproximateNumberOfMessagesNotVisible"], + ) + attrs = resp["Attributes"] + return int(attrs["ApproximateNumberOfMessages"]) + int( + attrs["ApproximateNumberOfMessagesNotVisible"] + ) + + +async def _drain_aws_queue(queue_url: str, region: str, *, timeout: float = _DRAIN_TIMEOUT) -> int: + """Drains an SQS queue (purge, then receive-and-delete) and returns the count deleted. + + Declares the queue empty only after several consecutive long-poll timeouts, so staggered + SNS->SQS delivery (or messages still becoming visible) does not cause a premature stop. + """ + deleted = 0 + empty_polls = 0 + loop = asyncio.get_running_loop() + deadline = loop.time() + timeout + async with get_boto_session().create_client("sqs", region_name=region) as client: + with contextlib.suppress(Exception): + await client.purge_queue(QueueUrl=queue_url) + while loop.time() < deadline: + resp = await client.receive_message( + QueueUrl=queue_url, + MaxNumberOfMessages=10, + WaitTimeSeconds=_DRAIN_POLL_WAIT, + ) + messages = resp.get("Messages") or [] + if messages: + await client.delete_message_batch( + QueueUrl=queue_url, + Entries=[ + {"Id": str(i), "ReceiptHandle": msg["ReceiptHandle"]} + for i, msg in enumerate(messages) + ], + ) + deleted += len(messages) + empty_polls = 0 + elif empty_polls + 1 >= _DRAIN_EMPTY_POLLS: + return deleted + else: + empty_polls += 1 + return deleted + + +async def _assert_aws_queue_empty(queue_url: str, region: str) -> None: + """Retries the depth check to absorb SQS approximate-count eventual consistency.""" + depth = await _aws_queue_depth(queue_url, region) + for _ in range(_DEPTH_CHECK_RETRIES): + if depth == 0: + return + await asyncio.sleep(_DEPTH_CHECK_INTERVAL) + depth = await _aws_queue_depth(queue_url, region) + pytest.fail(f"SQS queue {queue_url} not empty after drain (depth={depth})") + + +# --------------------------------------------------------------------------- +# GCP Pub/Sub helpers +# --------------------------------------------------------------------------- + + +def _drain_gcp_subscription(project_id: str, subscription_id: str, *, timeout: float = 2.0) -> int: + """Pulls and acknowledges all available messages on the subscription; returns the count. + + Loops until a pull returns nothing. An empty subscription either returns no messages or + raises ``DeadlineExceeded`` once the pull window elapses; both are treated as "drained". + Only ``pubsub.subscriptions.pull``/``acknowledge`` are used (the subscriber role), so + no ``pubsub.subscriptions.get`` permission is needed. + """ + sub_path = f"projects/{project_id}/subscriptions/{subscription_id}" + subscriber = pubsub_v1.SubscriberClient() + acked = 0 + try: + while True: + try: + response = subscriber.pull( + request={"subscription": sub_path, "max_messages": 10}, timeout=timeout + ) + except Exception: # DeadlineExceeded -> no messages available right now. + return acked + ack_ids = [msg.ack_id for msg in response.received_messages] + if not ack_ids: + return acked + subscriber.acknowledge(request={"subscription": sub_path, "ack_ids": ack_ids}) + acked += len(ack_ids) + finally: + subscriber.close() + + +# --------------------------------------------------------------------------- +# Hygiene orchestrator +# --------------------------------------------------------------------------- + + +async def _clean_broker( + *, + pre_test: bool, + queue_url: _t.Optional[str] = None, + region: _t.Optional[str] = None, + project_id: _t.Optional[str] = None, + subscription_id: _t.Optional[str] = None, +) -> None: + """Ensures a broker is empty. + + In ``pre_test`` mode a non-empty broker emits a warning and is drained, then the empty + state is asserted (a leftover that cannot be drained fails the test). In post-test mode + it is drained unconditionally so a failed run leaves nothing for the next one. + """ + if queue_url and region: + depth = await _aws_queue_depth(queue_url, region) + if depth: + message = f"SQS queue {queue_url} was not empty ({depth} message(s)); draining" + _LOG.warning(message) if pre_test else _LOG.info(message) + if pre_test: + warnings.warn(message, stacklevel=2) + await _drain_aws_queue(queue_url, region) + if pre_test: + await _assert_aws_queue_empty(queue_url, region) + + if project_id and subscription_id: + drained = _drain_gcp_subscription(project_id, subscription_id) + if drained: + message = ( + f"PubSub subscription {subscription_id} had {drained} leftover message(s); drained" + ) + if pre_test: + warnings.warn(message, stacklevel=2) + (_LOG.warning if pre_test else _LOG.info)(message) + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + @requires_broker async def test_aws_sns_sqs_round_trip(broker_env: None) -> None: """Writes records to the SNS topic and reads them back from the linked SQS queue.""" @@ -67,6 +234,8 @@ async def test_aws_sns_sqs_round_trip(broker_env: None) -> None: topic_arn = os.environ["AWS_SNS_TOPIC_ARN"] region = os.environ["AWS_REGION"] + await _clean_broker(pre_test=True, queue_url=queue_url, region=region) + writer = AWSSNSDataWriter( name="aws-sns-writer", field_names=FIELD_NAMES, @@ -88,12 +257,16 @@ async def test_aws_sns_sqs_round_trip(broker_env: None) -> None: await writer.init() try: await write_records(writer, RECORDS) + received = await read_records(reader, len(RECORDS)) finally: await writer.destroy() await reader.destroy() - received = await read_records(reader, len(RECORDS)) - assert sorted(map(tuple, received)) == sorted(map(tuple, RECORDS)) + assert sorted((record["x"], record["y"]) for record in received) == EXPECTED + + # Reading acked every record, so the queue should be clean; drain + assert as a safety. + await _clean_broker(pre_test=False, queue_url=queue_url, region=region) + await _assert_aws_queue_empty(queue_url, region) @requires_broker @@ -103,6 +276,8 @@ async def test_gcp_pubsub_round_trip(broker_env: None) -> None: topic_id = os.environ["GCP_TOPIC_ID"] subscription_id = os.environ["GCP_SUBSCRIPTION_ID"] + await _clean_broker(pre_test=True, project_id=project_id, subscription_id=subscription_id) + writer = GCPPubSubDataWriter( name="gcp-pubsub-writer", field_names=FIELD_NAMES, @@ -124,12 +299,20 @@ async def test_gcp_pubsub_round_trip(broker_env: None) -> None: await writer.init() try: await write_records(writer, RECORDS) + received = await read_records(reader, len(RECORDS)) finally: await writer.destroy() await reader.destroy() - received = await read_records(reader, len(RECORDS)) - assert sorted(map(tuple, received)) == sorted(map(tuple, RECORDS)) + assert sorted((record["x"], record["y"]) for record in received) == EXPECTED + + # Reading acked every record, so the subscription should be clean; drain as a safety. + await _clean_broker(pre_test=False, project_id=project_id, subscription_id=subscription_id) + + +# --------------------------------------------------------------------------- +# Driving helpers +# --------------------------------------------------------------------------- async def write_records( @@ -150,12 +333,12 @@ async def write_records( async def read_records( reader: AWSSQSDataReader | GCPPubSubDataReader, expected: int, - timeout: float = 60.0, + timeout: float = _READ_TIMEOUT, ) -> list[dict[str, _t.Any]]: """Steps the reader until ``expected`` records are collected or ``timeout`` elapses. - The reader waits for the next message rather than closing on an empty poll, so an - absent message surfaces as a timeout here instead of a silent pass. + A reader waits for the next message rather than closing on an empty poll, so an absent + message surfaces as a timeout here instead of a silent pass. """ collected: list[dict[str, _t.Any]] = [] loop = asyncio.get_running_loop() From 127c6157e3a13f196b584ac6afa8828af9cfa2e8 Mon Sep 17 00:00:00 2001 From: Chris Knight Date: Sun, 18 Oct 2026 23:09:38 +0200 Subject: [PATCH 15/21] fix: address PR #284 re-review findings Blocking: - B1: MessageDataReader.step now returns after io.close() (not raise), so an exhausted reader completes instead of being marked FAILED. - B2: GCPPubSubDataReader._receive returns [] on a DeadlineExceeded idle pull instead of treating it as a broker error. - B3: RetryPolicy.run now reconnects at the start of each retry attempt and skips that attempt when the reconnect fails, so the operation never runs on a disconnected client; on exhaustion the last real broker error is raised rather than a synthetic "not connected" error. - B4: Removed unnecessary _connect() calls from pure _convert unit tests so they never reach a real cloud client. - B5: Serialised the test-integration job with a concurrency group and gated PLUGBOARD_MESSAGE_BROKER_TESTS to the 3.12 matrix leg so the shared SQS queue and GCP subscription are not raced by concurrent jobs. Warnings: - W1: MessageDataReader._ack_pending retries acknowledgments through the policy. - W2: gcp_pubsub_io._wrap_error maps NotFound to NoMoreDataException only when the reader is exhausted, else a permanent error. - W3: GCP _decode_payload falls back to a "data" field for non-UTF-8, non-JSON, and JSON scalar/array payloads. - W4: MessageDataReader/WriterArgsDict now inherit the DataReader/DataWriter specs; topic is NotRequired with an empty default. - W5: Moved retry execution onto RetryPolicy.run and dropped the module-level helpers that exceeded the parameter limits. - W6: resolve_argument takes keyword-only names and rejects empty resolutions. - W7: Fixed the Raises sections on _fetch/_save and removed unimplemented MessageBrokerConnectionError claims from _connect. - W8: Export RetryPolicy from plugboard.utils and plugboard.library; removed the internal resolve_argument from the public exports. - W9: Assert the broker is empty right after reading and drain it in finally. - W10: Drive the integration writer through its public Component lifecycle (AsyncioConnector inputs + step + io.close + run) and assert region / subscription path via recorded client calls instead of private attributes. - W11: mock_s3_server now scopes its env vars with a MonkeyPatch context and keeps its temp file in a pytest temp dir. - W12: Added unit tests for the _decode_body / _decode_payload fallbacks and the _ack / _disconnect error-wrapping branches. - W13: Fixed the docs example to use AsyncioConnector-wrapped connectors, async with process, a transform that uses both inputs, and the project env var. - W14: Condensed the message-data guide to user-facing bullets (including the writer chunk_size default) and linked the config table instead of duplicating. Low: - L1: parse_json is keyword-only in the shared encoders and broker constructors. - L2: Extracted _parse_json_or_raw to de-nest _decode_body. - L3: Removed the never-read project/subscription/topic attributes on GCP. - L4: Trimmed the remaining rationale-narrating docstrings. - L5: Shared the backoff_delays fixture in conftest so the writer retry tests run without real sleeps. - L6: _drain_gcp_subscription catches DeadlineExceeded specifically. - L7: Extracted a _log_leftover helper in the integration test. - L8: Fixed the "evironment" typo and the /latest/ docs link in the README. - L9: Removed the section-banner comments from the new test modules. - L10: Deleted the superseded local proposal doc and its .gitignore entry. --- .github/workflows/lint-test.yaml | 10 +- .gitignore | 1 - README.md | 2 +- docs/usage/configuration.md | 2 +- docs/usage/message-data.md | 105 +++++------ plugboard/library/__init__.py | 2 + plugboard/library/aws_messaging_io.py | 28 +-- plugboard/library/gcp_pubsub_io.py | 59 ++++-- plugboard/library/message_reader.py | 58 +++--- plugboard/library/message_writer.py | 42 ++--- plugboard/utils/__init__.py | 5 +- plugboard/utils/retry.py | 140 +++++++------- plugboard/utils/settings.py | 3 +- tests/conftest.py | 19 ++ .../integration/test_message_data_brokers.py | 171 +++++++++++------- tests/unit/test_aws_messaging_io.py | 113 +++++++----- tests/unit/test_file_reader_writer.py | 30 +-- tests/unit/test_gcp_pubsub_io.py | 107 ++++++++--- tests/unit/test_message_data_reader.py | 115 +++++++----- tests/unit/test_message_data_writer.py | 52 +++--- 20 files changed, 622 insertions(+), 442 deletions(-) diff --git a/.github/workflows/lint-test.yaml b/.github/workflows/lint-test.yaml index 4f2dacf7..f2647270 100644 --- a/.github/workflows/lint-test.yaml +++ b/.github/workflows/lint-test.yaml @@ -115,6 +115,12 @@ jobs: name: Tests - integration runs-on: ubuntu-latest timeout-minutes: 10 + # The message broker integration tests publish to a single shared SQS queue and + # GCP subscription, so concurrent runs must not collide on them. Serializing the + # job across runs keeps one set of broker tests at a time. + concurrency: + group: test-integration + cancel-in-progress: false strategy: matrix: python_version: [3.12, 3.13, 3.14] @@ -156,7 +162,9 @@ jobs: RAY_ENABLE_UV_RUN_RUNTIME_ENV: 0 PLUGBOARD_IO_READ_TIMEOUT: 5.0 # Run the message broker integration tests (AWS SQS/SNS + GCP Pub/Sub). - PLUGBOARD_MESSAGE_BROKER_TESTS: ${{ vars.PLUGBOARD_MESSAGE_BROKER_TESTS }} + # Gated to a single Python leg: the tests share one SQS queue and GCP + # subscription, so the matrix legs must not run them concurrently. + PLUGBOARD_MESSAGE_BROKER_TESTS: ${{ vars.PLUGBOARD_MESSAGE_BROKER_TESTS == 'true' && matrix.python_version == 3.12 && 'true' || 'false' }} # AWS SQS/SNS identifiers. AWS_REGION: ${{ vars.AWS_REGION }} AWS_SQS_QUEUE_URL: ${{ vars.AWS_SQS_QUEUE_URL }} diff --git a/.gitignore b/.gitignore index 8b7a80bc..9107d228 100644 --- a/.gitignore +++ b/.gitignore @@ -170,4 +170,3 @@ cython_debug/ # Generated version files */_version.py -docs/message-data-reader-writer-proposal.md diff --git a/README.md b/README.md index dc82f7ba..8f9b933c 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, `plugboard[websockets]` for WebSocket I/O, and `plugboard[gcp-pubsub]` or `plugboard[aws-messaging]` for [message data components](https://docs.plugboard.dev/usage/message-data/). +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]` or `plugboard[aws-messaging]` for [message data components](https://docs.plugboard.dev/latest/usage/message-data/). ## ⚡ Quickstart with AI diff --git a/docs/usage/configuration.md b/docs/usage/configuration.md index a8e18ed0..1acd473d 100644 --- a/docs/usage/configuration.md +++ b/docs/usage/configuration.md @@ -33,7 +33,7 @@ through an external broker as part of the model itself, see 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: +environment below when they are not supplied: | Option Name | Description | Default Value | |---------------------------|----------------------------------------|---------------| diff --git a/docs/usage/message-data.md b/docs/usage/message-data.md index fe24a8e0..bcdde1a9 100644 --- a/docs/usage/message-data.md +++ b/docs/usage/message-data.md @@ -38,6 +38,7 @@ Messages are expected to carry a JSON object per record, whose keys match import asyncio from plugboard.component import Component, IOController as IO +from plugboard.connector import AsyncioConnector from plugboard.library import AWSSNSDataWriter, GCPPubSubDataReader from plugboard.process import LocalProcess from plugboard.schemas import ConnectorSpec @@ -49,14 +50,14 @@ class Bucket(Component): io = IO(inputs=["x", "y"], outputs=["bucket"]) async def step(self) -> None: - self.bucket = "high" if self.y > 50 else "low" + self.bucket = "high" if self.x + 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 + # project_id is optional; it defaults to the GCP_PUBSUB_PROJECT_ID variable field_names=["x", "y"], chunk_size=20, ) @@ -73,12 +74,13 @@ async def main() -> None: 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"), + AsyncioConnector(spec=ConnectorSpec(source="reader.x", target="bucket.x")), + AsyncioConnector(spec=ConnectorSpec(source="reader.y", target="bucket.y")), + AsyncioConnector(spec=ConnectorSpec(source="bucket.bucket", target="writer.bucket")), ], ) - await process.run() + async with process: + await process.run() asyncio.run(main()) @@ -89,62 +91,41 @@ 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. - -**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. - -**SNS fan-out is transparent to the SQS reader.** When you subscribe an SQS queue to an -SNS topic (the usual `AWSSNSDataWriter` -> `AWSSQSDataReader` topology), SQS delivers the -SNS notification envelope rather than the raw payload. `AWSSQSDataReader` unwraps the -envelope automatically, so the reader still yields the published fields. A body that is -not a JSON object (a raw string, or a JSON scalar) is exposed under a `data` field. - -**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 GCPPubSubDataReader -from plugboard.utils.retry import RetryPolicy - -reader = GCPPubSubDataReader( - name="reader", - subscription_id="measurements-pull", - project_id="plugboard-dev", - field_names=["x", "y"], - 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. +- **A reader waits; it does not finish.** An empty poll means "nothing yet", not "no + more data" (unlike a file). A reader started before its producer keeps running and + picks up messages as they arrive; to end a run that would otherwise wait forever, stop + the process. A reader only stops on its own when the broker reports the source as gone + (for example a deleted subscription or queue), at which point it closes its stream and + the run ends. +- **Messages are acknowledged after they are processed.** A batch is acknowledged only + once every record in it has been published downstream, so a crash mid-batch leaves the + rest to be redelivered (at-least-once). For SQS, "acknowledge" means deleting the + messages. +- **Writers batch by `chunk_size`.** A writer flushes to the broker when its input + buffer reaches `chunk_size` records, and flushes whatever is left over when the run + ends. `chunk_size` defaults to `None`, so a writer with no `chunk_size` buffers every + record in memory and publishes them all at the end rather than streaming during the + run — set `chunk_size` (for example `chunk_size=1` for live per-record publishing) to + stream. For SQS the effective size is capped at 10 messages per send, the broker + limit. +- **Failures are retried with backoff.** Transient broker errors trigger a reconnect and + retry; permanent errors (such as `AccessDenied`) fail immediately. Override the policy + with the `retry_policy` argument, e.g. + `retry_policy=RetryPolicy(max_retries=5, base_delay=0.5, max_delay=30.0)`; an empty + poll pauses for `idle_poll_delay` seconds. +- **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. +- **SNS fan-out is transparent.** Subscribing an SQS queue to an SNS topic delivers the + notification envelope, which `AWSSQSDataReader` unwraps automatically; a non-object + body is exposed under `data`. ## 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` | - -Credentials themselves are never configured through Plugboard: the clients use their -normal provider mechanisms (Application Default Credentials for PubSub, and the standard -AWS credential chain for SQS/SNS). +Connection details are passed explicitly or read from the environment (explicit arguments +always win), so the same model can target a different account or region without code +changes. The message-data components read the `GCP_PUBSUB_PROJECT_ID` and `AWS_REGION` +environment variables; see [Configuration](configuration.md#message-data-brokers) for +how these are resolved. Credentials are never configured through Plugboard — the clients +use their normal provider mechanisms (Application Default Credentials for PubSub, and the +standard AWS credential chain for SQS/SNS). diff --git a/plugboard/library/__init__.py b/plugboard/library/__init__.py index 22999240..d674c903 100644 --- a/plugboard/library/__init__.py +++ b/plugboard/library/__init__.py @@ -1,5 +1,6 @@ """Provides implementations of Plugboard objects for use in user models.""" +from plugboard.utils.retry import RetryPolicy from .aws_messaging_io import AWSSNSDataWriter, AWSSQSDataReader from .data_reader import DataReader from .data_writer import DataWriter @@ -25,6 +26,7 @@ "LLMImageProcessor", "MessageDataReader", "MessageDataWriter", + "RetryPolicy", "SQLReader", "SQLWriter", "WebsocketBase", diff --git a/plugboard/library/aws_messaging_io.py b/plugboard/library/aws_messaging_io.py index 22d6c448..ae73e2a2 100644 --- a/plugboard/library/aws_messaging_io.py +++ b/plugboard/library/aws_messaging_io.py @@ -80,6 +80,14 @@ ) +def _parse_json_or_raw(text: str) -> _t.Any: + """Parses `text` as JSON, returning the raw text unchanged if it is not valid JSON.""" + try: + return json.loads(text) + except (json.JSONDecodeError, TypeError): + return text + + def _decode_body(body: str, *, parse_json: bool) -> dict[str, _t.Any]: """Decodes an SQS message body into a record dict. @@ -97,17 +105,11 @@ def _decode_body(body: str, *, parse_json: bool) -> dict[str, _t.Any]: """ if not parse_json: return {"data": body} - try: - record = json.loads(body) - except (json.JSONDecodeError, TypeError): - return {"data": body} + record = _parse_json_or_raw(body) if isinstance(record, dict) and record.get("Type") == "Notification" and "Message" in record: - record = record["Message"] - if isinstance(record, str): - try: - record = json.loads(record) - except (json.JSONDecodeError, TypeError): - return {"data": record} + message = record["Message"] + if isinstance(message, str): + record = _parse_json_or_raw(message) if not isinstance(record, dict): return {"data": record} return record @@ -177,7 +179,7 @@ def __init__( [`MessageDataReader`][plugboard.library.MessageDataReader]. """ region = resolve_argument( - region, DI.settings.resolve_sync().aws.region, "region", "AWS_REGION" + region, DI.settings.resolve_sync().aws.region, arg_name="region", env_var="AWS_REGION" ) kwargs.setdefault("topic", queue_url) super().__init__(**kwargs) @@ -310,7 +312,7 @@ def __init__( [`MessageDataWriter`][plugboard.library.MessageDataWriter]. """ region = resolve_argument( - region, DI.settings.resolve_sync().aws.region, "region", "AWS_REGION" + region, DI.settings.resolve_sync().aws.region, arg_name="region", env_var="AWS_REGION" ) kwargs.setdefault("topic", topic_arn) super().__init__(**kwargs) @@ -362,4 +364,4 @@ async def _convert(self, data: dict[str, deque]) -> list[_t.Any]: Returns: A list of message strings ready to publish. """ - return encode_records(data, self._parse_json) + return encode_records(data, parse_json=self._parse_json) diff --git a/plugboard/library/gcp_pubsub_io.py b/plugboard/library/gcp_pubsub_io.py index 1dcdb795..1b62e13a 100644 --- a/plugboard/library/gcp_pubsub_io.py +++ b/plugboard/library/gcp_pubsub_io.py @@ -56,20 +56,23 @@ _NOT_FOUND_ERRORS = () -def _wrap_error(error: Exception) -> Exception: +def _wrap_error(error: Exception, *, exhausted: bool = False) -> Exception: """Maps a PubSub SDK error onto the Plugboard message broker exception hierarchy. Args: error: The error raised by the PubSub client. + exhausted: Whether a missing topic/subscription 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 topic/subscription, a + A `NoMoreDataException` for a deleted topic/subscription (readers only), a `MessageBrokerPermanentError` for a rejection that will not change on retry, a `MessageBrokerTransientError` for a failure worth retrying, otherwise the original error. """ if isinstance(error, _NOT_FOUND_ERRORS): - return NoMoreDataException() + return NoMoreDataException() if exhausted else MessageBrokerPermanentError(str(error)) if isinstance(error, _PERMANENT_ERRORS): return MessageBrokerPermanentError(str(error)) if isinstance(error, _TRANSIENT_ERRORS): @@ -77,6 +80,30 @@ def _wrap_error(error: Exception) -> Exception: return error +def _decode_payload(payload: bytes, *, parse_json: bool) -> dict[str, _t.Any]: + """Decodes a PubSub message payload into a record dict. + + A payload that is not valid UTF-8, is not JSON, or is a JSON scalar or array cannot + be read as a record, so it is exposed under the ``data`` field rather than raising. + + Args: + payload: The raw PubSub message data. + parse_json: Whether the record is JSON. + + Returns: + The decoded record as a dict. + """ + if not parse_json: + return {"data": payload} + try: + record = json.loads(payload.decode("utf-8")) + except (UnicodeDecodeError, json.JSONDecodeError): + return {"data": payload} + if not isinstance(record, dict): + return {"data": record} + return record + + class GCPPubSubDataReader(MessageDataReader): """Reads data from a Google Cloud PubSub subscription. @@ -111,13 +138,11 @@ def __init__( project_id = resolve_argument( project_id, DI.settings.resolve_sync().gcp_pubsub.project_id, - "project_id", - "GCP_PUBSUB_PROJECT_ID", + arg_name="project_id", + env_var="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 @@ -154,8 +179,12 @@ async def _receive(self) -> list[_t.Any]: response = await asyncio.to_thread( self._subscriber.pull, request=request, timeout=self._pull_timeout ) + except gcp_exceptions.DeadlineExceeded: + # The pull window elapsed on an empty subscription: this is "nothing yet", + # not a broker error, so it must not trigger a reconnect or fail the reader. + return [] except Exception as error: # noqa: BLE001 - raise _wrap_error(error) from error + raise _wrap_error(error, exhausted=True) from error return list(response.received_messages) async def _convert(self, data: list[_t.Any]) -> dict[str, deque]: @@ -169,11 +198,7 @@ async def _convert(self, data: list[_t.Any]) -> dict[str, deque]: """ converted: dict[str, deque] = {field: deque() for field in self.io.outputs} for msg_wrapper in data: - payload = msg_wrapper.message.data - if self._parse_json: - record = json.loads(payload.decode("utf-8")) - else: - record = {"data": payload} + record = _decode_payload(msg_wrapper.message.data, parse_json=self._parse_json) for field in self.io.outputs: converted[field].append(record.get(field)) return converted @@ -229,13 +254,11 @@ def __init__( project_id = resolve_argument( project_id, DI.settings.resolve_sync().gcp_pubsub.project_id, - "project_id", - "GCP_PUBSUB_PROJECT_ID", + arg_name="project_id", + env_var="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 @@ -294,4 +317,4 @@ async def _convert(self, data: dict[str, deque]) -> list[_t.Any]: Returns: A list of bytes objects ready to publish. """ - return encode_records_bytes(data, self._parse_json) + return encode_records_bytes(data, parse_json=self._parse_json) diff --git a/plugboard/library/message_reader.py b/plugboard/library/message_reader.py index 1b22ec06..4cc2e74a 100644 --- a/plugboard/library/message_reader.py +++ b/plugboard/library/message_reader.py @@ -9,27 +9,25 @@ import typing as _t from plugboard.component import IOController -from plugboard.exceptions import IOStreamClosedError, NoMoreDataException -from plugboard.library.data_reader import DataReader +from plugboard.exceptions import NoMoreDataException +from plugboard.library.data_reader import DataReader, DataReaderArgsSpec from plugboard.schemas import ComponentArgsDict -from plugboard.utils.retry import RetryPolicy, attempt_reconnect, with_retry +from plugboard.utils.retry import RetryPolicy, attempt_reconnect -class MessageDataReaderArgsDict(ComponentArgsDict): +class MessageDataReaderArgsDict(DataReaderArgsSpec): """Specification of the `MessageDataReader` constructor arguments. + Adds the message-broker-specific arguments to the `DataReader` ones. + 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. 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: str - chunk_size: _t.NotRequired[int | None] + topic: _t.NotRequired[str] retry_policy: _t.NotRequired[RetryPolicy] idle_poll_delay: _t.NotRequired[float] @@ -57,7 +55,7 @@ class MessageDataReader(DataReader): def __init__( self, field_names: list[str], - topic: str, + topic: str = "", chunk_size: _t.Optional[int] = None, retry_policy: RetryPolicy = RetryPolicy(), idle_poll_delay: float = 0.1, @@ -83,11 +81,7 @@ def __init__( @abstractmethod async def _connect(self) -> None: - """Establishes connection to the message broker. - - Raises: - MessageBrokerConnectionError: If connection cannot be established. - """ + """Establishes connection to the message broker.""" pass @abstractmethod @@ -142,12 +136,12 @@ async def _fetch(self) -> list[_t.Any]: Raises: NoMoreDataException: If the source is exhausted. - MessageBrokerConnectionError: If all retries are exhausted. + MessageBrokerPermanentError: If the failure cannot be recovered by retrying. + Exception: The error from the final attempt, if retries are exhausted. """ - return await with_retry( + return await self._retry_policy.run( self._receive, self._reconnect, - policy=self._retry_policy, logger=self._logger, description="receiving messages", ) @@ -196,15 +190,26 @@ def _records_available(self) -> int: return min(len(values) for values in self._buffer.values()) async def _ack_pending(self) -> None: - """Acknowledges the batch whose records have all been consumed. + """Acknowledges the batch whose records have all been consumed, retrying transient failures. - Acknowledgment waits for the connection lock, so it never runs against a client - that a concurrent reconnect is replacing. + The acknowledgment and any reconnect run under the connection lock so neither can + use a client that a concurrent reconnect is replacing; the lock is held inside the + retried operation so a mid-retry reconnect does not deadlock on it. """ if not self._pending_ack: return - async with self._connection_lock: - await self._ack(self._pending_ack) + pending = self._pending_ack + + async def ack() -> None: + async with self._connection_lock: + await self._ack(pending) + + await self._retry_policy.run( + ack, + self._reconnect, + logger=self._logger, + description="acknowledging messages", + ) self._pending_ack = [] async def init(self) -> None: @@ -224,8 +229,9 @@ async def step(self) -> None: 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 the source is exhausted. + An exhausted source closes the IO stream and returns (rather than raising): the + closed stream ends the run through `io.write`, exactly as it does for + `DataReader`, so the component completes instead of being marked failed. """ while not self._records_available: try: @@ -233,7 +239,7 @@ async def step(self) -> None: except NoMoreDataException: await self._ack_pending() await self.io.close() - raise IOStreamClosedError("No more messages from broker") from None + return self._consume_record() if not self._records_available: await self._ack_pending() diff --git a/plugboard/library/message_writer.py b/plugboard/library/message_writer.py index 6309b497..6462e5ba 100644 --- a/plugboard/library/message_writer.py +++ b/plugboard/library/message_writer.py @@ -9,34 +9,28 @@ import typing as _t from plugboard.component import IOController -from plugboard.library.data_writer import DataWriter +from plugboard.library.data_writer import DataWriter, DataWriterArgsSpec from plugboard.schemas import ComponentArgsDict -from plugboard.utils.retry import RetryPolicy, attempt_reconnect, with_retry +from plugboard.utils.retry import RetryPolicy, attempt_reconnect -class MessageDataWriterArgsDict(ComponentArgsDict): +class MessageDataWriterArgsDict(DataWriterArgsSpec): """Specification of the `MessageDataWriter` constructor arguments. + Adds the message-broker-specific arguments to the `DataWriter` ones. + 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 each send. retry_policy: Optional; The backoff policy for retrying failed broker calls. """ - field_names: list[str] - topic: str - chunk_size: _t.NotRequired[int | None] + topic: _t.NotRequired[str] retry_policy: _t.NotRequired[RetryPolicy] 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. @@ -66,7 +60,7 @@ def _raw_bytes(record: dict[str, _t.Any]) -> bytes: 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]: +def encode_records(data: dict[str, deque], *, parse_json: bool) -> list[str]: """Encodes a field buffer as one text message per record. Args: @@ -82,7 +76,7 @@ def encode_records(data: dict[str, deque], parse_json: bool) -> list[str]: ] -def encode_records_bytes(data: dict[str, deque], parse_json: bool) -> list[bytes]: +def encode_records_bytes(data: dict[str, deque], *, parse_json: bool) -> list[bytes]: """Encodes a field buffer as one binary message per record. Args: @@ -116,7 +110,7 @@ class MessageDataWriter(DataWriter): def __init__( self, field_names: list[str], - topic: str, + topic: str = "", chunk_size: _t.Optional[int] = None, retry_policy: RetryPolicy = RetryPolicy(), **kwargs: _t.Unpack[ComponentArgsDict], @@ -136,11 +130,7 @@ def __init__( @abstractmethod async def _connect(self) -> None: - """Establishes connection to the message broker. - - Raises: - MessageBrokerConnectionError: If connection cannot be established. - """ + """Establishes connection to the message broker.""" pass @abstractmethod @@ -181,22 +171,18 @@ async def _save(self, data: list[_t.Any]) -> None: data: The converted messages to send, as produced by `_convert`. Raises: - MessageBrokerConnectionError: If all retries are exhausted. + MessageBrokerPermanentError: If the failure cannot be recovered by retrying. + Exception: The error from the final attempt, if retries are exhausted. """ - await with_retry( + await self._retry_policy.run( 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. - - No lock is needed here: unlike a reader, a writer never acknowledges - concurrently, and sends are serialised through `self._task`. - """ + """Re-establishes the broker connection, ignoring a failure to close the old one.""" await attempt_reconnect( self._connect, self._disconnect, diff --git a/plugboard/utils/__init__.py b/plugboard/utils/__init__.py index e8d398aa..7cb168f4 100644 --- a/plugboard/utils/__init__.py +++ b/plugboard/utils/__init__.py @@ -9,7 +9,8 @@ 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, resolve_argument +from plugboard.utils.retry import RetryPolicy +from plugboard.utils.settings import Settings __all__ = [ @@ -19,12 +20,12 @@ "gather_except", "gen_rand_str", "is_on_ray_worker", - "resolve_argument", "run_coro_sync", "ClassRegistry", "DI", "EntityIdGen", "Exportable", "ExportMixin", + "RetryPolicy", "Settings", ] diff --git a/plugboard/utils/retry.py b/plugboard/utils/retry.py index 112d9515..2d6929a2 100644 --- a/plugboard/utils/retry.py +++ b/plugboard/utils/retry.py @@ -41,6 +41,74 @@ def delay_for(self, attempt: int) -> float: """ return min(self.base_delay * (2**attempt), self.max_delay) + async def run( + self, + operation: Callable[[], Awaitable[_T]], + reconnect: Callable[[], Awaitable[None]], + *, + 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. Every retry attempt re-establishes the connection before it + runs, so the operation never relies on a connection that a failed attempt may have + torn down. If the connection cannot be re-established, that attempt is skipped and + the retries continue, so the connection is retried on every attempt. When all + attempts are used up the last error is raised, so the caller sees the broker's own + error (or the reconnect failure) rather than a synthetic "not connected" one. + + Args: + operation: The coroutine function to run. + reconnect: The coroutine function that re-establishes the connection. + 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 last error encountered, if every attempt failed. + """ + last_error: Exception | None = None + for attempt in range(self.max_retries + 1): + if attempt: + delay = self.delay_for(attempt - 1) + logger.warning( + f"Transient error {description}, retrying", + attempt=attempt, + delay=delay, + ) + await asyncio.sleep(delay) + try: + await reconnect() + except Exception as reconnect_error: # noqa: BLE001 + # Without a connection the operation would fail with a synthetic + # "not connected" error, so it is skipped: the retries continue and the + # connection is tried again next attempt. `last_error` keeps the prior + # broker error, so what is eventually raised is always a real one. + logger.warning( + f"Reconnection for {description} failed, skipping attempt", + attempt=attempt, + error=str(reconnect_error), + ) + continue + try: + return await operation() + except (NoMoreDataException, MessageBrokerPermanentError): + raise + except Exception as error: # noqa: BLE001 + last_error = error + # Unreachable: the range is never empty, so at least one attempt ran and either + # returned, raised a non-retryable error, or recorded the last error. + if last_error is None: # pragma: no cover + raise RuntimeError("retry loop ended without recording an error") + raise last_error + async def attempt_reconnect( connect: Callable[[], Awaitable[None]], @@ -68,75 +136,3 @@ async def attempt_reconnect( 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 42d0cb49..dfbea070 100644 --- a/plugboard/utils/settings.py +++ b/plugboard/utils/settings.py @@ -116,6 +116,7 @@ class Settings(BaseSettings): def resolve_argument( value: _t.Optional[_ArgumentT], fallback: _t.Optional[_ArgumentT], + *, arg_name: str, env_var: str, ) -> _ArgumentT: @@ -138,7 +139,7 @@ def resolve_argument( ValueError: If neither an explicit value nor a configured value is available. """ resolved = value if value else fallback - if resolved is None: + if not resolved: raise ValueError( f"'{arg_name}' is required: pass it explicitly or set the '{env_var}' " f"environment variable." diff --git a/tests/conftest.py b/tests/conftest.py index 3f4c63fa..66197b71 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -4,6 +4,7 @@ import asyncio from contextlib import contextmanager import multiprocessing +from types import SimpleNamespace import typing as _t from unittest.mock import patch @@ -17,10 +18,28 @@ from plugboard.component.io_controller import IOStreamClosedError from plugboard.connector import Connector from plugboard.schemas import Status +from plugboard.utils import retry as retry_utils from plugboard.utils.di import DI from plugboard.utils.settings import Settings +@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 + + @contextmanager def override_settings(settings: Settings) -> _t.Iterator[None]: """Temporarily override DI settings for a test and always reset the override.""" diff --git a/tests/integration/test_message_data_brokers.py b/tests/integration/test_message_data_brokers.py index 238fbf82..a0123cb5 100644 --- a/tests/integration/test_message_data_brokers.py +++ b/tests/integration/test_message_data_brokers.py @@ -22,7 +22,6 @@ from __future__ import annotations import asyncio -from collections import deque import contextlib import logging import os @@ -30,15 +29,18 @@ import warnings from aiobotocore.session import get_session as get_boto_session +from google.api_core import exceptions as gcp_exceptions from google.cloud import pubsub_v1 import pytest +from plugboard.connector import AsyncioConnector from plugboard.library import ( AWSSNSDataWriter, AWSSQSDataReader, GCPPubSubDataReader, GCPPubSubDataWriter, ) +from plugboard.schemas import ConnectorSpec # Gate the whole module behind an opt-in env var (default: skip). The message broker @@ -82,11 +84,6 @@ def broker_env() -> _t.Iterator[None]: yield -# --------------------------------------------------------------------------- -# AWS SQS helpers -# --------------------------------------------------------------------------- - - async def _aws_queue_depth(queue_url: str, region: str) -> int: """Reads the number of messages in the queue (visible + in-flight).""" async with get_boto_session().create_client("sqs", region_name=region) as client: @@ -148,42 +145,76 @@ async def _assert_aws_queue_empty(queue_url: str, region: str) -> None: pytest.fail(f"SQS queue {queue_url} not empty after drain (depth={depth})") -# --------------------------------------------------------------------------- -# GCP Pub/Sub helpers -# --------------------------------------------------------------------------- +def _gcp_sub_path(project_id: str, subscription_id: str) -> str: + """The fully-qualified PubSub subscription path.""" + return f"projects/{project_id}/subscriptions/{subscription_id}" -def _drain_gcp_subscription(project_id: str, subscription_id: str, *, timeout: float = 2.0) -> int: - """Pulls and acknowledges all available messages on the subscription; returns the count. +def _gcp_pull_one(project_id: str, subscription_id: str, *, timeout: float) -> list[object]: + """Pulls up to one batch of messages, returning an empty list if the window elapses. - Loops until a pull returns nothing. An empty subscription either returns no messages or - raises ``DeadlineExceeded`` once the pull window elapses; both are treated as "drained". - Only ``pubsub.subscriptions.pull``/``acknowledge`` are used (the subscriber role), so - no ``pubsub.subscriptions.get`` permission is needed. + Uses only ``pubsub.subscriptions.pull`` (the subscriber role). An empty subscription + raises ``DeadlineExceeded`` once the pull window elapses; that is reported as "none + available" rather than a broker error, while any other error propagates so an + auth/permission problem is not mistaken for an empty broker during setup. """ - sub_path = f"projects/{project_id}/subscriptions/{subscription_id}" + subscriber = pubsub_v1.SubscriberClient() + try: + response = subscriber.pull( + request={ + "subscription": _gcp_sub_path(project_id, subscription_id), + "max_messages": 10, + }, + timeout=timeout, + ) + return list(response.received_messages) + except gcp_exceptions.DeadlineExceeded: + return [] + finally: + subscriber.close() + + +def _gcp_subscription_undelivered(project_id: str, subscription_id: str) -> bool: + """Returns True if the subscription has any available (undelivered) messages.""" + return len(_gcp_pull_one(project_id, subscription_id, timeout=2.0)) > 0 + + +def _drain_gcp_subscription(project_id: str, subscription_id: str, *, timeout: float = 2.0) -> int: + """Pulls and acknowledges all available messages on the subscription; returns the count.""" subscriber = pubsub_v1.SubscriberClient() acked = 0 try: while True: - try: - response = subscriber.pull( - request={"subscription": sub_path, "max_messages": 10}, timeout=timeout - ) - except Exception: # DeadlineExceeded -> no messages available right now. - return acked + response = subscriber.pull( + request={ + "subscription": _gcp_sub_path(project_id, subscription_id), + "max_messages": 10, + }, + timeout=timeout, + ) ack_ids = [msg.ack_id for msg in response.received_messages] if not ack_ids: return acked - subscriber.acknowledge(request={"subscription": sub_path, "ack_ids": ack_ids}) + subscriber.acknowledge( + request={ + "subscription": _gcp_sub_path(project_id, subscription_id), + "ack_ids": ack_ids, + } + ) acked += len(ack_ids) + except gcp_exceptions.DeadlineExceeded: + return acked finally: subscriber.close() -# --------------------------------------------------------------------------- -# Hygiene orchestrator -# --------------------------------------------------------------------------- +def _log_leftover(message: str, *, pre_test: bool) -> None: + """Reports a leftover message, escalating to a warning (and a pytest warning) pre-test.""" + if pre_test: + warnings.warn(message, stacklevel=3) + _LOG.warning(message) + else: + _LOG.info(message) async def _clean_broker( @@ -203,10 +234,10 @@ async def _clean_broker( if queue_url and region: depth = await _aws_queue_depth(queue_url, region) if depth: - message = f"SQS queue {queue_url} was not empty ({depth} message(s)); draining" - _LOG.warning(message) if pre_test else _LOG.info(message) - if pre_test: - warnings.warn(message, stacklevel=2) + _log_leftover( + f"SQS queue {queue_url} was not empty ({depth} message(s)); draining", + pre_test=pre_test, + ) await _drain_aws_queue(queue_url, region) if pre_test: await _assert_aws_queue_empty(queue_url, region) @@ -214,17 +245,10 @@ async def _clean_broker( if project_id and subscription_id: drained = _drain_gcp_subscription(project_id, subscription_id) if drained: - message = ( - f"PubSub subscription {subscription_id} had {drained} leftover message(s); drained" + _log_leftover( + f"PubSub subscription {subscription_id} had {drained} leftover message(s); drained", + pre_test=pre_test, ) - if pre_test: - warnings.warn(message, stacklevel=2) - (_LOG.warning if pre_test else _LOG.info)(message) - - -# --------------------------------------------------------------------------- -# Tests -# --------------------------------------------------------------------------- @requires_broker @@ -253,20 +277,22 @@ async def test_aws_sns_sqs_round_trip(broker_env: None) -> None: chunk_size=10, ) + writer_channels = await connect_writer(writer) await reader.init() await writer.init() + received: list[dict[str, _t.Any]] = [] try: - await write_records(writer, RECORDS) + await write_records(writer, writer_channels, RECORDS) received = await read_records(reader, len(RECORDS)) + assert sorted((record["x"], record["y"]) for record in received) == EXPECTED + # The reader acknowledged every record, so nothing should remain on the queue. + await _assert_aws_queue_empty(queue_url, region) finally: await writer.destroy() await reader.destroy() - - assert sorted((record["x"], record["y"]) for record in received) == EXPECTED - - # Reading acked every record, so the queue should be clean; drain + assert as a safety. - await _clean_broker(pre_test=False, queue_url=queue_url, region=region) - await _assert_aws_queue_empty(queue_url, region) + # Safety net: if the read or assertion failed mid-way, drain any unacked + # leftovers so they do not interfere with the next run. + await _clean_broker(pre_test=False, queue_url=queue_url, region=region) @requires_broker @@ -295,39 +321,60 @@ async def test_gcp_pubsub_round_trip(broker_env: None) -> None: pull_timeout=3, ) + writer_channels = await connect_writer(writer) await reader.init() await writer.init() + received: list[dict[str, _t.Any]] = [] try: - await write_records(writer, RECORDS) + await write_records(writer, writer_channels, RECORDS) received = await read_records(reader, len(RECORDS)) + assert sorted((record["x"], record["y"]) for record in received) == EXPECTED finally: await writer.destroy() await reader.destroy() + # Clean up so nothing interferes with the next run, then confirm the + # subscription is genuinely empty (the reader should have acked everything). + await _clean_broker(pre_test=False, project_id=project_id, subscription_id=subscription_id) + assert _gcp_subscription_undelivered(project_id, subscription_id) == 0, ( + "PubSub subscription not empty after the round trip" + ) - assert sorted((record["x"], record["y"]) for record in received) == EXPECTED - - # Reading acked every record, so the subscription should be clean; drain as a safety. - await _clean_broker(pre_test=False, project_id=project_id, subscription_id=subscription_id) +async def connect_writer( + writer: AWSSNSDataWriter | GCPPubSubDataWriter, +) -> dict[str, AsyncioConnector]: + """Connects an input channel to each of the writer's fields. -# --------------------------------------------------------------------------- -# Driving helpers -# --------------------------------------------------------------------------- + Inputs must be connected before the writer is initialised, mirroring the unit tests. + """ + channels = { + field: AsyncioConnector( + spec=ConnectorSpec(source="none.none", target=f"{writer.name}.{field}") + ) + for field in writer.io.inputs + } + await writer.io.connect(list(channels.values())) + return channels async def write_records( writer: AWSSNSDataWriter | GCPPubSubDataWriter, + channels: dict[str, AsyncioConnector], records: list[dict[str, _t.Any]], ) -> None: - """Publishes records through the writer's own convert + send path. + """Publishes records through the writer's public `Component` lifecycle. - Building the field buffer as one ``deque`` per field and calling ``_convert`` then - ``_send`` is exactly how the writer reaches the broker, so this validates the real - message encoding without depending on the Component input-IO machinery. + Each record is sent to the connected input channels and the writer stepped; closing + and running the writer then flushes the remaining buffer to the broker, so the real + ``step``/``run``/``_save`` path (not the private broker methods) is exercised. """ - buffer = {field: deque(record[field] for record in records) for field in writer.io.inputs} - messages = await writer._convert(buffer) - await writer._send(messages) + 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() + await writer.io.close() + await writer.run() async def read_records( diff --git a/tests/unit/test_aws_messaging_io.py b/tests/unit/test_aws_messaging_io.py index d1329f87..8acecc09 100644 --- a/tests/unit/test_aws_messaging_io.py +++ b/tests/unit/test_aws_messaging_io.py @@ -91,11 +91,6 @@ def sqs_message(body: str, handle: str) -> dict[str, str]: return {"Body": body, "ReceiptHandle": handle, "MessageId": handle} -# --------------------------------------------------------------------------- -# Reader: connection and receiving -# --------------------------------------------------------------------------- - - 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() @@ -147,15 +142,9 @@ async def test_reader_empty_poll_returns_no_messages(aws_client: MagicMock) -> N assert await reader._receive() == [] -# --------------------------------------------------------------------------- -# Reader: conversion and acknowledgment -# --------------------------------------------------------------------------- - - 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"])} @@ -163,7 +152,6 @@ async def test_reader_converts_json_bodies() -> None: 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"])} @@ -171,7 +159,6 @@ async def test_reader_non_json_body_uses_data_field() -> None: async def test_reader_unwraps_sns_envelope() -> None: """Tests that an SNS notification envelope delivers the nested record.""" reader = make_reader() - await reader._connect() envelope = json.dumps( { "Type": "Notification", @@ -184,6 +171,37 @@ async def test_reader_unwraps_sns_envelope() -> None: assert await reader._convert(batch) == {"x": deque([1]), "y": deque(["a"])} +def test_decode_body_raw_string_is_exposed_under_data() -> None: + """Tests that a non-JSON body is exposed under `data` rather than raising.""" + decoded = aws_messaging_io._decode_body("just some text", parse_json=True) + assert decoded == {"data": "just some text"} + + +def test_decode_body_envelope_with_non_json_message_is_exposed_under_data() -> None: + """Tests that an SNS envelope whose Message is a plain string is unwrapped to `data`.""" + envelope = json.dumps({"Type": "Notification", "Message": "plain payload"}) + assert aws_messaging_io._decode_body(envelope, parse_json=True) == {"data": "plain payload"} + + +@pytest.mark.parametrize( + ("body", "value"), + [ + ("42", 42), + ('"hello"', "hello"), + ("true", True), + ("[1, 2, 3]", [1, 2, 3]), + ], +) +def test_decode_body_json_scalar_or_array_is_exposed_under_data(body: str, value: _t.Any) -> None: + """Tests that a non-object JSON payload is wrapped under `data`.""" + assert aws_messaging_io._decode_body(body, parse_json=True) == {"data": value} + + +def test_decode_body_without_parse_json_is_raw() -> None: + """Tests that `parse_json=False` exposes the raw body under `data`.""" + assert aws_messaging_io._decode_body('{"x": 1}', parse_json=False) == {"data": '{"x": 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() @@ -234,6 +252,23 @@ async def test_reader_reports_failed_deletes_as_transient(aws_client: MagicMock) await reader._ack([sqs_message('{"x": 1}', "h1")]) +async def test_reader_ack_without_connection_raises_connection_error(aws_client: MagicMock) -> None: + """Tests that acking before connecting fails loudly.""" + reader = make_reader() + with pytest.raises(MessageBrokerConnectionError): + await reader._ack([sqs_message('{"x": 1}', "h1")]) + + +async def test_reader_ack_wraps_transient_delete_errors(aws_client: MagicMock) -> None: + """Tests that a throttled delete is wrapped for retry rather than leaking the SDK error.""" + aws_client.delete_message_batch.side_effect = client_error("ThrottlingException") + reader = make_reader() + await reader._connect() + + with pytest.raises(MessageBrokerTransientError): + await reader._ack([sqs_message('{"x": 1}', "h1")]) + + async def test_reader_without_connection_raises_connection_error(aws_client: MagicMock) -> None: """Tests that receiving before connecting fails loudly.""" reader = make_reader() @@ -241,11 +276,6 @@ async def test_reader_without_connection_raises_connection_error(aws_client: Mag await reader._receive() -# --------------------------------------------------------------------------- -# Reader: error mapping -# --------------------------------------------------------------------------- - - @pytest.mark.parametrize( ("code", "expected"), [ @@ -288,11 +318,6 @@ async def test_reader_unclassified_error_propagates(aws_client: MagicMock) -> No await reader._receive() -# --------------------------------------------------------------------------- -# Writer -# --------------------------------------------------------------------------- - - 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() @@ -316,7 +341,6 @@ async def test_writer_publishes_each_message(aws_client: MagicMock) -> None: 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() converted = await writer._convert({"x": deque([1, 2]), "y": deque(["a", "b"])}) assert converted == ['{"x": 1, "y": "a"}', '{"x": 2, "y": "b"}'] @@ -324,7 +348,6 @@ async def test_writer_converts_fields_to_json_strings() -> 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"] @@ -353,22 +376,19 @@ async def test_writer_without_connection_raises_connection_error(aws_client: Mag await writer._send(["one"]) -# --------------------------------------------------------------------------- -# Settings resolution -# --------------------------------------------------------------------------- - - -async def test_region_falls_back_to_settings() -> None: - """Tests that `AWS_REGION` supplies the region when none is passed.""" +async def test_region_falls_back_to_settings(aws_client: MagicMock) -> None: + """Tests that `AWS_REGION` supplies the client's 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" + await reader._connect() + assert aws_client.created == [("sqs", {"region_name": "ap-south-1"})] + await reader._disconnect() -def test_explicit_region_wins_over_settings() -> None: +async def test_explicit_region_wins_over_settings(aws_client: MagicMock) -> 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): @@ -378,7 +398,9 @@ def test_explicit_region_wins_over_settings() -> None: field_names=["x"], region="us-west-2", ) - assert reader._region == "us-west-2" + await reader._connect() + assert aws_client.created == [("sqs", {"region_name": "us-west-2"})] + await reader._disconnect() def test_missing_region_names_the_environment_variable() -> None: @@ -393,13 +415,22 @@ def test_missing_region_names_the_environment_variable() -> None: ) -def test_reader_topic_defaults_to_queue_url() -> None: - """Tests that the queue URL identifies the reader's topic.""" +async def test_reader_receives_from_the_configured_queue(aws_client: MagicMock) -> None: + """Tests that the reader receives from the queue it was configured with.""" + aws_client.receive_message.return_value = {"Messages": []} reader = make_reader() - assert reader._topic == "https://sqs.test/1/messages" + await reader._connect() + await reader._receive() + aws_client.receive_message.assert_awaited_once() + assert aws_client.receive_message.await_args.kwargs["QueueUrl"] == "https://sqs.test/1/messages" + await reader._disconnect() -def test_writer_topic_defaults_to_topic_arn() -> None: - """Tests that the topic ARN identifies the writer's topic.""" +async def test_writer_publishes_to_the_configured_topic(aws_client: MagicMock) -> None: + """Tests that the writer publishes to the topic it was configured with.""" writer = make_writer() - assert writer._topic == "arn:aws:sns:eu-west-1:1:messages" + await writer._connect() + await writer._send(['{"x": 1}']) + aws_client.publish.assert_awaited_once() + assert aws_client.publish.await_args.kwargs["TopicArn"] == "arn:aws:sns:eu-west-1:1:messages" + await writer._disconnect() diff --git a/tests/unit/test_file_reader_writer.py b/tests/unit/test_file_reader_writer.py index 60617ec6..b8ebc9f0 100644 --- a/tests/unit/test_file_reader_writer.py +++ b/tests/unit/test_file_reader_writer.py @@ -22,28 +22,28 @@ @pytest.fixture(scope="module", autouse=True) -def mock_s3_server() -> _t.Generator[None, None, None]: +def mock_s3_server(tmp_path_factory: pytest.TempPathFactory) -> _t.Generator[None, None, None]: """Provides a mock S3 server for testing. Required for async tests until this is resolved: https://github.com/aio-libs/aiobotocore/issues/755 + + The AWS env vars are set on a `MonkeyPatch` context so they are restored on + teardown, and the empty config file lives in a pytest temp dir so it is cleaned up. """ 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: - os.environ["AWS_ACCESS_KEY_ID"] = "test-key-id" - if "AWS_ENDPOINT_URL" not in os.environ: - os.environ["AWS_ENDPOINT_URL"] = S3_URL - - yield + # 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 = tmp_path_factory.mktemp("aws") / "config" + empty_config.touch() + with pytest.MonkeyPatch.context() as monkeypatch: + monkeypatch.setenv("AWS_CONFIG_FILE", str(empty_config)) + monkeypatch.setenv("AWS_SHARED_CREDENTIALS_FILE", str(empty_config)) + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "test-access-key") # noqa: S105 + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "test-key-id") + monkeypatch.setenv("AWS_ENDPOINT_URL", S3_URL) + yield server.stop() diff --git a/tests/unit/test_gcp_pubsub_io.py b/tests/unit/test_gcp_pubsub_io.py index b037c954..bdacdd05 100644 --- a/tests/unit/test_gcp_pubsub_io.py +++ b/tests/unit/test_gcp_pubsub_io.py @@ -3,6 +3,7 @@ from __future__ import annotations from collections import deque +import json import threading from types import SimpleNamespace import typing as _t @@ -84,11 +85,6 @@ def make_writer(**kwargs: _t.Any) -> GCPPubSubDataWriter: ) -# --------------------------------------------------------------------------- -# Reader -# --------------------------------------------------------------------------- - - async def test_reader_connect_builds_client_off_the_event_loop( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -162,10 +158,9 @@ def _pull(request: dict[str, _t.Any], timeout: float) -> MagicMock: assert caller_threads[0] != threading.current_thread().name -async def test_reader_converts_json_messages(subscriber: MagicMock) -> None: +async def test_reader_converts_json_messages() -> None: """Tests that JSON payloads become field buffers keyed by the declared fields.""" reader = make_reader() - await reader._connect() 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"])} @@ -173,11 +168,31 @@ async def test_reader_converts_json_messages(subscriber: MagicMock) -> None: 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"])} +def test_decode_payload_non_utf8_is_exposed_under_data() -> None: + """Tests that a payload which is not valid UTF-8 is exposed under `data`.""" + assert gcp_pubsub_io._decode_payload(b"\xff\xfe", parse_json=True) == {"data": b"\xff\xfe"} + + +def test_decode_payload_non_json_is_exposed_under_data() -> None: + """Tests that a non-JSON payload is exposed under `data`.""" + assert gcp_pubsub_io._decode_payload(b"not json", parse_json=True) == {"data": b"not json"} + + +@pytest.mark.parametrize("payload", [b"42", b'"hello"', b"true", b"[1, 2]"]) +def test_decode_payload_json_scalar_or_array_is_exposed_under_data(payload: bytes) -> None: + """Tests that a non-object JSON payload is wrapped under `data`.""" + assert gcp_pubsub_io._decode_payload(payload, parse_json=True) == {"data": json.loads(payload)} + + +def test_decode_payload_without_parse_json_is_raw() -> None: + """Tests that `parse_json=False` exposes the raw payload under `data`.""" + assert gcp_pubsub_io._decode_payload(b'{"x": 1}', parse_json=False) == {"data": b'{"x": 1}'} + + async def test_reader_ack_sends_ack_ids(subscriber: MagicMock) -> None: """Tests that acknowledging a batch sends exactly its ack ids.""" reader = make_reader() @@ -194,6 +209,23 @@ async def test_reader_ack_sends_ack_ids(subscriber: MagicMock) -> None: ) +async def test_reader_ack_without_connection_raises_connection_error() -> None: + """Tests that acking before connecting fails loudly.""" + reader = make_reader() + with pytest.raises(MessageBrokerConnectionError): + await reader._ack([make_message(b'{"x": 1}')]) + + +async def test_reader_ack_wraps_failed_acknowledgments(subscriber: MagicMock) -> None: + """Tests that a failed acknowledgment is wrapped for retry rather than leaking gRPC.""" + subscriber.acknowledge.side_effect = gcp_exceptions.DeadlineExceeded("deadline") + reader = make_reader() + await reader._connect() + + with pytest.raises(MessageBrokerTransientError): + await reader._ack([make_message(b'{"x": 1}')]) + + async def test_reader_disconnect_closes_subscriber(subscriber: MagicMock) -> None: """Tests that teardown closes the client, releasing its channel.""" reader = make_reader() @@ -209,7 +241,6 @@ async def test_reader_disconnect_closes_subscriber(subscriber: MagicMock) -> Non (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( @@ -224,6 +255,18 @@ async def test_reader_maps_broker_errors( await reader._receive() +async def test_reader_idle_pull_timeout_returns_no_messages(subscriber: MagicMock) -> None: + """Tests that a pull window elapsing on an empty subscription reads as 'nothing yet'. + + An idle pull raises `DeadlineExceeded`; treating it as a broker error would trigger a + needless reconnect and eventually fail the reader on an empty subscription. + """ + subscriber.pull.side_effect = gcp_exceptions.DeadlineExceeded("late") + reader = make_reader() + await reader._connect() + assert await reader._receive() == [] + + 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") @@ -253,8 +296,8 @@ async def test_reader_missing_project_id(monkeypatch: pytest.MonkeyPatch) -> Non ) -async def test_reader_project_id_from_settings() -> None: - """Tests that `GCP_PUBSUB_PROJECT_ID` supplies the project when not passed.""" +async def test_reader_project_id_from_settings(subscriber: MagicMock) -> None: + """Tests that `GCP_PUBSUB_PROJECT_ID` supplies the pull's project when not passed.""" settings = Settings.model_validate({"gcp_pubsub": {"project_id": "settings-project"}}) with conftest.override_settings(settings): reader = GCPPubSubDataReader( @@ -262,12 +305,13 @@ async def test_reader_project_id_from_settings() -> None: subscription_id="test-subscription", field_names=["x"], ) - assert reader._subscription_path == "projects/settings-project/subscriptions/test-subscription" - - -# --------------------------------------------------------------------------- -# Writer -# --------------------------------------------------------------------------- + await reader._connect() + await reader._receive() + subscriber.pull.assert_called_once() + assert subscriber.pull.call_args.kwargs["request"]["subscription"] == ( + "projects/settings-project/subscriptions/test-subscription" + ) + await reader._disconnect() async def test_writer_publishes_encoded_records(publisher: MagicMock) -> None: @@ -309,18 +353,16 @@ def result(self, timeout: float) -> None: assert all(name != threading.current_thread().name for name in caller_threads) -async def test_writer_converts_fields_to_json_bytes(publisher: MagicMock) -> None: +async def test_writer_converts_fields_to_json_bytes() -> 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"}'] -async def test_writer_non_json_sends_first_field(publisher: MagicMock) -> 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(["hello"])}) == [b"hello"] @@ -332,6 +374,15 @@ async def test_writer_disconnect_stops_publisher(publisher: MagicMock) -> None: publisher.stop.assert_called_once_with() +async def test_writer_disconnect_wraps_stop_errors(publisher: MagicMock) -> None: + """Tests that a failed publisher stop is wrapped rather than leaking gRPC.""" + publisher.stop.side_effect = gcp_exceptions.FailedPrecondition("not stopped") + writer = make_writer() + await writer._connect() + with pytest.raises(MessageBrokerPermanentError): + await writer._disconnect() + + 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") @@ -342,6 +393,20 @@ async def test_writer_publish_failure_maps_error(publisher: MagicMock) -> None: await writer._send([b"one"]) +async def test_writer_publish_to_missing_topic_is_permanent(publisher: MagicMock) -> None: + """Tests that publishing to a deleted topic is permanent, not an exhausted source. + + `NoMoreDataException` is reserved for exhausted read sources; a writer cannot recover + from a missing topic, so it must fail as permanent. + """ + publisher.publish.side_effect = gcp_exceptions.NotFound("gone") + writer = make_writer() + await writer._connect() + + with pytest.raises(MessageBrokerPermanentError): + await writer._send([b"one"]) + + async def test_writer_without_connection_raises_connection_error(publisher: MagicMock) -> None: """Tests that sending before connecting fails loudly.""" writer = make_writer() diff --git a/tests/unit/test_message_data_reader.py b/tests/unit/test_message_data_reader.py index 27befb9a..2d9265aa 100644 --- a/tests/unit/test_message_data_reader.py +++ b/tests/unit/test_message_data_reader.py @@ -3,19 +3,18 @@ from __future__ import annotations from collections import deque -from types import SimpleNamespace import typing as _t import pytest from plugboard.exceptions import ( IOStreamClosedError, + MessageBrokerConnectionError, MessageBrokerPermanentError, MessageBrokerTransientError, NoMoreDataException, ) from plugboard.library.message_reader import MessageDataReader -from plugboard.utils import retry as retry_utils from plugboard.utils.retry import RetryPolicy @@ -111,28 +110,6 @@ async def drain(reader: ScriptedReader, limit: int = 100) -> list[dict[str, _t.A 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 - - -# --------------------------------------------------------------------------- -# Lifecycle -# --------------------------------------------------------------------------- - - 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]]) @@ -163,11 +140,6 @@ async def test_topic_is_retained() -> None: 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()]) @@ -216,11 +188,6 @@ async def test_single_field_reader() -> None: await reader.destroy() -# --------------------------------------------------------------------------- -# Empty polls are not the end of the stream -# --------------------------------------------------------------------------- - - 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. @@ -259,11 +226,6 @@ async def test_no_dangling_task_when_source_exhausted() -> None: await reader.destroy() -# --------------------------------------------------------------------------- -# Acknowledgment timing -# --------------------------------------------------------------------------- - - 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. @@ -303,11 +265,6 @@ async def test_unconsumed_messages_are_not_acknowledged() -> None: assert reader.acked == [] -# --------------------------------------------------------------------------- -# Retry and reconnection -# --------------------------------------------------------------------------- - - async def test_transient_receive_failure_reconnects_and_recovers( backoff_delays: list[float], ) -> None: @@ -406,10 +363,11 @@ async def test_failing_disconnect_still_reconnects() -> None: 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. + """Tests that a failing reconnect skips that attempt but the retries continue. - 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. + During a network partition the reconnect can fail; rather than running the read and + failing with a synthetic "not connected" error, that attempt is skipped, the retries + keep reconnecting, and the error eventually raised is the broker's own. """ reader = UnreconnectableReader( name="test-reader", @@ -424,7 +382,66 @@ async def test_reconnect_failure_does_not_abort_retries( 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 reader.polls == 1 # Only the initial read ran; the retries were skipped. + assert reader.connects == 3 # Initial connect plus one reconnect per retry. + assert backoff_delays == [1.0, 2.0] + await reader.destroy() + + +class _FlakyConnectionReader(ScriptedReader): + """Models a connection that drops and cannot be re-established. + + The first read fails with a real broker error; the link is then torn down and every + reconnect fails, leaving the reader disconnected (where a real receive would raise + `MessageBrokerConnectionError`). + """ + + def __init__(self, *args: _t.Any, script: list[_t.Any], **kwargs: _t.Any) -> None: + super().__init__(*args, script=script, **kwargs) + self.connected = True + + async def _connect(self) -> None: + self.connected = True + self.connects += 1 + + async def _reconnect(self) -> None: + # The link is down and cannot be re-established: the client is left disconnected. + self.connected = False + raise OSError("broker unreachable") + + async def _receive(self) -> list[_t.Any]: + if not self.connected: + raise MessageBrokerConnectionError("client is not connected") + 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 test_retries_exhausted_after_failed_reconnect_raise_broker_error( + backoff_delays: list[float], +) -> None: + """Tests that total failure reports the broker's error, not a synthetic one. + + When every reconnect fails the reader is left disconnected, so running a read would + surface `MessageBrokerConnectionError("client is not connected")`. The retries must + instead surface the last real broker error from the attempt that last ran. + """ + reader = _FlakyConnectionReader( + name="test-reader", + field_names=["x", "y"], + topic="test-topic", + idle_poll_delay=0.0, + script=[ConnectionError("broker timeout")], + retry_policy=RetryPolicy(max_retries=2), + ) + await reader.init() + + with pytest.raises(ConnectionError, match="broker timeout") as excinfo: + await reader.step() + + assert "not connected" not in str(excinfo.value) # Never the synthetic connection error. + assert reader.connects == 1 # Only `init` connected; the reconnects all failed. 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 a1926e10..c4370263 100644 --- a/tests/unit/test_message_data_writer.py +++ b/tests/unit/test_message_data_writer.py @@ -126,11 +126,6 @@ def delivered(writer: RecordingWriter) -> list[_t.Any]: SINGLE = [{"x": 1}] # For the single-field writers used by the retry tests. -# --------------------------------------------------------------------------- -# Lifecycle -# --------------------------------------------------------------------------- - - async def test_init_connects_to_broker() -> None: """Tests that `init` opens the broker connection.""" writer = make_writer(["x"]) @@ -175,11 +170,6 @@ async def test_destroy_waits_for_in_flight_send() -> None: @pytest.mark.parametrize("chunk_size", [1, 2, 3]) -# --------------------------------------------------------------------------- -# Batching -# --------------------------------------------------------------------------- - - 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) @@ -208,12 +198,9 @@ async def test_run_flushes_remaining_buffer() -> None: assert delivered(writer) == RECORDS -# --------------------------------------------------------------------------- -# Retry and reconnection -# --------------------------------------------------------------------------- - - -async def test_transient_send_failure_reconnects_and_recovers() -> None: +async def test_transient_send_failure_reconnects_and_recovers( + backoff_delays: list[float], +) -> 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) @@ -227,6 +214,7 @@ async def test_transient_send_failure_reconnects_and_recovers() -> None: assert writer.attempts == 2 # Failed once, then succeeded. assert writer.connects == 2 # Initial connect plus one reconnect. assert writer.disconnects == 1 + assert backoff_delays == [1.0] # Base delay on the single retry. async def test_permanent_send_failure_is_not_retried() -> None: @@ -252,7 +240,9 @@ async def test_permanent_send_failure_is_not_retried() -> None: await writer.destroy() -async def test_retry_exhaustion_raises_last_send_error() -> None: +async def test_retry_exhaustion_raises_last_send_error( + backoff_delays: list[float], +) -> 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) @@ -268,12 +258,20 @@ async def test_retry_exhaustion_raises_last_send_error() -> None: assert writer.attempts == 3 # Initial attempt plus two retries. assert writer.connects == 3 # Initial connect plus one reconnect per retry. + assert backoff_delays == [1.0, 2.0] # Backoff doubles on consecutive attempts. writer._task = None await writer.destroy() -async def test_reconnect_failure_does_not_abort_retries() -> None: - """Tests that a failing reconnect costs an attempt instead of ending the retrying.""" +async def test_reconnect_failure_does_not_abort_retries( + backoff_delays: list[float], +) -> None: + """Tests that a failing reconnect skips that attempt but the retries continue. + + Without a connection a send would fail with a synthetic "not connected" error, so a + failed reconnect skips the attempt; the retries keep reconnecting and the error that + is eventually raised is still the broker's own. + """ writer = make_writer( ["x"], script=[ConnectionError("send failed")] * 3, @@ -290,17 +288,12 @@ async def test_reconnect_failure_does_not_abort_retries() -> 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 + assert writer.attempts == 1 # Only the initial send ran; the retries were skipped. + assert writer.connects == 3 # Initial connect plus one reconnect per retry. writer._task = None await writer.destroy() -# --------------------------------------------------------------------------- -# Shared record encoding helpers -# --------------------------------------------------------------------------- - - 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"])} @@ -337,7 +330,10 @@ def test_iter_records_is_linear_in_row_count() -> None: ) 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 + assert ( + encode_records({"x": deque([1, 2]), "y": deque(["a", "b"])}, parse_json=parse_json) + == expected + ) @pytest.mark.parametrize( @@ -349,7 +345,7 @@ def test_encode_records_as_text(parse_json: bool, expected: list[str]) -> None: ) 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 + assert encode_records_bytes({"x": deque([1, 2])}, parse_json=parse_json) == expected def test_encode_records_bytes_preserves_binary_payloads() -> None: From 99b2b2fa3aa1ae62f5d01806885eb69f58c63d5a Mon Sep 17 00:00:00 2001 From: Chris Knight Date: Mon, 19 Oct 2026 10:19:38 +0200 Subject: [PATCH 16/21] fix: run broker tests on every Python version and fix GCP drain check MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The broker tests must run on every supported Python version, but they share a single SQS queue and GCP subscription, so they cannot run at the same time. - Split the broker tests out of the parallel test-integration matrix into a dedicated test-message-brokers job that runs them for 3.12, 3.13 and 3.14 in turn (a per-version uv loop in one job). A single job per run contributes at most one member to its concurrency group, so it is serialized across runs without ever cancelling a waiting request — the failure mode the previous shared-group matrix leg hit. - Drop the concurrency group and the 3.12-only gate from test-integration, and exclude the broker test file so the matrix legs no longer touch the brokers. - Move the GCP round-trip emptiness assertion out of finally and before the post-test drain, so a leftover unacked message is caught instead of being silently removed (and an assert no longer masks an exception from the try). --- .github/workflows/lint-test.yaml | 82 ++++++++++++++----- .../integration/test_message_data_brokers.py | 12 +-- 2 files changed, 67 insertions(+), 27 deletions(-) diff --git a/.github/workflows/lint-test.yaml b/.github/workflows/lint-test.yaml index f2647270..f515a0f7 100644 --- a/.github/workflows/lint-test.yaml +++ b/.github/workflows/lint-test.yaml @@ -115,12 +115,6 @@ jobs: name: Tests - integration runs-on: ubuntu-latest timeout-minutes: 10 - # The message broker integration tests publish to a single shared SQS queue and - # GCP subscription, so concurrent runs must not collide on them. Serializing the - # job across runs keeps one set of broker tests at a time. - concurrency: - group: test-integration - cancel-in-progress: false strategy: matrix: python_version: [3.12, 3.13, 3.14] @@ -147,6 +141,50 @@ jobs: docker compose up -d sleep 10 # Wait for services to start + - name: Run integration tests + # The message broker tests (AWS SQS/SNS + GCP Pub/Sub) run in the + # test-message-brokers job, which is serialized on the shared SQS queue and GCP + # subscription; they are excluded here so the matrix legs cannot race on them. + run: COVERAGE_FILE=.coverage.py${{ matrix.python_version }}.integration uv run coverage run -m pytest ./tests/integration/ --ignore=./tests/integration/test_message_data_brokers.py -m "not tuner" + env: + RABBITMQ_URL: amqp://user:password@localhost:5672/ + REDIS_URL: redis://default:password@localhost:6379/ + RAY_ENABLE_UV_RUN_RUNTIME_ENV: 0 + PLUGBOARD_IO_READ_TIMEOUT: 5.0 + + - name: Upload integration test coverage + uses: actions/upload-artifact@v4 + with: + name: coverage-integration-py${{ matrix.python_version }} + include-hidden-files: true + path: .coverage.py${{ matrix.python_version }}.integration* + + test-message-brokers: + name: Tests - message brokers + runs-on: ubuntu-latest + # The broker tests run on every supported Python version, each takes ~15 s against the + # real brokers, plus install time per version. + timeout-minutes: 20 + # The broker tests publish to a single shared SQS queue and GCP subscription, so they + # must never run at the same time. A dedicated concurrency group serializes the job + # across concurrent workflow runs; because this is a single job (not a matrix), a run + # contributes at most one member to the group, so at most one is running and one queued + # and no leg is ever cancelled as a waiting request. The Python loop below runs the + # tests for each supported version in turn, keeping them on every version while + # contending for the shared brokers one at a time. + concurrency: + group: message-brokers + cancel-in-progress: false + steps: + - name: Checkout + uses: actions/checkout@v4 + + - name: Install uv + uses: astral-sh/setup-uv@v4 + with: + enable-cache: true + cache-dependency-glob: "uv.lock" + - name: Configure GCP Pub/Sub credentials if: ${{ vars.PLUGBOARD_MESSAGE_BROKER_TESTS == 'true' }} env: @@ -154,17 +192,24 @@ jobs: run: | echo "$GCP_SA_KEY_JSON" > /tmp/gcp-sa-key.json - - name: Run integration tests - run: COVERAGE_FILE=.coverage.py${{ matrix.python_version }}.integration uv run coverage run -m pytest ./tests/integration/ -m "not tuner" + - name: Run message broker tests on every supported Python + # Runs on every supported Python version (the loop) but one version at a time, so + # the shared SQS queue and GCP subscription are never contended. This is a plain + # test gate (no coverage): the broker code it exercises is already measured by the + # unit tests. When PLUGBOARD_MESSAGE_BROKER_TESTS is off, both tests skip and the + # job passes. + if: ${{ vars.PLUGBOARD_MESSAGE_BROKER_TESTS == 'true' }} + run: | + for version in 3.12 3.13 3.14; do + echo "::group::Python $version" + uv python install "$version" + uv sync --group test --python "$version" + uv run --python "$version" pytest ./tests/integration/test_message_data_brokers.py + echo "::endgroup::" + done env: - RABBITMQ_URL: amqp://user:password@localhost:5672/ - REDIS_URL: redis://default:password@localhost:6379/ - RAY_ENABLE_UV_RUN_RUNTIME_ENV: 0 + PLUGBOARD_MESSAGE_BROKER_TESTS: "true" PLUGBOARD_IO_READ_TIMEOUT: 5.0 - # Run the message broker integration tests (AWS SQS/SNS + GCP Pub/Sub). - # Gated to a single Python leg: the tests share one SQS queue and GCP - # subscription, so the matrix legs must not run them concurrently. - PLUGBOARD_MESSAGE_BROKER_TESTS: ${{ vars.PLUGBOARD_MESSAGE_BROKER_TESTS == 'true' && matrix.python_version == 3.12 && 'true' || 'false' }} # AWS SQS/SNS identifiers. AWS_REGION: ${{ vars.AWS_REGION }} AWS_SQS_QUEUE_URL: ${{ vars.AWS_SQS_QUEUE_URL }} @@ -179,13 +224,6 @@ jobs: AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }} GOOGLE_APPLICATION_CREDENTIALS: /tmp/gcp-sa-key.json - - name: Upload integration test coverage - uses: actions/upload-artifact@v4 - with: - name: coverage-integration-py${{ matrix.python_version }} - include-hidden-files: true - path: .coverage.py${{ matrix.python_version }}.integration* - test-integration-tuner: name: Tests - integration:tuner runs-on: ubuntu-latest diff --git a/tests/integration/test_message_data_brokers.py b/tests/integration/test_message_data_brokers.py index a0123cb5..a9f05db4 100644 --- a/tests/integration/test_message_data_brokers.py +++ b/tests/integration/test_message_data_brokers.py @@ -329,15 +329,17 @@ async def test_gcp_pubsub_round_trip(broker_env: None) -> None: await write_records(writer, writer_channels, RECORDS) received = await read_records(reader, len(RECORDS)) assert sorted((record["x"], record["y"]) for record in received) == EXPECTED + # The reader acknowledged every record, so the subscription should already be + # empty. Checked before the finally cleanup drains it, so leftover messages are + # caught rather than silently removed. + assert _gcp_subscription_undelivered(project_id, subscription_id) == 0, ( + "PubSub subscription not empty after the round trip" + ) finally: await writer.destroy() await reader.destroy() - # Clean up so nothing interferes with the next run, then confirm the - # subscription is genuinely empty (the reader should have acked everything). + # Clean up so nothing interferes with the next run. await _clean_broker(pre_test=False, project_id=project_id, subscription_id=subscription_id) - assert _gcp_subscription_undelivered(project_id, subscription_id) == 0, ( - "PubSub subscription not empty after the round trip" - ) async def connect_writer( From f058e36578f961ac968dce17f18697e4714df062 Mon Sep 17 00:00:00 2001 From: Chris Knight Date: Mon, 19 Oct 2026 10:36:46 +0200 Subject: [PATCH 17/21] ci: run broker tests concurrently using per-python-version lanes The message broker tests must run on every Python version, but they share a single SQS queue and GCP subscription, so they cannot run in parallel. Instead of forcing them into one serial job, each matrix leg now uses its own dedicated broker lane (a separate SNS+SQS pair and Pub/Sub topic+subscription per Python version, provisioned in plugboard-infra). - Drop the serial test-message-brokers job and the --ignore of the broker test file; the broker round trips run again in the normal parallel test-integration matrix, one leg per version. - Each leg resolves its lane from the new BROKER_LANES_JSON repository variable (keyed by python version) with jq, exporting only its own lane's queue/topic identifiers, so legs never contend on a shared broker. - If the lanes variable is absent (e.g. a fork pull request) the broker tests self-skip; a present-but-missing lane for a version fails loudly. --- .github/workflows/lint-test.yaml | 110 ++++++++++++------------------- 1 file changed, 41 insertions(+), 69 deletions(-) diff --git a/.github/workflows/lint-test.yaml b/.github/workflows/lint-test.yaml index f515a0f7..556d9dac 100644 --- a/.github/workflows/lint-test.yaml +++ b/.github/workflows/lint-test.yaml @@ -142,15 +142,52 @@ jobs: sleep 10 # Wait for services to start - name: Run integration tests - # The message broker tests (AWS SQS/SNS + GCP Pub/Sub) run in the - # test-message-brokers job, which is serialized on the shared SQS queue and GCP - # subscription; they are excluded here so the matrix legs cannot race on them. - run: COVERAGE_FILE=.coverage.py${{ matrix.python_version }}.integration uv run coverage run -m pytest ./tests/integration/ --ignore=./tests/integration/test_message_data_brokers.py -m "not tuner" + # The message broker round trips are part of this suite and run on every Python + # leg. Each leg resolves its own dedicated broker lane (one SNS+SQS pair and + # Pub/Sub topic+subscription per Python version, provisioned in plugboard-infra + # and stored in the BROKER_LANES_JSON variable), so the legs run concurrently + # without contending on a shared queue or subscription. When broker tests are off + # the lane resolution is skipped and the broker tests self-skip. + run: | + # Resolve this leg's broker lane. The broker tests self-skip unless these env vars + # are set, so when broker tests are off, or lanes have not been provisioned for this + # run (e.g. a fork pull request where vars/secrets are unavailable), they simply skip. + if [ "${PLUGBOARD_MESSAGE_BROKER_TESTS}" = "true" ] && [ -n "${BROKER_LANES_JSON:-}" ]; then + # A lanes file is present but has no entry for this Python version: that is a + # configuration mismatch (version in the matrix but not in terraform), so fail loud. + if ! jq -e --arg v "$PYTHON_VERSION" 'has($v)' <<< "$BROKER_LANES_JSON" > /dev/null; then + echo "::error::No broker lane for Python $PYTHON_VERSION in BROKER_LANES_JSON" + exit 1 + fi + printf '%s' "$GCP_SA_KEY_JSON" > /tmp/gcp-sa-key.json + # The lane values are queue URLs, topic ARNs and resource ids (no shell metacharacters), + # so the generated `export` lines are safe to eval. + eval "$(jq -r --arg v "$PYTHON_VERSION" \ + '.[$v] | "export AWS_SQS_QUEUE_URL=\(.aws_sqs_queue_url)\nexport AWS_SNS_TOPIC_ARN=\(.aws_sns_topic_arn)\nexport GCP_TOPIC_ID=\(.gcp_topic_id)\nexport GCP_SUBSCRIPTION_ID=\(.gcp_subscription_id)"' \ + <<< "$BROKER_LANES_JSON")" + export GOOGLE_APPLICATION_CREDENTIALS=/tmp/gcp-sa-key.json + echo "Broker lane for Python $PYTHON_VERSION:" + echo " AWS_SQS_QUEUE_URL=$AWS_SQS_QUEUE_URL" + echo " GCP_TOPIC_ID=$GCP_TOPIC_ID" + fi + COVERAGE_FILE=.coverage.py${{ matrix.python_version }}.integration \ + uv run coverage run -m pytest ./tests/integration/ -m "not tuner" env: RABBITMQ_URL: amqp://user:password@localhost:5672/ REDIS_URL: redis://default:password@localhost:6379/ RAY_ENABLE_UV_RUN_RUNTIME_ENV: 0 PLUGBOARD_IO_READ_TIMEOUT: 5.0 + # Message broker test inputs. The per-lane identifiers are resolved from + # BROKER_LANES_JSON in the step above (only when the tests are enabled). + PYTHON_VERSION: ${{ matrix.python_version }} + PLUGBOARD_MESSAGE_BROKER_TESTS: ${{ vars.PLUGBOARD_MESSAGE_BROKER_TESTS }} + BROKER_LANES_JSON: ${{ vars.BROKER_LANES_JSON }} + GCP_SA_KEY_JSON: ${{ secrets.GCP_SA_KEY_JSON }} + AWS_REGION: ${{ vars.AWS_REGION }} + GCP_PROJECT_ID: ${{ vars.GCP_PROJECT_ID }} + GCP_PUBSUB_PROJECT_ID: ${{ vars.GCP_PUBSUB_PROJECT_ID }} + AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }} + AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }} - name: Upload integration test coverage uses: actions/upload-artifact@v4 @@ -159,71 +196,6 @@ jobs: include-hidden-files: true path: .coverage.py${{ matrix.python_version }}.integration* - test-message-brokers: - name: Tests - message brokers - runs-on: ubuntu-latest - # The broker tests run on every supported Python version, each takes ~15 s against the - # real brokers, plus install time per version. - timeout-minutes: 20 - # The broker tests publish to a single shared SQS queue and GCP subscription, so they - # must never run at the same time. A dedicated concurrency group serializes the job - # across concurrent workflow runs; because this is a single job (not a matrix), a run - # contributes at most one member to the group, so at most one is running and one queued - # and no leg is ever cancelled as a waiting request. The Python loop below runs the - # tests for each supported version in turn, keeping them on every version while - # contending for the shared brokers one at a time. - concurrency: - group: message-brokers - cancel-in-progress: false - steps: - - name: Checkout - uses: actions/checkout@v4 - - - name: Install uv - uses: astral-sh/setup-uv@v4 - with: - enable-cache: true - cache-dependency-glob: "uv.lock" - - - name: Configure GCP Pub/Sub credentials - if: ${{ vars.PLUGBOARD_MESSAGE_BROKER_TESTS == 'true' }} - env: - GCP_SA_KEY_JSON: ${{ secrets.GCP_SA_KEY_JSON }} - run: | - echo "$GCP_SA_KEY_JSON" > /tmp/gcp-sa-key.json - - - name: Run message broker tests on every supported Python - # Runs on every supported Python version (the loop) but one version at a time, so - # the shared SQS queue and GCP subscription are never contended. This is a plain - # test gate (no coverage): the broker code it exercises is already measured by the - # unit tests. When PLUGBOARD_MESSAGE_BROKER_TESTS is off, both tests skip and the - # job passes. - if: ${{ vars.PLUGBOARD_MESSAGE_BROKER_TESTS == 'true' }} - run: | - for version in 3.12 3.13 3.14; do - echo "::group::Python $version" - uv python install "$version" - uv sync --group test --python "$version" - uv run --python "$version" pytest ./tests/integration/test_message_data_brokers.py - echo "::endgroup::" - done - env: - PLUGBOARD_MESSAGE_BROKER_TESTS: "true" - PLUGBOARD_IO_READ_TIMEOUT: 5.0 - # AWS SQS/SNS identifiers. - AWS_REGION: ${{ vars.AWS_REGION }} - AWS_SQS_QUEUE_URL: ${{ vars.AWS_SQS_QUEUE_URL }} - AWS_SNS_TOPIC_ARN: ${{ vars.AWS_SNS_TOPIC_ARN }} - # GCP Pub/Sub identifiers. - GCP_PROJECT_ID: ${{ vars.GCP_PROJECT_ID }} - GCP_PUBSUB_PROJECT_ID: ${{ vars.GCP_PUBSUB_PROJECT_ID }} - GCP_TOPIC_ID: ${{ vars.GCP_TOPIC_ID }} - GCP_SUBSCRIPTION_ID: ${{ vars.GCP_SUBSCRIPTION_ID }} - # Broker credentials (AWS credential chain + GCP Application Default Credentials). - AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }} - AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }} - GOOGLE_APPLICATION_CREDENTIALS: /tmp/gcp-sa-key.json - test-integration-tuner: name: Tests - integration:tuner runs-on: ubuntu-latest From 1a27a002dc65497d98d359fade713e791e8729e4 Mon Sep 17 00:00:00 2001 From: Chris Knight Date: Wed, 4 Nov 2026 13:05:35 +0100 Subject: [PATCH 18/21] feat: filter message broker readers and writers by job_id and queue CI lanes - Add `job_id` to `MessageDataReader` and `MessageDataWriter` (with precedence: explicit `job_id` argument -> connected `StateBackend.job_id` -> active `DI.job_id` context -> generated job ID). - Tag every message published by `AWSSNSDataWriter` and `GCPPubSubDataWriter` with the `job_id` message attribute. - Filter incoming messages in `AWSSQSDataReader` and `GCPPubSubDataReader` by `job_id` (acknowledging and discarding non-matching messages). - Add a per-version `concurrency` group (`test-integration-${{ matrix.python_version }}`) with `cancel-in-progress: false` and `queue: max` on `test-integration` so concurrent pipeline runs queue up to 100 pending jobs per Python version in FIFO order without cancelling queued runs or blocking across versions. --- .github/workflows/lint-test.yaml | 8 ++ docs/usage/message-data.md | 4 + plugboard/library/aws_messaging_io.py | 46 +++++++++- plugboard/library/gcp_pubsub_io.py | 24 ++++- plugboard/library/message_reader.py | 28 ++++++ plugboard/library/message_writer.py | 29 ++++++ .../integration/test_message_data_brokers.py | 31 ++++++- tests/unit/test_aws_messaging_io.py | 90 +++++++++++++++++-- tests/unit/test_gcp_pubsub_io.py | 42 ++++++++- tests/unit/test_message_data_reader.py | 28 ++++++ tests/unit/test_message_data_writer.py | 27 ++++++ 11 files changed, 344 insertions(+), 13 deletions(-) diff --git a/.github/workflows/lint-test.yaml b/.github/workflows/lint-test.yaml index 556d9dac..bb3658f0 100644 --- a/.github/workflows/lint-test.yaml +++ b/.github/workflows/lint-test.yaml @@ -118,6 +118,14 @@ jobs: strategy: matrix: python_version: [3.12, 3.13, 3.14] + # Each Python version has its own dedicated broker lane in plugboard-infra, so the + # matrix legs within a run execute concurrently. A per-version concurrency group with + # `queue: max` serialises same-version legs across concurrent pipeline runs (queueing + # up to 100 pending jobs in FIFO order rather than cancelling earlier pending ones). + concurrency: + group: test-integration-${{ matrix.python_version }} + cancel-in-progress: false + queue: max steps: - name: Checkout uses: actions/checkout@v4 diff --git a/docs/usage/message-data.md b/docs/usage/message-data.md index bcdde1a9..392ed812 100644 --- a/docs/usage/message-data.md +++ b/docs/usage/message-data.md @@ -119,6 +119,10 @@ any connector. - **SNS fan-out is transparent.** Subscribing an SQS queue to an SNS topic delivers the notification envelope, which `AWSSQSDataReader` unwraps automatically; a non-object body is exposed under `data`. +- **Messages are scoped by `job_id`.** Writers attach the active Plugboard `job_id` as a + `job_id` message attribute, and readers filter incoming messages by `job_id` + (acknowledging and discarding non-matching messages). Pass `job_id=...` to override the + active process job ID when reading messages produced by another job. ## Configuration diff --git a/plugboard/library/aws_messaging_io.py b/plugboard/library/aws_messaging_io.py index ae73e2a2..f303fbf5 100644 --- a/plugboard/library/aws_messaging_io.py +++ b/plugboard/library/aws_messaging_io.py @@ -115,6 +115,31 @@ def _decode_body(body: str, *, parse_json: bool) -> dict[str, _t.Any]: return record +def _attr_string(attrs: _t.Any, key: str) -> _t.Optional[str]: + """Reads a string message attribute from an SQS or SNS `MessageAttributes` mapping.""" + if not isinstance(attrs, dict): + return None + entry = attrs.get(key) + if not isinstance(entry, dict): + return None + value = entry.get("StringValue") or entry.get("Value") + return value if isinstance(value, str) else None + + +def _extract_job_id(msg: dict[str, _t.Any]) -> _t.Optional[str]: + """Extracts the `job_id` attribute from an SQS message or its SNS envelope.""" + direct = _attr_string(msg.get("MessageAttributes"), "job_id") + if direct is not None: + return direct + body = msg.get("Body") + if not isinstance(body, str): + return None + parsed = _parse_json_or_raw(body) + if isinstance(parsed, dict) and parsed.get("Type") == "Notification": + return _attr_string(parsed.get("MessageAttributes"), "job_id") + return None + + 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): @@ -227,12 +252,19 @@ async def _receive(self) -> list[_t.Any]: "QueueUrl": self._queue_url, "MaxNumberOfMessages": min(self._chunk_size or _SQS_MAX_MESSAGES, _SQS_MAX_MESSAGES), "WaitTimeSeconds": self._wait_time_seconds, + "MessageAttributeNames": ["job_id"], } try: 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", []) + messages = response.get("Messages", []) + target_job_id = self.job_id + matching = [msg for msg in messages if _extract_job_id(msg) == target_job_id] + unmatched = [msg for msg in messages if _extract_job_id(msg) != target_job_id] + if unmatched: + await self._ack(unmatched) + return matching async def _convert(self, data: list[_t.Any]) -> dict[str, deque]: """Converts SQS messages to a field buffer. @@ -349,9 +381,19 @@ async def _send(self, messages: list[_t.Any]) -> None: """ if self._client is None: raise MessageBrokerConnectionError("SNS client is not connected") + message_attributes = { + "job_id": { + "DataType": "String", + "StringValue": self.job_id, + } + } for msg_data in messages: try: - await self._client.publish(TopicArn=self._topic_arn, Message=msg_data) + await self._client.publish( + TopicArn=self._topic_arn, + Message=msg_data, + MessageAttributes=message_attributes, + ) except Exception as error: # noqa: BLE001 raise _wrap_error(error) from error diff --git a/plugboard/library/gcp_pubsub_io.py b/plugboard/library/gcp_pubsub_io.py index 1b62e13a..36f7035c 100644 --- a/plugboard/library/gcp_pubsub_io.py +++ b/plugboard/library/gcp_pubsub_io.py @@ -104,6 +104,16 @@ def _decode_payload(payload: bytes, *, parse_json: bool) -> dict[str, _t.Any]: return record +def _extract_job_id(msg_wrapper: _t.Any) -> _t.Optional[str]: + """Extracts the `job_id` attribute from a PubSub `ReceivedMessage`, if present.""" + message = getattr(msg_wrapper, "message", None) + attributes = getattr(message, "attributes", None) + if not isinstance(attributes, _t.Mapping): + return None + value = attributes.get("job_id") + return value if isinstance(value, str) else None + + class GCPPubSubDataReader(MessageDataReader): """Reads data from a Google Cloud PubSub subscription. @@ -185,7 +195,13 @@ async def _receive(self) -> list[_t.Any]: return [] except Exception as error: # noqa: BLE001 raise _wrap_error(error, exhausted=True) from error - return list(response.received_messages) + messages = list(response.received_messages) + target_job_id = self.job_id + matching = [msg for msg in messages if _extract_job_id(msg) == target_job_id] + unmatched = [msg for msg in messages if _extract_job_id(msg) != target_job_id] + if unmatched: + await self._ack(unmatched) + return matching async def _convert(self, data: list[_t.Any]) -> dict[str, deque]: """Converts PubSub messages to a field buffer. @@ -295,9 +311,13 @@ async def _send(self, messages: list[_t.Any]) -> None: if self._publisher is None: raise MessageBrokerConnectionError("Publisher client is not connected") publisher = self._publisher + job_id = self.job_id def _publish() -> None: - futures = [publisher.publish(self._topic_path, data=msg_data) for msg_data in messages] + futures = [ + publisher.publish(self._topic_path, data=msg_data, job_id=job_id) + for msg_data in messages + ] # The publisher batches internally; each future resolves once the broker # has accepted that message. for future in futures: diff --git a/plugboard/library/message_reader.py b/plugboard/library/message_reader.py index 4cc2e74a..d28d1ccd 100644 --- a/plugboard/library/message_reader.py +++ b/plugboard/library/message_reader.py @@ -12,6 +12,7 @@ from plugboard.exceptions import NoMoreDataException from plugboard.library.data_reader import DataReader, DataReaderArgsSpec from plugboard.schemas import ComponentArgsDict +from plugboard.utils import DI from plugboard.utils.retry import RetryPolicy, attempt_reconnect @@ -25,11 +26,14 @@ class MessageDataReaderArgsDict(DataReaderArgsSpec): retry_policy: Optional; The backoff policy for retrying failed broker calls. idle_poll_delay: Optional; Delay in seconds between polls that return no messages. + job_id: Optional; The job ID to filter incoming messages by. Defaults to the + active Plugboard job ID. """ topic: _t.NotRequired[str] retry_policy: _t.NotRequired[RetryPolicy] idle_poll_delay: _t.NotRequired[float] + job_id: _t.NotRequired[_t.Optional[str]] class MessageDataReader(DataReader): @@ -59,6 +63,7 @@ def __init__( chunk_size: _t.Optional[int] = None, retry_policy: RetryPolicy = RetryPolicy(), idle_poll_delay: float = 0.1, + job_id: _t.Optional[str] = None, **kwargs: _t.Unpack[ComponentArgsDict], ) -> None: """Instantiates the `MessageDataReader`. @@ -70,15 +75,38 @@ def __init__( 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. + job_id: Optional; The job ID to filter incoming messages by. Defaults to the + active Plugboard job ID. **kwargs: Additional keyword arguments for [`Component`][plugboard.component.Component]. """ super().__init__(field_names=field_names, chunk_size=chunk_size, **kwargs) self._topic = topic self._retry_policy = retry_policy self._idle_poll_delay = idle_poll_delay + self._explicit_job_id = job_id + self._job_id: _t.Optional[str] = job_id + if self._job_id is None: + with contextlib.suppress(RuntimeError): + self._job_id = DI.job_id.resolve_sync() self._pending_ack: list[_t.Any] = [] self._connection_lock = asyncio.Lock() + @property + def job_id(self) -> str: + """The job ID used to filter incoming messages.""" + if self._explicit_job_id is not None: + return self._explicit_job_id + if self._state is not None and self._state.job_id is not None: + return self._state.job_id + if self._job_id is not None: + return self._job_id + try: + self._job_id = DI.job_id.resolve_sync() + except RuntimeError: + with self._job_id_ctx(): + self._job_id = DI.job_id.resolve_sync() + return self._job_id + @abstractmethod async def _connect(self) -> None: """Establishes connection to the message broker.""" diff --git a/plugboard/library/message_writer.py b/plugboard/library/message_writer.py index 6462e5ba..3b6cbc88 100644 --- a/plugboard/library/message_writer.py +++ b/plugboard/library/message_writer.py @@ -5,12 +5,14 @@ from abc import abstractmethod import asyncio from collections import deque +import contextlib import json import typing as _t from plugboard.component import IOController from plugboard.library.data_writer import DataWriter, DataWriterArgsSpec from plugboard.schemas import ComponentArgsDict +from plugboard.utils import DI from plugboard.utils.retry import RetryPolicy, attempt_reconnect @@ -22,10 +24,13 @@ class MessageDataWriterArgsDict(DataWriterArgsSpec): Attributes: topic: The topic/queue to write to. retry_policy: Optional; The backoff policy for retrying failed broker calls. + job_id: Optional; The job ID to attach to outgoing messages. Defaults to the + active Plugboard job ID. """ topic: _t.NotRequired[str] retry_policy: _t.NotRequired[RetryPolicy] + job_id: _t.NotRequired[_t.Optional[str]] def iter_records(data: dict[str, deque]) -> _t.Iterator[dict[str, _t.Any]]: @@ -113,6 +118,7 @@ def __init__( topic: str = "", chunk_size: _t.Optional[int] = None, retry_policy: RetryPolicy = RetryPolicy(), + job_id: _t.Optional[str] = None, **kwargs: _t.Unpack[ComponentArgsDict], ) -> None: """Instantiates the `MessageDataWriter`. @@ -122,11 +128,34 @@ def __init__( topic: The topic/queue to write to. chunk_size: Optional; The number of records to batch into a single send. retry_policy: The backoff policy for retrying failed broker calls. + job_id: Optional; The job ID to attach to outgoing messages. Defaults to the + active Plugboard job ID. **kwargs: Additional keyword arguments for [`Component`][plugboard.component.Component]. """ super().__init__(field_names=field_names, chunk_size=chunk_size, **kwargs) self._topic = topic self._retry_policy = retry_policy + self._explicit_job_id = job_id + self._job_id: _t.Optional[str] = job_id + if self._job_id is None: + with contextlib.suppress(RuntimeError): + self._job_id = DI.job_id.resolve_sync() + + @property + def job_id(self) -> str: + """The job ID attached to outgoing messages.""" + if self._explicit_job_id is not None: + return self._explicit_job_id + if self._state is not None and self._state.job_id is not None: + return self._state.job_id + if self._job_id is not None: + return self._job_id + try: + self._job_id = DI.job_id.resolve_sync() + except RuntimeError: + with self._job_id_ctx(): + self._job_id = DI.job_id.resolve_sync() + return self._job_id @abstractmethod async def _connect(self) -> None: diff --git a/tests/integration/test_message_data_brokers.py b/tests/integration/test_message_data_brokers.py index a9f05db4..46f88dd1 100644 --- a/tests/integration/test_message_data_brokers.py +++ b/tests/integration/test_message_data_brokers.py @@ -53,6 +53,7 @@ FIELD_NAMES = ["x", "y"] # Enough to verify batch delivery, but small so a broken broker fails fast. RECORDS = [{"x": i, "y": f"value-{i}"} for i in range(1, 6)] +OTHER_JOB_RECORDS = [{"x": 999, "y": "other-job"}] EXPECTED = sorted((record["x"], record["y"]) for record in RECORDS) _LOG = logging.getLogger(__name__) @@ -252,7 +253,7 @@ async def _clean_broker( @requires_broker -async def test_aws_sns_sqs_round_trip(broker_env: None) -> None: +async def test_aws_sns_sqs_round_trip(broker_env: None, job_id_ctx: str) -> None: """Writes records to the SNS topic and reads them back from the linked SQS queue.""" queue_url = os.environ["AWS_SQS_QUEUE_URL"] topic_arn = os.environ["AWS_SNS_TOPIC_ARN"] @@ -267,6 +268,14 @@ async def test_aws_sns_sqs_round_trip(broker_env: None) -> None: region=region, parse_json=True, ) + other_writer = AWSSNSDataWriter( + name="aws-sns-other-writer", + field_names=FIELD_NAMES, + topic_arn=topic_arn, + region=region, + parse_json=True, + job_id="Job_other00000000", + ) reader = AWSSQSDataReader( name="aws-sqs-reader", field_names=FIELD_NAMES, @@ -276,18 +285,23 @@ async def test_aws_sns_sqs_round_trip(broker_env: None) -> None: wait_time_seconds=2, chunk_size=10, ) + assert writer.job_id == reader.job_id == job_id_ctx writer_channels = await connect_writer(writer) + other_channels = await connect_writer(other_writer) await reader.init() await writer.init() + await other_writer.init() received: list[dict[str, _t.Any]] = [] try: + await write_records(other_writer, other_channels, OTHER_JOB_RECORDS) await write_records(writer, writer_channels, RECORDS) received = await read_records(reader, len(RECORDS)) assert sorted((record["x"], record["y"]) for record in received) == EXPECTED # The reader acknowledged every record, so nothing should remain on the queue. await _assert_aws_queue_empty(queue_url, region) finally: + await other_writer.destroy() await writer.destroy() await reader.destroy() # Safety net: if the read or assertion failed mid-way, drain any unacked @@ -296,7 +310,7 @@ async def test_aws_sns_sqs_round_trip(broker_env: None) -> None: @requires_broker -async def test_gcp_pubsub_round_trip(broker_env: None) -> None: +async def test_gcp_pubsub_round_trip(broker_env: None, job_id_ctx: str) -> None: """Writes records to the PubSub topic and reads them back from the linked subscription.""" project_id = os.environ["GCP_PUBSUB_PROJECT_ID"] topic_id = os.environ["GCP_TOPIC_ID"] @@ -311,6 +325,14 @@ async def test_gcp_pubsub_round_trip(broker_env: None) -> None: project_id=project_id, parse_json=True, ) + other_writer = GCPPubSubDataWriter( + name="gcp-pubsub-other-writer", + field_names=FIELD_NAMES, + topic_id=topic_id, + project_id=project_id, + parse_json=True, + job_id="Job_other00000000", + ) reader = GCPPubSubDataReader( name="gcp-pubsub-reader", field_names=FIELD_NAMES, @@ -320,12 +342,16 @@ async def test_gcp_pubsub_round_trip(broker_env: None) -> None: chunk_size=10, pull_timeout=3, ) + assert writer.job_id == reader.job_id == job_id_ctx writer_channels = await connect_writer(writer) + other_channels = await connect_writer(other_writer) await reader.init() await writer.init() + await other_writer.init() received: list[dict[str, _t.Any]] = [] try: + await write_records(other_writer, other_channels, OTHER_JOB_RECORDS) await write_records(writer, writer_channels, RECORDS) received = await read_records(reader, len(RECORDS)) assert sorted((record["x"], record["y"]) for record in received) == EXPECTED @@ -336,6 +362,7 @@ async def test_gcp_pubsub_round_trip(broker_env: None) -> None: "PubSub subscription not empty after the round trip" ) finally: + await other_writer.destroy() await writer.destroy() await reader.destroy() # Clean up so nothing interferes with the next run. diff --git a/tests/unit/test_aws_messaging_io.py b/tests/unit/test_aws_messaging_io.py index 8acecc09..fc7f14f4 100644 --- a/tests/unit/test_aws_messaging_io.py +++ b/tests/unit/test_aws_messaging_io.py @@ -64,9 +64,13 @@ def _create_client(service_name: str, **kwargs: _t.Any) -> _t.Any: yield client +DEFAULT_JOB_ID = "Job_test12345678" + + def make_reader(**kwargs: _t.Any) -> AWSSQSDataReader: """Builds a reader for the fixed test queue.""" kwargs.setdefault("field_names", ["x", "y"]) + kwargs.setdefault("job_id", DEFAULT_JOB_ID) return AWSSQSDataReader( name="sqs-reader", queue_url="https://sqs.test/1/messages", @@ -78,6 +82,7 @@ def make_reader(**kwargs: _t.Any) -> AWSSQSDataReader: def make_writer(**kwargs: _t.Any) -> AWSSNSDataWriter: """Builds a writer for the fixed test topic.""" kwargs.setdefault("field_names", ["x", "y"]) + kwargs.setdefault("job_id", DEFAULT_JOB_ID) return AWSSNSDataWriter( name="sns-writer", topic_arn="arn:aws:sns:eu-west-1:1:messages", @@ -86,9 +91,14 @@ def make_writer(**kwargs: _t.Any) -> AWSSNSDataWriter: ) -def sqs_message(body: str, handle: str) -> dict[str, str]: +def sqs_message( + body: str, handle: str, *, job_id: _t.Optional[str] = DEFAULT_JOB_ID +) -> dict[str, _t.Any]: """Builds an SQS message dict as returned by `ReceiveMessage`.""" - return {"Body": body, "ReceiptHandle": handle, "MessageId": handle} + msg: dict[str, _t.Any] = {"Body": body, "ReceiptHandle": handle, "MessageId": handle} + if job_id is not None: + msg["MessageAttributes"] = {"job_id": {"DataType": "String", "StringValue": job_id}} + return msg async def test_reader_connects_sqs_client_in_region(aws_client: MagicMock) -> None: @@ -115,11 +125,73 @@ async def test_reader_receives_with_long_polling(aws_client: MagicMock) -> None: messages = await reader._receive() aws_client.receive_message.assert_awaited_once_with( - QueueUrl="https://sqs.test/1/messages", MaxNumberOfMessages=5, WaitTimeSeconds=7 + QueueUrl="https://sqs.test/1/messages", + MaxNumberOfMessages=5, + WaitTimeSeconds=7, + MessageAttributeNames=["job_id"], ) assert [msg["ReceiptHandle"] for msg in messages] == ["h1"] +async def test_reader_filters_messages_by_job_id_and_acks_unmatched( + aws_client: MagicMock, +) -> None: + """Tests that only messages matching `job_id` are returned and others are discarded.""" + sns_matching = { + "Body": json.dumps( + { + "Type": "Notification", + "Message": '{"x": 2}', + "MessageAttributes": {"job_id": {"Type": "String", "Value": DEFAULT_JOB_ID}}, + } + ), + "ReceiptHandle": "h-sns-match", + "MessageId": "h-sns-match", + } + sns_wrong = { + "Body": json.dumps( + { + "Type": "Notification", + "Message": '{"x": 99}', + "MessageAttributes": {"job_id": {"Type": "String", "Value": "Job_other"}}, + } + ), + "ReceiptHandle": "h-sns-wrong", + "MessageId": "h-sns-wrong", + } + aws_client.receive_message.return_value = { + "Messages": [ + sqs_message('{"x": 1}', "h-direct-match", job_id=DEFAULT_JOB_ID), + sqs_message('{"x": 98}', "h-direct-wrong", job_id="Job_other"), + sqs_message('{"x": 97}', "h-missing", job_id=None), + sns_matching, + sns_wrong, + ] + } + reader = make_reader() + await reader._connect() + + messages = await reader._receive() + + assert [msg["ReceiptHandle"] for msg in messages] == ["h-direct-match", "h-sns-match"] + aws_client.delete_message_batch.assert_awaited_once_with( + QueueUrl="https://sqs.test/1/messages", + Entries=[ + {"Id": "0", "ReceiptHandle": "h-direct-wrong"}, + {"Id": "1", "ReceiptHandle": "h-missing"}, + {"Id": "2", "ReceiptHandle": "h-sns-wrong"}, + ], + ) + + +def test_extract_job_id_handles_malformed_attributes() -> None: + """Tests that malformed message attributes or bodies resolve to `None` cleanly.""" + assert aws_messaging_io._extract_job_id({}) is None + assert aws_messaging_io._extract_job_id({"MessageAttributes": {"job_id": "bad"}}) is None + assert aws_messaging_io._extract_job_id({"Body": 123}) is None + assert aws_messaging_io._extract_job_id({"Body": '{"Type": "Notification"}'}) is None + + 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) @@ -333,8 +405,16 @@ async def test_writer_publishes_each_message(aws_client: MagicMock) -> None: 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}'), + call( + TopicArn="arn:aws:sns:eu-west-1:1:messages", + Message='{"x": 1}', + MessageAttributes={"job_id": {"DataType": "String", "StringValue": DEFAULT_JOB_ID}}, + ), + call( + TopicArn="arn:aws:sns:eu-west-1:1:messages", + Message='{"x": 2}', + MessageAttributes={"job_id": {"DataType": "String", "StringValue": DEFAULT_JOB_ID}}, + ), ] diff --git a/tests/unit/test_gcp_pubsub_io.py b/tests/unit/test_gcp_pubsub_io.py index bdacdd05..6c8ebb39 100644 --- a/tests/unit/test_gcp_pubsub_io.py +++ b/tests/unit/test_gcp_pubsub_io.py @@ -24,11 +24,15 @@ from tests import conftest -def make_message(data: bytes) -> MagicMock: +DEFAULT_JOB_ID = "Job_test12345678" + + +def make_message(data: bytes, *, job_id: _t.Optional[str] = DEFAULT_JOB_ID) -> MagicMock: """Builds a PubSub `ReceivedMessage` double carrying `data`.""" wrapper = MagicMock() wrapper.message.data = data - wrapper.ack_id = f"ack-{data!r}" + wrapper.message.attributes = {"job_id": job_id} if job_id is not None else {} + wrapper.ack_id = f"ack-{data!r}-{job_id}" return wrapper @@ -66,6 +70,7 @@ def publisher(monkeypatch: pytest.MonkeyPatch) -> MagicMock: def make_reader(**kwargs: _t.Any) -> GCPPubSubDataReader: """Builds a reader for the fixed test project/subscription.""" kwargs.setdefault("field_names", ["x", "y"]) + kwargs.setdefault("job_id", DEFAULT_JOB_ID) return GCPPubSubDataReader( name="gcp-reader", project_id="test-project", @@ -77,6 +82,7 @@ def make_reader(**kwargs: _t.Any) -> GCPPubSubDataReader: def make_writer(**kwargs: _t.Any) -> GCPPubSubDataWriter: """Builds a writer for the fixed test project/topic.""" kwargs.setdefault("field_names", ["x", "y"]) + kwargs.setdefault("job_id", DEFAULT_JOB_ID) return GCPPubSubDataWriter( name="gcp-writer", project_id="test-project", @@ -129,6 +135,34 @@ async def test_reader_pulls_from_subscription_path(subscriber: MagicMock) -> Non assert len(messages) == 1 +async def test_reader_filters_messages_by_job_id_and_acks_unmatched( + subscriber: MagicMock, +) -> None: + """Tests that only messages matching `job_id` are returned and others are acknowledged.""" + matched = make_message(b'{"x": 1}', job_id=DEFAULT_JOB_ID) + wrong_job = make_message(b'{"x": 2}', job_id="Job_other") + no_job = make_message(b'{"x": 3}', job_id=None) + subscriber.pull.return_value = MagicMock(received_messages=[matched, wrong_job, no_job]) + reader = make_reader() + await reader._connect() + + messages = await reader._receive() + + assert messages == [matched] + subscriber.acknowledge.assert_called_once_with( + request={ + "subscription": "projects/test-project/subscriptions/test-subscription", + "ack_ids": [wrong_job.ack_id, no_job.ack_id], + } + ) + + +def test_extract_job_id_handles_missing_attributes() -> None: + """Tests that a message wrapper with non-mapping attributes returns `None`.""" + wrapper = SimpleNamespace(message=SimpleNamespace(attributes=None)) + assert gcp_pubsub_io._extract_job_id(wrapper) is None + + 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=[]) @@ -327,6 +361,10 @@ async def test_writer_publishes_encoded_records(publisher: MagicMock) -> None: b'{"x": 1}', b'{"x": 2}', ] + assert [published.kwargs["job_id"] for published in publisher.publish.call_args_list] == [ + DEFAULT_JOB_ID, + DEFAULT_JOB_ID, + ] assert [published.args for published in publisher.publish.call_args_list] == [ ("projects/test-project/topics/test-topic",) ] * 2 diff --git a/tests/unit/test_message_data_reader.py b/tests/unit/test_message_data_reader.py index 2d9265aa..5a4b4e5c 100644 --- a/tests/unit/test_message_data_reader.py +++ b/tests/unit/test_message_data_reader.py @@ -15,6 +15,8 @@ NoMoreDataException, ) from plugboard.library.message_reader import MessageDataReader +from plugboard.state import DictStateBackend +from plugboard.utils.entities import EntityIdGen from plugboard.utils.retry import RetryPolicy @@ -445,3 +447,29 @@ async def test_retries_exhausted_after_failed_reconnect_raise_broker_error( assert reader.connects == 1 # Only `init` connected; the reconnects all failed. assert backoff_delays == [1.0, 2.0] await reader.destroy() + + +async def test_reader_job_id_resolution(job_id_ctx: str) -> None: + """Tests `job_id` precedence: explicit > state > DI context > generated fallback.""" + from_ctx = make_reader([]) + assert from_ctx.job_id == job_id_ctx + + explicit = make_reader([], job_id="Job_explicit123") + assert explicit.job_id == "Job_explicit123" + + state = DictStateBackend(job_id="Job_state1234567") + await state.init() + from_state = make_reader([], state=state) + assert from_state.job_id == "Job_state1234567" + + +def test_reader_job_id_standalone_fallback() -> None: + """Tests that a standalone reader outside a DI context generates a stable `job_id`.""" + reader = make_reader([]) + resolved = reader.job_id + assert EntityIdGen.is_job_id(resolved) + assert reader.job_id == resolved + # Re-resolving inside the component's own context hits the active DI context path. + reader._job_id = None + with reader._job_id_ctx(): + assert EntityIdGen.is_job_id(reader.job_id) diff --git a/tests/unit/test_message_data_writer.py b/tests/unit/test_message_data_writer.py index c4370263..9b598528 100644 --- a/tests/unit/test_message_data_writer.py +++ b/tests/unit/test_message_data_writer.py @@ -19,6 +19,8 @@ iter_records, ) from plugboard.schemas import ConnectorSpec +from plugboard.state import DictStateBackend +from plugboard.utils.entities import EntityIdGen from plugboard.utils.retry import RetryPolicy @@ -352,3 +354,28 @@ 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] + + +async def test_writer_job_id_resolution(job_id_ctx: str) -> None: + """Tests `job_id` precedence: explicit > state > DI context > generated fallback.""" + from_ctx = make_writer(["x"]) + assert from_ctx.job_id == job_id_ctx + + explicit = make_writer(["x"], job_id="Job_explicit123") + assert explicit.job_id == "Job_explicit123" + + state = DictStateBackend(job_id="Job_state1234567") + await state.init() + from_state = make_writer(["x"], state=state) + assert from_state.job_id == "Job_state1234567" + + +def test_writer_job_id_standalone_fallback() -> None: + """Tests that a standalone writer outside a DI context generates a stable `job_id`.""" + writer = make_writer(["x"]) + resolved = writer.job_id + assert EntityIdGen.is_job_id(resolved) + assert writer.job_id == resolved + writer._job_id = None + with writer._job_id_ctx(): + assert EntityIdGen.is_job_id(writer.job_id) From eba84515162153aa03e3293fd92c3339e658251a Mon Sep 17 00:00:00 2001 From: Chris Knight Date: Wed, 4 Nov 2026 13:20:23 +0100 Subject: [PATCH 19/21] fix: lock broker-side discards and cover the discard/retry contracts in tests - Take the connection lock around the in-`_receive` discard of non-matching messages, matching `_ack_pending`: without it a concurrent reconnect could swap the client mid-delete/mid-ack. Also filter in a single pass so the SQS/SNS body is parsed once per message. - Add tests that a failed discard surfaces as a transient (retryable) broker error on both SQS and PubSub, and a test locking in the frozen/shared default `RetryPolicy` contract that makes the constructor default safe. --- plugboard/library/aws_messaging_io.py | 11 ++++++++--- plugboard/library/gcp_pubsub_io.py | 11 ++++++++--- tests/unit/test_aws_messaging_io.py | 13 +++++++++++++ tests/unit/test_gcp_pubsub_io.py | 13 +++++++++++++ tests/unit/test_message_data_writer.py | 14 ++++++++++++++ 5 files changed, 56 insertions(+), 6 deletions(-) diff --git a/plugboard/library/aws_messaging_io.py b/plugboard/library/aws_messaging_io.py index f303fbf5..37de35c2 100644 --- a/plugboard/library/aws_messaging_io.py +++ b/plugboard/library/aws_messaging_io.py @@ -260,10 +260,15 @@ async def _receive(self) -> list[_t.Any]: raise _wrap_error(error, exhausted=True) from error messages = response.get("Messages", []) target_job_id = self.job_id - matching = [msg for msg in messages if _extract_job_id(msg) == target_job_id] - unmatched = [msg for msg in messages if _extract_job_id(msg) != target_job_id] + matching: list[_t.Any] = [] + unmatched: list[_t.Any] = [] + for msg in messages: + (matching if _extract_job_id(msg) == target_job_id else unmatched).append(msg) if unmatched: - await self._ack(unmatched) + # Discarded under the connection lock, like every other acknowledgment, so a + # concurrent reconnect cannot swap the client mid-delete. + async with self._connection_lock: + await self._ack(unmatched) return matching async def _convert(self, data: list[_t.Any]) -> dict[str, deque]: diff --git a/plugboard/library/gcp_pubsub_io.py b/plugboard/library/gcp_pubsub_io.py index 36f7035c..7b3c99c9 100644 --- a/plugboard/library/gcp_pubsub_io.py +++ b/plugboard/library/gcp_pubsub_io.py @@ -197,10 +197,15 @@ async def _receive(self) -> list[_t.Any]: raise _wrap_error(error, exhausted=True) from error messages = list(response.received_messages) target_job_id = self.job_id - matching = [msg for msg in messages if _extract_job_id(msg) == target_job_id] - unmatched = [msg for msg in messages if _extract_job_id(msg) != target_job_id] + matching: list[_t.Any] = [] + unmatched: list[_t.Any] = [] + for msg in messages: + (matching if _extract_job_id(msg) == target_job_id else unmatched).append(msg) if unmatched: - await self._ack(unmatched) + # Discarded under the connection lock, like every other acknowledgment, so a + # concurrent reconnect cannot swap the client mid-ack. + async with self._connection_lock: + await self._ack(unmatched) return matching async def _convert(self, data: list[_t.Any]) -> dict[str, deque]: diff --git a/tests/unit/test_aws_messaging_io.py b/tests/unit/test_aws_messaging_io.py index fc7f14f4..7d6662da 100644 --- a/tests/unit/test_aws_messaging_io.py +++ b/tests/unit/test_aws_messaging_io.py @@ -184,6 +184,19 @@ async def test_reader_filters_messages_by_job_id_and_acks_unmatched( ) +async def test_reader_discard_failure_wraps_transient_error(aws_client: MagicMock) -> None: + """Tests that a failed discard of an unmatched message surfaces as retryable.""" + aws_client.receive_message.return_value = { + "Messages": [sqs_message('{"x": 99}', "h-wrong", job_id="Job_other")] + } + aws_client.delete_message_batch.side_effect = client_error("ThrottlingException") + reader = make_reader() + await reader._connect() + + with pytest.raises(MessageBrokerTransientError): + await reader._receive() + + def test_extract_job_id_handles_malformed_attributes() -> None: """Tests that malformed message attributes or bodies resolve to `None` cleanly.""" assert aws_messaging_io._extract_job_id({}) is None diff --git a/tests/unit/test_gcp_pubsub_io.py b/tests/unit/test_gcp_pubsub_io.py index 6c8ebb39..f6397266 100644 --- a/tests/unit/test_gcp_pubsub_io.py +++ b/tests/unit/test_gcp_pubsub_io.py @@ -157,6 +157,19 @@ async def test_reader_filters_messages_by_job_id_and_acks_unmatched( ) +async def test_reader_discard_failure_wraps_transient_error(subscriber: MagicMock) -> None: + """Tests that a failed discard of an unmatched message surfaces as retryable.""" + subscriber.pull.return_value = MagicMock( + received_messages=[make_message(b'{"x": 99}', job_id="Job_other")] + ) + subscriber.acknowledge.side_effect = gcp_exceptions.DeadlineExceeded("deadline") + reader = make_reader() + await reader._connect() + + with pytest.raises(MessageBrokerTransientError): + await reader._receive() + + def test_extract_job_id_handles_missing_attributes() -> None: """Tests that a message wrapper with non-mapping attributes returns `None`.""" wrapper = SimpleNamespace(message=SimpleNamespace(attributes=None)) diff --git a/tests/unit/test_message_data_writer.py b/tests/unit/test_message_data_writer.py index 9b598528..b74350f5 100644 --- a/tests/unit/test_message_data_writer.py +++ b/tests/unit/test_message_data_writer.py @@ -356,6 +356,20 @@ def test_encode_records_bytes_preserves_binary_payloads() -> None: assert encode_records_bytes({"x": deque([payload])}, parse_json=False) == [payload] +def test_default_retry_policy_is_shared_and_immutable() -> None: + """Tests why instantiating `RetryPolicy` in the signature default is safe. + + The default is evaluated once, so every writer that does not override it shares one + instance. That is safe only while the policy stays immutable, so this test locks in + the frozen contract that makes the shared default correct. + """ + first = RecordingWriter(name="first", field_names=["x"], topic="test-topic") + second = RecordingWriter(name="second", field_names=["x"], topic="test-topic") + assert first._retry_policy is second._retry_policy + with pytest.raises(AttributeError, match="immutable"): + setattr(first._retry_policy, "max_retries", 10) + + async def test_writer_job_id_resolution(job_id_ctx: str) -> None: """Tests `job_id` precedence: explicit > state > DI context > generated fallback.""" from_ctx = make_writer(["x"]) From b4294116c7c8a7b607d0c1e6b65a21e28bc84bbb Mon Sep 17 00:00:00 2001 From: Chris Knight Date: Wed, 4 Nov 2026 13:33:04 +0100 Subject: [PATCH 20/21] docs: document broker credentials and classify auth-expiry failures - Add a Credentials section to the message-data guide: the mechanisms each provider picks up (AWS credential chain incl. AWS_PROFILE/SSO and instance roles; GCP Application Default Credentials), that Plugboard never handles authentication itself, and how long-running jobs interact with expiring sessions (SSO sessions expire and need interactive re-login; refreshable credentials renew automatically; retries re-run credential resolution). - Classify google-auth RefreshError as permanent on the PubSub side so a dead refresh token fails fast instead of burning the retry budget; GCP UNAUTHENTICATED was already covered via its Unauthorized superclass. - Test both paths: TokenRetrievalError (SSO/credential expiry) stays transient so reconnects re-run credential resolution, and RefreshError maps to a permanent error. --- docs/usage/message-data.md | 43 +++++++++++++++++++++++++++-- plugboard/library/gcp_pubsub_io.py | 6 +++- tests/unit/test_aws_messaging_io.py | 21 +++++++++++++- tests/unit/test_gcp_pubsub_io.py | 2 ++ 4 files changed, 67 insertions(+), 5 deletions(-) diff --git a/docs/usage/message-data.md b/docs/usage/message-data.md index 392ed812..f86f5ed6 100644 --- a/docs/usage/message-data.md +++ b/docs/usage/message-data.md @@ -130,6 +130,43 @@ Connection details are passed explicitly or read from the environment (explicit always win), so the same model can target a different account or region without code changes. The message-data components read the `GCP_PUBSUB_PROJECT_ID` and `AWS_REGION` environment variables; see [Configuration](configuration.md#message-data-brokers) for -how these are resolved. Credentials are never configured through Plugboard — the clients -use their normal provider mechanisms (Application Default Credentials for PubSub, and the -standard AWS credential chain for SQS/SNS). +how these are resolved. Credentials are never configured through Plugboard — see +[Credentials](#credentials) for the mechanisms each provider picks up. + +## Credentials + +Plugboard never configures, stores or refreshes credentials itself: the broker client +libraries resolve them from your environment, so authenticate **before** starting a run. +The mechanisms below are picked up in the provider's own precedence order — see the +provider documentation for full details. + +**AWS (SQS/SNS)** uses the +[standard credential chain](https://docs.aws.amazon.com/sdkref/latest/guide/standardized-credentials.html): + +- Static keys in the environment: `AWS_ACCESS_KEY_ID` + `AWS_SECRET_ACCESS_KEY` + (plus `AWS_SESSION_TOKEN` for temporary credentials). +- Shared credentials/config files, optionally selected with `AWS_PROFILE` — including + IAM Identity Center (SSO) and assume-role profiles. +- The instance or container role when running on EC2/ECS/EKS. + +**GCP (PubSub)** uses +[Application Default Credentials](https://cloud.google.com/docs/authentication/application-default-credentials): + +- A service account key file pointed at by `GOOGLE_APPLICATION_CREDENTIALS`. +- User credentials from `gcloud auth application-default login`. +- The metadata server / Workload Identity when running on GCP infrastructure. + +### Long-running jobs and expiring sessions + +Credentials that renew themselves — instance/container roles, service account keys, and +ADC refresh tokens — keep working for the life of a long-running job: the client +libraries refresh them automatically as access tokens expire. + +IAM Identity Center (SSO) sessions are the exception: the login token is valid only for +your organisation's configured session length (1-12 hours, commonly 8) and can only be +renewed by an interactive `aws sso login`. If a job outlives the session, the next broker +call fails; the component retries with backoff — each reconnect re-runs credential +resolution, which heals any provider that can refresh — and then fails with the +provider's own error rather than a synthetic one. Prefer a non-interactive credential +source (static keys, an instance role, or a service account) for jobs that may run +longer than a login session. diff --git a/plugboard/library/gcp_pubsub_io.py b/plugboard/library/gcp_pubsub_io.py index 7b3c99c9..7c185c11 100644 --- a/plugboard/library/gcp_pubsub_io.py +++ b/plugboard/library/gcp_pubsub_io.py @@ -25,16 +25,20 @@ try: from google.api_core import exceptions as gcp_exceptions + from google.auth import exceptions as auth_exceptions from google.cloud import pubsub_v1 # PubSub rejections that retrying cannot resolve, so the retry loop should stop - # rather than back off repeatedly against a guaranteed failure. + # rather than back off repeatedly against a guaranteed failure. `RefreshError` is + # the google-auth failure to renew an expired access token (a revoked or dead + # refresh token): reconnecting cannot fix it, so it fails fast like the others. _PERMANENT_ERRORS: tuple[type[Exception], ...] = ( gcp_exceptions.PermissionDenied, gcp_exceptions.Unauthorized, gcp_exceptions.InvalidArgument, gcp_exceptions.MethodNotImplemented, gcp_exceptions.FailedPrecondition, + auth_exceptions.RefreshError, ) # PubSub failures that are expected to clear on their own, so they are worth retrying. diff --git a/tests/unit/test_aws_messaging_io.py b/tests/unit/test_aws_messaging_io.py index 7d6662da..0268df72 100644 --- a/tests/unit/test_aws_messaging_io.py +++ b/tests/unit/test_aws_messaging_io.py @@ -8,7 +8,7 @@ import typing as _t from unittest.mock import AsyncMock, MagicMock, call -from botocore.exceptions import BotoCoreError, ClientError +from botocore.exceptions import BotoCoreError, ClientError, TokenRetrievalError import pytest from plugboard.exceptions import ( @@ -393,6 +393,25 @@ async def test_reader_transport_error_is_transient(aws_client: MagicMock) -> Non await reader._receive() +async def test_reader_credential_expiry_is_transient_for_reconnect(aws_client: MagicMock) -> None: + """Tests that a credential-expiry error stays retryable so reconnects can heal it. + + botocore raises `TokenRetrievalError` when a profile's credentials cannot be + refreshed, for example an expired IAM Identity Center (SSO) session. Reconnects + re-run credential resolution, so the error must remain retryable: a refreshable + provider recovers, and a genuinely dead session exhausts the retries and surfaces + botocore's own message. + """ + aws_client.receive_message.side_effect = TokenRetrievalError( + provider="sso", error_msg="SSO session expired" + ) + reader = make_reader() + await reader._connect() + + with pytest.raises(MessageBrokerTransientError): + await reader._receive() + + 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") diff --git a/tests/unit/test_gcp_pubsub_io.py b/tests/unit/test_gcp_pubsub_io.py index f6397266..7cb2cbdb 100644 --- a/tests/unit/test_gcp_pubsub_io.py +++ b/tests/unit/test_gcp_pubsub_io.py @@ -10,6 +10,7 @@ from unittest.mock import MagicMock, call from google.api_core import exceptions as gcp_exceptions +from google.auth.exceptions import RefreshError import pytest from plugboard.exceptions import ( @@ -288,6 +289,7 @@ async def test_reader_disconnect_closes_subscriber(subscriber: MagicMock) -> Non (gcp_exceptions.PermissionDenied("nope"), MessageBrokerPermanentError), (gcp_exceptions.InvalidArgument("bad"), MessageBrokerPermanentError), (gcp_exceptions.ServiceUnavailable("busy"), MessageBrokerTransientError), + (RefreshError("refresh token revoked"), MessageBrokerPermanentError), ], ) async def test_reader_maps_broker_errors( From 8b3584293d649b9409820e613105dc66b9cc71c9 Mon Sep 17 00:00:00 2001 From: Chris Knight Date: Wed, 4 Nov 2026 14:00:46 +0100 Subject: [PATCH 21/21] chore: merge main and bump the message broker dependencies to latest Merge main (numpy channel fix, Ray TaskGroup fail-fast, deferred external resources, dependency bumps) and refresh the lockfile for the newly added broker dependencies and their subtrees: - aiobotocore 3.9.2 -> 3.10.0 (botocore/boto3 1.43.106 -> 1.43.117) - google-cloud-pubsub 2.42.0 -> 2.43.0 - google-auth 2.49.1 -> 2.62.1 - google-api-core 2.30.2 -> 2.43.2 - grpcio / grpcio-status 1.80.0 -> 1.85.1 (clears the PQC FutureWarning: google-api-core and google-auth now require grpcio >= 1.83.0 by April 2027) Major-version constraints in pyproject are unchanged: both new extras are on their current majors (aiobotocore 3.x, google-cloud-pubsub 2.x). Verified against the upgraded deps: all non-Ray unit tests pass, and the real SNS->SQS and Pub/Sub round trips pass against live infrastructure. --- uv.lock | 165 ++++++++++++++++++++++++++------------------------------ 1 file changed, 77 insertions(+), 88 deletions(-) diff --git a/uv.lock b/uv.lock index dca043e8..7b2ded23 100644 --- a/uv.lock +++ b/uv.lock @@ -52,7 +52,7 @@ wheels = [ [[package]] name = "aiobotocore" -version = "3.9.2" +version = "3.10.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "aiohttp" }, @@ -63,9 +63,9 @@ dependencies = [ { name = "python-dateutil" }, { name = "wrapt" }, ] -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" } +sdist = { url = "https://files.pythonhosted.org/packages/16/40/659f2ac3915d329e3e8e5bf674d35ce1774f3117c6f60339d5e05e505c36/aiobotocore-3.10.0.tar.gz", hash = "sha256:35480ef11c51e23087273b85d97f80da2f356a46ce5f3fb88e24483cd343ffdb", size = 530781, upload-time = "2026-10-23T10:44:50.843Z" } wheels = [ - { 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" }, + { url = "https://files.pythonhosted.org/packages/7f/bc/3bb8965204ba51bbf4d021bc89d1357cd2179a17ef71b5fc429e21b511a2/aiobotocore-3.10.0-py3-none-any.whl", hash = "sha256:de0ccbcb82b6b0009f1b6667f48442e9d0432f9b8bfe6172fa51777cbea89d03", size = 102759, upload-time = "2026-10-23T10:44:49.109Z" }, ] [[package]] @@ -662,30 +662,30 @@ wheels = [ [[package]] name = "boto3" -version = "1.43.106" +version = "1.43.117" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "botocore" }, { name = "jmespath" }, { name = "s3transfer" }, ] -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" } +sdist = { url = "https://files.pythonhosted.org/packages/94/a7/61f0295f79dfdcebb5aa01f86c614fac7e1294faabbd66259d3609c50539/boto3-1.43.117.tar.gz", hash = "sha256:f3d314908d866b577ee6b74ed432c21e565c4c4dbd640ef3fd5e0da30fa13a98", size = 112677, upload-time = "2026-10-21T19:30:03.849Z" } wheels = [ - { 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" }, + { url = "https://files.pythonhosted.org/packages/f3/3b/22b05a545fb95983733d84681364c6b0a36afb0710d832f77732d1fc4c0b/boto3-1.43.117-py3-none-any.whl", hash = "sha256:a5b6cd7e67653987fcf0e212187c16f7b58af2ae192e368f7212ddab7c05b25d", size = 140043, upload-time = "2026-10-21T19:30:02.114Z" }, ] [[package]] name = "botocore" -version = "1.43.106" +version = "1.43.117" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "jmespath" }, { name = "python-dateutil" }, { name = "urllib3" }, ] -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" } +sdist = { url = "https://files.pythonhosted.org/packages/6e/a0/d0a5c44ad6a01f195681ddcb4ce203f06c712811daf1d15f8f8785e8a9f8/botocore-1.43.117.tar.gz", hash = "sha256:71993702832cc37c2495048b9945f2929d306fccb3fb767e9964e437e637fa6c", size = 16366888, upload-time = "2026-10-21T19:29:59.274Z" } wheels = [ - { 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" }, + { url = "https://files.pythonhosted.org/packages/10/94/050e35c9febf6dd2fce4a76be587ae4a323964865700467d9d7b082f7133/botocore-1.43.117-py3-none-any.whl", hash = "sha256:5920ef9b3002e1361846549078549d9154a6e0004c22853c01ce7b279d05d495", size = 16061765, upload-time = "2026-10-21T19:29:56.001Z" }, ] [[package]] @@ -1437,18 +1437,19 @@ wheels = [ [[package]] name = "google-api-core" -version = "2.30.2" +version = "2.43.2" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "google-auth" }, { name = "googleapis-common-protos" }, + { name = "opentelemetry-api" }, { name = "proto-plus" }, { name = "protobuf" }, { name = "requests" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/1a/2e/83ca41eb400eb228f9279ec14ed66f6475218b59af4c6daec2d5a509fe83/google_api_core-2.30.2.tar.gz", hash = "sha256:9a8113e1a88bdc09a7ff629707f2214d98d61c7f6ceb0ea38c42a095d02dc0f9", size = 176862, upload-time = "2026-04-02T21:23:44.876Z" } +sdist = { url = "https://files.pythonhosted.org/packages/a4/74/307919972c4dc84598f4128a8d6ae05ac3d9e6ca73fe60a7bb93c670931c/google_api_core-2.43.2.tar.gz", hash = "sha256:5dfc5e9d52982bcc03bd6af380a893365b5fdf907dc527f63d3a372833c33329", size = 286250, upload-time = "2026-10-29T17:20:40.193Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/84/e1/ebd5100cbb202e561c0c8b59e485ef3bd63fa9beb610f3fdcaea443f0288/google_api_core-2.30.2-py3-none-any.whl", hash = "sha256:a4c226766d6af2580577db1f1a51bf53cd262f722b49731ce7414c43068a9594", size = 173236, upload-time = "2026-04-02T21:23:06.395Z" }, + { url = "https://files.pythonhosted.org/packages/b0/1d/a0120c7d937597293fce74510c80fe0cec8e45b1d702a816b567b6754c25/google_api_core-2.43.2-py3-none-any.whl", hash = "sha256:5a9d21baab97fd3c0547ccfc186a9e64f4227d8d7026f393835316400b29a6d4", size = 224304, upload-time = "2026-10-29T17:20:17.733Z" }, ] [package.optional-dependencies] @@ -1459,15 +1460,15 @@ grpc = [ [[package]] name = "google-auth" -version = "2.49.1" +version = "2.62.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "cryptography" }, { name = "pyasn1-modules" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/ea/80/6a696a07d3d3b0a92488933532f03dbefa4a24ab80fb231395b9a2a1be77/google_auth-2.49.1.tar.gz", hash = "sha256:16d40da1c3c5a0533f57d268fe72e0ebb0ae1cc3b567024122651c045d879b64", size = 333825, upload-time = "2026-03-12T19:30:58.135Z" } +sdist = { url = "https://files.pythonhosted.org/packages/c4/38/d79ab09dcfc1e1638d887c8f65657619d3ebc454165c558ef79355971ce3/google_auth-2.62.1.tar.gz", hash = "sha256:3c8e1ce2553784ae17fd220e6f88df3281d86ab98019278f436d06dffab17a9e", size = 398519, upload-time = "2026-11-03T03:06:03.008Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/e9/eb/c6c2478d8a8d633460be40e2a8a6f8f429171997a35a96f81d3b680dec83/google_auth-2.49.1-py3-none-any.whl", hash = "sha256:195ebe3dca18eddd1b3db5edc5189b76c13e96f29e73043b923ebcf3f1a860f7", size = 240737, upload-time = "2026-03-12T19:30:53.159Z" }, + { url = "https://files.pythonhosted.org/packages/14/dd/2e11f0976ecb3ce48ef3f324837460e9ea50071fa4eb01f1e49df87d3cee/google_auth-2.62.1-py3-none-any.whl", hash = "sha256:015399a5abde39033c69404cc1c079e6d4c2de2993e164c9f4eac6e5376ad72b", size = 267975, upload-time = "2026-11-03T03:06:01.422Z" }, ] [[package]] @@ -1498,7 +1499,7 @@ wheels = [ [[package]] name = "google-cloud-pubsub" -version = "2.42.0" +version = "2.43.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "google-api-core", extra = ["grpc"] }, @@ -1511,9 +1512,9 @@ dependencies = [ { name = "proto-plus" }, { name = "protobuf" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/41/60/21e7de2ae629d8a025838b4e53f197aea2b45d878c6e37f222187b2ecf84/google_cloud_pubsub-2.42.0.tar.gz", hash = "sha256:92ed998c41ce7d4f6b5b7c950df8291abd27c4d50116714e5c46fac47313e74d", size = 397642, upload-time = "2026-10-01T18:17:44.845Z" } +sdist = { url = "https://files.pythonhosted.org/packages/25/ce/210febb5edd160e9490f1a7a8190fd870f0997fa8be558b2c902461a27d7/google_cloud_pubsub-2.43.0.tar.gz", hash = "sha256:9a0e01876024ecd5e101a3c41cd67fa50ef0c4e39aa56f53b81524599f90e346", size = 398611, upload-time = "2026-10-15T17:49:51.415Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/79/7d/a3db64849bc41a52b9c8625f5c76fbdbba02ad2d89b91d38c80a304ade51/google_cloud_pubsub-2.42.0-py3-none-any.whl", hash = "sha256:6eff5d2155b5f6967885b0e7f3da6fda91d2cfb15227fffeaef39b7339044a3d", size = 322310, upload-time = "2026-10-01T18:11:47.708Z" }, + { url = "https://files.pythonhosted.org/packages/ea/fe/05daee3bd9ca21bf90ce7cf0f62b95bae07c397e107dc1b6762a809a0098/google_cloud_pubsub-2.43.0-py3-none-any.whl", hash = "sha256:db5436dd2d99774171d13c1bf1fb960b6bf3aa23c15694c50090aa569316d19a", size = 323338, upload-time = "2026-10-15T17:49:22.528Z" }, ] [[package]] @@ -1709,57 +1710,67 @@ wheels = [ [[package]] name = "grpcio" -version = "1.80.0" +version = "1.85.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "typing-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/b7/48/af6173dbca4454f4637a4678b67f52ca7e0c1ed7d5894d89d434fecede05/grpcio-1.80.0.tar.gz", hash = "sha256:29aca15edd0688c22ba01d7cc01cb000d72b2033f4a3c72a81a19b56fd143257", size = 12978905, upload-time = "2026-03-30T08:49:10.502Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/5c/e8/a2b749265eb3415abc94f2e619bbd9e9707bebdda787e61c593004ec927a/grpcio-1.80.0-cp312-cp312-linux_armv7l.whl", hash = "sha256:c624cc9f1008361014378c9d776de7182b11fe8b2e5a81bc69f23a295f2a1ad0", size = 6015616, upload-time = "2026-03-30T08:47:13.428Z" }, - { url = "https://files.pythonhosted.org/packages/3e/97/b1282161a15d699d1e90c360df18d19165a045ce1c343c7f313f5e8a0b77/grpcio-1.80.0-cp312-cp312-macosx_11_0_universal2.whl", hash = "sha256:f49eddcac43c3bf350c0385366a58f36bed8cc2c0ec35ef7b74b49e56552c0c2", size = 12014204, upload-time = "2026-03-30T08:47:15.873Z" }, - { url = "https://files.pythonhosted.org/packages/6e/5e/d319c6e997b50c155ac5a8cb12f5173d5b42677510e886d250d50264949d/grpcio-1.80.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:d334591df610ab94714048e0d5b4f3dd5ad1bee74dfec11eee344220077a79de", size = 6563866, upload-time = "2026-03-30T08:47:18.588Z" }, - { url = "https://files.pythonhosted.org/packages/ae/f6/fdd975a2cb4d78eb67769a7b3b3830970bfa2e919f1decf724ae4445f42c/grpcio-1.80.0-cp312-cp312-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:0cb517eb1d0d0aaf1d87af7cc5b801d686557c1d88b2619f5e31fab3c2315921", size = 7273060, upload-time = "2026-03-30T08:47:21.113Z" }, - { url = "https://files.pythonhosted.org/packages/db/f0/a3deb5feba60d9538a962913e37bd2e69a195f1c3376a3dd44fe0427e996/grpcio-1.80.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:4e78c4ac0d97dc2e569b2f4bcbbb447491167cb358d1a389fc4af71ab6f70411", size = 6782121, upload-time = "2026-03-30T08:47:23.827Z" }, - { url = "https://files.pythonhosted.org/packages/ca/84/36c6dcfddc093e108141f757c407902a05085e0c328007cb090d56646cdf/grpcio-1.80.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:2ed770b4c06984f3b47eb0517b1c69ad0b84ef3f40128f51448433be904634cd", size = 7383811, upload-time = "2026-03-30T08:47:26.517Z" }, - { url = "https://files.pythonhosted.org/packages/7c/ef/f3a77e3dc5b471a0ec86c564c98d6adfa3510d38f8ee99010410858d591e/grpcio-1.80.0-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:256507e2f524092f1473071a05e65a5b10d84b82e3ff24c5b571513cfaa61e2f", size = 8393860, upload-time = "2026-03-30T08:47:29.439Z" }, - { url = "https://files.pythonhosted.org/packages/9b/8d/9d4d27ed7f33d109c50d6b5ce578a9914aa68edab75d65869a17e630a8d1/grpcio-1.80.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:9a6284a5d907c37db53350645567c522be314bac859a64a7a5ca63b77bb7958f", size = 7830132, upload-time = "2026-03-30T08:47:33.254Z" }, - { url = "https://files.pythonhosted.org/packages/14/e4/9990b41c6d7a44e1e9dee8ac11d7a9802ba1378b40d77468a7761d1ad288/grpcio-1.80.0-cp312-cp312-win32.whl", hash = "sha256:c71309cfce2f22be26aa4a847357c502db6c621f1a49825ae98aa0907595b193", size = 4140904, upload-time = "2026-03-30T08:47:35.319Z" }, - { url = "https://files.pythonhosted.org/packages/2f/2c/296f6138caca1f4b92a31ace4ae1b87dab692fc16a7a3417af3bb3c805bf/grpcio-1.80.0-cp312-cp312-win_amd64.whl", hash = "sha256:9fe648599c0e37594c4809d81a9e77bd138cc82eb8baa71b6a86af65426723ff", size = 4880944, upload-time = "2026-03-30T08:47:37.831Z" }, - { url = "https://files.pythonhosted.org/packages/2f/3a/7c3c25789e3f069e581dc342e03613c5b1cb012c4e8c7d9d5cf960a75856/grpcio-1.80.0-cp313-cp313-linux_armv7l.whl", hash = "sha256:e9e408fc016dffd20661f0126c53d8a31c2821b5c13c5d67a0f5ed5de93319ad", size = 6017243, upload-time = "2026-03-30T08:47:40.075Z" }, - { url = "https://files.pythonhosted.org/packages/04/19/21a9806eb8240e174fd1ab0cd5b9aa948bb0e05c2f2f55f9d5d7405e6d08/grpcio-1.80.0-cp313-cp313-macosx_11_0_universal2.whl", hash = "sha256:92d787312e613754d4d8b9ca6d3297e69994a7912a32fa38c4c4e01c272974b0", size = 12010840, upload-time = "2026-03-30T08:47:43.11Z" }, - { url = "https://files.pythonhosted.org/packages/18/3a/23347d35f76f639e807fb7a36fad3068aed100996849a33809591f26eca6/grpcio-1.80.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:8ac393b58aa16991a2f1144ec578084d544038c12242da3a215966b512904d0f", size = 6567644, upload-time = "2026-03-30T08:47:46.806Z" }, - { url = "https://files.pythonhosted.org/packages/ff/40/96e07ecb604a6a67ae6ab151e3e35b132875d98bc68ec65f3e5ab3e781d7/grpcio-1.80.0-cp313-cp313-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:68e5851ac4b9afe07e7f84483803ad167852570d65326b34d54ca560bfa53fb6", size = 7277830, upload-time = "2026-03-30T08:47:49.643Z" }, - { url = "https://files.pythonhosted.org/packages/9b/e2/da1506ecea1f34a5e365964644b35edef53803052b763ca214ba3870c856/grpcio-1.80.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:873ff5d17d68992ef6605330127425d2fc4e77e612fa3c3e0ed4e668685e3140", size = 6783216, upload-time = "2026-03-30T08:47:52.817Z" }, - { url = "https://files.pythonhosted.org/packages/44/83/3b20ff58d0c3b7f6caaa3af9a4174d4023701df40a3f39f7f1c8e7c48f9d/grpcio-1.80.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:2bea16af2750fd0a899bf1abd9022244418b55d1f37da2202249ba4ba673838d", size = 7385866, upload-time = "2026-03-30T08:47:55.687Z" }, - { url = "https://files.pythonhosted.org/packages/47/45/55c507599c5520416de5eefecc927d6a0d7af55e91cfffb2e410607e5744/grpcio-1.80.0-cp313-cp313-musllinux_1_2_i686.whl", hash = "sha256:ba0db34f7e1d803a878284cd70e4c63cb6ae2510ba51937bf8f45ba997cefcf7", size = 8391602, upload-time = "2026-03-30T08:47:58.303Z" }, - { url = "https://files.pythonhosted.org/packages/10/bb/dd06f4c24c01db9cf11341b547d0a016b2c90ed7dbbb086a5710df7dd1d7/grpcio-1.80.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:8eb613f02d34721f1acf3626dfdb3545bd3c8505b0e52bf8b5710a28d02e8aa7", size = 7826752, upload-time = "2026-03-30T08:48:01.311Z" }, - { url = "https://files.pythonhosted.org/packages/f9/1e/9d67992ba23371fd63d4527096eb8c6b76d74d52b500df992a3343fd7251/grpcio-1.80.0-cp313-cp313-win32.whl", hash = "sha256:93b6f823810720912fd131f561f91f5fed0fda372b6b7028a2681b8194d5d294", size = 4142310, upload-time = "2026-03-30T08:48:04.594Z" }, - { url = "https://files.pythonhosted.org/packages/cf/e6/283326a27da9e2c3038bc93eeea36fb118ce0b2d03922a9cda6688f53c5b/grpcio-1.80.0-cp313-cp313-win_amd64.whl", hash = "sha256:e172cf795a3ba5246d3529e4d34c53db70e888fa582a8ffebd2e6e48bc0cba50", size = 4882833, upload-time = "2026-03-30T08:48:07.363Z" }, - { url = "https://files.pythonhosted.org/packages/c5/6d/e65307ce20f5a09244ba9e9d8476e99fb039de7154f37fb85f26978b59c3/grpcio-1.80.0-cp314-cp314-linux_armv7l.whl", hash = "sha256:3d4147a97c8344d065d01bbf8b6acec2cf86fb0400d40696c8bdad34a64ffc0e", size = 6017376, upload-time = "2026-03-30T08:48:10.005Z" }, - { url = "https://files.pythonhosted.org/packages/69/10/9cef5d9650c72625a699c549940f0abb3c4bfdb5ed45a5ce431f92f31806/grpcio-1.80.0-cp314-cp314-macosx_11_0_universal2.whl", hash = "sha256:d8e11f167935b3eb089ac9038e1a063e6d7dbe995c0bb4a661e614583352e76f", size = 12018133, upload-time = "2026-03-30T08:48:12.927Z" }, - { url = "https://files.pythonhosted.org/packages/04/82/983aabaad82ba26113caceeb9091706a0696b25da004fe3defb5b346e15b/grpcio-1.80.0-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:f14b618fc30de822681ee986cfdcc2d9327229dc4c98aed16896761cacd468b9", size = 6574748, upload-time = "2026-03-30T08:48:16.386Z" }, - { url = "https://files.pythonhosted.org/packages/07/d7/031666ef155aa0bf399ed7e19439656c38bbd143779ae0861b038ce82abd/grpcio-1.80.0-cp314-cp314-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:4ed39fbdcf9b87370f6e8df4e39ca7b38b3e5e9d1b0013c7b6be9639d6578d14", size = 7277711, upload-time = "2026-03-30T08:48:19.627Z" }, - { url = "https://files.pythonhosted.org/packages/e8/43/f437a78f7f4f1d311804189e8f11fb311a01049b2e08557c1068d470cb2e/grpcio-1.80.0-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:2dcc70e9f0ba987526e8e8603a610fb4f460e42899e74e7a518bf3c68fe1bf05", size = 6785372, upload-time = "2026-03-30T08:48:22.373Z" }, - { url = "https://files.pythonhosted.org/packages/93/3d/f6558e9c6296cb4227faa5c43c54a34c68d32654b829f53288313d16a86e/grpcio-1.80.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:448c884b668b868562b1bda833c5fce6272d26e1926ec46747cda05741d302c1", size = 7395268, upload-time = "2026-03-30T08:48:25.638Z" }, - { url = "https://files.pythonhosted.org/packages/06/21/0fdd77e84720b08843c371a2efa6f2e19dbebf56adc72df73d891f5506f0/grpcio-1.80.0-cp314-cp314-musllinux_1_2_i686.whl", hash = "sha256:a1dc80fe55685b4a543555e6eef975303b36c8db1023b1599b094b92aa77965f", size = 8392000, upload-time = "2026-03-30T08:48:28.974Z" }, - { url = "https://files.pythonhosted.org/packages/f5/68/67f4947ed55d2e69f2cc199ab9fd85e0a0034d813bbeef84df6d2ba4d4b7/grpcio-1.80.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:31b9ac4ad1aa28ffee5503821fafd09e4da0a261ce1c1281c6c8da0423c83b6e", size = 7828477, upload-time = "2026-03-30T08:48:32.054Z" }, - { url = "https://files.pythonhosted.org/packages/44/b6/8d4096691b2e385e8271911a0de4f35f0a6c7d05aff7098e296c3de86939/grpcio-1.80.0-cp314-cp314-win32.whl", hash = "sha256:367ce30ba67d05e0592470428f0ec1c31714cab9ef19b8f2e37be1f4c7d32fae", size = 4218563, upload-time = "2026-03-30T08:48:34.538Z" }, - { url = "https://files.pythonhosted.org/packages/e5/8c/bbe6baf2557262834f2070cf668515fa308b2d38a4bbf771f8f7872a7036/grpcio-1.80.0-cp314-cp314-win_amd64.whl", hash = "sha256:3b01e1f5464c583d2f567b2e46ff0d516ef979978f72091fd81f5ab7fa6e2e7f", size = 5019457, upload-time = "2026-03-30T08:48:37.308Z" }, +sdist = { url = "https://files.pythonhosted.org/packages/f0/03/f9e805772d4d01da2a1754a5f63c58601eacaf8f1e878729217188de7d16/grpcio-1.85.1.tar.gz", hash = "sha256:5f530287d5de5f7f66089ead81de44f6787162944dab3922efa65dfdff108a1d", size = 14091456, upload-time = "2026-10-30T16:01:16.909Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/31/8c/b6d9ea0bf02a6b47ba34c5f936da2921d76f5e40cbb411e9abb246b70d2d/grpcio-1.85.1-cp312-cp312-linux_armv7l.whl", hash = "sha256:1a10d337967ec2be049c6a2345320f6a59f71bf9f54febfd80e5a966cffaa945", size = 6797195, upload-time = "2026-10-30T15:59:35.364Z" }, + { url = "https://files.pythonhosted.org/packages/c6/ee/f60cfacd679fe811944be391c899955c3d9c7e07db922b47e535ddda550c/grpcio-1.85.1-cp312-cp312-macosx_11_0_universal2.whl", hash = "sha256:658a9d89c5736bf4aa3ef4ecbc79528e26d8e23c5e1c83417d2870a1a06063e5", size = 12633025, upload-time = "2026-10-30T15:59:37.66Z" }, + { url = "https://files.pythonhosted.org/packages/f1/0d/5dbe225f0ceb208538dc8c16c5f52b0c402b91611f981604b5e280439dfb/grpcio-1.85.1-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:843aa281a48188621f50880496d35defed51298bcb66d5345d5f282fd3f87354", size = 7170616, upload-time = "2026-10-30T15:59:40.138Z" }, + { url = "https://files.pythonhosted.org/packages/54/35/ce972a4e5bb721ba85f5c9e20a0029253ae0b2deb5276dafa3383e025133/grpcio-1.85.1-cp312-cp312-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:e768556a2346d1d2921f511f3145e2462e4fe68095c31e4db407558bcd3f1d4f", size = 8180786, upload-time = "2026-10-30T15:59:42.7Z" }, + { url = "https://files.pythonhosted.org/packages/00/8d/74a8649d2787dae0a661855cdf09a8e2238fa741689bb368ca1b4f6d5b6f/grpcio-1.85.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:a1a72d19c0ab84b4b73b6c47eeaade76d84b1973a51593ff1160a73df33d03d9", size = 7362744, upload-time = "2026-10-30T15:59:45.076Z" }, + { url = "https://files.pythonhosted.org/packages/03/72/87ef13aba6bd2c5c6d4c5117238b17b1e3c1c1efa26c66cb377053124eb7/grpcio-1.85.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:0e07f477e1d1a05bdbe71d63add964c5457e7ab2e3f5fd32859639ed8ecb976e", size = 7929643, upload-time = "2026-10-30T15:59:47.294Z" }, + { url = "https://files.pythonhosted.org/packages/09/f1/2e75b2906c9643e0be81a74bf23ae10c29661e0199c8f3a6e2c23707e882/grpcio-1.85.1-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:b2846c70ed58c8e149e8e5d867af40805a13b151d84157d8272952d711a77c5c", size = 9181161, upload-time = "2026-10-30T15:59:49.399Z" }, + { url = "https://files.pythonhosted.org/packages/c1/d9/6b9127e8fa3ef14a3ec199436bec6c1a54a0bb429ca261b1a2fb518938d1/grpcio-1.85.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:9998f03ac0dfed226eead6e93cee1b547d7a32e47e1915dee94eff0e113b167e", size = 8318829, upload-time = "2026-10-30T15:59:52.09Z" }, + { url = "https://files.pythonhosted.org/packages/21/46/ff3374d33efee490a0e50e3c3fc59280b7f7fc2ed3ab1ee8c4e9873dc83d/grpcio-1.85.1-cp312-cp312-win32.whl", hash = "sha256:2670c03d154f6bec66a0fef97c607116761b032183187dfdfe229b5975820a4f", size = 4768056, upload-time = "2026-10-30T15:59:54.132Z" }, + { url = "https://files.pythonhosted.org/packages/dd/48/4030a63540089d00a36cea15282623bcf084a5ea4c0ffc4685a42ef0f342/grpcio-1.85.1-cp312-cp312-win_amd64.whl", hash = "sha256:eaad7e6aa9ac3aded4f6bc86a18bf2cff3de356f972bf32a5558df9788085b85", size = 5607218, upload-time = "2026-10-30T15:59:56.721Z" }, + { url = "https://files.pythonhosted.org/packages/c1/05/2d40be7dbfd8e66e8e5675ec695fafacc69823dcd2ef5d387ad7431c7af7/grpcio-1.85.1-cp313-cp313-linux_armv7l.whl", hash = "sha256:e117b074d707d418ff411e0ad6fc86e97e83047b742488b41d7a5b7a62953cd3", size = 6801759, upload-time = "2026-10-30T15:59:58.834Z" }, + { url = "https://files.pythonhosted.org/packages/b0/98/530af74d704dc8e79c79b26888af44ccbdd1827bfabc9cb405b3d4cc184e/grpcio-1.85.1-cp313-cp313-macosx_11_0_universal2.whl", hash = "sha256:6bb4bbb5cb46010687db552fd88852c405595814fb7d435b072124bd6b93f827", size = 12629341, upload-time = "2026-10-30T16:00:01.159Z" }, + { url = "https://files.pythonhosted.org/packages/21/cb/bac12e6c7ab3b1c6b50816fdcbf7f97aec1bd0992cc0d193e332aff716a5/grpcio-1.85.1-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:0bde4264ca71aaf3bc56c587380c426b1c959d1db0ab0a3f632bc66515f9726f", size = 7174387, upload-time = "2026-10-30T16:00:04.847Z" }, + { url = "https://files.pythonhosted.org/packages/90/0a/7770a652508a3a6c7d16fbb81db0f0b2be8a73a469af88a6cdfdc3d29cf2/grpcio-1.85.1-cp313-cp313-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:7b931335c163510fb054c30007fe3556ef3bf354a1c37b5da403f2e52c6bbd95", size = 8185778, upload-time = "2026-10-30T16:00:07.009Z" }, + { url = "https://files.pythonhosted.org/packages/bf/cb/52417fb272b7603d13ec8688e843dfd8df64584da1951ad954b9120d7c80/grpcio-1.85.1-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:de520b1555ae445cdc4f8f1cc7ccba12a19992bfafa7c292ebe7fe1fe5da6c53", size = 7364062, upload-time = "2026-10-30T16:00:09.777Z" }, + { url = "https://files.pythonhosted.org/packages/ba/86/9329799fa9defe9327084b9ca7b9a9745e83e7174bb90f2bfb2f37063967/grpcio-1.85.1-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:2b20ba567ca818b8393fd4668878b5688cab32a46e8e7a4d1a58641fc0072df9", size = 7932775, upload-time = "2026-10-30T16:00:11.956Z" }, + { url = "https://files.pythonhosted.org/packages/48/e4/826e0ae8f89a3d2ab3a15b91ee4ef486bdd7a09fd1ff72e12d01505341dc/grpcio-1.85.1-cp313-cp313-musllinux_1_2_i686.whl", hash = "sha256:f0ebfbd8ba42497e0ff7229c8c5e166e07c2b348462c1481dbd90faa978ad8af", size = 9173461, upload-time = "2026-10-30T16:00:15.696Z" }, + { url = "https://files.pythonhosted.org/packages/88/2d/d4baf37bf1a9ba8222dcb133c6aeb09119b183bd7b471061762a1198461c/grpcio-1.85.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:000999ea9bba80217b1625782d638e7a688790f6154e845e18826f7da859bfc7", size = 8311104, upload-time = "2026-10-30T16:00:18.533Z" }, + { url = "https://files.pythonhosted.org/packages/e4/02/b07af232a583727fa84665d11165bbd7be6ddd31d84c2b8cb62a7b69062f/grpcio-1.85.1-cp313-cp313-win32.whl", hash = "sha256:b3722f4e01820f4d2422679a567fe33f31b6f14eba4352657a1078f4b4ed323e", size = 4769351, upload-time = "2026-10-30T16:00:21.165Z" }, + { url = "https://files.pythonhosted.org/packages/f9/f0/d2448e860d884e8c89dd9135cd9373df064c10227d90973d822c2125428d/grpcio-1.85.1-cp313-cp313-win_amd64.whl", hash = "sha256:c17d61af28b322c48ab3881c685e5f6e7e367c56844cf589ee7707be9477b969", size = 5606831, upload-time = "2026-10-30T16:00:23.621Z" }, + { url = "https://files.pythonhosted.org/packages/39/cc/da49c22b9c571b5315acc507e2ff6568e5ff9d4c99895285910283bf87e3/grpcio-1.85.1-cp314-cp314-linux_armv7l.whl", hash = "sha256:2a378eb60ea42996be5a8c7167302fb30d968fcfcf9edd7b3fa6c985906c4500", size = 6803626, upload-time = "2026-10-30T16:00:26.587Z" }, + { url = "https://files.pythonhosted.org/packages/12/92/35bcee92bcfc2b0f66290a5dd6dd23ced7b9423fe269cf4495136e73d55b/grpcio-1.85.1-cp314-cp314-macosx_11_0_universal2.whl", hash = "sha256:90a7e0f26de39d15f5c741b8848bb2065c9ab7351c57d4fc973ff5cbda0cb9c7", size = 12640307, upload-time = "2026-10-30T16:00:28.969Z" }, + { url = "https://files.pythonhosted.org/packages/f2/b5/469b69dadb2c413706f16e0c3e7a2749cc70abaab10d718873603bec29b2/grpcio-1.85.1-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:421925365f9a665830eded01b8053033abbc2fbabc9bb4e9c59c3b5b12acbab9", size = 7187308, upload-time = "2026-10-30T16:00:31.652Z" }, + { url = "https://files.pythonhosted.org/packages/c4/df/2fcb1567d37e91ce26bddfca30f39fe6c39cbc5a2f4a210cfb1246ced1b2/grpcio-1.85.1-cp314-cp314-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:8a921a98ed6142a29651b4a6f4c7b8e051d115f1e78c31655a9642dcf6356dd9", size = 8188002, upload-time = "2026-10-30T16:00:33.893Z" }, + { url = "https://files.pythonhosted.org/packages/7b/b2/c67c1d8c7b70a094d404b28233458887bc1ca1319daa95d8952621216717/grpcio-1.85.1-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:e0e6d5829b50a2ffcfcd2dbf08d02ac0d6919051a2e97c2ee27c99bc9916a245", size = 7368045, upload-time = "2026-10-30T16:00:36.675Z" }, + { url = "https://files.pythonhosted.org/packages/80/48/5cf6c704be4b059d508e8f6c981154223533595ba7a3ee531211ef8be62e/grpcio-1.85.1-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:58b5ebe509da292bec5a1913c1e4c14d842a33f6cb330a87d3713c5089dd0f55", size = 7941330, upload-time = "2026-10-30T16:00:39.173Z" }, + { url = "https://files.pythonhosted.org/packages/67/96/5ef14edb784ea2053610753993601e98a76a18b541a8a74103dc390540ab/grpcio-1.85.1-cp314-cp314-musllinux_1_2_i686.whl", hash = "sha256:016982ebd21d2b48e955aa808e9db8d628708a35c6eba852cd0e70765826ab06", size = 9176909, upload-time = "2026-10-30T16:00:41.335Z" }, + { url = "https://files.pythonhosted.org/packages/50/01/4264ae95364bdeede4c8a5f0d83df83a21b477935f9ab8c829ae5606e1d4/grpcio-1.85.1-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:0fa3d7b8b95e8a609967991c3b0a302b2f9d510fe417efadfa7c7d2d044a99ba", size = 8316226, upload-time = "2026-10-30T16:00:44.199Z" }, + { url = "https://files.pythonhosted.org/packages/75/b8/79dace38eb4bb904751b804cfdc71609a3840f34084af7475169059dce20/grpcio-1.85.1-cp314-cp314-win32.whl", hash = "sha256:39b1965266b579f3758db20d74f9075b130d1b1edcf0464cda316e403a15da9b", size = 4853260, upload-time = "2026-10-30T16:00:47.098Z" }, + { url = "https://files.pythonhosted.org/packages/48/c9/68a6bbd2450f08942c3056acd8995ac28af2b4eeb0574c0a8e59fa59b140/grpcio-1.85.1-cp314-cp314-win_amd64.whl", hash = "sha256:5f015600190124813032da0560ec950c13c257f77c1307abd77c393c0994f76f", size = 5758958, upload-time = "2026-10-30T16:00:49.145Z" }, + { url = "https://files.pythonhosted.org/packages/1b/bd/1a2481b354ab4528eb8b3462937904c9206e2a54748f3f74a2576205ef3e/grpcio-1.85.1-cp315-cp315-linux_armv7l.whl", hash = "sha256:07bfef5df75875d8481d3403f864785e30096594162e8041c53a84fb11ee4f8e", size = 6803392, upload-time = "2026-10-30T16:00:51.782Z" }, + { url = "https://files.pythonhosted.org/packages/40/45/1868c12c18058d778fcecee0e04c867891a82d983d2dd0de2639158aa960/grpcio-1.85.1-cp315-cp315-macosx_10_15_universal2.whl", hash = "sha256:e909e0f9145a555d3968dd2a448616b0d3cb58748b9e2053de3e7974b3e429bc", size = 12657793, upload-time = "2026-10-30T16:00:55.335Z" }, + { url = "https://files.pythonhosted.org/packages/ed/af/c283032a42df5217913b98bd0113738b68dd94b91cea16b2209565fefe67/grpcio-1.85.1-cp315-cp315-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:1f3064e209167e5fcb75a72cbc18918c06027b92b29fa06eaff0126f9a2aab5a", size = 7209146, upload-time = "2026-10-30T16:00:58.279Z" }, + { url = "https://files.pythonhosted.org/packages/45/64/8a72fd6c20c206cff07eaa49b2ca63bb15d7ed9b6bf7c653ea10ec578e4b/grpcio-1.85.1-cp315-cp315-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:07675c635a2da0c772e58ac736a7fb10c7c7c1613157766441e6866ad97d25b4", size = 7511067, upload-time = "2026-10-30T16:01:00.806Z" }, + { url = "https://files.pythonhosted.org/packages/5d/b8/32afa404288c7eea4c6034d031f1a07eff772a8cdbfc7751314382e7b724/grpcio-1.85.1-cp315-cp315-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:759f7df78954e3ac0ed4c87612eff9125021b9271be37be8c88167aa533e29ee", size = 7410227, upload-time = "2026-10-30T16:01:02.966Z" }, + { url = "https://files.pythonhosted.org/packages/89/a4/1c51e2c2d0d9e4d37e8e8a7be412c5dfb3fd5c80ab100dd531fd81b44d43/grpcio-1.85.1-cp315-cp315-musllinux_1_2_aarch64.whl", hash = "sha256:7cd08bd5ad063ad9639dbebf95d3a6b8c8aa2c9d49049bc0e8f47e73ecf1c1a2", size = 7964886, upload-time = "2026-10-30T16:01:05.208Z" }, + { url = "https://files.pythonhosted.org/packages/07/f3/e26f570129c162d8f15a02af6edb49e592ac1c88a63490cd3ec9142d8c67/grpcio-1.85.1-cp315-cp315-musllinux_1_2_i686.whl", hash = "sha256:2198269c8b69cb1e9d5d20628029245e40b7227d5f6fc6f298a420701d8fb652", size = 8504060, upload-time = "2026-10-30T16:01:07.82Z" }, + { url = "https://files.pythonhosted.org/packages/16/e1/e334cf7442c01892ed813563b21b2a8ab5b042e34f33ad76cf64526d1de8/grpcio-1.85.1-cp315-cp315-musllinux_1_2_x86_64.whl", hash = "sha256:7bc5de50374416223bc35332aa8185567ec065f86190ff0ca6e13176ed39a3b4", size = 8346563, upload-time = "2026-10-30T16:01:10.394Z" }, + { url = "https://files.pythonhosted.org/packages/66/bb/cf72913ce0baae1fed95445f1ba8ac309184ac16b4f20fb5f88219eff72c/grpcio-1.85.1-cp315-cp315-win32.whl", hash = "sha256:5740aedf1e6acde538d7b80a2507eb6ee0b8983c577a36890d27a0e2b62cc1eb", size = 4853404, upload-time = "2026-10-30T16:01:12.239Z" }, + { url = "https://files.pythonhosted.org/packages/ea/b8/b7cff7b9e1ff7b6a8d36a19ce8b9c575963e71fcbb3878b276cff6748b05/grpcio-1.85.1-cp315-cp315-win_amd64.whl", hash = "sha256:b32f39c56d22d2e8722eececd4376074d94e01595a29f4f944ae0b6cea1caa70", size = 5758584, upload-time = "2026-10-30T16:01:14.325Z" }, ] [[package]] name = "grpcio-status" -version = "1.80.0" +version = "1.85.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "googleapis-common-protos" }, { name = "grpcio" }, { name = "protobuf" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/b1/ed/105f619bdd00cb47a49aa2feea6232ea2bbb04199d52a22cc6a7d603b5cb/grpcio_status-1.80.0.tar.gz", hash = "sha256:df73802a4c89a3ea88aa2aff971e886fccce162bc2e6511408b3d67a144381cd", size = 13901, upload-time = "2026-03-30T08:54:34.784Z" } +sdist = { url = "https://files.pythonhosted.org/packages/fe/16/0a161d0016f6e318ebf217780a91e007a67d9d244a06f79caeef689f5f02/grpcio_status-1.85.1.tar.gz", hash = "sha256:3f6ba6090d794b7148b23f03355addd56bee2b78174a12784c4d9b4e28c8928e", size = 13982, upload-time = "2026-10-30T16:05:25.631Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/76/80/58cd2dfc19a07d022abe44bde7c365627f6c7cb6f692ada6c65ca437d09a/grpcio_status-1.80.0-py3-none-any.whl", hash = "sha256:4b56990363af50dbf2c2ebb80f1967185c07d87aa25aa2bea45ddb75fc181dbe", size = 14638, upload-time = "2026-03-30T08:54:01.569Z" }, + { url = "https://files.pythonhosted.org/packages/f1/e9/5e34e2543b586ffddf90721ab242e7b662c063484ee3978bf61ca90d7579/grpcio_status-1.85.1-py3-none-any.whl", hash = "sha256:ef84cb262c7a05fd7c27e61bb70c9f63bce4484ce8b3366585a381fed738fcb0", size = 14640, upload-time = "2026-10-30T16:05:11.252Z" }, ] [[package]] @@ -1926,18 +1937,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/1e/5e/d4e9f1a599fb8e573b7b87160658329fbf28d19eac2718f51fc3def3aa5a/idna-3.18-py3-none-any.whl", hash = "sha256:7f952cbe720b688055e3f87de14f5c3e5fdaa8bc3928985c4077ca689de849a2", size = 65455, upload-time = "2026-06-02T14:34:06.319Z" }, ] -[[package]] -name = "importlib-metadata" -version = "8.7.1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "zipp" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/f3/49/3b30cad09e7771a4982d9975a8cbf64f00d4a1ececb53297f1d9a7be1b10/importlib_metadata-8.7.1.tar.gz", hash = "sha256:49fef1ae6440c182052f407c8d34a68f72efc36db9ca90dc0113398f2fdde8bb", size = 57107, upload-time = "2025-12-21T10:00:19.278Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/fa/5e/f8e9a1d23b9c20a551a8a02ea3637b4642e22c2626e3a13a9a29cdea99eb/importlib_metadata-8.7.1-py3-none-any.whl", hash = "sha256:5a1f80bf1daa489495071efbb095d75a634cf28a8bc299581244063b53176151", size = 27865, upload-time = "2025-12-21T10:00:18.329Z" }, -] - [[package]] name = "iniconfig" version = "2.3.0" @@ -3558,29 +3557,28 @@ wheels = [ [[package]] name = "opentelemetry-api" -version = "1.41.0" +version = "1.45.2" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "importlib-metadata" }, { name = "typing-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/47/8e/3778a7e87801d994869a9396b9fc2a289e5f9be91ff54a27d41eace494b0/opentelemetry_api-1.41.0.tar.gz", hash = "sha256:9421d911326ec12dee8bc933f7839090cad7a3f13fcfb0f9e82f8174dc003c09", size = 71416, upload-time = "2026-04-09T14:38:34.544Z" } +sdist = { url = "https://files.pythonhosted.org/packages/19/14/cf3a6b6448849656e1e5f4e2101f3988dc6789724050096aba1bc5ccf592/opentelemetry_api-1.45.2.tar.gz", hash = "sha256:d7b0008e21e5fd61cbe7a7f34ca3b6905a78ab677d9bc395e299359c75cda6b4", size = 73277, upload-time = "2026-10-29T17:29:37.888Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/58/ee/99ab786653b3bda9c37ade7e24a7b607a1b1f696063172768417539d876d/opentelemetry_api-1.41.0-py3-none-any.whl", hash = "sha256:0e77c806e6a89c9e4f8d372034622f3e1418a11bdbe1c80a50b3d3397ad0fa4f", size = 69007, upload-time = "2026-04-09T14:38:11.833Z" }, + { url = "https://files.pythonhosted.org/packages/a9/2d/610db864d7513158cebb8f36cc88e67e73bd6c7aa1905ef18d89e1281c81/opentelemetry_api-1.45.2-py3-none-any.whl", hash = "sha256:2cbd737377e710ee740d261c94e93fa0f08e6d037c18f17513cc0a2d94479a34", size = 60335, upload-time = "2026-10-29T17:29:06.955Z" }, ] [[package]] name = "opentelemetry-exporter-prometheus" -version = "0.62b0" +version = "0.66b2" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "opentelemetry-api" }, { name = "opentelemetry-sdk" }, { name = "prometheus-client" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/0f/ec/fa8a722199dc2e75dc582779d62207b00b0bdb014b5635594afa0cf3ee43/opentelemetry_exporter_prometheus-0.62b0.tar.gz", hash = "sha256:4d1106566a9b3e8dff028e69e9f2dc90723e6b431c900ff8c72982fcf11dbae5", size = 15441, upload-time = "2026-04-09T14:38:38.934Z" } +sdist = { url = "https://files.pythonhosted.org/packages/e8/c5/d49a08ebd180fb1471ae44ccaa265abf56e95cb5164ed5208e1663584b06/opentelemetry_exporter_prometheus-0.66b2.tar.gz", hash = "sha256:ada6b2507a31e10405f686e6a8498cb463a6931b5791337b04cfe14288db497f", size = 17337, upload-time = "2026-10-29T17:29:46.903Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/d1/1e/43645fadd561471af2aec95906a3dd54af1a8e7782322310e802a810ad3a/opentelemetry_exporter_prometheus-0.62b0-py3-none-any.whl", hash = "sha256:cd7e8acae3be5f425ffa2e0864eea474fa7a40706f786de7a2d23846573d8f75", size = 13278, upload-time = "2026-04-09T14:38:19.367Z" }, + { url = "https://files.pythonhosted.org/packages/b3/ec/91617f3e9bbe50487fd899a9efee4d111eabc3c61b66c2ef3ecf0aaba68e/opentelemetry_exporter_prometheus-0.66b2-py3-none-any.whl", hash = "sha256:b73998788aa30bec3bfe20fae312650878adb59f6563ed8bf37c188cb450a189", size = 13371, upload-time = "2026-10-29T17:29:21.28Z" }, ] [[package]] @@ -3597,29 +3595,29 @@ wheels = [ [[package]] name = "opentelemetry-sdk" -version = "1.41.0" +version = "1.45.2" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "opentelemetry-api" }, { name = "opentelemetry-semantic-conventions" }, { name = "typing-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/f8/0e/a586df1186f9f56b5a0879d52653effc40357b8e88fc50fe300038c3c08b/opentelemetry_sdk-1.41.0.tar.gz", hash = "sha256:7bddf3961131b318fc2d158947971a8e37e38b1cd23470cfb72b624e7cc108bd", size = 230181, upload-time = "2026-04-09T14:38:47.225Z" } +sdist = { url = "https://files.pythonhosted.org/packages/7b/1a/74a21d4d6969fb8b432ac397e66112a8167a6deb267ea4559dbaf2fe91ef/opentelemetry_sdk-1.45.2.tar.gz", hash = "sha256:7eec4de13dd21ff3a1670dcb0e02079e692736f94d733dfdce5af461d884248d", size = 221705, upload-time = "2026-10-29T17:29:54.685Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/2c/13/a7825118208cb32e6a4edcd0a99f925cbef81e77b3b0aedfd9125583c543/opentelemetry_sdk-1.41.0-py3-none-any.whl", hash = "sha256:a596f5687964a3e0d7f8edfdcf5b79cbca9c93c7025ebf5fb00f398a9443b0bd", size = 180214, upload-time = "2026-04-09T14:38:30.657Z" }, + { url = "https://files.pythonhosted.org/packages/9f/a5/9a3b793ef48245f2336a16beb2324c2397dbb3f987cd68942f77cca06156/opentelemetry_sdk-1.45.2-py3-none-any.whl", hash = "sha256:ff974ea70ee4f1d51e2ec152db37c2b626c7bf82bb3f7d86114fd8704e7199c6", size = 141756, upload-time = "2026-10-29T17:29:33.722Z" }, ] [[package]] name = "opentelemetry-semantic-conventions" -version = "0.62b0" +version = "0.66b2" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "opentelemetry-api" }, { name = "typing-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/a3/b0/c14f723e86c049b7bf8ff431160d982519b97a7be2857ed2247377397a24/opentelemetry_semantic_conventions-0.62b0.tar.gz", hash = "sha256:cbfb3c8fc259575cf68a6e1b94083cc35adc4a6b06e8cf431efa0d62606c0097", size = 145753, upload-time = "2026-04-09T14:38:48.274Z" } +sdist = { url = "https://files.pythonhosted.org/packages/37/3f/50f95972b50e8e353f5b9342034c9cd17730f1956f009ce6a3fbbd89f9ce/opentelemetry_semantic_conventions-0.66b2.tar.gz", hash = "sha256:7da5f1a5dc4383eb66711d4d3d3e71824d9403ac1559fea021cc103c1a4c4976", size = 150246, upload-time = "2026-10-29T17:29:55.682Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/58/6c/5e86fa1759a525ef91c2d8b79d668574760ff3f900d114297765eb8786cb/opentelemetry_semantic_conventions-0.62b0-py3-none-any.whl", hash = "sha256:0ddac1ce59eaf1a827d9987ab60d9315fb27aea23304144242d1fcad9e16b489", size = 231619, upload-time = "2026-04-09T14:38:32.394Z" }, + { url = "https://files.pythonhosted.org/packages/67/bd/6bf33ed31521c565931c037e9e81bac47c1acb000e850c6e24220e2a0c00/opentelemetry_semantic_conventions-0.66b2-py3-none-any.whl", hash = "sha256:4cbe61e5f9275d57ab31708446e7938ef4291ebbe6a7fca6f7a629c5acae87d8", size = 206280, upload-time = "2026-10-29T17:29:35.062Z" }, ] [[package]] @@ -6404,12 +6402,3 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/51/47/3fa2286c3cb162c71cdb34c4224d5745a1ceceb391b2bd9b19b668a8d724/yarl-1.23.0-cp314-cp314t-win_arm64.whl", hash = "sha256:44bb7bef4ea409384e3f8bc36c063d77ea1b8d4a5b2706956c0d6695f07dcc25", size = 86041, upload-time = "2026-03-01T22:07:49.026Z" }, { url = "https://files.pythonhosted.org/packages/69/68/c8739671f5699c7dc470580a4f821ef37c32c4cb0b047ce223a7f115757f/yarl-1.23.0-py3-none-any.whl", hash = "sha256:a2df6afe50dea8ae15fa34c9f824a3ee958d785fd5d089063d960bae1daa0a3f", size = 48288, upload-time = "2026-03-01T22:07:51.388Z" }, ] - -[[package]] -name = "zipp" -version = "3.23.1" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/30/21/093488dfc7cc8964ded15ab726fad40f25fd3d788fd741cc1c5a17d78ee8/zipp-3.23.1.tar.gz", hash = "sha256:32120e378d32cd9714ad503c1d024619063ec28aad2248dc6672ad13edfa5110", size = 25965, upload-time = "2026-04-13T23:21:46.6Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/08/8a/0861bec20485572fbddf3dfba2910e38fe249796cb73ecdeb74e07eeb8d3/zipp-3.23.1-py3-none-any.whl", hash = "sha256:0b3596c50a5c700c9cb40ba8d86d9f2cc4807e9bedb06bcdf7fac85633e444dc", size = 10378, upload-time = "2026-04-13T23:21:45.386Z" }, -]