Merge pull request #199 from azryve/grpcio-service-reflection

feature: GRPCService reflection flag
aiokitchen · Mar 7, 2024 · 12bfae8 · 12bfae8
2 parents 22ae4a0 + f68e4e7
commit 12bfae8
Show file tree

Hide file tree

Showing 4 changed files with 721 additions and 641 deletions.
diff --git a/aiomisc/service/grpc_server.py b/aiomisc/service/grpc_server.py
@@ -11,6 +11,7 @@
 
 try:
     import grpc.aio
+    from grpc_reflection.v1alpha import reflection
 except ImportError as e:
     raise ImportError(
         "You must install 'grpcio' manually or using extras 'aiomisc[grpc]'",
@@ -39,11 +40,12 @@ class GRPCService(Service):
     def __init__(
         self, *,
         migration_thread_pool: Optional[Executor] = None,
-        handlers: Optional[Sequence[grpc.GenericRpcHandler]] = None,
+        handlers: Optional[Sequence[grpc.ServiceRpcHandler]] = None,
         interceptors: Optional[Sequence[Any]] = None,
         options: Optional[Sequence[Tuple[str, Any]]] = None,
         maximum_concurrent_rpcs: Optional[int] = None,
         compression: Optional[grpc.Compression] = None,
+        reflection: bool = False,
         **kwds: Any,
     ):
         self._server_args = MappingProxyType({
@@ -54,9 +56,10 @@ def __init__(
             "migration_thread_pool": migration_thread_pool,
             "options": options,
         })
-        self._services: Set[grpc.GenericRpcHandler] = set()
+        self._services: Set[grpc.ServiceRpcHandler] = set()
         self._insecure_ports = set()
         self._secure_ports = set()
+        self._reflection = reflection
         super().__init__(**kwds)
 
     @classmethod
@@ -82,14 +85,19 @@ async def start(self) -> None:
             future.set_result(port)
             self._log_port("Listening secure address", address, port)
 
+        if self._reflection:
+            service_names = [x.service_name() for x in self._services]
+            service_names.append(reflection.SERVICE_NAME)
+            reflection.enable_server_reflection(service_names, self._server)
+
         self._server.add_generic_rpc_handlers(tuple(self._services))
         await self._server.start()
 
     async def stop(self, exception: Optional[Exception] = None) -> None:
         await self._server.stop(self.GRACEFUL_STOP_TIME)
 
     def add_generic_rpc_handlers(
-        self, generic_rpc_handlers: Sequence[grpc.GenericRpcHandler],
+        self, generic_rpc_handlers: Sequence[grpc.ServiceRpcHandler],
     ) -> None:
         for service in generic_rpc_handlers:
             self._services.add(service)

diff --git a/docs/source/services.rst b/docs/source/services.rst
@@ -800,6 +800,13 @@ Service initialization example:
         main()
 
 
+To enable reflection for the service you use reflection flag:
+
+.. code-block:: python
+
+    GRPCService(reflection=True)
+
+
 .. _memory-tracer:
 
 Memory Tracer