From 861fa1d63e05f207f8a543ec5d4ab3e79d6f5a78 Mon Sep 17 00:00:00 2001 From: Chris Knight Date: Sun, 18 Oct 2026 17:06:37 +0200 Subject: [PATCH] feat: Kafka message data components Adds `KafkaDataReader` and `KafkaDataWriter` on the `MessageDataReader` / `MessageDataWriter` bases: - Consumer uses a consumer group with auto-commit disabled; acknowledging a batch commits the offsets following the processed records, highest per partition, so records still buffered remain eligible for redelivery. - Producer submits a batch's records together and flushes once, keeping the client's internal batching instead of one round trip per record. - `bootstrap_servers` falls back to `KAFKA_BOOTSTRAP_SERVERS` and accepts a string or a list; aiokafka failures are mapped onto the broker exception hierarchy (authorization / record-too-large are permanent, connection and rebalance errors are retried). - The `kafka` extra declares aiokafka, and aiokafka is in the test dependency group so the unit tests can import the real error and offset types. 372 lines of tests cover offset commits (including out-of-order records across partitions), empty polls returning no messages, multi-partition flattening, batched sends with a single flush, error mapping, settings resolution and JSON vs raw encoding. Split out of #284, where it landed with the PubSub and SQS/SNS implementations. --- README.md | 2 +- docs/usage/configuration.md | 1 + docs/usage/message-data.md | 22 +- plugboard/library/__init__.py | 3 + plugboard/library/kafka_io.py | 297 +++++++++++++++++++++++++++ plugboard/utils/settings.py | 14 ++ pyproject.toml | 6 +- tests/unit/test_kafka_io.py | 372 ++++++++++++++++++++++++++++++++++ uv.lock | 50 ++++- 9 files changed, 755 insertions(+), 12 deletions(-) create mode 100644 plugboard/library/kafka_io.py create mode 100644 tests/unit/test_kafka_io.py diff --git a/README.md b/README.md index dc82f7ba..1efbf135 100644 --- a/README.md +++ b/README.md @@ -69,7 +69,7 @@ Optional integrations for different cloud providers can be installed using `plug Support for parallelisation and hyperparameter optimisation can be installed using `plugboard[ray]`. -Additional optional extras: `plugboard[llm]` for LLM components, `plugboard[redis]` for Redis-based connectors, `plugboard[omq]` for the pyomq backend for ZMQ connectors, `plugboard[websockets]` for WebSocket I/O, and `plugboard[gcp-pubsub]` or `plugboard[aws-messaging]` for [message data components](https://docs.plugboard.dev/usage/message-data/). +Additional optional extras: `plugboard[llm]` for LLM components, `plugboard[redis]` for Redis-based connectors, `plugboard[omq]` for the pyomq backend for ZMQ connectors, `plugboard[websockets]` for WebSocket I/O, and `plugboard[gcp-pubsub]`, `plugboard[aws-messaging]` or `plugboard[kafka]` for [message data components](https://docs.plugboard.dev/usage/message-data/). ## ⚡ Quickstart with AI diff --git a/docs/usage/configuration.md b/docs/usage/configuration.md index a8e18ed0..5430c2eb 100644 --- a/docs/usage/configuration.md +++ b/docs/usage/configuration.md @@ -39,6 +39,7 @@ evironment below when they are not supplied: |---------------------------|----------------------------------------|---------------| | `GCP_PUBSUB_PROJECT_ID` | GCP project for PubSub topics and subscriptions | | | `AWS_REGION` | AWS region for SQS and SNS clients | | +| `KAFKA_BOOTSTRAP_SERVERS` | Kafka broker address(es) | | Note that `AWS_REGION` is also read by the AWS SDK itself, so setting it affects both Plugboard's defaults and the credential/endpoint resolution of the underlying client. diff --git a/docs/usage/message-data.md b/docs/usage/message-data.md index da7a2171..039db62a 100644 --- a/docs/usage/message-data.md +++ b/docs/usage/message-data.md @@ -12,8 +12,10 @@ component's outputs (for a reader) or inputs (for a writer). | [`GCPPubSubDataWriter`][plugboard.library.GCPPubSubDataWriter] | Google Cloud PubSub | | a topic | | [`AWSSQSDataReader`][plugboard.library.AWSSQSDataReader] | AWS SQS | a queue | | | [`AWSSNSDataWriter`][plugboard.library.AWSSNSDataWriter] | AWS SNS | | a topic | +| [`KafkaDataReader`][plugboard.library.KafkaDataReader] | Apache Kafka | a topic + consumer group | | +| [`KafkaDataWriter`][plugboard.library.KafkaDataWriter] | Apache Kafka | | a topic | -All four build on the [`MessageDataReader`][plugboard.library.MessageDataReader] and +All six build on the [`MessageDataReader`][plugboard.library.MessageDataReader] and [`MessageDataWriter`][plugboard.library.MessageDataWriter] base classes, which handle connecting, reconnection with backoff, batching and acknowledgment. @@ -24,6 +26,7 @@ Each broker uses its own client library, supplied as an extra: ```shell pip install "plugboard[gcp-pubsub]" # google-cloud-pubsub pip install "plugboard[aws-messaging]" # aiobotocore +pip install "plugboard[kafka]" # aiokafka ``` Constructing a component without its extra raises an `ImportError` naming the extra to @@ -99,7 +102,8 @@ run that would otherwise wait forever. **Messages are acknowledged after they are processed.** A batch fetched from the broker is acknowledged only once every record in it has been published downstream, so a crash mid-batch leaves the unread messages to be redelivered (at-least-once). For SQS, -"acknowledge" means deleting the messages. +"acknowledge" means deleting the messages; for Kafka it means committing the offsets of +the processed records, per partition. **Batching is bounded by the broker.** `chunk_size` sets how many messages are fetched per poll (and per send for writers). SQS delivers at most 10 messages per call, so a @@ -111,14 +115,15 @@ deleted topic — fail immediately instead of burning the retries. Both the clas and the policy are handled by the base classes: ```python -from plugboard.library import GCPPubSubDataReader +from plugboard.library import KafkaDataReader from plugboard.utils.retry import RetryPolicy -reader = GCPPubSubDataReader( +reader = KafkaDataReader( name="reader", - subscription_id="measurements-pull", - project_id="plugboard-dev", + topic="measurements", + group_id="plugboard", field_names=["x", "y"], + bootstrap_servers="localhost:9092", # Optional; defaults to KAFKA_BOOTSTRAP_SERVERS retry_policy=RetryPolicy(max_retries=5, base_delay=0.5, max_delay=30.0), idle_poll_delay=0.5, # Pause between empty polls, if the broker returns immediately ) @@ -138,7 +143,8 @@ arguments always win. |---------------------------|------------------------------------------|---------| | `GCP_PUBSUB_PROJECT_ID` | Default GCP project for PubSub topics | `GCPPubSubDataReader`, `GCPPubSubDataWriter` | | `AWS_REGION` | Default AWS region for SQS/SNS clients | `AWSSQSDataReader`, `AWSSNSDataWriter` | +| `KAFKA_BOOTSTRAP_SERVERS` | Default Kafka broker address(es) | `KafkaDataReader`, `KafkaDataWriter` | Credentials themselves are never configured through Plugboard: the clients use their -normal provider mechanisms (Application Default Credentials for PubSub, and the standard -AWS credential chain for SQS/SNS). +normal provider mechanisms (Application Default Credentials for PubSub, the standard AWS +credential chain, SASL settings for Kafka). diff --git a/plugboard/library/__init__.py b/plugboard/library/__init__.py index 22999240..c4bd2ac2 100644 --- a/plugboard/library/__init__.py +++ b/plugboard/library/__init__.py @@ -5,6 +5,7 @@ from .data_writer import DataWriter from .file_io import FileReader, FileWriter from .gcp_pubsub_io import GCPPubSubDataReader, GCPPubSubDataWriter +from .kafka_io import KafkaDataReader, KafkaDataWriter from .llm import LLMChat, LLMImageProcessor from .message_reader import MessageDataReader from .message_writer import MessageDataWriter @@ -21,6 +22,8 @@ "FileWriter", "GCPPubSubDataReader", "GCPPubSubDataWriter", + "KafkaDataReader", + "KafkaDataWriter", "LLMChat", "LLMImageProcessor", "MessageDataReader", diff --git a/plugboard/library/kafka_io.py b/plugboard/library/kafka_io.py new file mode 100644 index 00000000..133a34b3 --- /dev/null +++ b/plugboard/library/kafka_io.py @@ -0,0 +1,297 @@ +"""Provides `KafkaDataReader` and `KafkaDataWriter` for Apache Kafka messaging.""" + +from __future__ import annotations + +import asyncio +from collections import deque +import json +import typing as _t + +from plugboard.exceptions import ( + MessageBrokerConnectionError, + MessageBrokerPermanentError, + MessageBrokerTransientError, +) +from plugboard.library.message_reader import MessageDataReader, MessageDataReaderArgsDict +from plugboard.library.message_writer import ( + MessageDataWriter, + MessageDataWriterArgsDict, + encode_records_bytes, +) +from plugboard.utils import DI, depends_on_optional +from plugboard.utils.settings import resolve_argument + + +try: + from aiokafka import AIOKafkaConsumer, AIOKafkaProducer + from aiokafka.errors import ( + CoordinatorNotAvailableError, + KafkaConnectionError, + NodeNotReadyError, + NotLeaderForPartitionError, + RecordTooLargeError, + RequestTimedOutError, + TopicAuthorizationFailedError, + UnknownTopicOrPartitionError, + ) + from aiokafka.structs import OffsetAndMetadata, TopicPartition + + # Kafka failures that retrying cannot resolve, such as missing authorization or a + # record the broker will always reject. + _PERMANENT_ERRORS: tuple[type[Exception], ...] = ( + TopicAuthorizationFailedError, + RecordTooLargeError, + ) + + # Kafka failures raised while the cluster is rebalancing or a broker is unreachable, + # which are expected to clear. + _TRANSIENT_ERRORS: tuple[type[Exception], ...] = ( + KafkaConnectionError, + NodeNotReadyError, + CoordinatorNotAvailableError, + NotLeaderForPartitionError, + RequestTimedOutError, + UnknownTopicOrPartitionError, + ) +except ImportError: # pragma: no cover + # With the extra not installed, `depends_on_optional` blocks construction and + # classification is moot: `isinstance(error, ())` never matches. + _PERMANENT_ERRORS = () + _TRANSIENT_ERRORS = () + + +def _wrap_error(error: Exception) -> Exception: + """Maps an aiokafka error onto the Plugboard message broker exception hierarchy. + + Args: + error: The error raised by the Kafka client. + + Returns: + A `MessageBrokerPermanentError` for a rejection that will not change on retry, + a `MessageBrokerTransientError` for a failure worth retrying, otherwise the + original error. + """ + if isinstance(error, _PERMANENT_ERRORS): + return MessageBrokerPermanentError(str(error)) + if isinstance(error, _TRANSIENT_ERRORS): + return MessageBrokerTransientError(str(error)) + return error + + +class KafkaDataReader(MessageDataReader): + """Reads data from an Apache Kafka topic. + + Messages are consumed from the topic using a consumer group. Acknowledging a batch + commits the offsets following the records that were consumed, so a failure before + that point leaves the remaining records for redelivery. + """ + + @depends_on_optional("aiokafka", extra="kafka") + def __init__( + self, + group_id: str, + bootstrap_servers: _t.Optional[str | list[str]] = None, + parse_json: bool = True, + poll_timeout_ms: int = 30_000, + **kwargs: _t.Unpack[MessageDataReaderArgsDict], + ) -> None: + """Instantiates the `KafkaDataReader`. + + Args: + group_id: Consumer group ID. + bootstrap_servers: Optional; Kafka broker address(es), for example + `"localhost:9092"`. Falls back to `KAFKA_BOOTSTRAP_SERVERS`. + parse_json: Whether to parse message values as JSON. + poll_timeout_ms: Milliseconds to wait for a batch of records. + **kwargs: Additional keyword arguments for + [`MessageDataReader`][plugboard.library.MessageDataReader]. + """ + bootstrap_servers = resolve_argument( + bootstrap_servers, + DI.settings.resolve_sync().kafka.bootstrap_servers, + "bootstrap_servers", + "KAFKA_BOOTSTRAP_SERVERS", + ) + super().__init__(**kwargs) + self._bootstrap_servers = ( + [bootstrap_servers] if isinstance(bootstrap_servers, str) else bootstrap_servers + ) + self._group_id = group_id + self._parse_json = parse_json + self._poll_timeout_ms = poll_timeout_ms + self._consumer: _t.Optional[AIOKafkaConsumer] = None + + async def _connect(self) -> None: + """Creates and starts a Kafka consumer.""" + consumer = AIOKafkaConsumer( + self._topic, + bootstrap_servers=self._bootstrap_servers, + group_id=self._group_id, + auto_offset_reset="earliest", + enable_auto_commit=False, + max_poll_records=self._chunk_size or 10, + ) + await consumer.start() + self._consumer = consumer + + async def _disconnect(self) -> None: + """Stops and closes the Kafka consumer.""" + if self._consumer is not None: + consumer = self._consumer + self._consumer = None + await consumer.stop() + + async def _receive(self) -> list[_t.Any]: + """Receives a batch of messages from the Kafka topic. + + Returns: + A list of Kafka `ConsumerRecord` objects, empty if the poll timed out. + + Raises: + MessageBrokerError: If the poll failed for a retryable or permanent reason. + """ + if self._consumer is None: + raise MessageBrokerConnectionError("Kafka consumer is not connected") + max_messages = self._chunk_size or 10 + try: + data = await self._consumer.getmany( + timeout_ms=self._poll_timeout_ms, max_records=max_messages + ) + except Exception as error: # noqa: BLE001 + raise _wrap_error(error) from error + return [record for records in data.values() for record in records] + + async def _convert(self, data: list[_t.Any]) -> dict[str, deque]: + """Converts Kafka messages to a field buffer. + + Args: + data: A list of `ConsumerRecord` objects. + + Returns: + A dictionary mapping field names to deques of field values. + """ + converted: dict[str, deque] = {field: deque() for field in self.io.outputs} + for record in data: + value = record.value + if isinstance(value, bytes): + value = value.decode("utf-8") + if self._parse_json: + record_data = json.loads(value) + else: + record_data = {"data": value} + for field in self.io.outputs: + converted[field].append(record_data.get(field)) + return converted + + async def _ack(self, messages: list[_t.Any]) -> None: + """Commits the offsets that follow the processed Kafka messages. + + Only the offsets of the records that were actually consumed are committed, so + records still waiting in the buffer stay eligible for redelivery. + + Args: + messages: The `ConsumerRecord` objects that have been processed. + + Raises: + MessageBrokerConnectionError: If the consumer is not connected. + MessageBrokerError: If the commit failed. + """ + if self._consumer is None: + raise MessageBrokerConnectionError("Kafka consumer is not connected") + offsets: dict[TopicPartition, OffsetAndMetadata] = {} + for record in messages: + partition = TopicPartition(record.topic, record.partition) + next_offset = OffsetAndMetadata(record.offset + 1, "") + current = offsets.get(partition) + if current is None or current.offset < next_offset.offset: + offsets[partition] = next_offset + if not offsets: + return + try: + await self._consumer.commit(offsets=offsets) + except Exception as error: # noqa: BLE001 + raise _wrap_error(error) from error + + +class KafkaDataWriter(MessageDataWriter): + """Writes data to an Apache Kafka topic. + + Field data is converted to JSON-encoded messages and produced to the specified + topic. Records are submitted together and flushed once, so the producer's internal + batching is used instead of one round trip per record. + """ + + @depends_on_optional("aiokafka", extra="kafka") + def __init__( + self, + bootstrap_servers: _t.Optional[str | list[str]] = None, + parse_json: bool = True, + **kwargs: _t.Unpack[MessageDataWriterArgsDict], + ) -> None: + """Instantiates the `KafkaDataWriter`. + + Args: + bootstrap_servers: Optional; Kafka broker address(es), for example + `"localhost:9092"`. Falls back to `KAFKA_BOOTSTRAP_SERVERS`. + parse_json: Whether to encode message values as JSON. + **kwargs: Additional keyword arguments for + [`MessageDataWriter`][plugboard.library.MessageDataWriter]. + """ + bootstrap_servers = resolve_argument( + bootstrap_servers, + DI.settings.resolve_sync().kafka.bootstrap_servers, + "bootstrap_servers", + "KAFKA_BOOTSTRAP_SERVERS", + ) + super().__init__(**kwargs) + self._bootstrap_servers = ( + [bootstrap_servers] if isinstance(bootstrap_servers, str) else bootstrap_servers + ) + self._parse_json = parse_json + self._producer: _t.Optional[AIOKafkaProducer] = None + + async def _connect(self) -> None: + """Creates and starts a Kafka producer.""" + producer = AIOKafkaProducer(bootstrap_servers=self._bootstrap_servers) + await producer.start() + self._producer = producer + + async def _disconnect(self) -> None: + """Stops and closes the Kafka producer.""" + if self._producer is not None: + producer = self._producer + self._producer = None + await producer.stop() + + async def _send(self, messages: list[_t.Any]) -> None: + """Sends messages to the Kafka topic. + + Args: + messages: A list of bytes objects to send. + + Raises: + MessageBrokerConnectionError: If the producer is not connected. + MessageBrokerError: If a send failed or the buffer could not be flushed. + """ + if self._producer is None: + raise MessageBrokerConnectionError("Kafka producer is not connected") + try: + # `send` returns a future that resolves once the broker acknowledges the + # record; awaiting them together keeps the producer's batching intact. + await asyncio.gather( + *(self._producer.send(self._topic, value=msg_data) for msg_data in messages) + ) + await self._producer.flush() + except Exception as error: # noqa: BLE001 + raise _wrap_error(error) from error + + async def _convert(self, data: dict[str, deque]) -> list[_t.Any]: + """Converts field buffer data to JSON-encoded bytes messages. + + Args: + data: A dictionary mapping field names to deques of field values. + + Returns: + A list of bytes objects ready to send. + """ + return encode_records_bytes(data, self._parse_json) diff --git a/plugboard/utils/settings.py b/plugboard/utils/settings.py index 42d0cb49..93c6c683 100644 --- a/plugboard/utils/settings.py +++ b/plugboard/utils/settings.py @@ -81,6 +81,18 @@ class _AWSSettings(BaseSettings): region: _t.Optional[str] = None +class _KafkaSettings(BaseSettings): + """Apache Kafka settings. + + Attributes: + bootstrap_servers: Kafka broker address(es). + """ + + model_config = SettingsConfigDict(env_prefix="KAFKA_") + + bootstrap_servers: _t.Optional[str] = None + + class Settings(BaseSettings): """Settings for Plugboard. @@ -95,6 +107,7 @@ class Settings(BaseSettings): redis: Redis settings. gcp_pubsub: Google Cloud PubSub settings. aws: AWS settings for SNS/SQS messaging. + kafka: Apache Kafka settings. """ model_config = SettingsConfigDict(env_prefix=_ENV_PREFIX) @@ -108,6 +121,7 @@ class Settings(BaseSettings): redis: _RedisSettings = Field(default_factory=_RedisSettings) gcp_pubsub: _GCPPubSubSettings = Field(default_factory=_GCPPubSubSettings) aws: _AWSSettings = Field(default_factory=_AWSSettings) + kafka: _KafkaSettings = Field(default_factory=_KafkaSettings) _ArgumentT = _t.TypeVar("_ArgumentT") diff --git a/pyproject.toml b/pyproject.toml index bff3ec42..b0ad7b72 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -47,6 +47,7 @@ aws-messaging = ["aiobotocore>=3.4,<4"] azure = ["adlfs>=2024.7.0"] gcp = ["gcsfs>=2024.9.0"] gcp-pubsub = ["google-cloud-pubsub>=2.25,<3"] +kafka = ["aiokafka>=0.11,<1"] llm = [ "llama-index-core>=0.12.30,<1", "llama-index-llms-openai>=0.3.33,<1", @@ -77,9 +78,10 @@ dev = [ test = [ "aiofile>=3.9,<4", "aiosqlite>=0.20,<1", - # Mirrors the optional broker extras so tests can import the real client classes - # and inject fakes, instead of leaving them unresolved to the type checker. + # Mirrors the optional broker extras so tests can import the real client classes and + # inject fakes, instead of leaving them unresolved to the type checker. "aiobotocore>=3.4,<4", + "aiokafka>=0.11,<1", "coverage>=7.7.1,<8", "google-cloud-pubsub>=2.25,<3", "llama-index-core>=0.12.30,<1", diff --git a/tests/unit/test_kafka_io.py b/tests/unit/test_kafka_io.py new file mode 100644 index 00000000..425ab396 --- /dev/null +++ b/tests/unit/test_kafka_io.py @@ -0,0 +1,372 @@ +"""Unit tests for the Apache Kafka message data implementations.""" + +from __future__ import annotations + +from collections import deque +import typing as _t +from unittest.mock import AsyncMock, MagicMock, call + +from aiokafka.errors import ( + KafkaConnectionError, + RecordTooLargeError, + TopicAuthorizationFailedError, +) +from aiokafka.structs import OffsetAndMetadata, TopicPartition +import pytest + +from plugboard.exceptions import ( + MessageBrokerConnectionError, + MessageBrokerPermanentError, + MessageBrokerTransientError, +) +from plugboard.library import kafka_io +from plugboard.library.kafka_io import KafkaDataReader, KafkaDataWriter +from plugboard.utils.settings import Settings +from tests import conftest + + +def consumer_record(topic: str, partition: int, offset: int, value: bytes) -> MagicMock: + """Builds a Kafka `ConsumerRecord` double.""" + record = MagicMock(name="consumer_record") + record.topic = topic + record.partition = partition + record.offset = offset + record.value = value + return record + + +@pytest.fixture +def consumer(monkeypatch: pytest.MonkeyPatch) -> MagicMock: + """Replaces `AIOKafkaConsumer` with a double that records its construction.""" + client = MagicMock(name="consumer") + client.start = AsyncMock() + client.stop = AsyncMock() + client.commit = AsyncMock() + client.getmany = AsyncMock(return_value={}) + factory = MagicMock(name="AIOKafkaConsumer", return_value=client) + monkeypatch.setattr(kafka_io, "AIOKafkaConsumer", factory) + client.factory = factory + return client + + +@pytest.fixture +def producer(monkeypatch: pytest.MonkeyPatch) -> MagicMock: + """Replaces `AIOKafkaProducer` with a double that records its construction.""" + client = MagicMock(name="producer") + client.start = AsyncMock() + client.stop = AsyncMock() + client.send = AsyncMock() + client.flush = AsyncMock() + client.send_and_wait = AsyncMock() + factory = MagicMock(name="AIOKafkaProducer", return_value=client) + monkeypatch.setattr(kafka_io, "AIOKafkaProducer", factory) + client.factory = factory + return client + + +def make_reader(**kwargs: _t.Any) -> KafkaDataReader: + """Builds a reader for the fixed test topic/group.""" + kwargs.setdefault("field_names", ["x", "y"]) + return KafkaDataReader( + name="kafka-reader", + topic="test-topic", + group_id="test-group", + bootstrap_servers="localhost:9092", + **kwargs, + ) + + +def make_writer(**kwargs: _t.Any) -> KafkaDataWriter: + """Builds a writer for the fixed test topic.""" + kwargs.setdefault("field_names", ["x", "y"]) + return KafkaDataWriter( + name="kafka-writer", + topic="test-topic", + bootstrap_servers="localhost:9092", + **kwargs, + ) + + +# --------------------------------------------------------------------------- +# Reader: connection and consuming +# --------------------------------------------------------------------------- + + +async def test_reader_starts_consumer_on_topic_group(consumer: MagicMock) -> None: + """Tests the consumer is built for the topic, group and batch size.""" + reader = make_reader(chunk_size=5) + await reader._connect() + + consumer.factory.assert_called_once_with( + "test-topic", + bootstrap_servers=["localhost:9092"], + group_id="test-group", + auto_offset_reset="earliest", + enable_auto_commit=False, + max_poll_records=5, + ) + consumer.start.assert_awaited_once_with() + + +async def test_reader_stops_consumer_on_disconnect(consumer: MagicMock) -> None: + """Tests that teardown stops the consumer, releasing its group membership.""" + reader = make_reader() + await reader._connect() + await reader._disconnect() + consumer.stop.assert_awaited_once_with() + + +async def test_reader_flattens_records_from_partitions(consumer: MagicMock) -> None: + """Tests that records from every assigned partition are returned.""" + consumer.getmany.return_value = { + TopicPartition("test-topic", 0): [consumer_record("test-topic", 0, 1, b"{}")], + TopicPartition("test-topic", 1): [ + consumer_record("test-topic", 1, 7, b"{}"), + consumer_record("test-topic", 1, 8, b"{}"), + ], + } + reader = make_reader() + await reader._connect() + + messages = await reader._receive() + + consumer.getmany.assert_awaited_once_with(timeout_ms=30_000, max_records=10) + assert [(r.partition, r.offset) for r in messages] == [(0, 1), (1, 7), (1, 8)] + + +async def test_reader_empty_poll_returns_no_messages(consumer: MagicMock) -> None: + """Tests that a timed-out poll reads as 'nothing yet', not as an exhausted topic. + + A consumer catching up between messages must keep running, so an empty poll must + not raise `NoMoreDataException`. + """ + consumer.getmany.return_value = {} + reader = make_reader() + await reader._connect() + assert await reader._receive() == [] + + +async def test_reader_converts_json_values(consumer: MagicMock) -> None: + """Tests that JSON record values become field buffers.""" + reader = make_reader() + await reader._connect() + batch = [ + consumer_record("test-topic", 0, 1, b'{"x": 1, "y": "a"}'), + consumer_record("test-topic", 0, 2, b'{"x": 2, "y": "b"}'), + ] + assert await reader._convert(batch) == {"x": deque([1, 2]), "y": deque(["a", "b"])} + + +async def test_reader_non_json_value_uses_data_field(consumer: MagicMock) -> None: + """Tests that a raw value is exposed through the `data` field.""" + reader = make_reader(field_names=["data"], parse_json=False) + await reader._connect() + converted = await reader._convert([consumer_record("test-topic", 0, 1, b"hello")]) + assert converted == {"data": deque(["hello"])} + + +# --------------------------------------------------------------------------- +# Reader: offset commits +# --------------------------------------------------------------------------- + + +async def test_reader_commits_offsets_of_processed_records(consumer: MagicMock) -> None: + """Tests that an ack commits the offset *after* each processed record. + + Committing the consumer position instead of these offsets would mark records that + were never processed as done, and they would never be redelivered. + """ + reader = make_reader() + await reader._connect() + records = [ + consumer_record("test-topic", 0, 3, b"{}"), + consumer_record("test-topic", 0, 4, b"{}"), + ] + + await reader._ack(records) + + consumer.commit.assert_awaited_once_with( + offsets={TopicPartition("test-topic", 0): OffsetAndMetadata(5, "")} + ) + + +async def test_reader_commits_highest_offset_per_partition(consumer: MagicMock) -> None: + """Tests that out-of-order records commit the furthest offset for each partition.""" + reader = make_reader() + await reader._connect() + records = [ + consumer_record("test-topic", 1, 2, b"{}"), + consumer_record("test-topic", 0, 10, b"{}"), + consumer_record("test-topic", 1, 5, b"{}"), + consumer_record("test-topic", 0, 8, b"{}"), + ] + + await reader._ack(records) + + committed = consumer.commit.await_args.kwargs["offsets"] + assert committed == { + TopicPartition("test-topic", 0): OffsetAndMetadata(11, ""), + TopicPartition("test-topic", 1): OffsetAndMetadata(6, ""), + } + + +async def test_reader_commits_nothing_without_records(consumer: MagicMock) -> None: + """Tests that an empty batch does not issue a commit.""" + reader = make_reader() + await reader._connect() + await reader._ack([]) + consumer.commit.assert_not_awaited() + + +async def test_reader_without_connection_raises_connection_error(consumer: MagicMock) -> None: + """Tests that consuming before connecting fails loudly.""" + reader = make_reader() + with pytest.raises(MessageBrokerConnectionError): + await reader._receive() + with pytest.raises(MessageBrokerConnectionError): + await reader._ack([]) + + +# --------------------------------------------------------------------------- +# Reader: error mapping +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + ("error", "expected"), + [ + (TopicAuthorizationFailedError("denied"), MessageBrokerPermanentError), + (RecordTooLargeError("too big"), MessageBrokerPermanentError), + (KafkaConnectionError("down"), MessageBrokerTransientError), + ], +) +async def test_reader_maps_kafka_errors( + consumer: MagicMock, error: Exception, expected: type[Exception] +) -> None: + """Tests that Kafka failures land on the right broker exception type.""" + consumer.getmany.side_effect = error + reader = make_reader() + await reader._connect() + + with pytest.raises(expected): + await reader._receive() + + +async def test_reader_unclassified_error_propagates(consumer: MagicMock) -> None: + """Tests that an unknown error is not rewritten into a broker exception.""" + consumer.getmany.side_effect = ValueError("unexpected") + reader = make_reader() + await reader._connect() + + with pytest.raises(ValueError, match="unexpected"): + await reader._receive() + + +# --------------------------------------------------------------------------- +# Writer +# --------------------------------------------------------------------------- + + +async def test_writer_starts_producer(producer: MagicMock) -> None: + """Tests that the producer is created and started.""" + writer = make_writer() + await writer._connect() + producer.factory.assert_called_once_with(bootstrap_servers=["localhost:9092"]) + producer.start.assert_awaited_once_with() + + +async def test_writer_stops_producer_on_disconnect(producer: MagicMock) -> None: + """Tests that teardown stops the producer.""" + writer = make_writer() + await writer._connect() + await writer._disconnect() + producer.stop.assert_awaited_once_with() + + +async def test_writer_sends_batch_then_flushes_once(producer: MagicMock) -> None: + """Tests that records are submitted together and the buffer flushed once. + + Awaiting each send individually (as `send_and_wait` does) costs one round trip per + record and defeats the producer's batching. + """ + writer = make_writer() + await writer._connect() + + await writer._send([b"one", b"two", b"three"]) + + assert producer.send.await_args_list == [ + call("test-topic", value=b"one"), + call("test-topic", value=b"two"), + call("test-topic", value=b"three"), + ] + producer.send_and_wait.assert_not_awaited() + producer.flush.assert_awaited_once_with() + + +async def test_writer_converts_fields_to_json_bytes(producer: MagicMock) -> None: + """Tests that buffered fields become one JSON message per record.""" + writer = make_writer() + await writer._connect() + converted = await writer._convert({"x": deque([1, 2]), "y": deque(["a", "b"])}) + assert converted == [b'{"x": 1, "y": "a"}', b'{"x": 2, "y": "b"}'] + + +async def test_writer_non_json_sends_first_field(producer: MagicMock) -> None: + """Tests that a raw writer publishes the first field's value.""" + writer = make_writer(field_names=["x"], parse_json=False) + await writer._connect() + assert await writer._convert({"x": deque(["hello"])}) == [b"hello"] + + +async def test_writer_send_failure_maps_error(producer: MagicMock) -> None: + """Tests that a failed send surfaces as a broker exception.""" + producer.send.side_effect = KafkaConnectionError("down") + writer = make_writer() + await writer._connect() + + with pytest.raises(MessageBrokerTransientError): + await writer._send([b"one"]) + + +async def test_writer_without_connection_raises_connection_error(producer: MagicMock) -> None: + """Tests that sending before connecting fails loudly.""" + writer = make_writer() + with pytest.raises(MessageBrokerConnectionError): + await writer._send([b"one"]) + + +# --------------------------------------------------------------------------- +# Settings resolution +# --------------------------------------------------------------------------- + + +async def test_bootstrap_servers_fall_back_to_settings() -> None: + """Tests that `KAFKA_BOOTSTRAP_SERVERS` supplies the brokers when not passed.""" + settings = Settings.model_validate({"kafka": {"bootstrap_servers": "kafka:9092"}}) + with conftest.override_settings(settings): + writer = KafkaDataWriter(name="kafka-writer", topic="test-topic", field_names=["x"]) + assert writer._bootstrap_servers == ["kafka:9092"] + + +def test_bootstrap_servers_accepts_a_list() -> None: + """Tests that an explicit list of brokers is kept as given.""" + reader = KafkaDataReader( + name="kafka-reader", + topic="test-topic", + group_id="test-group", + field_names=["x"], + bootstrap_servers=["a:9092", "b:9092"], + ) + assert reader._bootstrap_servers == ["a:9092", "b:9092"] + + +def test_missing_bootstrap_servers_names_the_environment_variable() -> None: + """Tests that unresolvable brokers report what to configure.""" + with conftest.override_settings(Settings.model_validate({})): + with pytest.raises(ValueError, match="KAFKA_BOOTSTRAP_SERVERS"): + KafkaDataWriter( + name="kafka-writer", + topic="test-topic", + field_names=["x"], + bootstrap_servers=None, + ) diff --git a/uv.lock b/uv.lock index 0c1a25de..93c2b18c 100644 --- a/uv.lock +++ b/uv.lock @@ -210,6 +210,37 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/10/a1/510b0a7fadc6f43a6ce50152e69dbd86415240835868bb0bd9b5b88b1e06/aioitertools-0.13.0-py3-none-any.whl", hash = "sha256:0be0292b856f08dfac90e31f4739432f4cb6d7520ab9eb73e143f4f2fa5259be", size = 24182, upload-time = "2025-11-06T22:17:06.502Z" }, ] +[[package]] +name = "aiokafka" +version = "0.14.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "async-timeout" }, + { name = "packaging" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/89/5f/dfc1180fd22d1acdc91949ec36e97199c43742dacb057cb8efed3679ed04/aiokafka-0.14.0.tar.gz", hash = "sha256:8ffdc945798ba4d3d132b705d4244d0a1f493925efb57c637a2ca88ee82794e1", size = 601374, upload-time = "2026-04-29T10:43:03.574Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f9/9d/3441db94829f9feb802a2f4052df61c0d1a01272accd174c351d7e9e1f6a/aiokafka-0.14.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:284a90d617584d7e42688a181aaa8c2a909d9c658ab9b69c6cf92f4df5c4b320", size = 348458, upload-time = "2026-04-29T10:42:37.243Z" }, + { url = "https://files.pythonhosted.org/packages/a4/10/7297589aac95654596af13301b31da2c9502c80e7e308530ee7a9bd5b9f1/aiokafka-0.14.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:b4f211d9e03a1fc83871a37eefcf307bc0943ee99adae25aa39bd1722e70747b", size = 351057, upload-time = "2026-04-29T10:42:38.69Z" }, + { url = "https://files.pythonhosted.org/packages/26/4e/5c0aa8db717fff0ffb8f3e16deece8f98ded6ca17c6a543b6b20cc9a7f84/aiokafka-0.14.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:be517b9b9513eba43ba19961dd770a6e26d08325743093feb47182770d235dd9", size = 1142238, upload-time = "2026-04-29T10:42:39.96Z" }, + { url = "https://files.pythonhosted.org/packages/88/78/322f797b9593a4cc8afd647342fa66b9ad732ee55098e5e084188c6202aa/aiokafka-0.14.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:219d2dc66b97b1aaea100697c928024b6a0348b7baa370b824900054bf86916e", size = 1131567, upload-time = "2026-04-29T10:42:41.542Z" }, + { url = "https://files.pythonhosted.org/packages/8d/0a/a45320778385142299a7fc3ae402152ec1f383537130b8aa8e8587742fad/aiokafka-0.14.0-cp312-cp312-win32.whl", hash = "sha256:1086b470f6c452471603a2d9c8d6933739230c75758d777d8d113ff8112bad68", size = 312160, upload-time = "2026-04-29T10:42:42.811Z" }, + { url = "https://files.pythonhosted.org/packages/a3/fb/7802a0ed69200e3e8e8791df06bd6daf9b00523839d045662de4ff061b18/aiokafka-0.14.0-cp312-cp312-win_amd64.whl", hash = "sha256:bcf3a8f6592d73f45965ca0750bfdfccf2555c8625358175c92f75f2cce1261a", size = 331897, upload-time = "2026-04-29T10:42:43.984Z" }, + { url = "https://files.pythonhosted.org/packages/30/b0/c9384541b2e4cc52a16402fc53fb9d44af0d78d37954cf8c7271c376ad47/aiokafka-0.14.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:db16e43fac4c1c5006131046c1bf370c580d6ac4495a10ac7778245710943179", size = 345859, upload-time = "2026-04-29T10:42:45.449Z" }, + { url = "https://files.pythonhosted.org/packages/7f/d1/fc266d9f4ffba4f197356c6ffdfbb0fe32e7cb874e240f299935d058ac06/aiokafka-0.14.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:32a8e91d88cf3ccf0778927715610d6579888c5f4748db4c2022cda25d628a48", size = 348284, upload-time = "2026-04-29T10:42:47.104Z" }, + { url = "https://files.pythonhosted.org/packages/b2/8e/0c4c270786dac79f3fca74c6166c3a25b61b0d26132be0d69f0d7f206f0a/aiokafka-0.14.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:aad4a575a506e7784e25e430f27026fe2f4378560b21b7f4e8c9a54f0d06eaee", size = 1117867, upload-time = "2026-04-29T10:42:48.394Z" }, + { url = "https://files.pythonhosted.org/packages/9d/7f/3b89fbd0a3be9edfd5b51e20bb5cd695c851219b63c501c051cf84367fa9/aiokafka-0.14.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:75e4a003502c9c3b5c705fa7c00d634ba146bf38fa5d525b80bb6ff6e3e779fe", size = 1108860, upload-time = "2026-04-29T10:42:50.249Z" }, + { url = "https://files.pythonhosted.org/packages/b3/59/849aba75cff93277bf6bf8b630de79e902949ff7ec48e4b12a64e6e32cae/aiokafka-0.14.0-cp313-cp313-win32.whl", hash = "sha256:a128e213cbc2bce0ea3db65a68920e52cebeeb8209bf001ac7aa022a8bd54d7d", size = 310889, upload-time = "2026-04-29T10:42:52.038Z" }, + { url = "https://files.pythonhosted.org/packages/c4/e5/52eab8f8515d23da7b5d90e2c5ba10eab9494a0314f749e3f73e003f4a50/aiokafka-0.14.0-cp313-cp313-win_amd64.whl", hash = "sha256:d6fa16bef3544be87bd1a7a8317b9d85e3da59f3202326d9ff22735ed052746e", size = 329470, upload-time = "2026-04-29T10:42:53.536Z" }, + { url = "https://files.pythonhosted.org/packages/50/9d/984803315fe2b883ea6e08b1d9c8a752bd5c16e966d8714bacc67c72c417/aiokafka-0.14.0-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:5d70615d1530ad19d0c4da8d87abaec0a12b9fdaabffdcd4e400efa0c50ef80c", size = 346672, upload-time = "2026-04-29T10:42:55.267Z" }, + { url = "https://files.pythonhosted.org/packages/49/df/da314966b7f3c3117bd78b082563cb03dbe3007848cb8f4b0932faf390a0/aiokafka-0.14.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:7e2392360c370b1ba6564c57d2889e154ecdb43157a8f7b7d7afe5e3c02fcc1a", size = 349594, upload-time = "2026-04-29T10:42:56.565Z" }, + { url = "https://files.pythonhosted.org/packages/57/7a/160516944ea0e0f68ea78e38f944c52f5248c7c7df26cba22a40b9f25709/aiokafka-0.14.0-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:201e38ecc595f9f65a945f1ef9085157ddf28f25cd2e482fd9efa1fcf4638213", size = 1114112, upload-time = "2026-04-29T10:42:57.869Z" }, + { url = "https://files.pythonhosted.org/packages/68/c4/9841118a2157e913e8ebfbc0a2b58f7b60f1f7202040c3e1df8925ed1184/aiokafka-0.14.0-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:1cd651e1f56571baae306fdd0b5509047ab9625797a24cd75902e139c5a20318", size = 1098571, upload-time = "2026-04-29T10:42:59.356Z" }, + { url = "https://files.pythonhosted.org/packages/d6/a1/0af8a37849a4108ae227f46c4c62f6beab31863cf66ba318fb73b0be5b26/aiokafka-0.14.0-cp314-cp314-win32.whl", hash = "sha256:128127eb96dab98150b636bb5f480c80e15f02f82a118eec206a521c8cf7cf7c", size = 314107, upload-time = "2026-04-29T10:43:01.111Z" }, + { url = "https://files.pythonhosted.org/packages/fa/18/fb46c65f758900c71d0f1c73b7802720f99cabcb1f4a11676573f9bc1b8f/aiokafka-0.14.0-cp314-cp314-win_amd64.whl", hash = "sha256:aa385039aa9b235359319bbdcf48c9c86a75d81c9c547d645056d00361238903", size = 333320, upload-time = "2026-04-29T10:43:02.424Z" }, +] + [[package]] name = "aiormq" version = "7.0.0" @@ -382,6 +413,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/e5/e2/c2e3abf398f80732e58b03be77bde9022550d221dd8781bf586bd4d97cc1/async_lru-2.3.0-py3-none-any.whl", hash = "sha256:eea27b01841909316f2cc739807acea1c623df2be8c5cfad7583286397bb8315", size = 8403, upload-time = "2026-03-19T01:04:30.883Z" }, ] +[[package]] +name = "async-timeout" +version = "5.0.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/a5/ae/136395dfbfe00dfc94da3f3e136d0b13f394cba8f4841120e34226265780/async_timeout-5.0.1.tar.gz", hash = "sha256:d9321a7a3d5a6a5e187e824d2fa0793ce379a202935782d555d6e9d2735677d3", size = 9274, upload-time = "2024-11-06T16:41:39.6Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fe/ba/e2081de779ca30d473f21f5b30e0e737c438205440784c7dfc81efc2b029/async_timeout-5.0.1-py3-none-any.whl", hash = "sha256:39e3809566ff85354557ec2398b55e096c8364bacac9405a7a1fa429e77fe76c", size = 6233, upload-time = "2024-11-06T16:41:37.9Z" }, +] + [[package]] name = "asyncpg" version = "0.31.0" @@ -3880,6 +3920,9 @@ gcp = [ gcp-pubsub = [ { name = "google-cloud-pubsub" }, ] +kafka = [ + { name = "aiokafka" }, +] llm = [ { name = "llama-index-core" }, { name = "llama-index-llms-openai" }, @@ -3903,6 +3946,7 @@ websockets = [ all = [ { name = "aiobotocore" }, { name = "aiofile" }, + { name = "aiokafka" }, { name = "aiosqlite" }, { name = "coverage" }, { name = "google-cloud-pubsub" }, @@ -3964,6 +4008,7 @@ docs = [ test = [ { name = "aiobotocore" }, { name = "aiofile" }, + { name = "aiokafka" }, { name = "aiosqlite" }, { name = "coverage" }, { name = "google-cloud-pubsub" }, @@ -3990,6 +4035,7 @@ requires-dist = [ { name = "adlfs", marker = "extra == 'azure'", specifier = ">=2024.7.0" }, { name = "aio-pika", specifier = ">=9.5.5,<11" }, { name = "aiobotocore", marker = "extra == 'aws-messaging'", specifier = ">=3.4,<4" }, + { name = "aiokafka", marker = "extra == 'kafka'", specifier = ">=0.11,<1" }, { name = "aiosqlite", specifier = ">=0.20,<1" }, { name = "annotated-types", specifier = ">=0.7,<1" }, { name = "asyncpg", specifier = ">=0.29.0,<1" }, @@ -4021,12 +4067,13 @@ requires-dist = [ { name = "uvloop", marker = "sys_platform != 'win32'", specifier = ">=0.21.0,<1" }, { name = "websockets", marker = "extra == 'websockets'", specifier = ">=14.2,<18" }, ] -provides-extras = ["aws", "aws-messaging", "azure", "gcp", "gcp-pubsub", "llm", "omq", "ray", "redis", "websockets"] +provides-extras = ["aws", "aws-messaging", "azure", "gcp", "gcp-pubsub", "kafka", "llm", "omq", "ray", "redis", "websockets"] [package.metadata.requires-dev] all = [ { name = "aiobotocore", specifier = ">=3.4,<4" }, { name = "aiofile", specifier = ">=3.9,<4" }, + { name = "aiokafka", specifier = ">=0.11,<1" }, { name = "aiosqlite", specifier = ">=0.20,<1" }, { name = "coverage", specifier = ">=7.7.1,<8" }, { name = "google-cloud-pubsub", specifier = ">=2.25,<3" }, @@ -4088,6 +4135,7 @@ docs = [ test = [ { name = "aiobotocore", specifier = ">=3.4,<4" }, { name = "aiofile", specifier = ">=3.9,<4" }, + { name = "aiokafka", specifier = ">=0.11,<1" }, { name = "aiosqlite", specifier = ">=0.20,<1" }, { name = "coverage", specifier = ">=7.7.1,<8" }, { name = "google-cloud-pubsub", specifier = ">=2.25,<3" },