Use source transaction for exactly-once output

faust/app/_attached.py

@@ -174,6 +174,12 @@ def put(
             value_serializer,
             callback,
         )
+        if self.app.in_transaction:
+            fut.message = fut.message._replace(
+                transactional_id=self.app.consumer.transactions.transactional_id_for_tp(
+                    message.tp
+                )
+            )
         # Note: Since FutureMessage have members that are unhashable
         # we wrap it in an Unordered object to stop heappush from crashing.
         # Unordered simply orders by random order, which is fine

faust/topics.py

@@ -232,6 +232,14 @@ def send_soon(
             callback=callback,
             eager_partitioning=eager_partitioning,
         )
+        if self.app.in_transaction:
+            event = current_event()
+            if event is not None:
+                fut.message = fut.message._replace(
+                    transactional_id=self.app.consumer.transactions.transactional_id_for_tp(
+                        event.message.tp
+                    )
+                )
         self.app.producer.send_soon(fut)
         return fut
 
@@ -409,6 +417,9 @@ async def publish_message(
         )
         assert topic is not None
         producer = await self._get_producer()
+        transaction_id_options: dict[str, str] = {}
+        if message.transactional_id is not None:
+            transaction_id_options["transactional_id"] = message.transactional_id
         state = app.sensors.on_send_initiated(
             producer,
             topic,
@@ -428,6 +439,7 @@ async def publish_message(
                 partition=partition,
                 timestamp=timestamp,
                 headers=headers,
+                **transaction_id_options,
             )
             app.sensors.on_send_completed(producer, state, ret)
             return await self._finalize_message(fut, ret)
@@ -441,6 +453,7 @@ async def publish_message(
                     partition=partition,
                     timestamp=timestamp,
                     headers=headers,
+                    **transaction_id_options,
                 ),
             )
             callback = partial(

faust/transport/consumer.py

@@ -294,11 +294,16 @@ async def send(
         transactional_id: Optional[str] = None,
     ) -> Awaitable[RecordMetadata]:
         """Schedule message to be sent by producer."""
-        group = transactional_id = None
+        group = None
         p = self.consumer.key_partition(topic, key, partition)
-        if p is not None:
-            group = self.app.assignor.group_for_topic(topic)
-            transactional_id = f"{self.app.conf.id}-{group}-{p}"
+        if transactional_id is None:
+            from faust.streams import current_event
+
+            event = current_event()
+            if event is not None:
+                transactional_id = self.transactional_id_for_tp(event.message.tp)
+            elif p is not None:
+                transactional_id = self.transactional_id_for_tp(TP(topic, p))
         return await self.producer.send(
             topic,
             key,
@@ -309,6 +314,10 @@ async def send(
             transactional_id=transactional_id,
         )
 
+    def transactional_id_for_tp(self, tp: TP) -> str:
+        group = self.app.assignor.group_for_topic(tp.topic)
+        return f"{self.app.conf.id}-{group}-{tp.partition}"
+
     def send_soon(self, fut: FutureMessage) -> None:
         raise NotImplementedError()
 
@@ -324,7 +333,15 @@ async def send_and_wait(
         transactional_id: Optional[str] = None,
     ) -> RecordMetadata:
         """Send message and wait for it to be transmitted."""
-        fut = await self.send(topic, key, value, partition, timestamp, headers)
+        fut = await self.send(
+            topic,
+            key,
+            value,
+            partition,
+            timestamp,
+            headers,
+            transactional_id=transactional_id,
+        )
         return await fut
 
     async def commit(

faust/types/tuples.py

@@ -69,6 +69,7 @@ class PendingMessage(NamedTuple):
     topic: Optional[str] = None
     offset: Optional[int] = None
     generation_id: Optional[int] = None
+    transactional_id: Optional[str] = None
 
 
 def _PendingMessage_to_Message(p: PendingMessage) -> "Message":

tests/unit/test_events.py

@@ -3,6 +3,8 @@
 import pytest
 
 from faust import Event
+from faust.app._attached import Attachments
+from faust.types import TP
 from tests.helpers import AsyncMock
 
 
@@ -172,6 +174,26 @@ def test_attach(self, *, event, app):
         )
         assert result is app._attachments.put()
 
+    def test_attach__retains_source_transaction_id(self, *, app, message):
+        app.in_transaction = True
+        app._attachments = Attachments(app)
+        app.consumer = Mock(
+            transactions=Mock(
+                transactional_id_for_tp=Mock(return_value="app-2-4")
+            )
+        )
+        message.tp = TP("input", 4)
+        message.offset = 20
+
+        future = app._attachments.put(
+            message, "output", key=b"key", value=b"value"
+        )
+
+        assert future.message.transactional_id == "app-2-4"
+        app.consumer.transactions.transactional_id_for_tp.assert_called_once_with(
+            message.tp
+        )
+
     @pytest.mark.asyncio
     async def test__send(self, *, event, app):
         app._attachments.maybe_put = AsyncMock(name="maybe_put")

tests/unit/test_topics.py

@@ -7,7 +7,7 @@
 import faust
 from faust import Event, Record
 from faust.exceptions import KeyDecodeError, ValueDecodeError
-from faust.types import FutureMessage, Message, PendingMessage
+from faust.types import FutureMessage, Message, PendingMessage, TP
 from tests.helpers import AsyncMock
 
 
@@ -162,6 +162,18 @@ async def test_publish_message__wait_enabled(self, *, topic, app):
         )
         callback.assert_called_once_with(fm)
 
+    @pytest.mark.asyncio
+    async def test_publish_message_forwards_pending_transaction_id(self, *, topic):
+        send_result = asyncio.get_running_loop().create_future()
+        producer = Mock(send=AsyncMock(return_value=send_result))
+        topic._get_producer = AsyncMock(return_value=producer)
+        future = topic.as_future_message(key="foo", value="bar")
+        future.message = future.message._replace(transactional_id="app-2-4")
+
+        await topic.publish_message(future)
+
+        assert producer.send.await_args.kwargs["transactional_id"] == "app-2-4"
+
     @pytest.mark.asyncio
     async def test_publish_message__non_bytes_key_does_not_crash(self, *, topic, app):
         # A raw, unserialized non-bytes key/value can reach publish_message
@@ -323,6 +335,19 @@ def test_send_soon(self, *, topic, app):
 
         app.producer.send_soon.assert_called_once_with(fut)
 
+    def test_send_soon__retains_source_transaction(self, *, topic, app):
+        app.in_transaction = True
+        app.consumer.transactions.transactional_id_for_tp = Mock(
+            return_value="app-2-4"
+        )
+        app.producer.send_soon = Mock(name="send_soon")
+        event = Mock(name="event", message=Mock(tp=TP("input", 4)))
+
+        with patch("faust.topics.current_event", return_value=event):
+            future = topic.send_soon(key="foo", value="bar")
+
+        assert future.message.transactional_id == "app-2-4"
+
     @pytest.mark.asyncio
     async def test_decode__decode_error_propagate(self, *, topic, message):
         exc = KeyDecodeError()

tests/unit/transport/test_consumer.py

@@ -254,6 +254,32 @@ async def test_send__topic_not_transactive(self, *, manager, producer):
             transactional_id=None,
         )
 
+    @pytest.mark.asyncio
+    async def test_send__uses_source_transaction_when_destination_unassigned(
+        self, *, manager, producer
+    ):
+        manager.app.assignor._topic_groups = {"input": 4}
+        manager.consumer.key_partition.return_value = None
+        event = Mock(name="event", message=Mock(tp=TP("input", 2)))
+
+        with patch("faust.streams.current_event", return_value=event):
+            await manager.send("output", "k", "v", partition=None, headers=None, timestamp=None)
+
+        manager.consumer.key_partition.assert_called_once_with("output", "k", None)
+        producer.send.assert_called_once_with(
+            "output",
+            "k",
+            "v",
+            None,
+            None,
+            None,
+            transactional_id="testid-4-2",
+        )
+
+    def test_transactional_id_for_tp(self, *, manager):
+        manager.app.assignor._topic_groups = {"input": 4}
+        assert manager.transactional_id_for_tp(TP("input", 2)) == "testid-4-2"
+
     def test_send_soon(self, *, manager):
         with pytest.raises(NotImplementedError):
             manager.send_soon(Mock(name="FutureMessage"))
@@ -276,6 +302,30 @@ async def send(*args, **kwargs):
             3,
             43.2,
             {},
+            transactional_id=None,
+        )
+
+    @pytest.mark.asyncio
+    async def test_send_and_wait__forwards_transactional_id(self, *, manager):
+        on_send = Mock()
+
+        async def send(*args, **kwargs):
+            on_send(*args, **kwargs)
+            return done_future()
+
+        manager.send = send
+
+        await manager.send_and_wait(
+            "t", "k", "v", 3, 43.2, {}, transactional_id="testid-2-4"
+        )
+        on_send.assert_called_once_with(
+            "t",
+            "k",
+            "v",
+            3,
+            43.2,
+            {},
+            transactional_id="testid-2-4",
         )
 
     @pytest.mark.asyncio