CommitOffset with readsessionid

Author: vgvoleg

tests/topics/test_topic_reader.py

@@ -74,6 +74,28 @@ async def test_commit_offset_works(self, driver, topic_with_messages, topic_cons
                     topic_with_messages, topic_consumer, message.partition_id, message.offset + 1
                 )
 
+    async def test_commit_offset_with_session_id_works(self, driver, topic_with_messages, topic_consumer):
+        async with driver.topic_client.reader(topic_with_messages, topic_consumer) as reader:
+            msg1 = await reader.receive_message()
+            assert msg1.seqno == 1
+            msg2 = await reader.receive_message()
+            assert msg2.seqno == 2
+
+            await driver.topic_client.commit_offset(
+                topic_with_messages,
+                topic_consumer,
+                msg1.partition_id,
+                msg1.offset + 1,
+                reader.read_session_id,
+            )
+
+            msg3 = await reader.receive_message()
+            assert msg3.seqno == 3
+
+        async with driver.topic_client.reader(topic_with_messages, topic_consumer) as reader:
+            msg2 = await reader.receive_message()
+            assert msg2.seqno == 2
+
     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"]:
@@ -213,6 +235,28 @@ def test_commit_offset_works(self, driver_sync, topic_with_messages, topic_consu
                     topic_with_messages, topic_consumer, message.partition_id, message.offset + 1
                 )
 
+    def test_commit_offset_with_session_id_works(self, driver_sync, topic_with_messages, topic_consumer):
+        with driver_sync.topic_client.reader(topic_with_messages, topic_consumer) as reader:
+            msg1 = reader.receive_message()
+            assert msg1.seqno == 1
+            msg2 = reader.receive_message()
+            assert msg2.seqno == 2
+
+            driver_sync.topic_client.commit_offset(
+                topic_with_messages,
+                topic_consumer,
+                msg1.partition_id,
+                msg1.offset + 1,
+                reader.read_session_id,
+            )
+
+            msg3 = reader.receive_message()
+            assert msg3.seqno == 3
+
+        with driver_sync.topic_client.reader(topic_with_messages, topic_consumer) as reader:
+            msg2 = reader.receive_message()
+            assert msg2.seqno == 2
+
     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"]:

ydb/_grpc/grpcwrapper/ydb_topic.py

@@ -143,13 +143,15 @@ class CommitOffsetRequest(IToProto):
     consumer: str
     partition_id: int
     offset: int
+    read_session_id: Optional[str]
 
     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,
+            read_session_id=self.read_session_id,
         )
 
 

ydb/_topic_reader/topic_reader_asyncio.py

@@ -190,6 +190,10 @@ async def close(self, flush: bool = True):
         self._closed = True
         await self._reconnector.close(flush)
 
+    @property
+    def read_session_id(self) -> Optional[str]:
+        return self._reconnector.read_session_id
+
 
 class ReaderReconnector:
     _static_reader_reconnector_counter = AtomicCounter()
@@ -373,6 +377,12 @@ def _set_first_error(self, err: issues.Error):
             # skip if already has result
             pass
 
+    @property
+    def read_session_id(self) -> Optional[str]:
+        if not self._stream_reader:
+            return None
+        return self._stream_reader._session_id
+
 
 class ReaderStream:
     _static_id_counter = AtomicCounter()

ydb/_topic_reader/topic_reader_sync.py

@@ -204,3 +204,7 @@ def close(self, *, flush: bool = True, timeout: TimeoutType = None):
     def _check_closed(self):
         if self._closed:
             raise TopicReaderClosedError()
+
+    @property
+    def read_session_id(self) -> Optional[str]:
+        return self._async_reader.read_session_id

ydb/topic.py

@@ -340,12 +340,15 @@ 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:
+    async def commit_offset(
+        self, path: str, consumer: str, partition_id: int, offset: int, read_session_id: Optional[str] = None
+    ) -> None:
         req = _ydb_topic.CommitOffsetRequest(
             path=path,
             consumer=consumer,
             partition_id=partition_id,
             offset=offset,
+            read_session_id=read_session_id,
         )
 
         await self._driver(
@@ -618,12 +621,15 @@ 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:
+    def commit_offset(
+        self, path: str, consumer: str, partition_id: int, offset: int, read_session_id: Optional[str] = None
+    ) -> None:
         req = _ydb_topic.CommitOffsetRequest(
             path=path,
             consumer=consumer,
             partition_id=partition_id,
             offset=offset,
+            read_session_id=read_session_id,
         )
 
         self._driver(