read message from real servre

Author: rekby

tests/conftest.py

@@ -100,11 +100,17 @@ async def driver(endpoint, database, event_loop):
 
 
 @pytest.fixture()
-def topic_path(endpoint) -> str:
+def topic_consumer():
+    return "fixture-consumer"
+
+
+@pytest.fixture()
+def topic_path(endpoint, topic_consumer) -> str:
     subprocess.run(
         """docker-compose exec -T ydb /ydb -e grpc://%s -d /local topic drop /local/test-topic"""
         % endpoint,
         shell=True,
+        capture_output=True,
     )
     res = subprocess.run(
         """docker-compose exec -T ydb /ydb -e grpc://%s -d /local topic create /local/test-topic"""
@@ -114,4 +120,24 @@ def topic_path(endpoint) -> str:
     )
     assert res.returncode == 0, res.stderr + res.stdout
 
+    res = subprocess.run(
+        """docker-compose exec -T ydb /ydb -e grpc://%s -d /local topic consumer add --consumer %s /local/test-topic"""
+        % (endpoint, topic_consumer),
+        shell=True,
+        capture_output=True,
+    )
+    assert res.returncode == 0, res.stderr + res.stdout
+
     return "/local/test-topic"
+
+
+@pytest.fixture()
+@pytest.mark.asyncio()
+async def topic_with_messages(driver, topic_path):
+    pass
+    writer = driver.topic_client.topic_writer(topic_path, producer_and_message_group_id="fixture-producer-id")
+    await writer.write_with_ack(
+        ydb.TopicWriterMessage(data="123".encode()),
+        ydb.TopicWriterMessage(data="456".encode()),
+    )
+    await writer.close()

tests/topics/test_topic_reader.py

@@ -0,0 +1,16 @@
+import pytest
+
+from ydb._topic_reader.topic_reader_asyncio import PublicAsyncIOReader
+from ydb import TopicReaderSettings
+
+
+@pytest.mark.asyncio
+class TestTopicWriterAsyncIO:
+    async def test_read_message(self, driver, topic_path, topic_with_messages, topic_consumer):
+        reader = PublicAsyncIOReader(driver, TopicReaderSettings(
+            consumer=topic_consumer,
+            topic=topic_path,
+        ))
+        await reader.wait_messages()
+
+        assert reader.receive_batch() is not None

ydb/_topic_reader/topic_reader_asyncio.py

@@ -32,6 +32,12 @@ def __init__(self, driver: Driver, settings: PublicReaderSettings):
         self._loop = asyncio.get_running_loop()
         self._reconnector = ReaderReconnector(driver, settings)
 
+    async def wait_messages(self):
+        await self._reconnector.wait_message()
+
+    def receive_batch(self):
+        return self._reconnector.receive_batch_nowait()
+
 
 class ReaderReconnector:
     _settings: PublicReaderSettings
@@ -58,6 +64,7 @@ async def wait_message(self):
         while True:
             if self._stream_reader is not None:
                 await self._stream_reader.wait_messages()
+                return
 
             await self._state_changed.wait()
             self._state_changed.clear()

ydb/_topic_wrapper/common.py

@@ -179,11 +179,14 @@ async def _start_sync_driver(self, driver: ydb.Driver, stub, method):
 
     async def receive(self) -> typing.Any:
         # todo handle grpc exceptions and convert it to internal exceptions
-        grpc_item = await self.from_server_grpc.__anext__()
-        return self.convert_server_grpc_to_wrapper(grpc_item)
+        grpc_message = await self.from_server_grpc.__anext__()
+        # print("rekby, grpc, received", grpc_message)
+        return self.convert_server_grpc_to_wrapper(grpc_message)
 
     def write(self, wrap_message: IToProto):
-        self.from_client_grpc.put_nowait(wrap_message.to_proto())
+        grpc_message=wrap_message.to_proto()
+        # print("rekby, grpc, send", grpc_message)
+        self.from_client_grpc.put_nowait(grpc_message)
 
 
 @dataclass(init=False)

ydb/_topic_wrapper/reader.py

@@ -227,8 +227,12 @@ def __init__(self, client_message: "ReaderMessagesFromClientToServer"):
 
         def to_proto(self) -> ydb_topic_pb2.StreamReadMessage.FromClient:
             res = ydb_topic_pb2.StreamReadMessage.FromClient()
-            if isinstance(self.client_message, StreamReadMessage.InitRequest):
+            if isinstance(self.client_message, StreamReadMessage.ReadRequest):
+                res.read_request.CopyFrom(self.client_message.to_proto())
+            elif isinstance(self.client_message, StreamReadMessage.InitRequest):
                 res.init_request.CopyFrom(self.client_message.to_proto())
+            elif isinstance(self.client_message, StreamReadMessage.StartPartitionSessionResponse):
+                res.start_partition_session_response.CopyFrom(self.client_message.to_proto())
             else:
                 raise NotImplementedError()
             return res
@@ -242,12 +246,16 @@ def from_proto(msg: ydb_topic_pb2.StreamReadMessage.FromServer) -> "StreamReadMe
             mess_type = msg.WhichOneof("server_message")
             if mess_type == "read_response":
                 return StreamReadMessage.FromServer(
-                    server_message=StreamReadMessage.ReadResponse.from_proto(msg.init_response)
+                    server_message=StreamReadMessage.ReadResponse.from_proto(msg.read_response)
                 )
             elif mess_type == "init_response":
                 return StreamReadMessage.FromServer(
                     server_message=StreamReadMessage.InitResponse.from_proto(msg.init_response),
                 )
+            elif mess_type == "start_partition_session_request":
+                return StreamReadMessage.FromServer(
+                    server_message=StreamReadMessage.StartPartitionSessionRequest.from_proto(msg.start_partition_session_request)
+                )
 
             # todo replace exception to log
             raise NotImplementedError()

ydb/topic.py

@@ -2,6 +2,7 @@
 
 from . import aio, Credentials
 from ._topic_reader.topic_reader import (
+    PublicReaderSettings as TopicReaderSettings,
     Reader as TopicReader,
     ReaderAsyncIO as TopicReaderAsyncIO,
     Selector as TopicSelector,
@@ -18,6 +19,7 @@
     RetryPolicy as TopicWriterRetryPolicy,
 )
 
+
 from ydb._topic_writer.topic_writer_asyncio import WriterAsyncIO as TopicWriterAsyncIO