No Consumer Reader

Author: vgvoleg

tests/topics/test_topic_reader.py

@@ -251,3 +251,127 @@ async def wait(fut):
 
         await reader0.close()
         await reader1.close()
+
+
+@pytest.mark.asyncio
+class TestTopicNoConsumerReaderAsyncIO:
+    async def test_reader_with_sync_lambda(self, driver, topic_with_messages):
+        def sync_lambda(partition_id: int):
+            assert partition_id == 0
+            return 1
+
+        reader = driver.topic_client.no_consumer_reader(topic_with_messages, [0], sync_lambda)
+        msg = await reader.receive_message()
+
+        assert msg.seqno == 2
+
+        await reader.close()
+
+    async def test_reader_with_async_lambda(self, driver, topic_with_messages):
+        async def async_lambda(partition_id: int) -> int:
+            assert partition_id == 0
+            return 1
+
+        reader = driver.topic_client.no_consumer_reader(topic_with_messages, [0], async_lambda)
+        msg = await reader.receive_message()
+
+        assert msg.seqno == 2
+
+        await reader.close()
+
+    async def test_commit_not_allowed(self, driver, topic_with_messages):
+        reader = driver.topic_client.no_consumer_reader(topic_with_messages, [0], lambda x: None)
+        batch = await reader.receive_batch()
+
+        with pytest.raises(ydb.Error):
+            reader.commit(batch)
+
+        with pytest.raises(ydb.Error):
+            await reader.commit_with_ack(batch)
+
+        await reader.close()
+
+    async def test_offsets_updated_after_reconnect(self, driver, topic_with_messages):
+        current_offset = 0
+
+        def get_start_offset_lambda(partition_id: int) -> int:
+            nonlocal current_offset
+            return current_offset
+
+        reader = driver.topic_client.no_consumer_reader(topic_with_messages, [0], get_start_offset_lambda)
+        msg = await reader.receive_message()
+
+        assert msg.seqno == current_offset + 1
+
+        current_offset += 2
+        reader._reconnector._stream_reader._set_first_error(ydb.Unavailable("some retriable error"))
+
+        await asyncio.sleep(0)
+
+        msg = await reader.receive_message()
+
+        assert msg.seqno == current_offset + 1
+
+        await reader.close()
+
+
+class TestTopicNoConsumerReader:
+    def test_reader_with_sync_lambda(self, driver_sync, topic_with_messages):
+        def sync_lambda(partition_id: int):
+            assert partition_id == 0
+            return 1
+
+        reader = driver_sync.topic_client.no_consumer_reader(topic_with_messages, [0], sync_lambda)
+        msg = reader.receive_message()
+
+        assert msg.seqno == 2
+
+        reader.close()
+
+    def test_reader_with_async_lambda(self, driver_sync, topic_with_messages):
+        async def async_lambda(partition_id: int) -> int:
+            assert partition_id == 0
+            return 1
+
+        reader = driver_sync.topic_client.no_consumer_reader(topic_with_messages, [0], async_lambda)
+        msg = reader.receive_message()
+
+        assert msg.seqno == 2
+
+        reader.close()
+
+    def test_commit_not_allowed(self, driver_sync, topic_with_messages):
+        reader = driver_sync.topic_client.no_consumer_reader(topic_with_messages, [0], lambda x: None)
+        batch = reader.receive_batch()
+
+        with pytest.raises(ydb.Error):
+            reader.commit(batch)
+
+        with pytest.raises(ydb.Error):
+            reader.commit_with_ack(batch)
+
+        with pytest.raises(ydb.Error):
+            reader.async_commit_with_ack(batch)
+
+        reader.close()
+
+    def test_offsets_updated_after_reconnect(self, driver_sync, topic_with_messages):
+        current_offset = 0
+
+        def get_start_offset_lambda(partition_id: int) -> int:
+            nonlocal current_offset
+            return current_offset
+
+        reader = driver_sync.topic_client.no_consumer_reader(topic_with_messages, [0], get_start_offset_lambda)
+        msg = reader.receive_message()
+
+        assert msg.seqno == current_offset + 1
+
+        current_offset += 2
+        reader._async_reader._reconnector._stream_reader._set_first_error(ydb.Unavailable("some retriable error"))
+
+        msg = reader.receive_message()
+
+        assert msg.seqno == current_offset + 1
+
+        reader.close()

ydb/_grpc/grpcwrapper/ydb_topic.py

@@ -439,12 +439,13 @@ def from_proto(
     @dataclass
     class InitRequest(IToProto):
         topics_read_settings: List["StreamReadMessage.InitRequest.TopicReadSettings"]
-        consumer: str
+        consumer: Optional[str]
         auto_partitioning_support: bool
 
         def to_proto(self) -> ydb_topic_pb2.StreamReadMessage.InitRequest:
             res = ydb_topic_pb2.StreamReadMessage.InitRequest()
-            res.consumer = self.consumer
+            if self.consumer is not None:
+                res.consumer = self.consumer
             for settings in self.topics_read_settings:
                 res.topics_read_settings.append(settings.to_proto())
             res.auto_partitioning_support = self.auto_partitioning_support

ydb/_topic_reader/topic_reader.py

@@ -3,6 +3,7 @@
 import datetime
 from dataclasses import dataclass
 from typing import (
+    Awaitable,
     Union,
     Optional,
     List,
@@ -42,7 +43,7 @@ def _to_topic_read_settings(self) -> StreamReadMessage.InitRequest.TopicReadSett
 
 @dataclass
 class PublicReaderSettings:
-    consumer: str
+    consumer: Optional[str]
     topic: TopicSelectorTypes
     buffer_size_bytes: int = 50 * 1024 * 1024
     auto_partitioning_support: bool = True
@@ -54,12 +55,15 @@ class PublicReaderSettings:
     decoder_executor: Optional[concurrent.futures.Executor] = None
     update_token_interval: Union[int, float] = 3600
 
+    partition_ids: Optional[List[int]] = None
+    get_start_offset_lambda: Optional[Union[Callable[[int], int], Callable[[int], Awaitable[int]]]] = None
+
     def __post_init__(self):
         # check possible create init message
         _ = self._init_message()
 
     def _init_message(self) -> StreamReadMessage.InitRequest:
-        if not isinstance(self.consumer, str):
+        if self.consumer is not None and not isinstance(self.consumer, str):
             raise TypeError("Unsupported type for customer field: '%s'" % type(self.consumer))
 
         if isinstance(self.topic, list):
@@ -69,7 +73,7 @@ def _init_message(self) -> StreamReadMessage.InitRequest:
 
         for index, selector in enumerate(selectors):
             if isinstance(selector, str):
-                selectors[index] = PublicTopicSelector(path=selector)
+                selectors[index] = PublicTopicSelector(path=selector, partitions=self.partition_ids)
             elif isinstance(selector, PublicTopicSelector):
                 pass
             else:

ydb/_topic_reader/topic_reader_asyncio.py

@@ -182,6 +182,14 @@ async def close(self, flush: bool = True):
         await self._reconnector.close(flush)
 
 
+class PublicAsyncIONoConsumerReader(PublicAsyncIOReader):
+    def commit(self, batch: typing.Union[datatypes.PublicMessage, datatypes.PublicBatch]):
+        raise issues.Error("Commit operations are not supported for topic reader without consumer.")
+
+    async def commit_with_ack(self, batch: typing.Union[datatypes.PublicMessage, datatypes.PublicBatch]):
+        raise issues.Error("Commit operations are not supported for topic reader without consumer.")
+
+
 class ReaderReconnector:
     _static_reader_reconnector_counter = AtomicCounter()
 
@@ -393,6 +401,7 @@ class ReaderStream:
     _update_token_interval: Union[int, float]
     _update_token_event: asyncio.Event
     _get_token_function: Callable[[], str]
+    _settings: topic_reader.PublicReaderSettings
 
     def __init__(
         self,
@@ -425,6 +434,8 @@ def __init__(
         self._get_token_function = get_token_function
         self._update_token_event = asyncio.Event()
 
+        self._settings = settings
+
     @staticmethod
     async def create(
         reader_reconnector_id: int,
@@ -615,7 +626,7 @@ async def _read_messages_loop(self):
                         message.server_message,
                         StreamReadMessage.StartPartitionSessionRequest,
                     ):
-                        self._on_start_partition_session(message.server_message)
+                        await self._on_start_partition_session(message.server_message)
 
                     elif isinstance(
                         message.server_message,
@@ -660,7 +671,7 @@ async def _update_token(self, token: str):
         finally:
             self._update_token_event.clear()
 
-    def _on_start_partition_session(self, message: StreamReadMessage.StartPartitionSessionRequest):
+    async def _on_start_partition_session(self, message: StreamReadMessage.StartPartitionSessionRequest):
         try:
             if message.partition_session.partition_session_id in self._partition_sessions:
                 raise TopicReaderError(
@@ -676,11 +687,19 @@ def _on_start_partition_session(self, message: StreamReadMessage.StartPartitionS
                 reader_reconnector_id=self._reader_reconnector_id,
                 reader_stream_id=self._id,
             )
+
+            read_offset = None
+            callee = self._settings.get_start_offset_lambda
+            if callee is not None:
+                read_offset = callee(message.partition_session.partition_id)
+                if asyncio.iscoroutinefunction(callee):
+                    read_offset = await read_offset
+
             self._stream.write(
                 StreamReadMessage.FromClient(
                     client_message=StreamReadMessage.StartPartitionSessionResponse(
                         partition_session_id=message.partition_session.partition_session_id,
-                        read_offset=None,
+                        read_offset=read_offset,
                         commit_offset=None,
                     )
                 ),

ydb/_topic_reader/topic_reader_sync.py

@@ -4,6 +4,7 @@
 import typing
 from typing import List, Union, Optional
 
+from ydb import issues
 from ydb._grpc.grpcwrapper.common_utils import SupportedDriverType
 from ydb._topic_common.common import (
     _get_shared_event_loop,
@@ -191,3 +192,16 @@ def close(self, *, flush: bool = True, timeout: TimeoutType = None):
     def _check_closed(self):
         if self._closed:
             raise TopicReaderClosedError()
+
+
+class TopicNoConsumerReaderSync(TopicReaderSync):
+    def commit(self, batch: typing.Union[datatypes.PublicMessage, datatypes.PublicBatch]):
+        raise issues.Error("Commit operations are not supported for topic reader without consumer.")
+
+    def commit_with_ack(self, batch: typing.Union[datatypes.PublicMessage, datatypes.PublicBatch]):
+        raise issues.Error("Commit operations are not supported for topic reader without consumer.")
+
+    def async_commit_with_ack(
+        self, mess: typing.Union[datatypes.PublicMessage, datatypes.PublicBatch]
+    ) -> concurrent.futures.Future:
+        raise issues.Error("Commit operations are not supported for topic reader without consumer.")

ydb/topic.py

@@ -14,6 +14,7 @@
     "TopicError",
     "TopicMeteringMode",
     "TopicReader",
+    "TopicNoConsumerReaderAsyncIO",
     "TopicReaderAsyncIO",
     "TopicReaderBatch",
     "TopicReaderMessage",
@@ -36,7 +37,7 @@
 import datetime
 from dataclasses import dataclass
 import logging
-from typing import List, Union, Mapping, Optional, Dict, Callable
+from typing import Awaitable, List, Union, Mapping, Optional, Dict, Callable
 
 from . import aio, Credentials, _apis, issues
 
@@ -52,10 +53,14 @@
     PublicTopicSelector as TopicReaderSelector,
 )
 
-from ._topic_reader.topic_reader_sync import TopicReaderSync as TopicReader
+from ._topic_reader.topic_reader_sync import (
+    TopicReaderSync as TopicReader,
+    TopicNoConsumerReaderSync as TopicNoConsumerReader,
+)
 
 from ._topic_reader.topic_reader_asyncio import (
     PublicAsyncIOReader as TopicReaderAsyncIO,
+    PublicAsyncIONoConsumerReader as TopicNoConsumerReaderAsyncIO,
     PublicTopicReaderPartitionExpiredError as TopicReaderPartitionExpiredError,
     PublicTopicReaderUnexpectedCodecError as TopicReaderUnexpectedCodecError,
 )
@@ -261,6 +266,31 @@ def reader(
 
         return TopicReaderAsyncIO(self._driver, settings, _parent=self)
 
+    def no_consumer_reader(
+        self,
+        topic: Union[str, TopicReaderSelector, List[Union[str, TopicReaderSelector]]],
+        partition_ids: List[int],
+        get_start_offset_lambda: Union[Callable[[int], int], Callable[[int], Awaitable[int]]],
+        buffer_size_bytes: int = 50 * 1024 * 1024,
+        # decoders: map[codec_code] func(encoded_bytes)->decoded_bytes
+        # the func will be called from multiply threads in parallel
+        decoders: Union[Mapping[int, Callable[[bytes], bytes]], None] = None,
+        # custom decoder executor for call builtin and custom decoders. If None - use shared executor pool.
+        # if max_worker in the executor is 1 - then decoders will be called from the thread without parallel
+        decoder_executor: Optional[concurrent.futures.Executor] = None,
+        auto_partitioning_support: Optional[bool] = True,  # Auto partitioning feature flag. Default - True.
+    ) -> TopicNoConsumerReaderAsyncIO:
+        if not decoder_executor:
+            decoder_executor = self._executor
+
+        args = locals().copy()
+        del args["self"]
+        args["consumer"] = None
+
+        settings = TopicReaderSettings(**args)
+
+        return TopicNoConsumerReaderAsyncIO(self._driver, settings, _parent=self)
+
     def writer(
         self,
         topic,
@@ -505,6 +535,31 @@ def reader(
 
         return TopicReader(self._driver, settings, _parent=self)
 
+    def no_consumer_reader(
+        self,
+        topic: Union[str, TopicReaderSelector, List[Union[str, TopicReaderSelector]]],
+        partition_ids: List[int],
+        get_start_offset_lambda: Union[Callable[[int], int], Callable[[int], Awaitable[int]]],
+        buffer_size_bytes: int = 50 * 1024 * 1024,
+        # decoders: map[codec_code] func(encoded_bytes)->decoded_bytes
+        # the func will be called from multiply threads in parallel
+        decoders: Union[Mapping[int, Callable[[bytes], bytes]], None] = None,
+        # custom decoder executor for call builtin and custom decoders. If None - use shared executor pool.
+        # if max_worker in the executor is 1 - then decoders will be called from the thread without parallel
+        decoder_executor: Optional[concurrent.futures.Executor] = None,
+        auto_partitioning_support: Optional[bool] = True,  # Auto partitioning feature flag. Default - True.
+    ) -> TopicNoConsumerReader:
+        if not decoder_executor:
+            decoder_executor = self._executor
+
+        args = locals().copy()
+        del args["self"]
+        args["consumer"] = None
+
+        settings = TopicReaderSettings(**args)
+
+        return TopicNoConsumerReader(self._driver, settings, _parent=self)
+
     def writer(
         self,
         topic,