parse stop partition message

handle unexpected message for topic reader

Author: rekby

ydb/_grpc/grpcwrapper/ydb_topic.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import datetime
 import enum
 import typing
@@ -642,11 +644,20 @@ def to_proto(
             return res
 
     @dataclass
-    class StopPartitionSessionRequest:
+    class StopPartitionSessionRequest(IFromProto):
         partition_session_id: int
         graceful: bool
         committed_offset: int
 
+        @staticmethod
+        def from_proto(msg: ydb_topic_pb2.StreamReadMessage.StopPartitionSessionRequest) -> StreamReadMessage.StopPartitionSessionRequest:
+            return StreamReadMessage.StopPartitionSessionRequest(
+                partition_session_id=msg.partition_session_id,
+                graceful=msg.graceful,
+                committed_offset=msg.committed_offset,
+            )
+
+
     @dataclass
     class StopPartitionSessionResponse:
         partition_session_id: int
@@ -707,6 +718,13 @@ def from_proto(
                         msg.start_partition_session_request,
                     ),
                 )
+            elif mess_type == "stop_partition_session_request":
+                return StreamReadMessage.FromServer(
+                    server_status=server_status,
+                    server_message=StreamReadMessage.StopPartitionSessionRequest.from_proto(
+                        msg.stop_partition_session_request
+                    )
+                )
             elif mess_type == "update_token_response":
                 return StreamReadMessage.FromServer(
                     server_status=server_status,
@@ -720,9 +738,6 @@ def from_proto(
                     )
                 )
 
-            # todo replace exception to log
-            raise NotImplementedError()
-
 
 ReaderMessagesFromClientToServer = Union[
     StreamReadMessage.InitRequest,

ydb/_topic_reader/topic_reader.py

@@ -1,6 +1,7 @@
 import concurrent.futures
 import enum
 import datetime
+import logging
 from dataclasses import dataclass
 from typing import (
     Union,
@@ -52,6 +53,7 @@ class PublicReaderSettings:
     # decoder_executor, must be set for handle non raw messages
     decoder_executor: Optional[concurrent.futures.Executor] = None
     update_token_interval: Union[int, float] = 3600
+    logger: Optional[logging.Logger] = None
 
     def __post_init__(self):
         # check possible create init message

ydb/_topic_reader/topic_reader_asyncio.py

@@ -2,6 +2,7 @@
 
 import asyncio
 import concurrent.futures
+import copy
 import gzip
 import typing
 from asyncio import Task
@@ -26,7 +27,9 @@
     Codec,
 )
 from .._errors import check_retriable_error
+import logging
 
+_module_logger = logging.getLogger(__name__)
 
 class TopicReaderError(YdbError):
     pass
@@ -62,6 +65,7 @@ class PublicAsyncIOReader:
     _closed: bool
     _reconnector: ReaderReconnector
     _parent: typing.Any  # need for prevent close parent client by GC
+    _logger: logging.Logger
 
     def __init__(
         self,
@@ -70,8 +74,14 @@ def __init__(
         *,
         _parent=None,
     ):
+        if settings.logger:
+            self._logger = settings.logger
+        else:
+            self._logger = _module_logger
+
         self._loop = asyncio.get_running_loop()
         self._closed = False
+        self._logger = settings.logger
         self._reconnector = ReaderReconnector(driver, settings)
         self._parent = _parent
 
@@ -139,14 +149,19 @@ class ReaderReconnector:
     _settings: topic_reader.PublicReaderSettings
     _driver: Driver
     _background_tasks: Set[Task]
+    _logger: logging.Logger
 
     _state_changed: asyncio.Event
     _stream_reader: Optional["ReaderStream"]
     _first_error: asyncio.Future[YdbError]
 
     def __init__(self, driver: Driver, settings: topic_reader.PublicReaderSettings):
-        self._id = self._static_reader_reconnector_counter.inc_and_get()
+        if settings.logger:
+            self._logger = settings.logger
+        else:
+            self._logger = _module_logger
 
+        self._id = self._static_reader_reconnector_counter.inc_and_get()
         self._settings = settings
         self._driver = driver
         self._background_tasks = set()
@@ -234,6 +249,7 @@ class ReaderStream:
     _buffer_size_bytes: int  # use for init request, then for debug purposes only
     _decode_executor: concurrent.futures.Executor
     _decoders: Dict[int, typing.Callable[[bytes], bytes]]  # dict[codec_code] func(encoded_bytes)->decoded_bytes
+    _logger: logging.Logger
 
     if typing.TYPE_CHECKING:
         _batches_to_decode: asyncio.Queue[datatypes.PublicBatch]
@@ -255,6 +271,11 @@ def __init__(
         settings: topic_reader.PublicReaderSettings,
         get_token_function: Optional[Callable[[], str]] = None,
     ):
+        if settings.logger:
+            self._logger = settings.logger
+        else:
+            self._logger = _module_logger
+
         self._loop = asyncio.get_running_loop()
         self._id = ReaderStream._static_id_counter.inc_and_get()
         self._reader_reconnector_id = reader_reconnector_id
@@ -395,34 +416,37 @@ async def _read_messages_loop(self):
                 )
             )
             while True:
-                message = await self._stream.receive()  # type: StreamReadMessage.FromServer
-                _process_response(message.server_status)
-
-                if isinstance(message.server_message, StreamReadMessage.ReadResponse):
-                    self._on_read_response(message.server_message)
-
-                elif isinstance(message.server_message, StreamReadMessage.CommitOffsetResponse):
-                    self._on_commit_response(message.server_message)
-
-                elif isinstance(
-                    message.server_message,
-                    StreamReadMessage.StartPartitionSessionRequest,
-                ):
-                    self._on_start_partition_session(message.server_message)
-
-                elif isinstance(
-                    message.server_message,
-                    StreamReadMessage.StopPartitionSessionRequest,
-                ):
-                    self._on_partition_session_stop(message.server_message)
-
-                elif isinstance(message.server_message, UpdateTokenResponse):
-                    self._update_token_event.set()
-
-                else:
-                    raise NotImplementedError(
-                        "Unexpected type of StreamReadMessage.FromServer message: %s" % message.server_message
-                    )
+                try:
+                    message = await self._stream.receive()  # type: StreamReadMessage.FromServer
+                    _process_response(message.server_status)
+
+                    if isinstance(message.server_message, StreamReadMessage.ReadResponse):
+                        self._on_read_response(message.server_message)
+
+                    elif isinstance(message.server_message, StreamReadMessage.CommitOffsetResponse):
+                        self._on_commit_response(message.server_message)
+
+                    elif isinstance(
+                        message.server_message,
+                        StreamReadMessage.StartPartitionSessionRequest,
+                    ):
+                        self._on_start_partition_session(message.server_message)
+
+                    elif isinstance(
+                        message.server_message,
+                        StreamReadMessage.StopPartitionSessionRequest,
+                    ):
+                        self._on_partition_session_stop(message.server_message)
+
+                    elif isinstance(message.server_message, UpdateTokenResponse):
+                        self._update_token_event.set()
+
+                    else:
+                        raise issues.UnexpectedGrpcMessage(
+                            "Unexpected message in _read_messages_loop: %s" % type(message.server_message)
+                        )
+                except issues.UnexpectedGrpcMessage:
+                    self._logger.exception("unexpected message in stream reader")
 
                 self._state_changed.set()
         except Exception as e:

ydb/_topic_reader/topic_reader_asyncio_test.py

@@ -1127,6 +1127,30 @@ async def test_update_token(self, stream):
 
         await reader.close()
 
+    async def test_read_unknown_message(self, stream, stream_reader):
+        class TestMessage:
+            pass
+
+        logged = asyncio.Event()
+
+        def set_logged(*args, **kwargs):
+            logged.set()
+
+        stream_reader._logger = mock.Mock()
+        stream_reader._logger.exception = mock.Mock(side_effect=set_logged)
+
+        # noinspection PyTypeChecker
+        stream.from_server.put_nowait(StreamReadMessage.FromServer(
+            server_status=ServerStatus(
+                status=issues.StatusCode.SUCCESS,
+                issues=[],
+            ),
+            server_message=TestMessage(),
+        ))
+        await wait_for_fast(logged.wait())
+
+        stream_reader._logger.exception.assert_called_once()
+
 
 @pytest.mark.asyncio
 class TestReaderReconnector:

ydb/issues.py

@@ -156,6 +156,11 @@ class SessionPoolEmpty(Error, queue.Empty):
     status = StatusCode.SESSION_POOL_EMPTY
 
 
+class UnexpectedGrpcMessage(Error):
+    def __init__(self, message: str):
+        super().__init__(message)
+
+
 def _format_issues(issues):
     if not issues:
         return ""