sync

Author: rekby

ydb/_topic_reader/topic_reader_asyncio.py

@@ -6,6 +6,8 @@
 from collections import deque
 from typing import Optional, Set, Dict
 
+import grpc
+
 from .. import _apis
 from ..aio import Driver
 from ..issues import Error as YdbError
@@ -145,6 +147,9 @@ async def wait_messages(self):
             raise TopicReaderStreamClosedError()
 
         while len(self._message_batches) == 0:
+            if self._first_error is not None:
+                raise self._first_error
+
             await self._state_changed.wait()
             self._state_changed.clear()
 
@@ -180,6 +185,8 @@ async def _read_messages_loop(self, stream: IGrpcWrapperAsyncIO):
                     )
 
                 self._state_changed.set()
+        except grpc.RpcError as e:
+
         except Exception as e:
             self._set_first_error(e)
             raise e

ydb/_topic_reader/topic_reader_asyncio_test.py

@@ -2,15 +2,17 @@
 import datetime
 from unittest import mock
 
+import grpc
 import pytest
 
 from ydb import aio
-from ydb._topic_reader.datatypes import PublicBatch, PublicMessage
-from ydb._topic_reader.topic_reader import PublicReaderSettings
-from ydb._topic_reader.topic_reader_asyncio import ReaderStream, PartitionSession
-from ydb._topic_wrapper.common import OffsetsRange, Codec
-from ydb._topic_wrapper.reader import StreamReadMessage
-from ydb._topic_wrapper.test_helpers import StreamMock, wait_condition, wait_for_fast
+from .datatypes import PublicBatch, PublicMessage
+from .topic_reader import PublicReaderSettings
+from .topic_reader_asyncio import ReaderStream, PartitionSession
+from .._topic_wrapper.common import OffsetsRange, Codec
+from .._topic_wrapper.reader import StreamReadMessage
+from .._topic_wrapper.test_helpers import StreamMock, wait_condition, wait_for_fast
+from ..issues import Unavailable
 
 
 @pytest.fixture()
@@ -167,6 +169,21 @@ def batch_count():
         )))
         await wait_condition(lambda: batch_count() > initial_batches)
 
+    async def test_convert_errors_to_ydb(self, stream, stream_reader):
+        class TestError(grpc.RpcError):
+            _code: grpc.StatusCode
+
+            def __init__(self, code: grpc.StatusCode):
+                self._code = code
+
+            def code(self):
+                return self._code
+
+        stream.from_server.put_nowait(TestError(grpc.StatusCode.UNAVAILABLE))
+
+        with pytest.raises(Unavailable):
+            await wait_for_fast(stream_reader.wait_messages())
+
     async def test_init_reader(self, stream, default_reader_settings):
         reader = ReaderStream(default_reader_settings)
         init_message = StreamReadMessage.InitRequest(
@@ -563,3 +580,9 @@ async def test_receive_batch_nowait(self, stream, stream_reader, partition_sessi
 
         with pytest.raises(asyncio.QueueEmpty):
             stream.from_client.get_nowait()
+
+@pytest.mark.asyncio
+class TestReaderReconnector:
+    async def test_start(self):
+        pass
+

ydb/_topic_wrapper/common.py

@@ -238,3 +238,13 @@ def from_proto(msg: ydb_topic_pb2.UpdateTokenResponse) -> typing.Any:
 
 
 TokenGetterFuncType = typing.Optional[typing.Callable[[], str]]
+
+
+def callback_from_asyncio(callback: typing.Union[typing.Callable, typing.Coroutine]) -> [asyncio.Future, asyncio.Task]:
+    loop = asyncio.get_running_loop()
+
+    if asyncio.iscoroutinefunction(callback):
+        return loop.create_task(callback())
+    else:
+        return loop.run_in_executor(None, callback)
+

ydb/_topic_wrapper/common_test.py

@@ -0,0 +1,36 @@
+import asyncio
+
+import pytest
+
+from .common import callback_from_asyncio
+
+
+@pytest.mark.asyncio
+class Test:
+    async def test_callback_from_asyncio(self):
+        class TestError(Exception):
+            pass
+
+        def sync_success():
+            return 1
+
+        assert await callback_from_asyncio(sync_success) == 1
+
+        def sync_failed():
+            raise TestError()
+
+        with pytest.raises(TestError):
+            await callback_from_asyncio(sync_failed)
+
+        async def async_success():
+            await asyncio.sleep(0)
+            return 1
+
+        assert await callback_from_asyncio(async_success) == 1
+
+        async def async_failed():
+            await asyncio.sleep(0)
+            raise TestError()
+
+        with pytest.raises(TestError):
+            await callback_from_asyncio(async_failed)