add: stream subscriber mapping

Author: GefMar

README.md

@@ -87,6 +87,39 @@ async def regular_task() -> None:
 await regular_task.kiq()
 ```
 
+## Stream topics
+
+You can subscribe an existing task to a Kafka topic with raw messages.
+Messages from subscribed stream topics are wrapped into regular taskiq
+messages before execution, so result backends and worker middlewares keep
+working as usual.
+
+```python
+import json
+
+from taskiq_aio_kafka import AioKafkaBroker
+
+broker = AioKafkaBroker(
+    bootstrap_servers="localhost",
+    kafka_topic="taskiq-topic",
+)
+
+
+@broker.task
+async def process_user_created(event: dict[str, object]) -> None:
+    print(event)
+
+
+broker.subscribe(
+    "users.created",
+    process_user_created,
+    decoder=json.loads,
+)
+```
+
+The decoded value is passed as the first task argument. Stream messages also
+receive the `taskiq-stream` label with the source topic name.
+
 ## Configuration
 
 AioKafkaBroker parameters:

taskiq_aio_kafka/broker.py

@@ -6,17 +6,21 @@
 from typing import Any, TypeVar, overload
 
 from aiokafka import AIOKafkaConsumer, AIOKafkaProducer
+from aiokafka.structs import ConsumerRecord
 from kafka.admin import KafkaAdminClient, NewTopic
 from kafka.coordinator.assignors.roundrobin import RoundRobinPartitionAssignor
 from kafka.partitioner.default import DefaultPartitioner
 from taskiq import AsyncResultBackend, BrokerMessage
 from taskiq.abc.broker import AsyncBroker
+from taskiq.decor import AsyncTaskiqDecoratedTask
+from taskiq.message import TaskiqMessage
 from typing_extensions import ParamSpec
 
-from .constants import TASK_TOPIC_LABEL
+from .constants import TASK_STREAM_LABEL, TASK_TOPIC_LABEL
 from .decorated_task import AioKafkaDecoratedTask
 from .exceptions import WrongAioKafkaBrokerParametersError
 from .models import KafkaConsumerParameters, KafkaProducerParameters
+from .subscriber import StreamDecoder, StreamSubscriber
 from .topic import Topic
 from .types import TopicType
 from .utils import get_topic_name
@@ -79,6 +83,7 @@ def __init__(  # noqa: PLR0913
                     get_topic_name(topic),
                     topic,
                 )
+        self._stream_subscribers: dict[str, StreamSubscriber] = {}
 
         self._aiokafka_producer_params: KafkaProducerParameters = (
             KafkaProducerParameters()
@@ -165,6 +170,30 @@ def configure_consumer(self, **consumer_parameters: Any) -> None:
             **consumer_parameters,
         )
 
+    @staticmethod
+    def _default_stream_decoder(message: bytes) -> bytes:
+        return message
+
+    def subscribe(
+        self,
+        topic: TopicType,
+        task: AsyncTaskiqDecoratedTask[Any, Any],
+        decoder: StreamDecoder | None = None,
+        **labels: Any,
+    ) -> None:
+        """Subscribe task to raw Kafka topic messages."""
+        topic_name = get_topic_name(topic)
+        if topic_name in self._stream_subscribers:
+            error_message = f"Topic {topic_name!r} is already subscribed."
+            raise ValueError(error_message)
+
+        self._kafka_topics.setdefault(topic_name, topic)
+        self._stream_subscribers[topic_name] = StreamSubscriber(
+            task_name=task.task_name,
+            decoder=decoder or self._default_stream_decoder,
+            labels=labels,
+        )
+
     @overload
     def task(
         self,
@@ -347,4 +376,30 @@ async def listen(
             raise ValueError("Please run startup before listening.")
 
         async for raw_kafka_message in self._aiokafka_consumer:
-            yield raw_kafka_message.value
+            subscriber = self._stream_subscribers.get(raw_kafka_message.topic)
+            if subscriber is None:
+                yield raw_kafka_message.value
+                continue
+
+            yield self._build_stream_message(raw_kafka_message, subscriber)
+
+    def _build_stream_message(
+        self,
+        raw_kafka_message: ConsumerRecord[Any, bytes],
+        subscriber: StreamSubscriber,
+    ) -> bytes:
+        raw_value = raw_kafka_message.value
+        decoded_value = subscriber.decoder(raw_value)
+        labels = {
+            TASK_STREAM_LABEL: raw_kafka_message.topic,
+            **subscriber.labels,
+        }
+        message = TaskiqMessage(
+            task_id=self.id_generator(),
+            task_name=subscriber.task_name,
+            labels=labels,
+            labels_types={},
+            args=[decoded_value],
+            kwargs={},
+        )
+        return self.formatter.dumps(message).message

taskiq_aio_kafka/constants.py

@@ -1,3 +1,4 @@
-__all__ = ("TASK_TOPIC_LABEL",)
+__all__ = ("TASK_STREAM_LABEL", "TASK_TOPIC_LABEL")
 
+TASK_STREAM_LABEL = "taskiq-stream"
 TASK_TOPIC_LABEL = "taskiq_aio_kafka_topic"

taskiq_aio_kafka/subscriber.py

@@ -0,0 +1,16 @@
+__all__ = ("StreamDecoder", "StreamSubscriber")
+
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from typing import Any
+
+StreamDecoder = Callable[[bytes], Any]
+
+
+@dataclass(frozen=True)
+class StreamSubscriber:
+    """Kafka stream subscriber bound to a taskiq task."""
+
+    task_name: str
+    decoder: StreamDecoder
+    labels: dict[str, Any] = field(default_factory=dict)

tests/test_broker_multi_topic.py

@@ -2,10 +2,12 @@
 from unittest.mock import Mock
 
 import pytest
+from aiokafka.structs import ConsumerRecord
 from kafka.admin import KafkaAdminClient, NewTopic
 from taskiq import BrokerMessage
 
 from taskiq_aio_kafka.broker import AioKafkaBroker
+from taskiq_aio_kafka.constants import TASK_STREAM_LABEL
 from taskiq_aio_kafka.topic import Topic
 
 
@@ -40,13 +42,53 @@ async def stop(self) -> None:
         """Stop consumer."""
 
 
+class _ConsumerMock:
+    """Kafka consumer mock."""
+
+    def __init__(self, messages: list[ConsumerRecord[None, bytes]]) -> None:
+        self._messages = iter(messages)
+
+    def __aiter__(self) -> "_ConsumerMock":
+        return self
+
+    async def __anext__(self) -> ConsumerRecord[None, bytes]:
+        """Return next message."""
+        try:
+            return next(self._messages)
+        except StopIteration as exc:
+            raise StopAsyncIteration from exc
+
+
 def get_admin_client_mock() -> KafkaAdminClient:
     """Get kafka admin client mock."""
     admin_client = Mock(spec=KafkaAdminClient)
     admin_client.list_topics.return_value = []
     return admin_client
 
 
+def build_consumer_record(topic: str, value: bytes) -> ConsumerRecord[None, bytes]:
+    """Build Kafka consumer record."""
+    return ConsumerRecord(
+        topic=topic,
+        partition=0,
+        offset=0,
+        timestamp=0,
+        timestamp_type=0,
+        key=None,
+        value=value,
+        checksum=None,
+        serialized_key_size=0,
+        serialized_value_size=len(value),
+        headers=[],
+    )
+
+
+async def get_first_task(broker: AioKafkaBroker) -> bytes:  # type: ignore[return]
+    """Get first message from the topic."""
+    async for message in broker.listen():
+        return message
+
+
 async def test_task_topic_is_used_for_kick() -> None:
     """Test that task is sent to its declared topic."""
     broker = AioKafkaBroker(
@@ -252,3 +294,99 @@ def create_consumer(
     await broker.shutdown()
 
     assert consumer_topics == ("default-topic", "extra-topic")
+
+
+async def test_subscribe_wraps_raw_topic_message() -> None:
+    """Test that raw Kafka messages are wrapped as taskiq messages."""
+    broker = AioKafkaBroker(
+        bootstrap_servers="localhost",
+        kafka_topic="default-topic",
+        kafka_admin_client=get_admin_client_mock(),
+    )
+
+    @broker.task
+    async def test_task(message: str) -> None:
+        assert message
+
+    broker.subscribe(
+        "stream-topic",
+        test_task,
+        decoder=lambda message: message.decode(),
+    )
+    broker._aiokafka_consumer = _ConsumerMock(
+        [build_consumer_record("stream-topic", b"raw-message")],
+    )
+    broker._is_consumer_started = True
+
+    received_message = await get_first_task(broker)
+    taskiq_message = broker.formatter.loads(received_message)
+
+    assert taskiq_message.task_name == test_task.task_name
+    assert taskiq_message.args == ["raw-message"]
+    assert taskiq_message.kwargs == {}
+    assert taskiq_message.labels[TASK_STREAM_LABEL] == "stream-topic"
+
+
+async def test_subscribe_accepts_topic_object_and_custom_labels() -> None:
+    """Test that stream subscriptions can use Topic objects."""
+    broker = AioKafkaBroker(
+        bootstrap_servers="localhost",
+        kafka_topic="default-topic",
+        kafka_admin_client=get_admin_client_mock(),
+    )
+    stream_topic = Topic("stream-topic")
+
+    @broker.task
+    async def test_task(message: bytes) -> None:
+        assert message
+
+    broker.subscribe(stream_topic, test_task, source="external")
+    broker._aiokafka_consumer = _ConsumerMock(
+        [build_consumer_record(stream_topic.name, b"raw-message")],
+    )
+    broker._is_consumer_started = True
+
+    received_message = await get_first_task(broker)
+    taskiq_message = broker.formatter.loads(received_message)
+
+    assert taskiq_message.task_name == test_task.task_name
+    assert taskiq_message.args == ["raw-message"]
+    assert taskiq_message.labels == {
+        TASK_STREAM_LABEL: stream_topic.name,
+        "source": "external",
+    }
+
+
+def test_subscribe_adds_topic_to_broker_topics() -> None:
+    """Test that subscribed topic is included in listened topics."""
+    broker = AioKafkaBroker(
+        bootstrap_servers="localhost",
+        kafka_topic="default-topic",
+        kafka_admin_client=get_admin_client_mock(),
+    )
+
+    @broker.task
+    async def test_task(message: bytes) -> None:
+        assert message
+
+    broker.subscribe("stream-topic", test_task)
+
+    assert set(broker._kafka_topics) == {"default-topic", "stream-topic"}
+
+
+def test_subscribe_rejects_duplicate_topics() -> None:
+    """Test that stream topic cannot be subscribed twice."""
+    broker = AioKafkaBroker(
+        bootstrap_servers="localhost",
+        kafka_topic="default-topic",
+        kafka_admin_client=get_admin_client_mock(),
+    )
+
+    @broker.task
+    async def test_task(message: bytes) -> None:
+        assert message
+
+    broker.subscribe("stream-topic", test_task)
+
+    with pytest.raises(ValueError, match="already subscribed"):
+        broker.subscribe("stream-topic", test_task)