tx reader and writer based on callbacks

Author: vgvoleg

.github/workflows/tests.yaml

@@ -18,21 +18,15 @@ jobs:
       fail-fast: false
       matrix:
         python-version: [3.8, 3.9]
-        environment: [py-proto5, py-tls-proto5, py-proto4, py-tls-proto4, py-proto3, py-tls-proto3]
-        folder: [ydb, tests --ignore=tests/topics, tests/topics]
+        environment: [py, py-tls, py-proto4, py-tls-proto4, py-proto3, py-tls-proto3]
+        folder: [ydb, tests]
         exclude:
-          - environment: py-tls-proto5
+          - environment: py-tls
             folder: ydb
           - environment: py-tls-proto4
             folder: ydb
           - environment: py-tls-proto3
             folder: ydb
-          - environment: py-tls-proto5
-            folder: tests/topics
-          - environment: py-tls-proto4
-            folder: tests/topics
-          - environment: py-tls-proto3
-            folder: tests/topics
 
     steps:
     - uses: actions/checkout@v1

tests/topics/test_topic_transactions.py

@@ -1,44 +1,69 @@
 import asyncio
 from asyncio import wait_for
 import pytest
+from unittest import mock
 import ydb
 
+DEFAULT_TIMEOUT = 0.1
+DEFAULT_RETRY_SETTINGS = ydb.RetrySettings(max_retries=1)
+
 
-@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):
-        async with driver.topic_client.reader(topic_with_messages, topic_consumer) as reader:
-            async with ydb.aio.QuerySessionPool(driver) as pool:
+        async with ydb.aio.QuerySessionPool(driver) as pool:
+            async with driver.topic_client.reader(topic_with_messages, topic_consumer) as reader:
 
                 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"
+                    batch = await wait_for(reader.receive_batch_with_tx(tx, max_messages=1), DEFAULT_TIMEOUT)
+                    assert len(batch.messages) == 1
+                    assert batch.messages[0].data.decode() == "123"
 
-                await pool.retry_tx_async(callee)
+                await pool.retry_tx_async(callee, retry_settings=DEFAULT_RETRY_SETTINGS)
 
-                msg = await wait_for(reader.receive_message(), 1)
+            async with driver.topic_client.reader(topic_with_messages, topic_consumer) as reader:
+                msg = await wait_for(reader.receive_message(), DEFAULT_TIMEOUT)
                 assert msg.data.decode() == "456"
 
     async def test_rollback(self, driver: ydb.aio.Driver, topic_with_messages, topic_consumer):
         async with driver.topic_client.reader(topic_with_messages, topic_consumer) as reader:
             async with ydb.aio.QuerySessionPool(driver) as pool:
 
                 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"
+                    batch = await wait_for(reader.receive_batch_with_tx(tx, max_messages=1), DEFAULT_TIMEOUT)
+                    assert len(batch.messages) == 1
+                    assert batch.messages[0].data.decode() == "123"
 
                     await tx.rollback()
 
-                await pool.retry_tx_async(callee)
+                await pool.retry_tx_async(callee, retry_settings=DEFAULT_RETRY_SETTINGS)
 
-                msg = await wait_for(reader.receive_message(), 1)
+                msg = await wait_for(reader.receive_message(), DEFAULT_TIMEOUT)
                 assert msg.data.decode() == "123"
 
+    async def test_tx_failed_if_update_offsets_call_failed(
+        self, driver: ydb.aio.Driver, topic_with_messages, topic_consumer
+    ):
+        async with driver.topic_client.reader(topic_with_messages, topic_consumer) as reader:
+            async with ydb.aio.QuerySessionPool(driver) as pool:
+                with mock.patch.object(
+                    reader._reconnector,
+                    "_do_commit_batches_with_tx_call",
+                    side_effect=ydb.Error("Update offsets in tx failed"),
+                ):
+
+                    async def callee(tx: ydb.aio.QueryTxContext):
+                        batch = await wait_for(reader.receive_batch_with_tx(tx, max_messages=1), DEFAULT_TIMEOUT)
+                        assert len(batch.messages) == 1
+                        assert batch.messages[0].data.decode() == "123"
+
+                    with pytest.raises(ydb.Error):
+                        await pool.retry_tx_async(callee, retry_settings=DEFAULT_RETRY_SETTINGS)
+
+                    msg = await wait_for(reader.receive_message(), DEFAULT_TIMEOUT)
+                    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:
@@ -47,7 +72,7 @@ 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, retry_settings=DEFAULT_RETRY_SETTINGS)
 
         msg = await wait_for(topic_reader.receive_message(), 0.1)
         assert msg.data.decode() == "123"
@@ -61,7 +86,92 @@ async def callee(tx: ydb.aio.QueryTxContext):
 
                 await tx.rollback()
 
-            await pool.retry_tx_async(callee)
+            await pool.retry_tx_async(callee, retry_settings=DEFAULT_RETRY_SETTINGS)
+
+        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 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, retry_settings=DEFAULT_RETRY_SETTINGS)
 
         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 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")
+
+            await pool.retry_tx_async(callee, retry_settings=DEFAULT_RETRY_SETTINGS)
+
+            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)
+
+
+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, retry_settings=ydb.RetrySettings(max_retries=1))
+
+        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, retry_settings=DEFAULT_RETRY_SETTINGS)
+
+        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 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, retry_settings=DEFAULT_RETRY_SETTINGS)
+
+        with pytest.raises(TimeoutError):
+            topic_reader_sync.receive_message(timeout=0.1)

tox.ini

@@ -1,5 +1,5 @@
 [tox]
-envlist = py-proto5,py-proto4,py-proto3,py-tls-proto5,py-tls-proto4,py-tls-proto3,style,pylint,black,protoc,py-cov-proto4
+envlist = py,py-proto4,py-proto3,py-tls,py-tls-proto4,py-tls-proto3,style,pylint,black,protoc,py-cov-proto4
 minversion = 4.2.6
 skipsdist = True
 ignore_basepython_conflict = true
@@ -30,7 +30,7 @@ deps =
     -r{toxinidir}/test-requirements.txt
     protobuf<4.0.0
 
-[testenv:py-proto5]
+[testenv:py]
 commands =
     pytest -v -m "not tls" --docker-compose-remove-volumes --docker-compose=docker-compose.yml {posargs}
 deps =
@@ -60,7 +60,7 @@ deps =
     -r{toxinidir}/test-requirements.txt
     protobuf<4.0.0
 
-[testenv:py-tls-proto5]
+[testenv:py-tls]
 commands =
     pytest -v -m tls --docker-compose-remove-volumes --docker-compose=docker-compose-tls.yml {posargs}
 deps =

ydb/_apis.py

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

ydb/_errors.py

@@ -5,6 +5,7 @@
 
 _errors_retriable_fast_backoff_types = [
     issues.Unavailable,
+    issues.ClientInternalError,
 ]
 _errors_retriable_slow_backoff_types = [
     issues.Aborted,

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/datatypes.py

@@ -108,6 +108,9 @@ def ack_notify(self, offset: int):
             waiter = self._ack_waiters.popleft()
             waiter._finish_ok()
 
+    def _update_last_commited_offset_if_needed(self, offset: int):
+        self.committed_offset = max(self.committed_offset, offset)
+
     def close(self):
         if self.closed:
             return
@@ -211,3 +214,9 @@ def _pop_batch(self, message_count: int) -> PublicBatch:
         self._bytes_size = self._bytes_size - new_batch._bytes_size
 
         return new_batch
+
+    def _update_partition_offsets(self, tx, exc=None):
+        if exc is not None:
+            return
+        offsets = self._commit_get_offsets_range()
+        self._partition_session._update_last_commited_offset_if_needed(offsets.end)