handle grpc errors and errors in server status

Author: rekby

ydb/_topic_reader/topic_reader_asyncio.py

@@ -8,7 +8,7 @@
 
 import grpc
 
-from .. import _apis
+from .. import _apis, issues
 from ..aio import Driver
 from ..issues import Error as YdbError
 from .datatypes import PartitionSession, PublicMessage, PublicBatch
@@ -143,19 +143,19 @@ async def _start(self, stream: IGrpcWrapperAsyncIO, init_message: StreamReadMess
         self._background_tasks.add(read_messages_task)
 
     async def wait_messages(self):
-        if self._closed:
-            raise TopicReaderStreamClosedError()
-
-        while len(self._message_batches) == 0:
+        while True:
             if self._first_error is not None:
                 raise self._first_error
 
+            if len(self._message_batches) > 0:
+                return
+
             await self._state_changed.wait()
             self._state_changed.clear()
 
     def receive_batch_nowait(self):
-        if self._closed:
-            raise TopicReaderStreamClosedError()
+        if self._first_error is not None:
+            raise self._first_error
 
         try:
             batch = self._message_batches.popleft()
@@ -185,8 +185,6 @@ 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
@@ -296,7 +294,7 @@ async def close(self):
             raise TopicReaderError(message="Double closed ReaderStream")
 
         self._closed = True
-        self._set_first_error(TopicReaderError("Reader closed"))
+        self._set_first_error(TopicReaderStreamClosedError())
         self._state_changed.set()
 
         for task in self._background_tasks:

ydb/_topic_reader/topic_reader_asyncio_test.py

@@ -5,15 +5,23 @@
 import grpc
 import pytest
 
+import ydb
 from ydb import aio
 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.common import OffsetsRange, Codec, ServerStatus
 from .._topic_wrapper.reader import StreamReadMessage
 from .._topic_wrapper.test_helpers import StreamMock, wait_condition, wait_for_fast
 from ..issues import Unavailable
 
+# Workaround for good autocomplete in IDE and universal import at runtime
+# noinspection PyUnreachableCode
+if False:
+    from .._grpc.v4.protos import ydb_status_codes_pb2
+else:
+    from .._grpc.common.protos import ydb_status_codes_pb2
+
 
 @pytest.fixture()
 def default_reader_settings():
@@ -58,7 +66,7 @@ def second_partition_session(self, default_reader_settings):
         )
 
     @pytest.fixture()
-    async def stream_reader(self, stream, default_reader_settings, partition_session,
+    async def stream_reader_started(self, stream, default_reader_settings, partition_session,
                             second_partition_session) -> ReaderStream:
         reader = ReaderStream(default_reader_settings)
         init_message = object()
@@ -67,7 +75,8 @@ async def stream_reader(self, stream, default_reader_settings, partition_session
         start = asyncio.create_task(reader._start(stream, init_message))
 
         stream.from_server.put_nowait(StreamReadMessage.FromServer(
-            StreamReadMessage.InitResponse(session_id="test-session")
+            server_status=ServerStatus(ydb_status_codes_pb2.StatusIds.SUCCESS, []),
+            server_message=StreamReadMessage.InitResponse(session_id="test-session"),
         ))
 
         init_request = await wait_for_fast(stream.from_client.get())
@@ -77,7 +86,9 @@ async def stream_reader(self, stream, default_reader_settings, partition_session
         assert isinstance(read_request.client_message, StreamReadMessage.ReadRequest)
 
         stream.from_server.put_nowait(
-            StreamReadMessage.FromServer(server_message=StreamReadMessage.StartPartitionSessionRequest(
+            StreamReadMessage.FromServer(
+                server_status=ServerStatus(ydb_status_codes_pb2.StatusIds.SUCCESS, []),
+                server_message=StreamReadMessage.StartPartitionSessionRequest(
                 partition_session=StreamReadMessage.PartitionSession(
                     partition_session_id=partition_session.id,
                     path=partition_session.topic_path,
@@ -96,7 +107,9 @@ async def stream_reader(self, stream, default_reader_settings, partition_session
         assert isinstance(start_partition_resp.client_message, StreamReadMessage.StartPartitionSessionResponse)
 
         stream.from_server.put_nowait(
-            StreamReadMessage.FromServer(server_message=StreamReadMessage.StartPartitionSessionRequest(
+            StreamReadMessage.FromServer(
+                server_status=ServerStatus(ydb_status_codes_pb2.StatusIds.SUCCESS, []),
+                server_message=StreamReadMessage.StartPartitionSessionRequest(
                 partition_session=StreamReadMessage.PartitionSession(
                     partition_session_id=second_partition_session.id,
                     path=second_partition_session.topic_path,
@@ -116,11 +129,22 @@ async def stream_reader(self, stream, default_reader_settings, partition_session
         with pytest.raises(asyncio.QueueEmpty):
             stream.from_client.get_nowait()
 
-        yield reader
+        return reader
 
-        assert reader._first_error is None
+    @pytest.fixture()
+    async def stream_reader(self, stream_reader_started: ReaderStream):
+        yield stream_reader_started
+
+        assert stream_reader_started._first_error is None
+        await stream_reader_started.close()
+
+    @pytest.fixture()
+    async def stream_reader_finish_with_error(self, stream_reader_started: ReaderStream):
+        yield stream_reader_started
+
+        assert stream_reader_started._first_error is not None
+        await stream_reader_started.close()
 
-        await reader.close()
 
     @staticmethod
     def create_message(partition_session: PartitionSession, seqno: int):
@@ -143,7 +167,9 @@ def batch_count():
         initial_batches = batch_count()
 
         stream = stream_reader._stream  # type: StreamMock
-        stream.from_server.put_nowait(StreamReadMessage.FromServer(server_message=StreamReadMessage.ReadResponse(
+        stream.from_server.put_nowait(StreamReadMessage.FromServer(
+            server_status=ServerStatus(ydb_status_codes_pb2.StatusIds.SUCCESS, []),
+            server_message=StreamReadMessage.ReadResponse(
             partition_data=[StreamReadMessage.ReadResponse.PartitionData(
                 partition_session_id=message._partition_session.id,
                 batches=[
@@ -169,20 +195,25 @@ 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
+    async def test_first_error(self, stream, stream_reader_finish_with_error):
+        class TestError(grpc.RpcError, grpc.Call):
+            def __init__(self):
+                pass
 
             def code(self):
-                return self._code
+                return grpc.StatusCode.UNAUTHENTICATED
 
-        stream.from_server.put_nowait(TestError(grpc.StatusCode.UNAVAILABLE))
+            def details(self):
+                return "test error"
 
-        with pytest.raises(Unavailable):
-            await wait_for_fast(stream_reader.wait_messages())
+        test_err = TestError()
+        stream.from_server.put_nowait(test_err)
+
+        with pytest.raises(TestError):
+            await wait_for_fast(stream_reader_finish_with_error.wait_messages())
+
+        with pytest.raises(TestError):
+            stream_reader_finish_with_error.receive_batch_nowait()
 
     async def test_init_reader(self, stream, default_reader_settings):
         reader = ReaderStream(default_reader_settings)
@@ -202,6 +233,7 @@ async def test_init_reader(self, stream, default_reader_settings):
         assert sent_message == expected_sent_init_message
 
         stream.from_server.put_nowait(StreamReadMessage.FromServer(
+            server_status=ServerStatus(ydb_status_codes_pb2.StatusIds.SUCCESS, []),
             server_message=StreamReadMessage.InitResponse(session_id="test"))
         )
 
@@ -231,6 +263,7 @@ def session_count():
         test_topic_path = default_reader_settings.topic + "-asd"
 
         stream.from_server.put_nowait(StreamReadMessage.FromServer(
+            server_status=ServerStatus(ydb_status_codes_pb2.StatusIds.SUCCESS, []),
             server_message=StreamReadMessage.StartPartitionSessionRequest(
                 partition_session=StreamReadMessage.PartitionSession(
                     partition_session_id=test_partition_session_id,
@@ -266,6 +299,7 @@ def session_count():
         initial_session_count = session_count()
 
         stream.from_server.put_nowait(StreamReadMessage.FromServer(
+            server_status=ServerStatus(ydb_status_codes_pb2.StatusIds.SUCCESS, []),
             server_message=StreamReadMessage.StopPartitionSessionRequest(
                 partition_session_id=partition_session.id,
                 graceful=False,
@@ -287,6 +321,7 @@ def session_count():
         initial_session_count = session_count()
 
         stream.from_server.put_nowait(StreamReadMessage.FromServer(
+            server_status=ServerStatus(ydb_status_codes_pb2.StatusIds.SUCCESS, []),
             server_message=StreamReadMessage.StopPartitionSessionRequest(
                 partition_session_id=partition_session.id,
                 graceful=True,
@@ -303,6 +338,7 @@ def session_count():
         )
 
         stream.from_server.put_nowait(StreamReadMessage.FromServer(
+            server_status=ServerStatus(ydb_status_codes_pb2.StatusIds.SUCCESS, []),
             server_message=StreamReadMessage.StopPartitionSessionRequest(
                 partition_session_id=partition_session.id,
                 graceful=False,
@@ -330,7 +366,9 @@ def reader_batch_count():
         session_meta = {"a": "b"}
         message_group_id = "test-message-group-id"
 
-        stream.from_server.put_nowait(StreamReadMessage.FromServer(server_message=StreamReadMessage.ReadResponse(
+        stream.from_server.put_nowait(StreamReadMessage.FromServer(
+            server_status=ServerStatus(ydb_status_codes_pb2.StatusIds.SUCCESS, []),
+            server_message=StreamReadMessage.ReadResponse(
             bytes_size=bytes_size,
             partition_data=[
                 StreamReadMessage.ReadResponse.PartitionData(

ydb/_topic_wrapper/common.py

@@ -4,11 +4,15 @@
 from dataclasses import dataclass
 from enum import IntEnum
 
+import grpc
 from google.protobuf.message import Message
 
 import ydb.aio
 
+from .. import issues, connection
+
 # Workaround for good autocomplete in IDE and universal import at runtime
+# noinspection PyUnreachableCode
 if False:
     from ydb._grpc.v4.protos import (
         ydb_status_codes_pb2,
@@ -147,16 +151,19 @@ class GrpcWrapperAsyncIO(IGrpcWrapperAsyncIO):
     from_client_grpc: asyncio.Queue
     from_server_grpc: typing.AsyncIterator
     convert_server_grpc_to_wrapper: typing.Callable[[typing.Any], typing.Any]
+    _connection_state: str
 
     def __init__(self, convert_server_grpc_to_wrapper):
         self.from_client_grpc = asyncio.Queue()
         self.convert_server_grpc_to_wrapper = convert_server_grpc_to_wrapper
+        self._connection_state = "new"
 
     async def start(self, driver: SupportedDriverType, stub, method):
         if asyncio.iscoroutinefunction(driver.__call__):
             await self._start_asyncio_driver(driver, stub, method)
         else:
             await self._start_sync_driver(driver, stub, method)
+        self._connection_state = "started"
 
     async def _start_asyncio_driver(self, driver: ydb.aio.Driver, stub, method):
         requests_iterator = QueueToIteratorAsyncIO(self.from_client_grpc)
@@ -179,37 +186,49 @@ 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_message = await self.from_server_grpc.__anext__()
+        try:
+            grpc_message = await self.from_server_grpc.__anext__()
+        except grpc.RpcError as e:
+            raise connection._rpc_error_handler(self._connection_state, e)
+
+        issues._process_response(grpc_message)
+
+        if self._connection_state != "has_received_messages":
+            self._connection_state = "has_received_messages"
+
         # print("rekby, grpc, received", grpc_message)
         return self.convert_server_grpc_to_wrapper(grpc_message)
 
     def write(self, wrap_message: IToProto):
-        grpc_message=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)
 class ServerStatus(IFromProto):
-    __slots__ = ("status", "_issues")
+    __slots__ = ("_grpc_status_code", "_issues")
 
     def __init__(
-        self,
-        status: ydb_status_codes_pb2.StatusIds.StatusCode,
-        issues: typing.Iterable[ydb_issue_message_pb2.IssueMessage],
+            self,
+            status_code: ydb_status_codes_pb2.StatusIds.StatusCode,
+            grpc_issues: typing.Iterable[ydb_issue_message_pb2.IssueMessage],
     ):
-        self.status = status
-        self._issues = issues
+        self._grpc_status_code = status_code
+        self._issues = grpc_issues
 
     def __str__(self):
         return self.__repr__()
 
     @staticmethod
-    def from_proto(msg: Message) -> "ServerStatus":
-        return ServerStatus(msg.status)
+    def from_proto(msg: typing.Union[
+        ydb_topic_pb2.StreamReadMessage.FromServer,
+        ydb_topic_pb2.StreamWriteMessage.FromServer,
+    ]) -> "ServerStatus":
+        return ServerStatus(msg.status, msg.issues)
 
     def is_success(self) -> bool:
-        return self.status == ydb_status_codes_pb2.StatusIds.SUCCESS
+        return self._grpc_status_code == ydb_status_codes_pb2.StatusIds.SUCCESS
 
     @classmethod
     def issue_to_str(cls, issue: ydb_issue_message_pb2.IssueMessage):
@@ -248,3 +267,6 @@ def callback_from_asyncio(callback: typing.Union[typing.Callable, typing.Corouti
     else:
         return loop.run_in_executor(None, callback)
 
+
+def ensure_success_or_raise_error(server_status: ServerStatus):
+    error = issues._process_response(server_status._grpc_status_code, server_status._issues)