diff --git a/sea/server/multiprocessing.py b/sea/server/multiprocessing.py index 623cd62..7b491b3 100644 --- a/sea/server/multiprocessing.py +++ b/sea/server/multiprocessing.py @@ -9,6 +9,7 @@ from typing import List import grpc +from grpc_reflection.v1alpha import reflection from sea import signals @@ -45,9 +46,9 @@ def _run_server(self, bind_address): ], ) self.server = server # set server in slave process - if self.app.config.get("GRPC_REFLECTION", True): - from grpc_reflection.v1alpha import reflection - reflection.enable_server_reflection((reflection.SERVICE_NAME,), server) + # register reflection service + if self.app.config.get("GRPC_REFLECTION_SERVICES"): + reflection.enable_server_reflection((reflection.SERVICE_NAME, *self.app.config["GRPC_REFLECTION_SERVICES"]), self.server) for _, (add_func, servicer) in self.app.servicers.items(): add_func(servicer(), server) diff --git a/sea/server/threading.py b/sea/server/threading.py index bfeea43..4fa29b5 100644 --- a/sea/server/threading.py +++ b/sea/server/threading.py @@ -3,6 +3,7 @@ from concurrent import futures import grpc +from grpc_reflection.v1alpha import reflection from sea import signals @@ -28,9 +29,9 @@ def run(self): from prometheus_client import start_http_server start_http_server(self.app.config["PROMETHEUS_PORT"]) - if self.app.config.get("GRPC_REFLECTION", True): - from grpc_reflection.v1alpha import reflection - reflection.enable_server_reflection((reflection.SERVICE_NAME,), self.server) + # register reflection service + if self.app.config.get("GRPC_REFLECTION_SERVICES"): + reflection.enable_server_reflection((reflection.SERVICE_NAME, *self.app.config["GRPC_REFLECTION_SERVICES"]), self.server) # run grpc server for name, (add_func, servicer) in self.app.servicers.items(): add_func(servicer(), self.server)