Polish and add AsyncIO examples

* Add type annotations * Add test for asyncio auth example * Use dynamic stubs to allow same code runs manually and under Bazel * Add grpcio-tools as a Bazel dependency * asyncio.run is not yet existed in 3.6 * Improve readability
dmeister · Dec 8, 2020 · 40a5a65 · 40a5a65
1 parent 8c12cc8
commit 40a5a65
Show file tree

Hide file tree

Showing 13 changed files with 374 additions and 35 deletions.
diff --git a/examples/python/auth/BUILD.bazel b/examples/python/auth/BUILD.bazel
@@ -33,25 +33,53 @@ py_binary(
     name = "customized_auth_client",
     testonly = 1,
     srcs = ["customized_auth_client.py"],
+    data = ["helloworld.proto"],
     python_version = "PY3",
     deps = [
         ":_credentials",
-        "//examples:helloworld_py_pb2",
-        "//examples:helloworld_py_pb2_grpc",
         "//src/python/grpcio/grpc:grpcio",
+        "//tools/distrib/python/grpcio_tools:grpc_tools",
     ],
 )
 
 py_binary(
     name = "customized_auth_server",
     testonly = 1,
     srcs = ["customized_auth_server.py"],
+    data = ["helloworld.proto"],
     python_version = "PY3",
     deps = [
         ":_credentials",
-        "//examples:helloworld_py_pb2",
-        "//examples:helloworld_py_pb2_grpc",
         "//src/python/grpcio/grpc:grpcio",
+        "//tools/distrib/python/grpcio_tools:grpc_tools",
+    ],
+)
+
+py_binary(
+    name = "async_customized_auth_client",
+    testonly = 1,
+    srcs = ["async_customized_auth_client.py"],
+    data = ["helloworld.proto"],
+    imports = ["."],
+    python_version = "PY3",
+    deps = [
+        ":_credentials",
+        "//src/python/grpcio/grpc:grpcio",
+        "//tools/distrib/python/grpcio_tools:grpc_tools",
+    ],
+)
+
+py_binary(
+    name = "async_customized_auth_server",
+    testonly = 1,
+    srcs = ["async_customized_auth_server.py"],
+    data = ["helloworld.proto"],
+    imports = ["."],
+    python_version = "PY3",
+    deps = [
+        ":_credentials",
+        "//src/python/grpcio/grpc:grpcio",
+        "//tools/distrib/python/grpcio_tools:grpc_tools",
     ],
 )
 
@@ -61,9 +89,11 @@ py_test(
     python_version = "PY3",
     deps = [
         ":_credentials",
+        ":async_customized_auth_client",
+        ":async_customized_auth_server",
         ":customized_auth_client",
         ":customized_auth_server",
-        "//examples:helloworld_py_pb2",
         "//src/python/grpcio/grpc:grpcio",
+        "//tools/distrib/python/grpcio_tools:grpc_tools",
     ],
 )
diff --git a/examples/python/auth/_credentials.py b/examples/python/auth/_credentials.py
@@ -13,10 +13,6 @@
 # limitations under the License.
 """Loading SSL credentials for gRPC Python authentication example."""
 
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
 import os
 
 

diff --git a/examples/python/auth/async_customized_auth_client.py b/examples/python/auth/async_customized_auth_client.py
@@ -0,0 +1,101 @@
+# Copyright 2020 The gRPC Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Client of the Python AsyncIO example of customizing authentication mechanism."""
+
+import argparse
+import asyncio
+import logging
+
+import grpc
+
+import _credentials
+
+helloworld_pb2, helloworld_pb2_grpc = grpc.protos_and_services(
+    "helloworld.proto")
+
+_LOGGER = logging.getLogger(__name__)
+_LOGGER.setLevel(logging.INFO)
+
+_SERVER_ADDR_TEMPLATE = 'localhost:%d'
+_SIGNATURE_HEADER_KEY = 'x-signature'
+
+
+class AuthGateway(grpc.AuthMetadataPlugin):
+
+    def __call__(self, context: grpc.AuthMetadataContext,
+                 callback: grpc.AuthMetadataPluginCallback) -> None:
+        """Implements authentication by passing metadata to a callback.
+
+        Implementations of this method must not block.
+
+        Args:
+          context: An AuthMetadataContext providing information on the RPC that
+            the plugin is being called to authenticate.
+          callback: An AuthMetadataPluginCallback to be invoked either
+            synchronously or asynchronously.
+        """
+        # Example AuthMetadataContext object:
+        # AuthMetadataContext(
+        #     service_url=u'https://localhost:50051/helloworld.Greeter',
+        #     method_name=u'SayHello')
+        signature = context.method_name[::-1]
+        callback(((_SIGNATURE_HEADER_KEY, signature),), None)
+
+
+def create_client_channel(addr: str) -> grpc.aio.Channel:
+    # Call credential object will be invoked for every single RPC
+    call_credentials = grpc.metadata_call_credentials(AuthGateway(),
+                                                      name='auth gateway')
+    # Channel credential will be valid for the entire channel
+    channel_credential = grpc.ssl_channel_credentials(
+        _credentials.ROOT_CERTIFICATE)
+    # Combining channel credentials and call credentials together
+    composite_credentials = grpc.composite_channel_credentials(
+        channel_credential,
+        call_credentials,
+    )
+    channel = grpc.aio.secure_channel(addr, composite_credentials)
+    return channel
+
+
+async def send_rpc(channel: grpc.aio.Channel) -> helloworld_pb2.HelloReply:
+    stub = helloworld_pb2_grpc.GreeterStub(channel)
+    request = helloworld_pb2.HelloRequest(name='you')
+    try:
+        response = await stub.SayHello(request)
+    except grpc.RpcError as rpc_error:
+        _LOGGER.error('Received error: %s', rpc_error)
+        return rpc_error
+    else:
+        _LOGGER.info('Received message: %s', response)
+        return response
+
+
+async def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--port',
+                        nargs='?',
+                        type=int,
+                        default=50051,
+                        help='the address of server')
+    args = parser.parse_args()
+
+    channel = create_client_channel(_SERVER_ADDR_TEMPLATE % args.port)
+    await send_rpc(channel)
+    await channel.close()
+
+
+if __name__ == '__main__':
+    logging.basicConfig(level=logging.INFO)
+    asyncio.run(main())
diff --git a/examples/python/auth/async_customized_auth_server.py b/examples/python/auth/async_customized_auth_server.py
@@ -0,0 +1,103 @@
+# Copyright 2020 The gRPC Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Server of the Python AsyncIO example of customizing authentication mechanism."""
+
+import argparse
+import asyncio
+import logging
+from typing import Awaitable, Callable, Tuple
+
+import grpc
+
+import _credentials
+
+helloworld_pb2, helloworld_pb2_grpc = grpc.protos_and_services(
+    "helloworld.proto")
+
+_LOGGER = logging.getLogger(__name__)
+_LOGGER.setLevel(logging.INFO)
+
+_LISTEN_ADDRESS_TEMPLATE = 'localhost:%d'
+_SIGNATURE_HEADER_KEY = 'x-signature'
+
+
+class SignatureValidationInterceptor(grpc.aio.ServerInterceptor):
+
+    def __init__(self):
+
+        def abort(ignored_request, context: grpc.aio.ServicerContext) -> None:
+            context.abort(grpc.StatusCode.UNAUTHENTICATED, 'Invalid signature')
+
+        self._abort_handler = grpc.unary_unary_rpc_method_handler(abort)
+
+    async def intercept_service(
+            self, continuation: Callable[[grpc.HandlerCallDetails], Awaitable[
+                grpc.RpcMethodHandler]],
+            handler_call_details: grpc.HandlerCallDetails
+    ) -> grpc.RpcMethodHandler:
+        # Example HandlerCallDetails object:
+        #     _HandlerCallDetails(
+        #       method=u'/helloworld.Greeter/SayHello',
+        #       invocation_metadata=...)
+        method_name = handler_call_details.method.split('/')[-1]
+        expected_metadata = (_SIGNATURE_HEADER_KEY, method_name[::-1])
+        if expected_metadata in handler_call_details.invocation_metadata:
+            return await continuation(handler_call_details)
+        else:
+            return self._abort_handler
+
+
+class SimpleGreeter(helloworld_pb2_grpc.GreeterServicer):
+
+    async def SayHello(self, request: helloworld_pb2.HelloRequest,
+                       unused_context) -> helloworld_pb2.HelloReply:
+        return helloworld_pb2.HelloReply(message='Hello, %s!' % request.name)
+
+
+async def run_server(port: int) -> Tuple[grpc.aio.Server, int]:
+    # Bind interceptor to server
+    server = grpc.aio.server(interceptors=(SignatureValidationInterceptor(),))
+    helloworld_pb2_grpc.add_GreeterServicer_to_server(SimpleGreeter(), server)
+
+    # Loading credentials
+    server_credentials = grpc.ssl_server_credentials(((
+        _credentials.SERVER_CERTIFICATE_KEY,
+        _credentials.SERVER_CERTIFICATE,
+    ),))
+
+    # Pass down credentials
+    port = server.add_secure_port(_LISTEN_ADDRESS_TEMPLATE % port,
+                                  server_credentials)
+
+    await server.start()
+    return server, port
+
+
+async def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--port',
+                        nargs='?',
+                        type=int,
+                        default=50051,
+                        help='the listening port')
+    args = parser.parse_args()
+
+    server, port = await run_server(args.port)
+    logging.info('Server is listening at port :%d', port)
+    await server.wait_for_termination()
+
+
+if __name__ == '__main__':
+    logging.basicConfig(level=logging.INFO)
+    asyncio.run(main())
diff --git a/examples/python/auth/customized_auth_client.py b/examples/python/auth/customized_auth_client.py
@@ -13,18 +13,15 @@
 # limitations under the License.
 """Client of the Python example of customizing authentication mechanism."""
 
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
 import argparse
 import contextlib
 import logging
 
 import grpc
-from examples import helloworld_pb2
-from examples import helloworld_pb2_grpc
-from examples.python.auth import _credentials
+import _credentials
+
+helloworld_pb2, helloworld_pb2_grpc = grpc.protos_and_services(
+    "helloworld.proto")
 
 _LOGGER = logging.getLogger(__name__)
 _LOGGER.setLevel(logging.INFO)

diff --git a/examples/python/auth/customized_auth_server.py b/examples/python/auth/customized_auth_server.py
@@ -13,19 +13,16 @@
 # limitations under the License.
 """Server of the Python example of customizing authentication mechanism."""
 
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
 import argparse
 import contextlib
 import logging
 from concurrent import futures
 
 import grpc
-from examples import helloworld_pb2
-from examples import helloworld_pb2_grpc
-from examples.python.auth import _credentials
+import _credentials
+
+helloworld_pb2, helloworld_pb2_grpc = grpc.protos_and_services(
+    "helloworld.proto")
 
 _LOGGER = logging.getLogger(__name__)
 _LOGGER.setLevel(logging.INFO)

diff --git a/examples/python/auth/helloworld.proto b/examples/python/auth/helloworld.proto
@@ -0,0 +1 @@
+../../protos/helloworld.proto
diff --git a/examples/python/auth/test/_auth_example_test.py b/examples/python/auth/test/_auth_example_test.py
@@ -13,16 +13,15 @@
 # limitations under the License.
 """Test for gRPC Python authentication example."""
 
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
+import asyncio
 import unittest
 
 import grpc
 from examples.python.auth import _credentials
 from examples.python.auth import customized_auth_client
 from examples.python.auth import customized_auth_server
+from examples.python.auth import async_customized_auth_client
+from examples.python.auth import async_customized_auth_server
 
 _SERVER_ADDR_TEMPLATE = 'localhost:%d'
 
@@ -51,6 +50,19 @@ def test_no_call_credential(self):
                 resp = customized_auth_client.send_rpc(channel)
                 self.assertEqual(resp.code(), grpc.StatusCode.UNAUTHENTICATED)
 
+    def test_successful_call_asyncio(self):
+
+        async def test_body():
+            server, port = await async_customized_auth_server.run_server(0)
+            channel = async_customized_auth_client.create_client_channel(
+                _SERVER_ADDR_TEMPLATE % port)
+            await async_customized_auth_client.send_rpc(channel)
+            await channel.close()
+            await server.stop(0)
+            # No unhandled exception raised, test passed!
+
+        asyncio.get_event_loop().run_until_complete(test_body())
+
 
 if __name__ == '__main__':
     unittest.main(verbosity=2)
diff --git a/examples/python/helloworld/async_greeter_client.py b/examples/python/helloworld/async_greeter_client.py
@@ -21,10 +21,7 @@
 import helloworld_pb2_grpc
 
 
-async def run():
-    # NOTE(gRPC Python Team): .close() is possible on a channel and should be
-    # used in circumstances in which the with statement does not fit the needs
-    # of the code.
+async def run() -> None:
     async with grpc.aio.insecure_channel('localhost:50051') as channel:
         stub = helloworld_pb2_grpc.GreeterStub(channel)
         response = await stub.SayHello(helloworld_pb2.HelloRequest(name='you'))