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" },