sync

Author: rekby

tests/__init__.py

@@ -160,7 +160,18 @@ async def topic_path(driver, topic_consumer, database) -> str:
 
 @pytest.fixture()
 @pytest.mark.asyncio()
-async def topic_with_messages(driver, topic_path):
+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 +186,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]
 
@@ -73,8 +73,8 @@ def decode(b: bytes):
 
 
 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 +86,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,39 @@
 from .._grpc.grpcwrapper.ydb_topic import StreamReadMessage, OffsetsRange
 
 
-class Selector:
+class PublicTopicSelector:
     path: str
     partitions: Union[None, int, List[int]]
-    read_from_timestamp_ms: Optional[int]
-    max_time_lag_ms: Optional[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
 
+    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,
+        )
+
+
+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
@@ -39,12 +57,21 @@ class PublicReaderSettings:
     update_token_interval: Union[int, float] = 3600
 
     def _init_message(self) -> StreamReadMessage.InitRequest:
+        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("Value of %s not supported as topic selector" % 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,
         )
 

tests/topics/test_topic_reader.py

@@ -576,7 +576,7 @@ async def test_init_reader(self, stream, default_reader_settings):
                 StreamReadMessage.InitRequest.TopicReadSettings(
                     path="/local/test-topic",
                     partition_ids=[],
-                    max_lag_seconds=None,
+                    max_lag=None,
                     read_from=None,
                 )
             ],