Non stream CommitOffset feature

Author: vgvoleg

tests/topics/test_topic_reader.py

@@ -183,6 +183,26 @@ def test_read_and_commit_with_ack(self, driver_sync, topic_with_messages, topic_
 
         assert message != batch.messages[0]
 
+    def test_reader_fine_with_no_stream_commits(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_offsets(
+                    topic_with_messages, topic_consumer, message.partition_id, message.offset_to_commit
+                )
+
+    def test_no_stream_commits_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_offsets(
+                    topic_with_messages, topic_consumer, message.partition_id, message.offset_to_commit
+                )
+
     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,14 @@ 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
+
+    @property
+    def offset_to_commit(self) -> int:
+        return self._commit_end_offset
+
 
 @dataclass
 class PartitionSession:
@@ -184,6 +192,14 @@ def empty(self) -> bool:
     def alive(self) -> bool:
         return not self._partition_session.closed
 
+    @property
+    def partition_id(self) -> int:
+        return self._partition_session.partition_id
+
+    @property
+    def offset_to_commit(self) -> int:
+        return self._commit_get_offsets_range().end
+
     def pop_message(self) -> PublicMessage:
         return self.messages.pop(0)
 

ydb/topic.py

@@ -317,6 +317,21 @@ def tx_writer(
 
         return TopicTxWriterAsyncIO(tx=tx, driver=self._driver, settings=settings, _client=self)
 
+    async def commit_offsets(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
@@ -563,6 +578,21 @@ def tx_writer(
 
         return TopicTxWriter(tx, self._driver, settings, _parent=self)
 
+    def commit_offsets(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