Merge pull request #277 custom topic settings

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/__init__.py

@@ -160,7 +160,36 @@ async def topic_path(driver, topic_consumer, database) -> str:
 
 @pytest.fixture()
 @pytest.mark.asyncio()
-async def topic_with_messages(driver, topic_path):
+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):
+    topic_path = database + "/test-topic-with-messages"
+    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],
+    )
+
     writer = driver.topic_client.writer(topic_path, producer_id="fixture-producer-id", codec=ydb.TopicCodec.RAW)
     await writer.write_with_ack(
         [
@@ -175,6 +204,7 @@ async def topic_with_messages(driver, topic_path):
         ]
     )
     await writer.close()
+    return topic_path
 
 
 @pytest.fixture()

tests/conftest.py

@@ -5,8 +5,8 @@
 
 @pytest.mark.asyncio
 class TestTopicReaderAsyncIO:
-    async def test_read_batch(self, driver, topic_path, topic_with_messages, topic_consumer):
-        reader = driver.topic_client.reader(topic_path, topic_consumer)
+    async def test_read_batch(self, driver, topic_with_messages, topic_consumer):
+        reader = driver.topic_client.reader(topic_with_messages, topic_consumer)
         batch = await reader.receive_batch()
 
         assert batch is not None
@@ -18,30 +18,30 @@ async def test_link_to_client(self, driver, topic_path, topic_consumer):
         reader = driver.topic_client.reader(topic_path, topic_consumer)
         assert reader._parent is driver.topic_client
 
-    async def test_read_message(self, driver, topic_path, topic_with_messages, topic_consumer):
-        reader = driver.topic_client.reader(topic_path, topic_consumer)
+    async def test_read_message(self, driver, topic_with_messages, topic_consumer):
+        reader = driver.topic_client.reader(topic_with_messages, topic_consumer)
         msg = await reader.receive_message()
 
         assert msg is not None
         assert msg.seqno
 
         await reader.close()
 
-    async def test_read_and_commit_with_close_reader(self, driver, topic_path, topic_with_messages, topic_consumer):
-        async with driver.topic_client.reader(topic_path, topic_consumer) as reader:
+    async def test_read_and_commit_with_close_reader(self, driver, topic_with_messages, topic_consumer):
+        async with driver.topic_client.reader(topic_with_messages, topic_consumer) as reader:
             message = await reader.receive_message()
             reader.commit(message)
 
-        async with driver.topic_client.reader(topic_path, topic_consumer) as reader:
+        async with driver.topic_client.reader(topic_with_messages, topic_consumer) as reader:
             message2 = await reader.receive_message()
             assert message != message2
 
-    async def test_read_and_commit_with_ack(self, driver, topic_path, topic_with_messages, topic_consumer):
-        reader = driver.topic_client.reader(topic_path, topic_consumer)
+    async def test_read_and_commit_with_ack(self, driver, topic_with_messages, topic_consumer):
+        reader = driver.topic_client.reader(topic_with_messages, topic_consumer)
         batch = await reader.receive_batch()
         await reader.commit_with_ack(batch)
 
-        reader = driver.topic_client.reader(topic_path, topic_consumer)
+        reader = driver.topic_client.reader(topic_with_messages, topic_consumer)
         batch2 = await reader.receive_batch()
         assert batch.messages[0] != batch2.messages[0]
 
@@ -71,10 +71,35 @@ 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_path, topic_with_messages, topic_consumer):
-        reader = driver_sync.topic_client.reader(topic_path, topic_consumer)
+    def test_read_batch(self, driver_sync, topic_with_messages, topic_consumer):
+        reader = driver_sync.topic_client.reader(topic_with_messages, topic_consumer)
         batch = reader.receive_batch()
 
         assert batch is not None
@@ -86,30 +111,30 @@ def test_link_to_client(self, driver_sync, topic_path, topic_consumer):
         reader = driver_sync.topic_client.reader(topic_path, topic_consumer)
         assert reader._parent is driver_sync.topic_client
 
-    def test_read_message(self, driver_sync, topic_path, topic_with_messages, topic_consumer):
-        reader = driver_sync.topic_client.reader(topic_path, topic_consumer)
+    def test_read_message(self, driver_sync, topic_with_messages, topic_consumer):
+        reader = driver_sync.topic_client.reader(topic_with_messages, topic_consumer)
         msg = reader.receive_message()
 
         assert msg is not None
         assert msg.seqno
 
         reader.close()
 
-    def test_read_and_commit_with_close_reader(self, driver_sync, topic_path, topic_with_messages, topic_consumer):
-        with driver_sync.topic_client.reader(topic_path, topic_consumer) as reader:
+    def test_read_and_commit_with_close_reader(self, driver_sync, topic_with_messages, topic_consumer):
+        with driver_sync.topic_client.reader(topic_with_messages, topic_consumer) as reader:
             message = reader.receive_message()
             reader.commit(message)
 
-        with driver_sync.topic_client.reader(topic_path, topic_consumer) as reader:
+        with driver_sync.topic_client.reader(topic_with_messages, topic_consumer) as reader:
             message2 = reader.receive_message()
             assert message != message2
 
-    def test_read_and_commit_with_ack(self, driver_sync, topic_path, topic_with_messages, topic_consumer):
-        reader = driver_sync.topic_client.reader(topic_path, topic_consumer)
+    def test_read_and_commit_with_ack(self, driver_sync, topic_with_messages, topic_consumer):
+        reader = driver_sync.topic_client.reader(topic_with_messages, topic_consumer)
         batch = reader.receive_batch()
         reader.commit_with_ack(batch)
 
-        reader = driver_sync.topic_client.reader(topic_path, topic_consumer)
+        reader = driver_sync.topic_client.reader(topic_with_messages, topic_consumer)
         batch2 = reader.receive_batch()
         assert batch.messages[0] != batch2.messages[0]
 

tests/ssl/__init__.py

@@ -274,14 +274,15 @@ async def to_thread(func, /, *args, **kwargs):
     return await loop.run_in_executor(None, func_call)
 
 
-def proto_duration_from_timedelta(t: Optional[datetime.timedelta]) -> ProtoDuration:
+def proto_duration_from_timedelta(t: Optional[datetime.timedelta]) -> Optional[ProtoDuration]:
     if t is None:
         return None
+
     res = ProtoDuration()
     res.FromTimedelta(t)
 
 
-def proto_timestamp_from_datetime(t: Optional[datetime.datetime]) -> ProtoTimeStamp:
+def proto_timestamp_from_datetime(t: Optional[datetime.datetime]) -> Optional[ProtoTimeStamp]:
     if t is None:
         return None
 

tests/table/__init__.py

@@ -410,17 +410,23 @@ def to_proto(self) -> ydb_topic_pb2.StreamReadMessage.InitRequest:
         class TopicReadSettings(IToProto):
             path: str
             partition_ids: List[int] = field(default_factory=list)
-            max_lag_seconds: Union[datetime.timedelta, None] = None
-            read_from: Union[int, float, datetime.datetime, None] = None
+            max_lag: Optional[datetime.timedelta] = None
+            read_from: Optional[datetime.datetime] = None
 
             def to_proto(
                 self,
             ) -> ydb_topic_pb2.StreamReadMessage.InitRequest.TopicReadSettings:
                 res = ydb_topic_pb2.StreamReadMessage.InitRequest.TopicReadSettings()
                 res.path = self.path
                 res.partition_ids.extend(self.partition_ids)
-                if self.max_lag_seconds is not None:
-                    res.max_lag = proto_duration_from_timedelta(self.max_lag_seconds)
+                max_lag = proto_duration_from_timedelta(self.max_lag)
+                if max_lag is not None:
+                    res.max_lag = max_lag
+
+                read_from = proto_timestamp_from_datetime(self.read_from)
+                if read_from is not None:
+                    res.read_from = read_from
+
                 return res
 
     @dataclass

tests/topics/__init__.py

@@ -14,21 +14,36 @@
 from .._grpc.grpcwrapper.ydb_topic import StreamReadMessage, OffsetsRange
 
 
-class Selector:
+@dataclass
+class PublicTopicSelector:
     path: str
-    partitions: Union[None, int, List[int]]
-    read_from_timestamp_ms: Optional[int]
-    max_time_lag_ms: Optional[int]
+    partitions: Optional[Union[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:
+        partitions = self.partitions
+        if partitions is None:
+            partitions = []
+
+        elif not isinstance(partitions, list):
+            partitions = [partitions]
+
+        return StreamReadMessage.InitRequest.TopicReadSettings(
+            path=self.path,
+            partition_ids=partitions,
+            max_lag=self.max_lag,
+            read_from=self.read_from,
+        )
+
 
-    def __init__(self, path, *, partitions: Union[None, int, List[int]] = None):
-        self.path = path
-        self.partitions = partitions
+TopicSelectorTypes = Union[str, PublicTopicSelector, List[Union[str, PublicTopicSelector]]]
 
 
 @dataclass
 class PublicReaderSettings:
     consumer: str
-    topic: str
+    topic: TopicSelectorTypes
     buffer_size_bytes: int = 50 * 1024 * 1024
 
     decoders: Union[Mapping[int, Callable[[bytes], bytes]], None] = None
@@ -43,12 +58,24 @@ def __post_init__(self):
         _ = self._init_message()
 
     def _init_message(self) -> StreamReadMessage.InitRequest:
+        if not isinstance(self.consumer, str):
+            raise TypeError("Unsupported type for customer field: '%s'" % type(self.consumer))
+
+        if isinstance(self.topic, list):
+            selectors = self.topic
+        else:
+            selectors = [self.topic]
+
+        for index, selector in enumerate(selectors):
+            if isinstance(selector, str):
+                selectors[index] = PublicTopicSelector(path=selector)
+            elif isinstance(selector, PublicTopicSelector):
+                pass
+            else:
+                raise TypeError("Unsupported type for topic field: '%s'" % type(selector))
+
         return StreamReadMessage.InitRequest(
-            topics_read_settings=[
-                StreamReadMessage.InitRequest.TopicReadSettings(
-                    path=self.topic,
-                )
-            ],
+            topics_read_settings=list(map(PublicTopicSelector._to_topic_read_settings, selectors)),  # type: ignore
             consumer=self.consumer,
         )