Revert "Refactor wrappers to split UnaryStream and StreamStream wrappers"

This reverts commit 0ffe545.

Author: vgvoleg

ydb/_grpc/grpcwrapper/common_utils.py

@@ -145,7 +145,7 @@ def close(self):
 SupportedDriverType = Union[ydb.Driver, ydb.aio.Driver]
 
 
-class AbstractGrpcWrapperAsyncIO(IGrpcWrapperAsyncIO, abc.ABC):
+class GrpcWrapperAsyncIO(IGrpcWrapperAsyncIO):
     from_client_grpc: asyncio.Queue
     from_server_grpc: AsyncIterator
     convert_server_grpc_to_wrapper: Callable[[Any], Any]
@@ -163,6 +163,13 @@ def __init__(self, convert_server_grpc_to_wrapper):
     def __del__(self):
         self._clean_executor(wait=False)
 
+    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"
+
     def close(self):
         self.from_client_grpc.put_nowait(_stop_grpc_connection_marker)
         if self._stream_call:
@@ -174,35 +181,6 @@ def _clean_executor(self, wait: bool):
         if self._wait_executor:
             self._wait_executor.shutdown(wait)
 
-    async def receive(self) -> Any:
-        # todo handle grpc exceptions and convert it to internal exceptions
-        try:
-            grpc_message = await self.from_server_grpc.__anext__()
-        except (grpc.RpcError, grpc.aio.AioRpcError) 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()
-        # print("rekby, grpc, send", grpc_message)
-        self.from_client_grpc.put_nowait(grpc_message)
-
-
-class GrpcWrapperStreamStreamAsyncIO(AbstractGrpcWrapperAsyncIO):
-    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)
         stream_call = await driver(
@@ -221,30 +199,25 @@ async def _start_sync_driver(self, driver: ydb.Driver, stub, method):
         self._stream_call = stream_call
         self.from_server_grpc = SyncToAsyncIterator(stream_call.__iter__(), self._wait_executor)
 
+    async def receive(self) -> Any:
+        # todo handle grpc exceptions and convert it to internal exceptions
+        try:
+            grpc_message = await self.from_server_grpc.__anext__()
+        except (grpc.RpcError, grpc.aio.AioRpcError) as e:
+            raise connection._rpc_error_handler(self._connection_state, e)
 
-class GrpcWrapperUnaryStreamAsyncIO(AbstractGrpcWrapperAsyncIO):
-    async def start(self, driver: SupportedDriverType, request, stub, method):
-        if asyncio.iscoroutinefunction(driver.__call__):
-            await self._start_asyncio_driver(driver, request, stub, method)
-        else:
-            await self._start_sync_driver(driver, request, stub, method)
-        self._connection_state = "started"
+        issues._process_response(grpc_message)
 
-    async def _start_asyncio_driver(self, driver: ydb.aio.Driver, request, stub, method):
-        stream_call = await driver(
-            request,
-            stub,
-            method,
-        )
-        self._stream_call = stream_call
-        self.from_server_grpc = stream_call.__aiter__()
+        if self._connection_state != "has_received_messages":
+            self._connection_state = "has_received_messages"
 
-    async def _start_sync_driver(self, driver: ydb.Driver, request, stub, method):
-        self._wait_executor = concurrent.futures.ThreadPoolExecutor(max_workers=1)
+        # print("rekby, grpc, received", grpc_message)
+        return self.convert_server_grpc_to_wrapper(grpc_message)
 
-        stream_call = await to_thread(driver, request, stub, method, executor=self._wait_executor)
-        self._stream_call = stream_call
-        self.from_server_grpc = SyncToAsyncIterator(stream_call.__iter__(), self._wait_executor)
+    def write(self, wrap_message: IToProto):
+        grpc_message = wrap_message.to_proto()
+        # print("rekby, grpc, send", grpc_message)
+        self.from_client_grpc.put_nowait(grpc_message)
 
 
 @dataclass(init=False)
@@ -283,19 +256,6 @@ def issue_to_str(cls, issue: ydb_issue_message_pb2.IssueMessage):
         return res
 
 
-ResultType = typing.TypeVar("ResultType", bound=IFromProtoWithProtoType)
-
-
-def create_result_wrapper(
-    result_type: typing.Type[ResultType],
-) -> typing.Callable[[typing.Any, typing.Any, typing.Any], ResultType]:
-    def wrapper(rpc_state, response_pb, driver=None):
-        # issues._process_response(response_pb.operation)
-        return result_type.from_proto(response_pb)
-
-    return wrapper
-
-
 def callback_from_asyncio(callback: Union[Callable, Coroutine]) -> [asyncio.Future, asyncio.Task]:
     loop = asyncio.get_running_loop()
 

ydb/_topic_common/common_test.py

@@ -8,7 +8,7 @@
 
 from .common import CallFromSyncToAsync
 from .._grpc.grpcwrapper.common_utils import (
-    GrpcWrapperStreamStreamAsyncIO,
+    GrpcWrapperAsyncIO,
     ServerStatus,
     callback_from_asyncio,
 )
@@ -77,7 +77,7 @@ async def async_failed():
 
 
 @pytest.mark.asyncio
-class TestGrpcWrapperStreamStreamAsyncIO:
+class TestGrpcWrapperAsyncIO:
     async def test_convert_grpc_errors_to_ydb(self):
         class TestError(grpc.RpcError, grpc.Call):
             def __init__(self):
@@ -93,7 +93,7 @@ class FromServerMock:
             async def __anext__(self):
                 raise TestError()
 
-        wrapper = GrpcWrapperStreamStreamAsyncIO(lambda: None)
+        wrapper = GrpcWrapperAsyncIO(lambda: None)
         wrapper.from_server_grpc = FromServerMock()
 
         with pytest.raises(issues.Unauthenticated):
@@ -107,7 +107,7 @@ async def __anext__(self):
                     issues=[],
                 )
 
-        wrapper = GrpcWrapperStreamStreamAsyncIO(lambda: None)
+        wrapper = GrpcWrapperAsyncIO(lambda: None)
         wrapper.from_server_grpc = FromServerMock()
 
         with pytest.raises(issues.Overloaded):

ydb/_topic_reader/topic_reader_asyncio.py

@@ -18,7 +18,7 @@
 from .._grpc.grpcwrapper.common_utils import (
     IGrpcWrapperAsyncIO,
     SupportedDriverType,
-    GrpcWrapperStreamStreamAsyncIO,
+    GrpcWrapperAsyncIO,
 )
 from .._grpc.grpcwrapper.ydb_topic import (
     StreamReadMessage,
@@ -308,7 +308,7 @@ async def create(
         driver: SupportedDriverType,
         settings: topic_reader.PublicReaderSettings,
     ) -> "ReaderStream":
-        stream = GrpcWrapperStreamStreamAsyncIO(StreamReadMessage.FromServer.from_proto)
+        stream = GrpcWrapperAsyncIO(StreamReadMessage.FromServer.from_proto)
 
         await stream.start(driver, _apis.TopicService.Stub, _apis.TopicService.StreamRead)
 

ydb/_topic_writer/topic_writer_asyncio.py

@@ -39,7 +39,7 @@
 from .._grpc.grpcwrapper.common_utils import (
     IGrpcWrapperAsyncIO,
     SupportedDriverType,
-    GrpcWrapperStreamStreamAsyncIO,
+    GrpcWrapperAsyncIO,
 )
 
 logger = logging.getLogger(__name__)
@@ -613,7 +613,7 @@ async def create(
         init_request: StreamWriteMessage.InitRequest,
         update_token_interval: Optional[Union[int, float]] = None,
     ) -> "WriterAsyncIOStream":
-        stream = GrpcWrapperStreamStreamAsyncIO(StreamWriteMessage.FromServer.from_proto)
+        stream = GrpcWrapperAsyncIO(StreamWriteMessage.FromServer.from_proto)
 
         await stream.start(driver, _apis.TopicService.Stub, _apis.TopicService.StreamWrite)