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

@@ -1,9 +1,11 @@
 from typing import AsyncContextManager
 from typing import AsyncIterator
+from typing import DefaultDict
 
 import asyncio
 import contextlib
 import dataclasses
+from collections import defaultdict
 from collections.abc import AsyncGenerator
 from contextlib import asynccontextmanager
 from datetime import datetime
@@ -109,6 +111,8 @@ async def message_context(
     ) -> AsyncIterator[TopicMessage]:
         context = contextlib.AsyncExitStack()
         message_args: list[dict] = []
+        ids: list[str] = []
+        tag_lists: DefaultDict[str, list[str]] = defaultdict(list)
 
         for message_context in batch:
             message: TopicMessage
@@ -118,5 +122,16 @@ async def message_context(
                 message = message_context
             message_args.append(message.args)
 
+            ids.append(message.id)
+            for tag, value in message.tags.items():
+                tag_lists[tag].append(value)
+
+        tags = {"batched_ids": ", ".join(ids)} | {
+            tag: ", ".join(values) for tag, values in tag_lists.items()
+        }
+
         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,42 @@ 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", "hi": "a"}},
+                        {"id": "3", "args": {}, "tags": {"hello": "3"}},
+                        {"id": "4", "args": {}, "tags": {"hello": "4"}},
+                        {"id": "5", "args": {}, "tags": {"hello": "5", "hi": "b"}},
+                    ],
+                },
+            ),
+            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 == {
+        "batched_ids": "1, 2, 3, 4, 5",
+        "hello": "1, 2, 3, 4, 5",
+        "hi": "a, b",
+    }