transaction identity

Author: vgvoleg

Diff for: examples/topic/topic_transactions_example.py

@@ -5,13 +5,13 @@ def writer_example(driver: ydb.Driver, topic: str):
     session_pool = ydb.QuerySessionPool(driver)
 
     def callee(tx: ydb.QueryTxContext):
-        tx_writer: ydb.TopicTxWriter = driver.topic_client.tx_writer(tx, topic) # <=======
+        tx_writer: ydb.TopicTxWriter = driver.topic_client.tx_writer(tx, topic)  # <=======
         # дефолт - без дедупликации, без ретраев и без producer_id.
 
         with tx.execute(query="select 1") as result_sets:
             messages = [result_set.rows[0] for result_set in result_sets]
 
-        tx_writer.write(messages) # вне зависимости от состояния вышестоящего стрима поведение должно быть одинаковое
+        tx_writer.write(messages)  # вне зависимости от состояния вышестоящего стрима поведение должно быть одинаковое
 
     session_pool.retry_tx_sync(callee)
 
@@ -20,7 +20,7 @@ def reader_example(driver: ydb.Driver, reader: ydb.TopicReader):
     session_pool = ydb.QuerySessionPool(driver)
 
     def callee(tx: ydb.QueryTxContext):
-        batch = reader.receive_batch_with_tx(tx, max_messages=5) # <=======
+        batch = reader.receive_batch_with_tx(tx, max_messages=5)  # <=======
 
         with tx.execute(query="INSERT INTO max_values(val) VALUES ($val)", parameters={"$val": max(batch)}) as _:
             pass

Diff for: tests/query/test_query_transaction.py

@@ -92,3 +92,15 @@ def test_execute_two_results(self, tx: QueryTxContext):
 
         assert res == [[1], [2]]
         assert counter == 2
+
+    def test_tx_identity_before_begin_raises(self, tx: QueryTxContext):
+        with pytest.raises(RuntimeError):
+            tx._tx_identity()
+
+    def test_tx_identity_after_begin_works(self, tx: QueryTxContext):
+        tx.begin()
+
+        identity = tx._tx_identity()
+
+        assert identity.tx_id == tx.tx_id
+        assert identity.session_id == tx.session_id

Diff for: ydb/_grpc/grpcwrapper/ydb_topic.py

@@ -142,6 +142,18 @@ def from_proto(msg: ydb_topic_pb2.UpdateTokenResponse) -> typing.Any:
 ########################################################################################################################
 
 
+@dataclass
+class TransactionIdentity(IToProto):
+    tx_id: str
+    session_id: str
+
+    def to_proto(self) -> ydb_topic_pb2.TransactionIdentity:
+        return ydb_topic_pb2.TransactionIdentity(
+            id=self.tx_id,
+            session=self.session_id,
+        )
+
+
 class StreamWriteMessage:
     @dataclass()
     class InitRequest(IToProto):
@@ -200,6 +212,7 @@ def from_proto(
     class WriteRequest(IToProto):
         messages: typing.List["StreamWriteMessage.WriteRequest.MessageData"]
         codec: int
+        tx_identity: Optional[TransactionIdentity]
 
         @dataclass
         class MessageData(IToProto):
@@ -238,6 +251,9 @@ def to_proto(self) -> ydb_topic_pb2.StreamWriteMessage.WriteRequest:
             proto = ydb_topic_pb2.StreamWriteMessage.WriteRequest()
             proto.codec = self.codec
 
+            if self.tx_identity is not None:
+                proto.tx = self.tx_identity.to_proto()
+
             for message in self.messages:
                 proto_mess = proto.messages.add()
                 proto_mess.CopyFrom(message.to_proto())

Diff for: ydb/_topic_writer/topic_writer.py

@@ -11,6 +11,7 @@
 
 import ydb.aio
 from .._grpc.grpcwrapper.ydb_topic import StreamWriteMessage
+from .._grpc.grpcwrapper.ydb_topic import TransactionIdentity
 from .._grpc.grpcwrapper.common_utils import IToProto
 from .._grpc.grpcwrapper.ydb_topic_public_types import PublicCodec
 from .. import connection
@@ -205,6 +206,7 @@ def default_serializer_message_content(data: Any) -> bytes:
 
 def messages_to_proto_requests(
     messages: List[InternalMessage],
+    tx_identity: Optional[TransactionIdentity],
 ) -> List[StreamWriteMessage.FromClient]:
 
     gropus = _slit_messages_for_send(messages)
@@ -215,6 +217,7 @@ def messages_to_proto_requests(
             StreamWriteMessage.WriteRequest(
                 messages=list(map(InternalMessage.to_message_data, group)),
                 codec=group[0].codec,
+                tx_identity=tx_identity,
             )
         )
         res.append(req)
@@ -239,6 +242,7 @@ def messages_to_proto_requests(
                 ),
             ],
             codec=20000,
+            tx_identity=None,
         )
     )
     .to_proto()

Diff for: ydb/_topic_writer/topic_writer_asyncio.py

@@ -35,6 +35,7 @@
     UpdateTokenRequest,
     UpdateTokenResponse,
     StreamWriteMessage,
+    TransactionIdentity,
     WriterMessagesFromServerToClient,
 )
 from .._grpc.grpcwrapper.common_utils import (
@@ -43,6 +44,9 @@
     GrpcWrapperAsyncIO,
 )
 
+if typing.TYPE_CHECKING:
+    from ..query.transaction import BaseQueryTxContext
+
 logger = logging.getLogger(__name__)
 
 
@@ -165,7 +169,20 @@ async def wait_init(self) -> PublicWriterInitInfo:
 
 
 class TxWriterAsyncIO(WriterAsyncIO):
-    ...
+    _tx: object
+
+    def __init__(
+        self,
+        tx,
+        driver: SupportedDriverType,
+        settings: PublicWriterSettings,
+        _client=None,
+    ):
+        self._tx = tx
+        self._loop = asyncio.get_running_loop()
+        self._closed = False
+        self._reconnector = WriterAsyncIOReconnector(driver=driver, settings=WriterSettings(settings), tx=self._tx)
+        self._parent = _client
 
 
 class WriterAsyncIOReconnector:
@@ -182,6 +199,7 @@ class WriterAsyncIOReconnector:
     _codec_selector_batch_num: int
     _codec_selector_last_codec: Optional[PublicCodec]
     _codec_selector_check_batches_interval: int
+    _tx: Optional["BaseQueryTxContext"]
 
     if typing.TYPE_CHECKING:
         _messages_for_encode: asyncio.Queue[List[InternalMessage]]
@@ -199,7 +217,9 @@ class WriterAsyncIOReconnector:
         _stop_reason: asyncio.Future
     _init_info: Optional[PublicWriterInitInfo]
 
-    def __init__(self, driver: SupportedDriverType, settings: WriterSettings):
+    def __init__(
+        self, driver: SupportedDriverType, settings: WriterSettings, tx: Optional["BaseQueryTxContext"] = None
+    ):
         self._closed = False
         self._loop = asyncio.get_running_loop()
         self._driver = driver
@@ -209,6 +229,7 @@ def __init__(self, driver: SupportedDriverType, settings: WriterSettings):
         self._init_info = None
         self._stream_connected = asyncio.Event()
         self._settings = settings
+        self._tx = tx
 
         self._codec_functions = {
             PublicCodec.RAW: lambda data: data,
@@ -358,10 +379,12 @@ async def _connection_loop(self):
             # noinspection PyBroadException
             stream_writer = None
             try:
+                tx_identity = None if self._tx is None else self._tx._tx_identity()
                 stream_writer = await WriterAsyncIOStream.create(
                     self._driver,
                     self._init_message,
                     self._settings.update_token_interval,
+                    tx_identity=tx_identity,
                 )
                 try:
                     if self._init_info is None:
@@ -601,10 +624,13 @@ class WriterAsyncIOStream:
     _update_token_event: asyncio.Event
     _get_token_function: Optional[Callable[[], str]]
 
+    _tx_identity: Optional[TransactionIdentity]
+
     def __init__(
         self,
         update_token_interval: Optional[Union[int, float]] = None,
         get_token_function: Optional[Callable[[], str]] = None,
+        tx_identity: Optional[TransactionIdentity] = None,
     ):
         self._closed = False
 
@@ -613,6 +639,8 @@ def __init__(
         self._update_token_event = asyncio.Event()
         self._update_token_task = None
 
+        self._tx_identity = tx_identity
+
     async def close(self):
         if self._closed:
             return
@@ -629,6 +657,7 @@ async def create(
         driver: SupportedDriverType,
         init_request: StreamWriteMessage.InitRequest,
         update_token_interval: Optional[Union[int, float]] = None,
+        tx_identity: Optional[TransactionIdentity] = None,
     ) -> "WriterAsyncIOStream":
         stream = GrpcWrapperAsyncIO(StreamWriteMessage.FromServer.from_proto)
 
@@ -638,6 +667,7 @@ async def create(
         writer = WriterAsyncIOStream(
             update_token_interval=update_token_interval,
             get_token_function=creds.get_auth_token if creds else lambda: "",
+            tx_identity=tx_identity,
         )
         await writer._start(stream, init_request)
         return writer
@@ -684,7 +714,7 @@ def write(self, messages: List[InternalMessage]):
         if self._closed:
             raise RuntimeError("Can not write on closed stream.")
 
-        for request in messages_to_proto_requests(messages):
+        for request in messages_to_proto_requests(messages, self._tx_identity):
             self._stream.write(request)
 
     async def _update_token_loop(self):

Diff for: ydb/_topic_writer/topic_writer_asyncio_test.py

@@ -18,6 +18,7 @@
 from .._grpc.grpcwrapper.ydb_topic import (
     Codec,
     StreamWriteMessage,
+    TransactionIdentity,
     UpdateTokenRequest,
     UpdateTokenResponse,
 )
@@ -43,6 +44,12 @@
 from ..credentials import AnonymousCredentials
 
 
+FAKE_TRANSACTION_IDENTITY = TransactionIdentity(
+    tx_id="transaction_id",
+    session_id="session_id",
+)
+
+
 @pytest.fixture
 def default_driver() -> aio.Driver:
     driver = mock.Mock(spec=aio.Driver)
@@ -148,6 +155,44 @@ async def test_write_a_message(self, writer_and_stream: WriterWithMockedStream):
         expected_message = StreamWriteMessage.FromClient(
             StreamWriteMessage.WriteRequest(
                 codec=Codec.CODEC_RAW,
+                tx_identity=None,
+                messages=[
+                    StreamWriteMessage.WriteRequest.MessageData(
+                        seq_no=1,
+                        created_at=now,
+                        data=data,
+                        metadata_items={},
+                        uncompressed_size=len(data),
+                        partitioning=None,
+                    )
+                ],
+            )
+        )
+
+        sent_message = await writer_and_stream.stream.from_client.get()
+        assert expected_message == sent_message
+
+    async def test_write_a_message_with_tx(self, writer_and_stream: WriterWithMockedStream):
+        writer_and_stream.writer._tx_identity = FAKE_TRANSACTION_IDENTITY
+
+        data = "123".encode()
+        now = datetime.datetime.now(datetime.timezone.utc)
+        writer_and_stream.writer.write(
+            [
+                InternalMessage(
+                    PublicMessage(
+                        seqno=1,
+                        created_at=now,
+                        data=data,
+                    )
+                )
+            ]
+        )
+
+        expected_message = StreamWriteMessage.FromClient(
+            StreamWriteMessage.WriteRequest(
+                codec=Codec.CODEC_RAW,
+                tx_identity=FAKE_TRANSACTION_IDENTITY,
                 messages=[
                     StreamWriteMessage.WriteRequest.MessageData(
                         seq_no=1,
@@ -264,7 +309,7 @@ def _create(self):
 
         res = DoubleQueueWriters()
 
-        async def async_create(driver, init_message, token_getter):
+        async def async_create(driver, init_message, token_getter, tx_identity):
             return res.get_first()
 
         monkeypatch.setattr(WriterAsyncIOStream, "create", async_create)

Diff for: ydb/query/transaction.py

@@ -11,6 +11,7 @@
     _apis,
     issues,
 )
+from .._grpc.grpcwrapper import ydb_topic as _ydb_topic
 from .._grpc.grpcwrapper import ydb_query as _ydb_query
 from ..connection import _RpcState as RpcState
 
@@ -215,6 +216,11 @@ def tx_id(self) -> Optional[str]:
         """
         return self._tx_state.tx_id
 
+    def _tx_identity(self) -> _ydb_topic.TransactionIdentity:
+        if not self.tx_id:
+            raise RuntimeError("Unable to get tx identity without started tx.")
+        return _ydb_topic.TransactionIdentity(self.tx_id, self.session_id)
+
     def _begin_call(self, settings: Optional[BaseRequestSettings]) -> "BaseQueryTxContext":
         self._tx_state._check_invalid_transition(QueryTxStateEnum.BEGINED)
 

Diff for: ydb/topic.py

@@ -65,6 +65,7 @@
     PublicWriteResult as TopicWriteResult,
 )
 
+from ydb._topic_writer.topic_writer_asyncio import TxWriterAsyncIO as TopicTxWriterAsyncIO
 from ydb._topic_writer.topic_writer_asyncio import WriterAsyncIO as TopicWriterAsyncIO
 from ._topic_writer.topic_writer_sync import WriterSync as TopicWriter
 
@@ -294,7 +295,15 @@ def tx_writer(
         # If max_worker in the executor is 1 - then encoders will be called from the thread without parallel.
         encoder_executor: Optional[concurrent.futures.Executor] = None,
     ) -> TopicTxWriterAsyncIO:
-        
+        args = locals().copy()
+        del args["self"]
+
+        settings = TopicWriterSettings(**args)
+
+        if not settings.encoder_executor:
+            settings.encoder_executor = self._executor
+
+        return TopicTxWriterAsyncIO(tx=tx, driver=self._driver, settings=settings, _client=self)
 
     def close(self):
         if self._closed: