another approach

Author: vgvoleg

tests/topics/test_topic_reader.py

@@ -4,6 +4,8 @@
 
 import ydb
 
+from ydb._topic_reader.events import OnPartitionGetStartOffsetResponse
+
 
 @pytest.mark.asyncio
 class TestTopicReaderAsyncIO:
@@ -253,6 +255,13 @@ async def wait(fut):
         await reader1.close()
 
 
+@pytest.fixture()
+def topic_selector(topic_with_messages):
+    return ydb.TopicReaderSelector(
+        path=topic_with_messages,
+        partitions=[0]
+    )
+
 @pytest.mark.asyncio
 class TestTopicNoConsumerReaderAsyncIO:
     async def test_reader_with_no_partition_ids_raises(self, driver, topic_with_messages):
@@ -262,57 +271,63 @@ async def test_reader_with_no_partition_ids_raises(self, driver, topic_with_mess
                 consumer=None,
             )
 
-    async def test_reader_with_default_lambda(self, driver, topic_with_messages):
-        reader = driver.topic_client.reader(
-            topic_with_messages,
-            consumer=None,
-            partition_ids=[0],
-        )
+    async def test_reader_with_no_partition_ids_selector_raises(self, driver, topic_selector):
+        topic_selector.partitions = None
+
+        with pytest.raises(ydb.Error):
+            driver.topic_client.reader(
+                topic_selector,
+                consumer=None,
+            )
+
+    async def test_reader_with_default_lambda(self, driver, topic_selector):
+        reader = driver.topic_client.reader(topic_selector, consumer=None)
         msg = await reader.receive_message()
 
         assert msg.seqno == 1
 
         await reader.close()
 
-    async def test_reader_with_sync_lambda(self, driver, topic_with_messages):
-        def sync_lambda(partition_id: int):
-            assert partition_id == 0
-            return 1
+    async def test_reader_with_sync_lambda(self, driver, topic_selector):
+        class CustomEventHandler(ydb.TopicReaderEventHandler):
+            def on_partition_get_start_offset(self, event):
+                assert event.partition_id == 0
+                return OnPartitionGetStartOffsetResponse(1)
 
         reader = driver.topic_client.reader(
-            topic_with_messages,
+            topic_selector,
             consumer=None,
-            partition_ids=[0],
-            get_start_offset_lambda=sync_lambda,
+            event_handler=CustomEventHandler(),
         )
+
         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
+    async def test_reader_with_async_lambda(self, driver, topic_selector):
+        class CustomEventHandler(ydb.TopicReaderEventHandler):
+            async def on_partition_get_start_offset(self, event):
+                assert event.partition_id == 0
+                return OnPartitionGetStartOffsetResponse(1)
 
         reader = driver.topic_client.reader(
-            topic_with_messages,
+            topic_selector,
             consumer=None,
-            partition_ids=[0],
-            get_start_offset_lambda=async_lambda,
+            event_handler=CustomEventHandler(),
         )
+
         msg = await reader.receive_message()
 
         assert msg.seqno == 2
 
         await reader.close()
 
-    async def test_commit_not_allowed(self, driver, topic_with_messages):
+    async def test_commit_not_allowed(self, driver, topic_selector):
         reader = driver.topic_client.reader(
-            topic_with_messages,
+            topic_selector,
             consumer=None,
-            partition_ids=[0],
         )
         batch = await reader.receive_batch()
 
@@ -324,18 +339,18 @@ async def test_commit_not_allowed(self, driver, topic_with_messages):
 
         await reader.close()
 
-    async def test_offsets_updated_after_reconnect(self, driver, topic_with_messages):
+    async def test_offsets_updated_after_reconnect(self, driver, topic_selector):
         current_offset = 0
 
-        def get_start_offset_lambda(partition_id: int) -> int:
-            nonlocal current_offset
-            return current_offset
+        class CustomEventHandler(ydb.TopicReaderEventHandler):
+            def on_partition_get_start_offset(self, event):
+                nonlocal current_offset
+                return OnPartitionGetStartOffsetResponse(current_offset)
 
         reader = driver.topic_client.reader(
-            topic_with_messages,
+            topic_selector,
             consumer=None,
-            partition_ids=[0],
-            get_start_offset_lambda=get_start_offset_lambda,
+            event_handler=CustomEventHandler(),
         )
         msg = await reader.receive_message()
 
@@ -361,57 +376,63 @@ def test_reader_with_no_partition_ids_raises(self, driver_sync, topic_with_messa
                 consumer=None,
             )
 
-    def test_reader_with_default_lambda(self, driver_sync, topic_with_messages):
-        reader = driver_sync.topic_client.reader(
-            topic_with_messages,
-            consumer=None,
-            partition_ids=[0],
-        )
+    def test_reader_with_no_partition_ids_selector_raises(self, driver_sync, topic_selector):
+        topic_selector.partitions = None
+
+        with pytest.raises(ydb.Error):
+            driver_sync.topic_client.reader(
+                topic_selector,
+                consumer=None,
+            )
+
+    def test_reader_with_default_lambda(self, driver_sync, topic_selector):
+        reader = driver_sync.topic_client.reader(topic_selector, consumer=None)
         msg = reader.receive_message()
 
         assert msg.seqno == 1
 
         reader.close()
 
-    def test_reader_with_sync_lambda(self, driver_sync, topic_with_messages):
-        def sync_lambda(partition_id: int):
-            assert partition_id == 0
-            return 1
+    def test_reader_with_sync_lambda(self, driver_sync, topic_selector):
+        class CustomEventHandler(ydb.TopicReaderEventHandler):
+            def on_partition_get_start_offset(self, event):
+                assert event.partition_id == 0
+                return OnPartitionGetStartOffsetResponse(1)
 
         reader = driver_sync.topic_client.reader(
-            topic_with_messages,
+            topic_selector,
             consumer=None,
-            partition_ids=[0],
-            get_start_offset_lambda=sync_lambda,
+            event_handler=CustomEventHandler(),
         )
+
         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
+    def test_reader_with_async_lambda(self, driver_sync, topic_selector):
+        class CustomEventHandler(ydb.TopicReaderEventHandler):
+            async def on_partition_get_start_offset(self, event):
+                assert event.partition_id == 0
+                return OnPartitionGetStartOffsetResponse(1)
 
         reader = driver_sync.topic_client.reader(
-            topic_with_messages,
+            topic_selector,
             consumer=None,
-            partition_ids=[0],
-            get_start_offset_lambda=async_lambda,
+            event_handler=CustomEventHandler(),
         )
+
         msg = reader.receive_message()
 
         assert msg.seqno == 2
 
         reader.close()
 
-    def test_commit_not_allowed(self, driver_sync, topic_with_messages):
+    def test_commit_not_allowed(self, driver_sync, topic_selector):
         reader = driver_sync.topic_client.reader(
-            topic_with_messages,
+            topic_selector,
             consumer=None,
-            partition_ids=[0],
         )
         batch = reader.receive_batch()
 
@@ -421,30 +442,27 @@ def test_commit_not_allowed(self, driver_sync, topic_with_messages):
         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):
+    def test_offsets_updated_after_reconnect(self, driver_sync, topic_selector):
         current_offset = 0
 
-        def get_start_offset_lambda(partition_id: int) -> int:
-            nonlocal current_offset
-            return current_offset
+        class CustomEventHandler(ydb.TopicReaderEventHandler):
+            def on_partition_get_start_offset(self, event):
+                nonlocal current_offset
+                return OnPartitionGetStartOffsetResponse(current_offset)
 
         reader = driver_sync.topic_client.reader(
-            topic_with_messages,
+            topic_selector,
             consumer=None,
-            partition_ids=[0],
-            get_start_offset_lambda=get_start_offset_lambda,
+            event_handler=CustomEventHandler(),
         )
         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"))
+        reader._reconnector._stream_reader._set_first_error(ydb.Unavailable("some retriable error"))
 
         msg = reader.receive_message()
 

ydb/_topic_reader/events.py

@@ -0,0 +1,72 @@
+import asyncio
+from dataclasses import dataclass
+from typing import Awaitable, Union
+
+from ..issues import ClientInternalError
+
+
+class ReaderEvent:
+    pass
+
+
+@dataclass
+class OnCommit(ReaderEvent):
+    topic: str
+    offset: int
+
+
+@dataclass
+class OnPartitionGetStartOffsetRequest(ReaderEvent):
+    topic: str
+    partition_id: int
+
+
+@dataclass
+class OnPartitionGetStartOffsetResponse:
+    start_offset: int
+
+
+class OnInitPartition(ReaderEvent):
+    pass
+
+
+class OnShutdownPartition:
+    pass
+
+
+TopicEventDispatchType = Union[OnPartitionGetStartOffsetResponse, None]
+
+
+class ReaderEventHandler:
+    def on_commit(self, event: OnCommit) -> Union[None, Awaitable[None]]:
+        pass
+
+    def on_partition_get_start_offset(
+        self,
+        event: OnPartitionGetStartOffsetRequest,
+    ) -> Union[OnPartitionGetStartOffsetResponse, Awaitable[OnPartitionGetStartOffsetResponse]]:
+        pass
+
+    def on_init_partition(self, event: OnInitPartition) -> Union[None, Awaitable[None]]:
+        pass
+
+    def on_shutdown_partition(self, event: OnShutdownPartition) -> Union[None, Awaitable[None]]:
+        pass
+
+    async def _dispatch(self, event: ReaderEvent) -> Awaitable[TopicEventDispatchType]:
+        f = None
+        if isinstance(event, OnCommit):
+            f = self.on_commit
+        elif isinstance(event, OnPartitionGetStartOffsetRequest):
+            f = self.on_partition_get_start_offset
+        elif isinstance(event, OnInitPartition):
+            f = self.on_init_partition
+        elif isinstance(event, OnShutdownPartition):
+            f = self.on_shutdown_partition
+        else:
+            raise ClientInternalError("Unsupported topic reader event")
+
+        if asyncio.iscoroutinefunction(f):
+            return await f(event)
+
+        return f(event)

ydb/_topic_reader/topic_reader.py

@@ -11,6 +11,7 @@
     Callable,
 )
 
+from .events import ReaderEventHandler
 from ..retries import RetrySettings
 from .._grpc.grpcwrapper.ydb_topic import StreamReadMessage, OffsetsRange
 
@@ -21,6 +22,7 @@ class PublicTopicSelector:
     partitions: Optional[Union[int, List[int]]] = None
     read_from: Optional[datetime.datetime] = None
     max_lag: Optional[datetime.timedelta] = None
+    read_offset: Optional[int] = None
 
     def _to_topic_read_settings(self) -> StreamReadMessage.InitRequest.TopicReadSettings:
         partitions = self.partitions
@@ -54,9 +56,7 @@ class PublicReaderSettings:
     # decoder_executor, must be set for handle non raw messages
     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
+    event_handler: Optional[ReaderEventHandler] = None
 
     def __post_init__(self):
         # check possible create init message
@@ -89,25 +89,6 @@ def _retry_settings(self) -> RetrySettings:
         return RetrySettings(idempotent=True)
 
 
-class Events:
-    class OnCommit:
-        topic: str
-        offset: int
-
-    class OnPartitionGetStartOffsetRequest:
-        topic: str
-        partition_id: int
-
-    class OnPartitionGetStartOffsetResponse:
-        start_offset: int
-
-    class OnInitPartition:
-        pass
-
-    class OnShutdownPatition:
-        pass
-
-
 class RetryPolicy:
     connection_timeout_sec: float
     overload_timeout_sec: float