Merge pull request #10216 from mehrdada/python-service-reflection

Python Service Reflection

src/python/grpcio_reflection/grpc_reflection/v1alpha/reflection.py

@@ -35,6 +35,7 @@ from google.protobuf import descriptor_pb2
 from google.protobuf import descriptor_pool
 
 from grpc_reflection.v1alpha import reflection_pb2
+from grpc_reflection.v1alpha import reflection_pb2_grpc
 
 _POOL = descriptor_pool.Default()
 
@@ -64,7 +65,7 @@ class ReflectionServicer(reflection_pb2.ServerReflectionServicer):
     Args:
       service_names: Iterable of fully-qualified service names available.
     """
-        self._service_names = list(service_names)
+        self._service_names = tuple(sorted(service_names))
         self._pool = _POOL if pool is None else pool
 
     def _file_by_filename(self, filename):
@@ -84,23 +85,32 @@ class ReflectionServicer(reflection_pb2.ServerReflectionServicer):
         else:
             return _file_descriptor_response(descriptor)
 
-    def _file_containing_extension(containing_type, extension_number):
-        # TODO(atash) Python protobuf currently doesn't support querying extensions.
-        # https://github.com/google/protobuf/issues/2248
-        return reflection_pb2.ServerReflectionResponse(
-            error_response=reflection_pb2.ErrorResponse(
-                error_code=grpc.StatusCode.UNIMPLEMENTED.value[0],
-                error_message=grpc.StatusCode.UNIMPLMENTED.value[1].encode(),))
-
-    def _extension_numbers_of_type(fully_qualified_name):
-        # TODO(atash) We're allowed to leave this unsupported according to the
-        # protocol, but we should still eventually implement it. Hits the same issue
-        # as `_file_containing_extension`, however.
-        # https://github.com/google/protobuf/issues/2248
-        return reflection_pb2.ServerReflectionResponse(
-            error_response=reflection_pb2.ErrorResponse(
-                error_code=grpc.StatusCode.UNIMPLEMENTED.value[0],
-                error_message=grpc.StatusCode.UNIMPLMENTED.value[1].encode(),))
+    def _file_containing_extension(self, containing_type, extension_number):
+        try:
+            message_descriptor = self._pool.FindMessageTypeByName(containing_type)
+            extension_descriptor = self._pool.FindExtensionByNumber(
+                message_descriptor, extension_number)
+            descriptor = self._pool.FindFileContainingSymbol(
+                extension_descriptor.full_name)
+        except KeyError:
+            return _not_found_error()
+        else:
+            return _file_descriptor_response(descriptor)
+
+    def _all_extension_numbers_of_type(self, containing_type):
+        try:
+            message_descriptor = self._pool.FindMessageTypeByName(containing_type)
+            extension_numbers = tuple(sorted(
+                extension.number
+                for extension in self._pool.FindAllExtensions(message_descriptor)))
+        except KeyError:
+            return _not_found_error()
+        else:
+            return reflection_pb2.ServerReflectionResponse(
+                all_extension_numbers_response=reflection_pb2.
+                ExtensionNumberResponse(
+                    base_type_name=message_descriptor.full_name,
+                    extension_number=extension_numbers))
 
     def _list_services(self):
         return reflection_pb2.ServerReflectionResponse(
@@ -121,7 +131,7 @@ class ReflectionServicer(reflection_pb2.ServerReflectionServicer):
                     request.file_containing_extension.containing_type,
                     request.file_containing_extension.extension_number)
             elif request.HasField('all_extension_numbers_of_type'):
-                yield _all_extension_numbers_of_type(
+                yield self._all_extension_numbers_of_type(
                     request.all_extension_numbers_of_type)
             elif request.HasField('list_services'):
                 yield self._list_services()
@@ -131,3 +141,14 @@ class ReflectionServicer(reflection_pb2.ServerReflectionServicer):
                         error_code=grpc.StatusCode.INVALID_ARGUMENT.value[0],
                         error_message=grpc.StatusCode.INVALID_ARGUMENT.value[1]
                         .encode(),))
+
+
+def enable_server_reflection(service_names, server):
+    """Enables server reflection on a server.
+
+    Args:
+      service_names: Iterable of fully-qualified service names available.
+      server: grpc.Server to which reflection service will be added.
+    """
+    reflection_pb2_grpc.add_ServerReflectionServicer_to_server(
+        ReflectionServicer(service_names), server)

src/python/grpcio_tests/tests/reflection/_reflection_servicer_test.py

@@ -39,14 +39,19 @@ from grpc_reflection.v1alpha import reflection_pb2_grpc
 from google.protobuf import descriptor_pool
 from google.protobuf import descriptor_pb2
 
-from src.proto.grpc.testing.proto2 import empty2_extensions_pb2
 from src.proto.grpc.testing import empty_pb2
+#empty2_pb2 is imported for import-consequent side-effects.
+from src.proto.grpc.testing.proto2 import empty2_pb2  # pylint: disable=unused-import
+from src.proto.grpc.testing.proto2 import empty2_extensions_pb2
+
 from tests.unit.framework.common import test_constants
 
 _EMPTY_PROTO_FILE_NAME = 'src/proto/grpc/testing/empty.proto'
 _EMPTY_PROTO_SYMBOL_NAME = 'grpc.testing.Empty'
 _SERVICE_NAMES = ('Angstrom', 'Bohr', 'Curie', 'Dyson', 'Einstein', 'Feynman',
                   'Galilei')
+_EMPTY_EXTENSIONS_SYMBOL_NAME = 'grpc.testing.proto2.EmptyWithExtensions'
+_EMPTY_EXTENSIONS_NUMBERS = (124, 125, 126, 127, 128,)
 
 
 def _file_descriptor_to_proto(descriptor):
@@ -110,12 +115,12 @@ class ReflectionServicerTest(unittest.TestCase):
         self.assertSequenceEqual(expected_responses, responses)
 
     @unittest.skip(
-        'TODO(atash): implement file-containing-extension reflection '
-        '(see https://github.com/google/protobuf/issues/2248)')
+        'TODO(mmx): enable when (pure) python protobuf issue is fixed'
+        '(see https://github.com/google/protobuf/issues/2882)')
     def testFileContainingExtension(self):
         requests = (reflection_pb2.ServerReflectionRequest(
             file_containing_extension=reflection_pb2.ExtensionRequest(
-                containing_type='grpc.testing.proto2.Empty',
+                containing_type=_EMPTY_EXTENSIONS_SYMBOL_NAME,
                 extension_number=125,),
         ), reflection_pb2.ServerReflectionRequest(
             file_containing_extension=reflection_pb2.ExtensionRequest(
@@ -127,7 +132,28 @@ class ReflectionServicerTest(unittest.TestCase):
                 valid_host='',
                 file_descriptor_response=reflection_pb2.FileDescriptorResponse(
                     file_descriptor_proto=(_file_descriptor_to_proto(
-                        empty_extensions_pb2.DESCRIPTOR),))),
+                        empty2_extensions_pb2.DESCRIPTOR),))),
+            reflection_pb2.ServerReflectionResponse(
+                valid_host='',
+                error_response=reflection_pb2.ErrorResponse(
+                    error_code=grpc.StatusCode.NOT_FOUND.value[0],
+                    error_message=grpc.StatusCode.NOT_FOUND.value[1].encode(),
+                )),)
+        self.assertSequenceEqual(expected_responses, responses)
+
+    def testExtensionNumbersOfType(self):
+        requests = (reflection_pb2.ServerReflectionRequest(
+            all_extension_numbers_of_type=_EMPTY_EXTENSIONS_SYMBOL_NAME
+        ), reflection_pb2.ServerReflectionRequest(
+            all_extension_numbers_of_type='i.donut.exist.co.uk.net.name.foo'),)
+        responses = tuple(self._stub.ServerReflectionInfo(iter(requests)))
+        expected_responses = (
+            reflection_pb2.ServerReflectionResponse(
+                valid_host='',
+                all_extension_numbers_response=reflection_pb2.
+                ExtensionNumberResponse(
+                    base_type_name=_EMPTY_EXTENSIONS_SYMBOL_NAME,
+                    extension_number=_EMPTY_EXTENSIONS_NUMBERS)),
             reflection_pb2.ServerReflectionResponse(
                 valid_host='',
                 error_response=reflection_pb2.ErrorResponse(