From d74d30a47ae2bea298b133c0a7b3cc567d6f0c5e Mon Sep 17 00:00:00 2001 From: Chris Knight Date: Sat, 13 Jun 2026 21:35:28 +0200 Subject: [PATCH 1/4] 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 2/4] 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 3/4] 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 4/4] 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()