CommitOffset feature

Author: vgvoleg

tests/topics/test_topic_reader.py

@@ -64,6 +64,26 @@ async def test_read_and_commit_with_ack(self, driver, topic_with_messages, topic
 
         assert message != batch.messages[0]
 
+    async def test_commit_offset_works(self, driver, topic_with_messages, topic_consumer):
+        for out in ["123", "456", "789", "0"]:
+            async with driver.topic_client.reader(topic_with_messages, topic_consumer) as reader:
+                message = await reader.receive_message()
+                assert message.data.decode() == out
+
+                await driver.topic_client.commit_offset(
+                    topic_with_messages, topic_consumer, message.partition_id, message.offset + 1
+                )
+
+    async def test_reader_reconnect_after_commit_offset(self, driver, topic_with_messages, topic_consumer):
+        async with driver.topic_client.reader(topic_with_messages, topic_consumer) as reader:
+            for out in ["123", "456", "789", "0"]:
+                message = await reader.receive_message()
+                assert message.data.decode() == out
+
+                await driver.topic_client.commit_offset(
+                    topic_with_messages, topic_consumer, message.partition_id, message.offset + 1
+                )
+
     async def test_read_compressed_messages(self, driver, topic_path, topic_consumer):
         async with driver.topic_client.writer(topic_path, codec=ydb.TopicCodec.GZIP) as writer:
             await writer.write("123")
@@ -183,6 +203,26 @@ def test_read_and_commit_with_ack(self, driver_sync, topic_with_messages, topic_
 
         assert message != batch.messages[0]
 
+    def test_commit_offset_works(self, driver_sync, topic_with_messages, topic_consumer):
+        for out in ["123", "456", "789", "0"]:
+            with driver_sync.topic_client.reader(topic_with_messages, topic_consumer) as reader:
+                message = reader.receive_message()
+                assert message.data.decode() == out
+
+                driver_sync.topic_client.commit_offset(
+                    topic_with_messages, topic_consumer, message.partition_id, message.offset + 1
+                )
+
+    def test_reader_reconnect_after_commit_offset(self, driver_sync, topic_with_messages, topic_consumer):
+        with driver_sync.topic_client.reader(topic_with_messages, topic_consumer) as reader:
+            for out in ["123", "456", "789", "0"]:
+                message = reader.receive_message()
+                assert message.data.decode() == out
+
+                driver_sync.topic_client.commit_offset(
+                    topic_with_messages, topic_consumer, message.partition_id, message.offset + 1
+                )
+
     def test_read_compressed_messages(self, driver_sync, topic_path, topic_consumer):
         with driver_sync.topic_client.writer(topic_path, codec=ydb.TopicCodec.GZIP) as writer:
             writer.write("123")

ydb/_apis.py

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

ydb/_grpc/grpcwrapper/ydb_topic.py

@@ -137,6 +137,22 @@ def from_proto(msg: ydb_topic_pb2.UpdateTokenResponse) -> typing.Any:
         return UpdateTokenResponse()
 
 
+@dataclass
+class CommitOffsetRequest(IToProto):
+    path: str
+    consumer: str
+    partition_id: int
+    offset: int
+
+    def to_proto(self) -> ydb_topic_pb2.CommitOffsetRequest:
+        return ydb_topic_pb2.CommitOffsetRequest(
+            path=self.path,
+            consumer=self.consumer,
+            partition_id=self.partition_id,
+            offset=self.offset,
+        )
+
+
 ########################################################################################################################
 #  StreamWrite
 ########################################################################################################################

ydb/_topic_reader/datatypes.py

@@ -56,6 +56,10 @@ def _commit_get_offsets_range(self) -> OffsetsRange:
     def alive(self) -> bool:
         return not self._partition_session.closed
 
+    @property
+    def partition_id(self) -> int:
+        return self._partition_session.partition_id
+
 
 @dataclass
 class PartitionSession:

ydb/topic.py

@@ -340,6 +340,21 @@ def tx_writer(
 
         return TopicTxWriterAsyncIO(tx=tx, driver=self._driver, settings=settings, _client=self)
 
+    async def commit_offset(self, path: str, consumer: str, partition_id: int, offset: int) -> None:
+        req = _ydb_topic.CommitOffsetRequest(
+            path=path,
+            consumer=consumer,
+            partition_id=partition_id,
+            offset=offset,
+        )
+
+        await self._driver(
+            req.to_proto(),
+            _apis.TopicService.Stub,
+            _apis.TopicService.CommitOffset,
+            _wrap_operation,
+        )
+
     def close(self):
         if self._closed:
             return
@@ -603,6 +618,21 @@ def tx_writer(
 
         return TopicTxWriter(tx, self._driver, settings, _parent=self)
 
+    def commit_offset(self, path: str, consumer: str, partition_id: int, offset: int) -> None:
+        req = _ydb_topic.CommitOffsetRequest(
+            path=path,
+            consumer=consumer,
+            partition_id=partition_id,
+            offset=offset,
+        )
+
+        self._driver(
+            req.to_proto(),
+            _apis.TopicService.Stub,
+            _apis.TopicService.CommitOffset,
+            _wrap_operation,
+        )
+
     def close(self):
         if self._closed:
             return