Non stream CommitOffset feature

Author: vgvoleg

tests/topics/test_topic_reader.py

@@ -183,6 +183,32 @@ 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,
+                    0,
+                    message._commit_end_offset
+                )
+
+    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,
+                    0,
+                    message._commit_end_offset
+                )
+
     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.py

@@ -317,6 +317,23 @@ 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 +580,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