transactional reader draft

Author: vgvoleg

tests/topics/test_topic_transactions.py

@@ -4,7 +4,7 @@
 import ydb
 
 
-@pytest.mark.skip("Not implemented yet.")
+# @pytest.mark.skip("Not implemented yet.")
 @pytest.mark.asyncio
 class TestTopicTransactionalReader:
     async def test_commit(self, driver: ydb.aio.Driver, topic_with_messages, topic_consumer):
@@ -13,8 +13,8 @@ async def test_commit(self, driver: ydb.aio.Driver, topic_with_messages, topic_c
 
                 async def callee(tx: ydb.aio.QueryTxContext):
                     batch = await wait_for(reader.receive_batch_with_tx(tx, max_messages=1), 1)
-                    assert len(batch) == 1
-                    assert batch[0].data.decode() == "123"
+                    assert len(batch.messages) == 1
+                    assert batch.messages[0].data.decode() == "123"
 
                 await pool.retry_tx_async(callee)
 
@@ -27,8 +27,8 @@ async def test_rollback(self, driver: ydb.aio.Driver, topic_with_messages, topic
 
                 async def callee(tx: ydb.aio.QueryTxContext):
                     batch = await wait_for(reader.receive_batch_with_tx(tx, max_messages=1), 1)
-                    assert len(batch) == 1
-                    assert batch[0].data.decode() == "123"
+                    assert len(batch.messages) == 1
+                    assert batch.messages[0].data.decode() == "123"
 
                     await tx.rollback()
 
@@ -38,114 +38,114 @@ async def callee(tx: ydb.aio.QueryTxContext):
                 assert msg.data.decode() == "123"
 
 
-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:
+# 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:
 
-            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()))
+#             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()))
 
-            await pool.retry_tx_async(callee)
+#             await pool.retry_tx_async(callee)
 
-        msg = await wait_for(topic_reader.receive_message(), 0.1)
-        assert msg.data.decode() == "123"
+#         msg = await wait_for(topic_reader.receive_message(), 0.1)
+#         assert msg.data.decode() == "123"
 
-    async def test_rollback(self, driver: ydb.aio.Driver, topic_path, topic_reader: ydb.TopicReaderAsyncIO):
-        async with ydb.aio.QuerySessionPool(driver) as pool:
+#     async def test_rollback(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()))
+#             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()))
 
-                await tx.rollback()
+#                 await tx.rollback()
 
-            await pool.retry_tx_async(callee)
+#             await pool.retry_tx_async(callee)
 
-        with pytest.raises(asyncio.TimeoutError):
-            await wait_for(topic_reader.receive_message(), 0.1)
+#         with pytest.raises(asyncio.TimeoutError):
+#             await wait_for(topic_reader.receive_message(), 0.1)
 
-    async def test_no_msg_written_in_error_case(
-        self, driver: ydb.aio.Driver, topic_path, topic_reader: ydb.TopicReaderAsyncIO
-    ):
-        async with ydb.aio.QuerySessionPool(driver) as pool:
+#     async def test_no_msg_written_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()))
+#             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")
+#                 raise BaseException("error")
 
-            with pytest.raises(BaseException):
-                await pool.retry_tx_async(callee)
+#             with pytest.raises(BaseException):
+#                 await pool.retry_tx_async(callee)
 
-        with pytest.raises(asyncio.TimeoutError):
-            await wait_for(topic_reader.receive_message(), 0.1)
+#         with pytest.raises(asyncio.TimeoutError):
+#             await wait_for(topic_reader.receive_message(), 0.1)
 
-    async def test_msg_written_exactly_once_with_retries(
-        self, driver: ydb.aio.Driver, topic_path, topic_reader: ydb.TopicReaderAsyncIO
-    ):
-        error_raised = False
-        async with ydb.aio.QuerySessionPool(driver) as pool:
+#     async def test_msg_written_exactly_once_with_retries(
+#         self, driver: ydb.aio.Driver, topic_path, topic_reader: ydb.TopicReaderAsyncIO
+#     ):
+#         error_raised = False
+#         async with ydb.aio.QuerySessionPool(driver) as pool:
 
-            async def callee(tx: ydb.aio.QueryTxContext):
-                nonlocal error_raised
-                tx_writer = driver.topic_client.tx_writer(tx, topic_path)
-                await tx_writer.write(ydb.TopicWriterMessage(data="123".encode()))
+#             async def callee(tx: ydb.aio.QueryTxContext):
+#                 nonlocal error_raised
+#                 tx_writer = driver.topic_client.tx_writer(tx, topic_path)
+#                 await tx_writer.write(ydb.TopicWriterMessage(data="123".encode()))
 
-                if not error_raised:
-                    error_raised = True
-                    raise ydb.issues.Unavailable("some retriable error")
+#                 if not error_raised:
+#                     error_raised = True
+#                     raise ydb.issues.Unavailable("some retriable error")
 
-            await pool.retry_tx_async(callee)
+#             await pool.retry_tx_async(callee)
 
-            msg = await wait_for(topic_reader.receive_message(), 0.1)
-            assert msg.data.decode() == "123"
+#             msg = await wait_for(topic_reader.receive_message(), 0.1)
+#             assert msg.data.decode() == "123"
 
-            with pytest.raises(asyncio.TimeoutError):
-                await wait_for(topic_reader.receive_message(), 0.1)
+#             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:
+# 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()))
+#             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)
+#             pool.retry_tx_sync(callee)
 
-        msg = topic_reader_sync.receive_message(timeout=0.1)
-        assert msg.data.decode() == "123"
+#         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 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()))
+#             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()
+#                 tx.rollback()
 
-            pool.retry_tx_sync(callee)
+#             pool.retry_tx_sync(callee)
 
-        with pytest.raises(TimeoutError):
-            topic_reader_sync.receive_message(timeout=0.1)
+#         with pytest.raises(TimeoutError):
+#             topic_reader_sync.receive_message(timeout=0.1)
 
-    def test_no_msg_written_in_error_case(
-        self, driver_sync: ydb.Driver, topic_path, topic_reader_sync: ydb.TopicReaderAsyncIO
-    ):
-        with ydb.QuerySessionPool(driver_sync) as pool:
+#     def test_no_msg_written_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()))
+#             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")
+#                 raise BaseException("error")
 
-            with pytest.raises(BaseException):
-                pool.retry_tx_sync(callee)
+#             with pytest.raises(BaseException):
+#                 pool.retry_tx_sync(callee)
 
-        with pytest.raises(TimeoutError):
-            topic_reader_sync.receive_message(timeout=0.1)
+#         with pytest.raises(TimeoutError):
+#             topic_reader_sync.receive_message(timeout=0.1)

ydb/_apis.py

@@ -116,6 +116,7 @@ class TopicService(object):
     DropTopic = "DropTopic"
     StreamRead = "StreamRead"
     StreamWrite = "StreamWrite"
+    UpdateOffsetsInTransaction = "UpdateOffsetsInTransaction"
 
 
 class QueryService(object):

ydb/_grpc/grpcwrapper/ydb_topic.py

@@ -1209,6 +1209,52 @@ def to_public(self) -> ydb_topic_public_types.PublicMeteringMode:
             return ydb_topic_public_types.PublicMeteringMode.UNSPECIFIED
 
 
+@dataclass
+class UpdateOffsetsInTransactionRequest(IToProto):
+    tx: TransactionIdentity
+    topics: List[UpdateOffsetsInTransactionRequest.TopicOffsets]
+    consumer: str
+
+    def to_proto(self):
+        return ydb_topic_pb2.UpdateOffsetsInTransactionRequest(
+            tx=self.tx.to_proto(),
+            consumer=self.consumer,
+            topics=list(
+                map(
+                    UpdateOffsetsInTransactionRequest.TopicOffsets.to_proto,
+                    self.topics,
+                )
+            ),
+        )
+
+    @dataclass
+    class TopicOffsets(IToProto):
+        path: str
+        partitions: List[UpdateOffsetsInTransactionRequest.TopicOffsets.PartitionOffsets]
+
+        def to_proto(self):
+            return ydb_topic_pb2.UpdateOffsetsInTransactionRequest.TopicOffsets(
+                path=self.path,
+                partitions=list(
+                    map(
+                        UpdateOffsetsInTransactionRequest.TopicOffsets.PartitionOffsets.to_proto,
+                        self.partitions,
+                    )
+                )
+            )
+
+        @dataclass
+        class PartitionOffsets(IToProto):
+            partition_id: int
+            partition_offsets: List[OffsetsRange]
+
+            def to_proto(self) -> ydb_topic_pb2.UpdateOffsetsInTransactionRequest.TopicOffsets.PartitionOffsets:
+                return ydb_topic_pb2.UpdateOffsetsInTransactionRequest.TopicOffsets.PartitionOffsets(
+                    partition_id=self.partition_id,
+                    partition_offsets=list(map(OffsetsRange.to_proto, self.partition_offsets)),
+                )
+
+
 @dataclass
 class CreateTopicRequest(IToProto, IFromPublic):
     path: str

ydb/_topic_reader/topic_reader_asyncio.py

@@ -25,11 +25,17 @@
     StreamReadMessage,
     UpdateTokenRequest,
     UpdateTokenResponse,
+    UpdateOffsetsInTransactionRequest,
     Codec,
 )
 from .._errors import check_retriable_error
 import logging
 
+from ..query.base import TxListenerAsyncIO
+
+if typing.TYPE_CHECKING:
+    from ..query.transaction import BaseQueryTxContext
+
 logger = logging.getLogger(__name__)
 
 
@@ -121,6 +127,23 @@ async def receive_message(self) -> typing.Optional[datatypes.PublicMessage]:
         await self._reconnector.wait_message()
         return self._reconnector.receive_message_nowait()
 
+    async def receive_batch_with_tx(
+        self,
+        tx: "BaseQueryTxContext",
+        max_messages: typing.Union[int, None] = None,
+    ) -> typing.Union[datatypes.PublicBatch, None]:
+        """
+        Get one messages batch from reader.
+        All messages in a batch from same partition.
+
+        use asyncio.wait_for for wait with timeout.
+        """
+        await self._reconnector.wait_message()
+        return await self._reconnector.receive_batch_with_tx_nowait(
+            tx,
+            max_messages=max_messages,
+        )
+
     def commit(self, batch: typing.Union[datatypes.PublicMessage, datatypes.PublicBatch]):
         """
         Write commit message to a buffer.
@@ -165,6 +188,7 @@ class ReaderReconnector:
     _state_changed: asyncio.Event
     _stream_reader: Optional["ReaderStream"]
     _first_error: asyncio.Future[YdbError]
+    _batches_to_commit: asyncio.Queue
 
     def __init__(self, driver: Driver, settings: topic_reader.PublicReaderSettings):
         self._id = self._static_reader_reconnector_counter.inc_and_get()
@@ -177,6 +201,8 @@ def __init__(self, driver: Driver, settings: topic_reader.PublicReaderSettings):
         self._background_tasks.add(asyncio.create_task(self._connection_loop()))
         self._first_error = asyncio.get_running_loop().create_future()
 
+        self._batches_to_commit = asyncio.Queue()
+
     async def _connection_loop(self):
         attempt = 0
         while True:
@@ -228,6 +254,51 @@ def receive_message_nowait(self):
     def commit(self, batch: datatypes.ICommittable) -> datatypes.PartitionSession.CommitAckWaiter:
         return self._stream_reader.commit(batch)
 
+    async def _commit_with_tx(self, tx: "BaseQueryTxContext", batch: datatypes.ICommittable) -> None:
+        pass
+
+    async def receive_batch_with_tx_nowait(self, tx: "BaseQueryTxContext", max_messages: Optional[int] = None):
+        batch = self.receive_batch_nowait(max_messages=max_messages)
+        tx._add_listener(batch)
+        await self._update_offsets_in_tx_call(self._driver, tx, batch)
+        return batch
+        # self._batches_to_commit.put_nowait((tx, batch))
+
+    async def _update_offsets_in_tx_loop(self):
+        while True:
+            tx, batch = self._batches_to_commit.get()
+            await self._update_offsets_in_tx_call(self._driver, tx, batch)
+
+    async def _update_offsets_in_tx_call(self, driver: SupportedDriverType, tx: "BaseQueryTxContext", batch: datatypes.ICommittable) -> None:
+        partition_session = batch._commit_get_partition_session()
+        request = UpdateOffsetsInTransactionRequest(
+            tx=tx._tx_identity(),
+            consumer=self._settings.consumer,
+            topics=[
+                UpdateOffsetsInTransactionRequest.TopicOffsets(
+                    path=partition_session.topic_path,
+                    partitions=[
+                        UpdateOffsetsInTransactionRequest.TopicOffsets.PartitionOffsets(
+                            partition_id=partition_session.partition_id,
+                            partition_offsets=[batch._commit_get_offsets_range()]
+                        )
+                    ],
+                )
+            ],
+        ).to_proto()
+
+        res = driver(
+            request,
+            _apis.TopicService.Stub,
+            _apis.TopicService.UpdateOffsetsInTransaction,
+            topic_common.wrap_operation,
+        )
+
+        if asyncio.iscoroutinefunction(driver.__call__):
+            res = await res
+
+        return res
+
     async def close(self, flush: bool):
         if self._stream_reader:
             await self._stream_reader.close(flush)