Batching: Keep ids and tags in BatchingTopic and BatchingInventory.

Author: Mathieu Gascon-Lefebvre

src/saturn_engine/worker/inventories/batching.py

@@ -41,4 +41,8 @@ async def iterate(self, after: Optional[str] = None) -> AsyncIterator[Item]:
                 return
 
             after = batch[-1].id
-            yield Item(id=after, args={"batch": [item.args for item in batch]})
+            yield Item(
+                id=after,
+                args={"batch": [item.args for item in batch]},
+                tags={"batched_ids": ", ".join([item.id for item in batch])},
+            )

src/saturn_engine/worker/topics/batching.py

@@ -4,6 +4,7 @@
 import asyncio
 import contextlib
 import dataclasses
+import json
 from collections.abc import AsyncGenerator
 from contextlib import asynccontextmanager
 from datetime import datetime
@@ -109,6 +110,7 @@ async def message_context(
     ) -> AsyncIterator[TopicMessage]:
         context = contextlib.AsyncExitStack()
         message_args: list[dict] = []
+        tags: dict[str, str] = {}
 
         for message_context in batch:
             message: TopicMessage
@@ -117,6 +119,7 @@ async def message_context(
             else:
                 message = message_context
             message_args.append(message.args)
+            tags[f"{message.id}_tags"] = json.dumps(message.tags)
 
         async with context:
-            yield TopicMessage(args={"batch": message_args})
+            yield TopicMessage(args={"batch": message_args}, tags=tags)

tests/worker/inventories/test_batching_inventory.py

@@ -32,10 +32,21 @@ async def test_batching_inventory() -> None:
             },
         ),
     ]
+    assert [i.tags for i in items] == [
+        {"batched_ids": "0, 1, 2"},
+        {"batched_ids": "3, 4, 5"},
+        {"batched_ids": "6, 7, 8"},
+        {"batched_ids": "9"},
+    ]
 
     items = await alib.list(inventory.iterate(after="4"))
 
     assert [(i.id, i.args) for i in items] == [
         ("7", {"batch": [{"a": "5"}, {"a": "6"}, {"a": "7"}]}),
         ("9", {"batch": [{"a": "8"}, {"a": "9"}]}),
     ]
+
+    assert [i.tags for i in items] == [
+        {"batched_ids": "5, 6, 7"},
+        {"batched_ids": "8, 9"},
+    ]

tests/worker/topics/test_batching_topic.py

@@ -149,3 +149,44 @@ async def test_batching_topic_context_manager(
 
     assert batch_number == 2
     await topic.close()
+
+
+@pytest.mark.asyncio
+async def test_batching_topic_tags() -> None:
+    BATCH_SIZE = 5
+
+    topic = BatchingTopic(
+        options=BatchingTopic.Options(
+            topic=TopicItem(
+                name="static-topic-with-tags",
+                type="StaticTopic",
+                options={
+                    "messages": [
+                        {"id": 1, "args": {}, "tags": {"hello": "1"}},
+                        {"id": 2, "args": {}, "tags": {"hello": "2"}},
+                        {"id": 3, "args": {}, "tags": {"hello": "3"}},
+                        {"id": 4, "args": {}, "tags": {"hello": "4"}},
+                        {"id": 5, "args": {}, "tags": {"hello": "5"}},
+                    ],
+                },
+            ),
+            batch_size=BATCH_SIZE,
+        ),
+        services=ServicesNamespace(strict=False),
+    )
+
+    async with alib.scoped_iter(topic.run()) as scoped_topic_iter:
+        context = await scoped_topic_iter.__anext__()
+        assert isinstance(context, AsyncContextManager)
+        async with context as message:
+            ...
+
+    await topic.close()
+
+    assert message.tags == {
+        "1_tags": '{"hello": "1"}',
+        "2_tags": '{"hello": "2"}',
+        "3_tags": '{"hello": "3"}',
+        "4_tags": '{"hello": "4"}',
+        "5_tags": '{"hello": "5"}',
+    }