sync writer

Author: vgvoleg

tests/topics/test_topic_transactions.py

@@ -38,7 +38,6 @@ async def callee(tx: ydb.aio.QueryTxContext):
                 assert msg.data.decode() == "123"
 
 
-# @pytest.mark.skip("Not implemented yet.")
 class TestTopicTransactionalWriter:
     async def test_commit(self, driver: ydb.aio.Driver, topic_path, topic_reader: ydb.TopicReaderAsyncIO):
         async with ydb.aio.QuerySessionPool(driver) as pool:
@@ -65,3 +64,65 @@ async def callee(tx: ydb.aio.QueryTxContext):
 
         with pytest.raises(asyncio.TimeoutError):
             await wait_for(topic_reader.receive_message(), 0.1)
+
+    async def test_no_msg_writter_in_error_case(
+        self, driver: ydb.aio.Driver, topic_path, topic_reader: ydb.TopicReaderAsyncIO
+    ):
+        async with ydb.aio.QuerySessionPool(driver) as pool:
+
+            async def callee(tx: ydb.aio.QueryTxContext):
+                tx_writer = driver.topic_client.tx_writer(tx, topic_path)
+                await tx_writer.write(ydb.TopicWriterMessage(data="123".encode()))
+
+                raise BaseException("error")
+
+            with pytest.raises(BaseException):
+                await pool.retry_tx_async(callee)
+
+        with pytest.raises(asyncio.TimeoutError):
+            await wait_for(topic_reader.receive_message(), 0.1)
+
+
+class TestTopicTransactionalWriterSync:
+    def test_commit(self, driver_sync: ydb.Driver, topic_path, topic_reader_sync: ydb.TopicReader):
+        with ydb.QuerySessionPool(driver_sync) as pool:
+
+            def callee(tx: ydb.QueryTxContext):
+                tx_writer = driver_sync.topic_client.tx_writer(tx, topic_path)
+                tx_writer.write(ydb.TopicWriterMessage(data="123".encode()))
+
+            pool.retry_tx_sync(callee)
+
+        msg = topic_reader_sync.receive_message(timeout=0.1)
+        assert msg.data.decode() == "123"
+
+    def test_rollback(self, driver_sync: ydb.aio.Driver, topic_path, topic_reader_sync: ydb.TopicReader):
+        with ydb.QuerySessionPool(driver_sync) as pool:
+
+            def callee(tx: ydb.QueryTxContext):
+                tx_writer = driver_sync.topic_client.tx_writer(tx, topic_path)
+                tx_writer.write(ydb.TopicWriterMessage(data="123".encode()))
+
+                tx.rollback()
+
+            pool.retry_tx_sync(callee)
+
+        with pytest.raises(TimeoutError):
+            topic_reader_sync.receive_message(timeout=0.1)
+
+    def test_no_msg_writter_in_error_case(
+        self, driver_sync: ydb.Driver, topic_path, topic_reader_sync: ydb.TopicReaderAsyncIO
+    ):
+        with ydb.QuerySessionPool(driver_sync) as pool:
+
+            def callee(tx: ydb.QueryTxContext):
+                tx_writer = driver_sync.topic_client.tx_writer(tx, topic_path)
+                tx_writer.write(ydb.TopicWriterMessage(data="123".encode()))
+
+                raise BaseException("error")
+
+            with pytest.raises(BaseException):
+                pool.retry_tx_sync(callee)
+
+        with pytest.raises(TimeoutError):
+            topic_reader_sync.receive_message(timeout=0.1)

ydb/_topic_writer/topic_writer_sync.py

@@ -14,13 +14,21 @@
     TopicWriterClosedError,
 )
 
-from .topic_writer_asyncio import WriterAsyncIO
+from ..query.base import TxListener
+
+from .topic_writer_asyncio import (
+    TxWriterAsyncIO,
+    WriterAsyncIO,
+)
 from .._topic_common.common import (
     _get_shared_event_loop,
     TimeoutType,
     CallFromSyncToAsync,
 )
 
+if typing.TYPE_CHECKING:
+    from ..query.transaction import BaseQueryTxContext
+
 
 class WriterSync:
     _caller: CallFromSyncToAsync
@@ -122,3 +130,38 @@ def write_with_ack(
         self._check_closed()
 
         return self._caller.unsafe_call_with_result(self._async_writer.write_with_ack(messages), timeout=timeout)
+
+
+class TxWriterSync(WriterSync, TxListener):
+    def __init__(
+        self,
+        tx: "BaseQueryTxContext",
+        driver: SupportedDriverType,
+        settings: PublicWriterSettings,
+        *,
+        eventloop: Optional[asyncio.AbstractEventLoop] = None,
+        _parent=None,
+    ):
+
+        self._closed = False
+
+        if eventloop:
+            loop = eventloop
+        else:
+            loop = _get_shared_event_loop()
+
+        self._caller = CallFromSyncToAsync(loop)
+
+        async def create_async_writer():
+            return TxWriterAsyncIO(tx, driver, settings)
+
+        self._async_writer = self._caller.safe_call_with_result(create_async_writer(), None)
+        self._parent = _parent
+
+        tx._add_listener(self)
+
+    def _on_before_commit(self):
+        self.close()
+
+    def _on_before_rollback(self):
+        self.close()

ydb/topic.py

@@ -68,6 +68,7 @@
 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
+from ._topic_writer.topic_writer_sync import TxWriterSync as TxTopicWriter
 
 from ._topic_common.common import (
     wrap_operation as _wrap_operation,
@@ -517,6 +518,36 @@ def writer(
 
         return TopicWriter(self._driver, settings, _parent=self)
 
+    def tx_writer(
+        self,
+        tx,
+        topic,
+        *,
+        producer_id: Optional[str] = None,  # default - random
+        session_metadata: Mapping[str, str] = None,
+        partition_id: Union[int, None] = None,
+        auto_seqno: bool = True,
+        auto_created_at: bool = True,
+        codec: Optional[TopicCodec] = None,  # default mean auto-select
+        # encoders: map[codec_code] func(encoded_bytes)->decoded_bytes
+        # the func will be called from multiply threads in parallel.
+        encoders: Optional[Mapping[_ydb_topic_public_types.PublicCodec, Callable[[bytes], bytes]]] = None,
+        # custom encoder executor for call builtin and custom decoders. If None - use shared executor pool.
+        # 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,  # default shared client executor pool
+    ) -> TopicWriter:
+        args = locals().copy()
+        del args["self"]
+        del args["tx"]
+        self._check_closed()
+
+        settings = TopicWriterSettings(**args)
+
+        if not settings.encoder_executor:
+            settings.encoder_executor = self._executor
+
+        return TxTopicWriter(tx, self._driver, settings, _parent=self)
+
     def close(self):
         if self._closed:
             return