add: stream data normalizer

Author: GefMar

README.md

@@ -117,8 +117,33 @@ broker.subscribe(
 )
 ```
 
-The decoded value is passed as the first task argument. Stream messages also
-receive the `taskiq-stream` label with the source topic name.
+The decoded value is passed as the first task argument. A decoder can also
+return `StreamMessage` to map one Kafka message into the exact `args` and
+`kwargs` expected by the task.
+
+```python
+import json
+
+from taskiq_aio_kafka import StreamMessage
+
+
+@broker.task
+async def process_user(user_id: int, email: str) -> None:
+    print(user_id, email)
+
+
+def decode_user_event(message: bytes) -> StreamMessage:
+    event = json.loads(message)
+    return StreamMessage(
+        args=(event["user"]["id"],),
+        kwargs={"email": event["user"]["email"]},
+    )
+
+
+broker.subscribe("users.created", process_user, decoder=decode_user_event)
+```
+
+Stream messages also receive the `taskiq-stream` label with the source topic name.
 
 ## Configuration
 

taskiq_aio_kafka/__init__.py

@@ -1,5 +1,6 @@
 """Taskiq integration with aiokafka."""
 
-__all__ = ("AioKafkaBroker",)
+__all__ = ("AioKafkaBroker", "StreamMessage")
 
 from taskiq_aio_kafka.broker import AioKafkaBroker
+from taskiq_aio_kafka.subscriber import StreamMessage

taskiq_aio_kafka/broker.py

@@ -20,7 +20,7 @@
 from .decorated_task import AioKafkaDecoratedTask
 from .exceptions import WrongAioKafkaBrokerParametersError
 from .models import KafkaConsumerParameters, KafkaProducerParameters
-from .subscriber import StreamDecoder, StreamSubscriber
+from .subscriber import StreamDecoder, StreamMessage, StreamSubscriber
 from .topic import Topic
 from .types import TopicType
 from .utils import get_topic_name
@@ -171,8 +171,8 @@ def configure_consumer(self, **consumer_parameters: Any) -> None:
         )
 
     @staticmethod
-    def _default_stream_decoder(message: bytes) -> bytes:
-        return message
+    def _default_stream_decoder(message: bytes) -> StreamMessage:
+        return StreamMessage(args=(message,))
 
     def subscribe(
         self,
@@ -390,6 +390,7 @@ def _build_stream_message(
     ) -> bytes:
         raw_value = raw_kafka_message.value
         decoded_value = subscriber.decoder(raw_value)
+        stream_message = self._normalize_stream_message(decoded_value)
         labels = {
             TASK_STREAM_LABEL: raw_kafka_message.topic,
             **subscriber.labels,
@@ -399,7 +400,13 @@ def _build_stream_message(
             task_name=subscriber.task_name,
             labels=labels,
             labels_types={},
-            args=[decoded_value],
-            kwargs={},
+            args=list(stream_message.args),
+            kwargs=stream_message.kwargs,
         )
         return self.formatter.dumps(message).message
+
+    @staticmethod
+    def _normalize_stream_message(message: Any | StreamMessage) -> StreamMessage:
+        if isinstance(message, StreamMessage):
+            return message
+        return StreamMessage(args=(message,))

taskiq_aio_kafka/subscriber.py

@@ -1,10 +1,19 @@
-__all__ = ("StreamDecoder", "StreamSubscriber")
+__all__ = ("StreamDecoder", "StreamMessage", "StreamSubscriber")
 
-from collections.abc import Callable
+from collections.abc import Callable, Sequence
 from dataclasses import dataclass, field
 from typing import Any
 
-StreamDecoder = Callable[[bytes], Any]
+
+@dataclass(frozen=True)
+class StreamMessage:
+    """Decoded stream message arguments for a taskiq task."""
+
+    args: Sequence[Any] = ()
+    kwargs: dict[str, Any] = field(default_factory=dict)
+
+
+StreamDecoder = Callable[[bytes], Any | StreamMessage]
 
 
 @dataclass(frozen=True)

tests/test_broker_multi_topic.py

@@ -8,6 +8,7 @@
 
 from taskiq_aio_kafka.broker import AioKafkaBroker
 from taskiq_aio_kafka.constants import TASK_STREAM_LABEL
+from taskiq_aio_kafka.subscriber import StreamMessage
 from taskiq_aio_kafka.topic import Topic
 
 
@@ -357,6 +358,47 @@ async def test_task(message: bytes) -> None:
     }
 
 
+async def test_subscribe_decoder_can_return_task_args_and_kwargs() -> None:
+    """Test that stream messages can be mapped to task args and kwargs."""
+    broker = AioKafkaBroker(
+        bootstrap_servers="localhost",
+        kafka_topic="default-topic",
+        kafka_admin_client=get_admin_client_mock(),
+    )
+
+    @broker.task
+    async def test_task(user_id: int, email: str, active: bool) -> None:
+        assert user_id
+        assert email
+        assert active
+
+    def decode_message(message: bytes) -> StreamMessage:
+        user_id, email, active = message.decode().split(":")
+        return StreamMessage(
+            args=(int(user_id),),
+            kwargs={
+                "email": email,
+                "active": active == "true",
+            },
+        )
+
+    broker.subscribe("stream-topic", test_task, decoder=decode_message)
+    broker._aiokafka_consumer = _ConsumerMock(
+        [build_consumer_record("stream-topic", b"1:user@example.com:true")],
+    )
+    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 == [1]
+    assert taskiq_message.kwargs == {
+        "email": "user@example.com",
+        "active": True,
+    }
+
+
 def test_subscribe_adds_topic_to_broker_topics() -> None:
     """Test that subscribed topic is included in listened topics."""
     broker = AioKafkaBroker(