diff --git a/faust/app/_attached.py b/faust/app/_attached.py index d7a94b317..b3f199c57 100644 --- a/faust/app/_attached.py +++ b/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 diff --git a/faust/topics.py b/faust/topics.py index 5b8fa5323..b47c13629 100644 --- a/faust/topics.py +++ b/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( diff --git a/faust/transport/consumer.py b/faust/transport/consumer.py index 9e49a9a8f..71009520c 100644 --- a/faust/transport/consumer.py +++ b/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( diff --git a/faust/types/tuples.py b/faust/types/tuples.py index c25898d8a..97c649c58 100644 --- a/faust/types/tuples.py +++ b/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": diff --git a/tests/unit/test_events.py b/tests/unit/test_events.py index 626329893..5fc4696cc 100644 --- a/tests/unit/test_events.py +++ b/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") diff --git a/tests/unit/test_topics.py b/tests/unit/test_topics.py index e880ce000..4c1b624d3 100644 --- a/tests/unit/test_topics.py +++ b/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() diff --git a/tests/unit/transport/test_consumer.py b/tests/unit/transport/test_consumer.py index d0261179a..5ce31ca71 100644 --- a/tests/unit/transport/test_consumer.py +++ b/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