added support many topics in one reader

Author: rekby

CHANGELOG.md

@@ -1,3 +1,5 @@
+* Added support to set many topics and topic reader settings for read in one reader
+
 ## 3.2.2 ##
 * Fixed set keep_in_cache algorithm
 

tests/conftest.py

@@ -158,6 +158,24 @@ async def topic_path(driver, topic_consumer, database) -> str:
     return topic_path
 
 
+@pytest.fixture()
+@pytest.mark.asyncio()
+async def topic2_path(driver, topic_consumer, database) -> str:
+    topic_path = database + "/test-topic2"
+
+    try:
+        await driver.topic_client.drop_topic(topic_path)
+    except issues.SchemeError:
+        pass
+
+    await driver.topic_client.create_topic(
+        path=topic_path,
+        consumers=[topic_consumer],
+    )
+
+    return topic_path
+
+
 @pytest.fixture()
 @pytest.mark.asyncio()
 async def topic_with_messages(driver, topic_consumer, database):

tests/topics/test_topic_reader.py

@@ -71,6 +71,31 @@ def decode(b: bytes):
             batch = await reader.receive_batch()
             assert batch.messages[0].data.decode() == "123"
 
+    async def test_read_from_two_topics(self, driver, topic_path, topic2_path, topic_consumer):
+        async with driver.topic_client.writer(topic_path) as writer:
+            await writer.write("1")
+            await writer.flush()
+
+        async with driver.topic_client.writer(topic2_path) as writer:
+            await writer.write("2")
+            await writer.flush()
+
+        messages = []
+        async with driver.topic_client.reader(
+            [
+                topic_path,
+                ydb.TopicReaderSelector(path=topic2_path),
+            ],
+            consumer=topic_consumer,
+        ) as reader:
+            for _ in range(2):
+                message = await reader.receive_message()
+                messages.append(message)
+
+        messages = [message.data.decode() for message in messages]
+        messages.sort()
+        assert messages == ["1", "2"]
+
 
 class TestTopicReaderSync:
     def test_read_batch(self, driver_sync, topic_with_messages, topic_consumer):

ydb/_topic_reader/topic_reader.py

@@ -14,17 +14,14 @@
 from .._grpc.grpcwrapper.ydb_topic import StreamReadMessage, OffsetsRange
 
 
+@dataclass
 class PublicTopicSelector:
     path: str
-    partitions: Union[None, int, List[int]]
-    read_from: Optional[datetime.datetime]
-    max_lag: Optional[datetime.timedelta]
-
-    def __init__(self, path, *, partitions: Union[None, int, List[int]] = None):
-        self.path = path
-        self.partitions = partitions
+    partitions: Union[None, int, List[int]] = None
+    read_from: Optional[datetime.datetime] = None
+    max_lag: Optional[datetime.timedelta] = None
 
-    def _to_topic_read_settings(self)->StreamReadMessage.InitRequest.TopicReadSettings:
+    def _to_topic_read_settings(self) -> StreamReadMessage.InitRequest.TopicReadSettings:
         partitions = self.partitions
         if partitions is None:
             partitions = []
@@ -71,7 +68,7 @@ def _init_message(self) -> StreamReadMessage.InitRequest:
                 raise TypeError("Value of %s not supported as topic selector" % type(selector))
 
         return StreamReadMessage.InitRequest(
-            topics_read_settings=list(map(PublicTopicSelector._to_topic_read_settings, selectors)), # type: ignore
+            topics_read_settings=list(map(PublicTopicSelector._to_topic_read_settings, selectors)),  # type: ignore
             consumer=self.consumer,
         )
 

ydb/_topic_reader/topic_reader_asyncio.py

@@ -164,7 +164,7 @@ async def _connection_loop(self):
                 attempt = 0
                 self._state_changed.set()
                 await self._stream_reader.wait_error()
-            except issues.Error as err:
+            except BaseException as err:
                 retry_info = check_retriable_error(err, self._settings._retry_settings(), attempt)
                 if not retry_info.is_retriable:
                     self._set_first_error(err)

ydb/topic.py

@@ -11,6 +11,7 @@
     "TopicMeteringMode",
     "TopicReader",
     "TopicReaderAsyncIO",
+    "TopicReaderSelector",
     "TopicReaderSettings",
     "TopicStatWindow",
     "TopicWriter",
@@ -30,6 +31,7 @@
 
 from ._topic_reader.topic_reader import (
     PublicReaderSettings as TopicReaderSettings,
+    PublicTopicSelector as TopicReaderSelector,
 )
 
 from ._topic_reader.topic_reader_sync import TopicReaderSync as TopicReader
@@ -151,7 +153,7 @@ async def drop_topic(self, path: str):
 
     def reader(
         self,
-        topic: str,
+        topic: Union[str, TopicReaderSelector, List[Union[str, TopicReaderSelector]]],
         consumer: str,
         buffer_size_bytes: int = 50 * 1024 * 1024,
         # decoders: map[codec_code] func(encoded_bytes)->decoded_bytes
@@ -300,7 +302,7 @@ def drop_topic(self, path: str):
 
     def reader(
         self,
-        topic: str,
+        topic: Union[str, TopicReaderSelector, List[Union[str, TopicReaderSelector]]],
         consumer: str,
         buffer_size_bytes: int = 50 * 1024 * 1024,
         # decoders: map[codec_code] func(encoded_bytes)->decoded_bytes