sync

rekby · rekby · commit 9cd569533dc2 · 2023-02-06T19:41:12.000+03:00
diff --git a/ydb/_topic_reader/topic_reader.py b/ydb/_topic_reader/topic_reader.py
@@ -17,6 +17,7 @@
 )
 
 from ydb._topic_wrapper.common import OffsetsRange, TokenGetterFuncType
+from ydb._topic_wrapper.reader import StreamReadMessage
 
 
 class Selector:
@@ -270,6 +271,16 @@ class PublicReaderSettings:
     # connection_timeout: Union[float, None] = None
     # retry_policy: Union["RetryPolicy", None] = None
 
+    def _init_message(self) -> StreamReadMessage.InitRequest:
+        return StreamReadMessage.InitRequest(
+            topics_read_settings=[
+                StreamReadMessage.InitRequest.TopicReadSettings(
+                    path=self.topic,
+                )
+            ],
+            consumer=self.consumer,
+        )
+
 
 class Events:
     class OnCommit:
diff --git a/ydb/_topic_reader/topic_reader_asyncio.py b/ydb/_topic_reader/topic_reader_asyncio.py
@@ -2,13 +2,16 @@
 
 import asyncio
 import typing
+from asyncio import Task
 from collections import deque
 from typing import Optional, Set, Dict
 
+from .. import _apis
+from ..aio import Driver
 from ..issues import Error as YdbError
 from .datatypes import PartitionSession, PublicMessage, PublicBatch
 from .topic_reader import PublicReaderSettings
-from .._topic_wrapper.common import TokenGetterFuncType, IGrpcWrapperAsyncIO
+from .._topic_wrapper.common import TokenGetterFuncType, IGrpcWrapperAsyncIO, SupportedDriverType, GrpcWrapperAsyncIO
 from .._topic_wrapper.reader import StreamReadMessage
 
 
@@ -22,11 +25,52 @@ def __init__(self):
 
 
 class PublicAsyncIOReader:
-    pass
+    _loop: asyncio.AbstractEventLoop
+    _reconnector: ReaderReconnector
+
+    def __init__(self, driver: Driver, settings: PublicReaderSettings):
+        self._loop = asyncio.get_running_loop()
+        self._reconnector = ReaderReconnector(driver, settings)
 
 
 class ReaderReconnector:
-    pass
+    _settings: PublicReaderSettings
+    _driver: Driver
+    _background_tasks: Set[Task]
+
+    _state_changed: asyncio.Event
+    _stream_reader: Optional["ReaderStream"]
+
+    def __init__(self, driver: Driver, settings: PublicReaderSettings):
+        self._settings = settings
+        self._driver = driver
+        self._background_tasks = set()
+
+        self._state_changed = asyncio.Event()
+        self._stream_reader = None
+        self._background_tasks.add(asyncio.create_task(self.start()))
+
+    async def start(self):
+        self._stream_reader = await ReaderStream.create(self._driver, self._settings)
+        self._state_changed.set()
+
+    async def wait_message(self):
+        while True:
+            if self._stream_reader is not None:
+                await self._stream_reader.wait_messages()
+
+            await self._state_changed.wait()
+            self._state_changed.clear()
+
+    def receive_batch_nowait(self):
+        return self._stream_reader.receive_batch_nowait()
+
+    async def close(self):
+        await self._stream_reader.close()
+        for task in self._background_tasks:
+            task.cancel()
+
+        await asyncio.wait(self._background_tasks)
 
 
 class ReaderStream:
@@ -57,7 +101,22 @@ def __init__(self, settings: PublicReaderSettings):
         self._first_error = None
         self._message_batches = deque()
 
-    async def start(self, stream: IGrpcWrapperAsyncIO, init_message: StreamReadMessage.InitRequest):
+    @staticmethod
+    async def create(
+        driver: SupportedDriverType,
+        settings: PublicReaderSettings,
+    ) -> "ReaderStream":
+        stream = GrpcWrapperAsyncIO(StreamReadMessage.FromServer.from_proto)
+
+        await stream.start(
+            driver, _apis.TopicService.Stub, _apis.TopicService.StreamRead
+        )
+
+        reader = ReaderStream(settings)
+        await reader._start(stream, settings._init_message())
+        return reader
+
+    async def _start(self, stream: IGrpcWrapperAsyncIO, init_message: StreamReadMessage.InitRequest):
         if self._started:
             raise TopicReaderError("Double start ReaderStream")
 
diff --git a/ydb/_topic_reader/topic_reader_asyncio_test.py b/ydb/_topic_reader/topic_reader_asyncio_test.py
@@ -62,7 +62,7 @@ async def stream_reader(self, stream, default_reader_settings, partition_session
         init_message = object()
 
         # noinspection PyTypeChecker
-        start = asyncio.create_task(reader.start(stream, init_message))
+        start = asyncio.create_task(reader._start(stream, init_message))
 
         stream.from_server.put_nowait(StreamReadMessage.FromServer(
             StreamReadMessage.InitResponse(session_id="test-session")
@@ -178,7 +178,7 @@ async def test_init_reader(self, stream, default_reader_settings):
                 read_from=None,
             )]
         )
-        start_task = asyncio.create_task(reader.start(stream, init_message))
+        start_task = asyncio.create_task(reader._start(stream, init_message))
 
         sent_message = await wait_for_fast(stream.from_client.get())
         expected_sent_init_message = StreamReadMessage.FromClient(client_message=init_message)
diff --git a/ydb/_topic_wrapper/common.py b/ydb/_topic_wrapper/common.py
@@ -33,10 +33,17 @@ class Codec(IntEnum):
 
 
 @dataclass
-class OffsetsRange:
+class OffsetsRange(IFromProto):
     start: int
     end: int
 
+    @staticmethod
+    def from_proto(msg: ydb_topic_pb2.OffsetsRange) -> "OffsetsRange":
+        return OffsetsRange(
+            start=msg.start,
+            end=msg.end,
+        )
+
 
 class IToProto(abc.ABC):
     @abc.abstractmethod
diff --git a/ydb/_topic_wrapper/reader.py b/ydb/_topic_wrapper/reader.py
@@ -1,7 +1,10 @@
 import datetime
+import typing
 from dataclasses import dataclass, field
 from typing import List, Union, Dict
 
+from google.protobuf.message import Message
+
 from ydb._topic_wrapper.common import OffsetsRange, IToProto, UpdateTokenRequest, UpdateTokenResponse, IFromProto
 from google.protobuf.duration_pb2 import Duration as ProtoDuration
 
@@ -14,11 +17,19 @@
 
 class StreamReadMessage:
     @dataclass
-    class PartitionSession:
+    class PartitionSession(IFromProto):
         partition_session_id: int
         path: str
         partition_id: int
 
+        @staticmethod
+        def from_proto(msg: ydb_topic_pb2.StreamReadMessage.PartitionSession) -> "StreamReadMessage.PartitionSession":
+            return StreamReadMessage.PartitionSession(
+                partition_session_id=msg.partition_session_id,
+                path=msg.path,
+                partition_id=msg.partition_id,
+            )
+
     @dataclass
     class InitRequest(IToProto):
         topics_read_settings: List["StreamReadMessage.InitRequest.TopicReadSettings"]
@@ -56,36 +67,90 @@ def from_proto(msg: ydb_topic_pb2.StreamReadMessage.InitResponse) -> "StreamRead
             return StreamReadMessage.InitResponse(session_id=msg.session_id)
 
     @dataclass
-    class ReadRequest:
+    class ReadRequest(IToProto):
         bytes_size: int
 
+        def to_proto(self) -> ydb_topic_pb2.StreamReadMessage.ReadRequest:
+            res = ydb_topic_pb2.StreamReadMessage.ReadRequest()
+            res.bytes_size = self.bytes_size
+            return res
+
     @dataclass
-    class ReadResponse:
+    class ReadResponse(IFromProto):
         partition_data: List["StreamReadMessage.ReadResponse.PartitionData"]
         bytes_size: int
 
+        @staticmethod
+        def from_proto(msg: ydb_topic_pb2.StreamReadMessage.ReadResponse) -> "StreamReadMessage.ReadResponse":
+            partition_data = []
+            for proto_partition_data in msg.partition_data:
+                partition_data.append(StreamReadMessage.ReadResponse.PartitionData.from_proto(proto_partition_data))
+            return StreamReadMessage.ReadResponse(
+                partition_data=partition_data,
+                bytes_size=msg.bytes_size,
+            )
+
         @dataclass
-        class MessageData:
+        class MessageData(IFromProto):
             offset: int
             seq_no: int
             created_at: datetime.datetime
             data: bytes
             uncompresed_size: int
             message_group_id: str
 
+            @staticmethod
+            def from_proto(msg: ydb_topic_pb2.StreamReadMessage.ReadResponse.MessageData) ->\
+                    "StreamReadMessage.ReadResponse.MessageData":
+                return StreamReadMessage.ReadResponse.MessageData(
+                    offset=msg.offset,
+                    seq_no=msg.seq_no,
+                    created_at=msg.created_at.ToDatetime(),
+                    data=msg.data,
+                    uncompresed_size=msg.uncompressed_size,
+                    message_group_id=msg.message_group_id
+                )
+
         @dataclass
-        class Batch:
+        class Batch(IFromProto):
             message_data: List["StreamReadMessage.ReadResponse.MessageData"]
             producer_id: str
             write_session_meta: Dict[str, str]
             codec: int
             written_at: datetime.datetime
 
+            @staticmethod
+            def from_proto(msg: ydb_topic_pb2.StreamReadMessage.ReadResponse.Batch) -> \
+                    "StreamReadMessage.ReadResponse.Batch":
+                message_data = []
+                for message in msg.message_data:
+                    message_data.append(StreamReadMessage.ReadResponse.MessageData.from_proto(message))
+                return StreamReadMessage.ReadResponse.Batch(
+                    message_data=message_data,
+                    producer_id=msg.producer_id,
+                    write_session_meta=dict(msg.write_session_meta),
+                    codec=msg.codec,
+                    written_at=msg.written_at.ToDatetime(),
+                )
+
+
         @dataclass
-        class PartitionData:
+        class PartitionData(IFromProto):
             partition_session_id: int
             batches: List["StreamReadMessage.ReadResponse.Batch"]
 
+            @staticmethod
+            def from_proto(msg: ydb_topic_pb2.StreamReadMessage.ReadResponse.PartitionData) ->\
+                    "StreamReadMessage.ReadResponse.PartitionData":
+                batches = []
+                for proto_batch in msg.batches:
+                    batches.append(StreamReadMessage.ReadResponse.Batch.from_proto(proto_batch))
+                return StreamReadMessage.ReadResponse.PartitionData(
+                    partition_session_id=msg.partition_session_id,
+                    batches=batches,
+                )
+
+
     @dataclass
     class CommitOffsetRequest:
         commit_offsets: List["PartitionCommitOffset"]
@@ -116,17 +181,33 @@ class PartitionSessionStatusResponse:
         write_time_high_watermark: float
 
     @dataclass
-    class StartPartitionSessionRequest:
+    class StartPartitionSessionRequest(IFromProto):
         partition_session: "StreamReadMessage.PartitionSession"
         committed_offset: int
         partition_offsets: OffsetsRange
 
+        @staticmethod
+        def from_proto(msg: ydb_topic_pb2.StreamReadMessage.StartPartitionSessionRequest) -> \
+                "StreamReadMessage.StartPartitionSessionRequest":
+            return StreamReadMessage.StartPartitionSessionRequest(
+                partition_session=StreamReadMessage.PartitionSession.from_proto(msg.partition_session),
+                committed_offset=msg.committed_offset,
+                partition_offsets=OffsetsRange.from_proto(msg.partition_offsets)
+            )
+
     @dataclass
-    class StartPartitionSessionResponse:
+    class StartPartitionSessionResponse(IToProto):
         partition_session_id: int
         read_offset: int
         commit_offset: int
 
+        def to_proto(self) -> ydb_topic_pb2.StreamReadMessage.StartPartitionSessionResponse:
+            res = ydb_topic_pb2.StreamReadMessage.StartPartitionSessionResponse()
+            res.partition_session_id = self.partition_session_id
+            res.read_offset = self.read_offset
+            res.commit_offset = self.commit_offset
+            return res
+
     @dataclass
     class StopPartitionSessionRequest:
         partition_session_id: int
@@ -159,7 +240,11 @@ class FromServer(IFromProto):
         @staticmethod
         def from_proto(msg: ydb_topic_pb2.StreamReadMessage.FromServer) -> "StreamReadMessage.FromServer":
             mess_type = msg.WhichOneof("server_message")
-            if mess_type == "init_response":
+            if mess_type == "read_response":
+                return StreamReadMessage.FromServer(
+                    server_message=StreamReadMessage.ReadResponse.from_proto(msg.init_response)
+                )
+            elif mess_type == "init_response":
                 return StreamReadMessage.FromServer(
                     server_message=StreamReadMessage.InitResponse.from_proto(msg.init_response),
                 )
