support ServerReflection with a custom DescriptorPool

grpclib/reflection/_deprecated.py

@@ -14,7 +14,7 @@
 #     See the License for the specific language governing permissions and
 #     limitations under the License.
 #
-from typing import Collection
+from typing import Any, Collection, Optional
 
 from google.protobuf.descriptor import FileDescriptor
 from google.protobuf.descriptor_pb2 import FileDescriptorProto
@@ -35,10 +35,14 @@ class ServerReflection(ServerReflectionBase):
     """
     Implements server reflection protocol.
     """
-    def __init__(self, *, _service_names: Collection[str]):
+    def __init__(
+        self, *,
+        _service_names: Collection[str],
+        _pool: Optional[Any] = None
+    ):
         self._service_names = _service_names
         # FIXME: DescriptorPool has incomplete typings
-        self._pool = Default()  # type: ignore
+        self._pool = _pool or Default()  # type: ignore
 
     def _not_found_response(self) -> ServerReflectionResponse:
         return ServerReflectionResponse(

grpclib/reflection/service.py

@@ -14,7 +14,7 @@
 #     See the License for the specific language governing permissions and
 #     limitations under the License.
 #
-from typing import TYPE_CHECKING, Collection, List
+from typing import TYPE_CHECKING, Any, Collection, List, Optional
 
 from google.protobuf.descriptor import FileDescriptor
 from google.protobuf.descriptor_pb2 import FileDescriptorProto
@@ -41,10 +41,14 @@ class ServerReflection(ServerReflectionBase):
     """
     Implements server reflection protocol.
     """
-    def __init__(self, *, _service_names: Collection[str]):
+    def __init__(
+        self, *,
+        _service_names: Collection[str],
+        _pool: Optional[Any] = None
+    ):
         self._service_names = _service_names
         # FIXME: DescriptorPool has incomplete typings
-        self._pool = Default()  # type: ignore
+        self._pool = _pool or Default()  # type: ignore
 
     def _not_found_response(self) -> ServerReflectionResponse:
         return ServerReflectionResponse(
@@ -161,7 +165,11 @@ async def ServerReflectionInfo(
             await stream.send_message(response)
 
     @classmethod
-    def extend(cls, services: 'Collection[IServable]') -> 'List[IServable]':
+    def extend(
+        cls, services: 'Collection[IServable]',
+        *,
+        pool: Optional[Any] = None
+    ) -> 'List[IServable]':
         """
         Extends services list with reflection service:
 
@@ -181,6 +189,7 @@ def extend(cls, services: 'Collection[IServable]') -> 'List[IServable]':
         for service in services:
             service_names.append(_service_name(service))
         services = list(services)
-        services.append(cls(_service_names=service_names))
-        services.append(_ServerReflectionV1Alpha(_service_names=service_names))
+        services.append(cls(_service_names=service_names, _pool=pool))
+        services.append(
+            _ServerReflectionV1Alpha(_service_names=service_names, _pool=pool))
         return services

tests/test_reflection.py

@@ -1,5 +1,7 @@
 import socket
 
+from google.protobuf.descriptor_pool import DescriptorPool
+
 import pytest
 import pytest_asyncio
 
@@ -108,3 +110,45 @@ async def test_list_services_response(channel):
 
     service, = r1.list_services_response.service
     assert service.name == DESCRIPTOR.services_by_name['DummyService'].full_name
+
+
+@pytest.mark.asyncio
+async def test_file_containing_symbol_response_custom_pool(port):
+    my_pool = DescriptorPool()
+    services = [DummyService()]
+    services = ServerReflection.extend(services, pool=my_pool)
+
+    server = Server(services)
+    await server.start(port=port)
+
+    channel = Channel(port=port)
+    try:
+        # because we use our own pool (my_pool), there's no descriptors to find.
+        req = ServerReflectionRequest(
+            file_containing_symbol=(
+                DESCRIPTOR.message_types_by_name['DummyRequest'].full_name
+            ),
+        )
+        resp, = await ServerReflectionStub(channel).ServerReflectionInfo([req])
+
+        assert resp == ServerReflectionResponse(
+            error_response=ErrorResponse(
+                error_code=5,
+                error_message='not found',
+            ),
+        )
+
+        # once we update the pool, we should find the descriptor.
+        my_pool.AddSerializedFile(DESCRIPTOR.serialized_pb)
+
+        resp, = await ServerReflectionStub(channel).ServerReflectionInfo([req])
+
+        proto_bytes, = resp.file_descriptor_response.file_descriptor_proto
+        dummy_proto = FileDescriptorProto()
+        dummy_proto.ParseFromString(proto_bytes)
+        assert dummy_proto.name == DESCRIPTOR.name
+        assert dummy_proto.package == DESCRIPTOR.package
+    finally:
+        channel.close()
+        server.close()
+        await server.wait_closed()