Use composition to implement AsyncClient

lidizheng · lidizheng · commit 3cc730c80165 · 2020-04-17T13:33:02.000-07:00
diff --git a/gapic/templates/%namespace/%name_%version/%sub/services/%service/async_client.py.j2 b/gapic/templates/%namespace/%name_%version/%sub/services/%service/async_client.py.j2
@@ -2,6 +2,7 @@
 
 {% block content %}
 from collections import OrderedDict
+import functools
 import re
 from typing import Dict, {% if service.any_server_streaming %}AsyncIterable, {% endif %}{% if service.any_client_streaming %}AsyncIterator, {% endif %}Sequence, Tuple, Type, Union
 import pkg_resources
@@ -21,111 +22,31 @@ from google.oauth2 import service_account              # type: ignore
 {% endfor -%}
 {% endfilter %}
 from .transports.base import {{ service.name }}Transport
-from .transports.grpc_asyncio import {{ service.name }}GrpcAsyncIOTransport
+from .transports.grpc_asyncio import {{ service.grpc_asyncio_transport_name }}
+from .client import {{ service.client_name }}
 
 
-class {{ service.async_client_name }}Meta(type):
-    """Metaclass for the {{ service.name }} client.
-
-    This provides class-level methods for building and retrieving
-    support objects (e.g. transport) without polluting the client instance
-    objects.
-    """
-    _transport_registry: Dict[str, Type[{{ service.name }}Transport]] = OrderedDict()
-    _transport_registry['grpc_asyncio'] = {{ service.name }}GrpcAsyncIOTransport
-
-    def get_transport_class(cls,
-            label: str = None,
-            ) -> Type[{{ service.name }}Transport]:
-        """Return an appropriate transport class.
-
-        Args:
-            label: The name of the desired transport. If none is
-                provided, then the first transport in the registry is used.
-
-        Returns:
-            The transport class to use.
-        """
-        # If a specific transport is requested, return that one.
-        if label:
-            return cls._transport_registry[label]
-
-        # No transport is requested; return the default (that is, the first one
-        # in the dictionary).
-        return next(iter(cls._transport_registry.values()))
-
-
-class {{ service.async_client_name }}(metaclass={{ service.async_client_name }}Meta):
+class {{ service.async_client_name }}:
     """{{ service.meta.doc|rst(width=72, indent=4) }}"""
 
-    @staticmethod
-    def _get_default_mtls_endpoint(api_endpoint):
-        """Convert api endpoint to mTLS endpoint.
-        Convert "*.sandbox.googleapis.com" and "*.googleapis.com" to
-        "*.mtls.sandbox.googleapis.com" and "*.mtls.googleapis.com" respectively.
-        Args:
-            api_endpoint (Optional[str]): the api endpoint to convert.
-        Returns:
-            str: converted mTLS api endpoint.
-        """
-        if not api_endpoint:
-            return api_endpoint
-
-        mtls_endpoint_re = re.compile(
-            r"(?P<name>[^.]+)(?P<mtls>\.mtls)?(?P<sandbox>\.sandbox)?(?P<googledomain>\.googleapis\.com)?"
-        )
-
-        m = mtls_endpoint_re.match(api_endpoint)
-        name, mtls, sandbox, googledomain = m.groups()
-        if mtls or not googledomain:
-            return api_endpoint
-
-        if sandbox:
-            return api_endpoint.replace(
-                "sandbox.googleapis.com", "mtls.sandbox.googleapis.com"
-            )
+    _client: {{ service.client_name }}
 
-        return api_endpoint.replace(".googleapis.com", ".mtls.googleapis.com")
+    DEFAULT_ENDPOINT = {{ service.client_name }}.DEFAULT_ENDPOINT
+    DEFAULT_MTLS_ENDPOINT = {{ service.client_name }}.DEFAULT_MTLS_ENDPOINT
 
-    DEFAULT_ENDPOINT = {% if service.host %}'{{ service.host }}'{% else %}None{% endif %}
-    DEFAULT_MTLS_ENDPOINT = _get_default_mtls_endpoint.__func__(  # type: ignore
-        DEFAULT_ENDPOINT
-    )
-    DEFAULT_MTLS_TRANSPORT = {{ service.grpc_asyncio_transport_name }}
-
-    @classmethod
-    def from_service_account_file(cls, filename: str, *args, **kwargs):
-        """Creates an instance of this client using the provided credentials
-        file.
-
-        Args:
-            filename (str): The path to the service account private key json
-                file.
-            args: Additional arguments to pass to the constructor.
-            kwargs: Additional arguments to pass to the constructor.
+    {% for message in service.resource_messages -%}
+    {{ message.resource_type|snake_case }}_path = staticmethod({{ service.client_name }}.{{ message.resource_type|snake_case }}_path)
 
-        Returns:
-            {@api.name}: The constructed client.
-        """
-        credentials = service_account.Credentials.from_service_account_file(
-            filename)
-        kwargs['credentials'] = credentials
-        return cls(*args, **kwargs)
+    {% endfor %}
 
+    from_service_account_file = {{ service.client_name }}.from_service_account_file
     from_service_account_json = from_service_account_file
 
-
-    {% for message in service.resource_messages -%}
-    @staticmethod
-    def {{ message.resource_type|snake_case }}_path({% for arg in message.resource_path_args %}{{ arg }}: str,{% endfor %}) -> str:
-        """Return a fully-qualified {{ message.resource_type|snake_case }} string."""
-        return "{{ message.resource_path }}".format({% for arg in message.resource_path_args %}{{ arg }}={{ arg }}, {% endfor %})
-
-    {% endfor %}
+    get_transport_class = functools.partial(type({{ service.client_name }}).get_transport_class, type({{ service.client_name }}))
 
     def __init__(self, *,
             credentials: credentials.Credentials = None,
-            transport: Union[str, {{ service.name }}Transport] = None,
+            transport: Union[str, {{ service.name }}Transport] = "grpc_asyncio",
             client_options: ClientOptions = None,
             ) -> None:
         """Instantiate the {{ (service.client_name|snake_case).replace('_', ' ') }}.
@@ -152,52 +73,12 @@ class {{ service.async_client_name }}(metaclass={{ service.async_client_name }}M
             google.auth.exceptions.MutualTlsChannelError: If mutual TLS transport
                 creation failed for any reason.
         """
-        if isinstance(client_options, dict):
-            client_options = ClientOptions.from_dict(client_options)
-
-        # Save or instantiate the transport.
-        # Ordinarily, we provide the transport, but allowing a custom transport
-        # instance provides an extensibility point for unusual situations.
-        if isinstance(transport, {{ service.name }}Transport):
-            # transport is a {{ service.name }}Transport instance.
-            if credentials:
-                raise ValueError('When providing a transport instance, '
-                                 'provide its credentials directly.')
-            self._transport = transport
-        elif client_options is None or (
-            client_options.api_endpoint == None
-            and client_options.client_cert_source is None
-        ):
-            # Don't trigger mTLS if we get an empty ClientOptions.
-            Transport = type(self).get_transport_class(transport)
-            self._transport = Transport(
-                credentials=credentials, host=self.DEFAULT_ENDPOINT
-            )
-        else:
-            # We have a non-empty ClientOptions. If client_cert_source is
-            # provided, trigger mTLS with user provided endpoint or the default
-            # mTLS endpoint.
-            if client_options.client_cert_source:
-                api_mtls_endpoint = (
-                    client_options.api_endpoint
-                    if client_options.api_endpoint
-                    else self.DEFAULT_MTLS_ENDPOINT
-                )
-            else:
-                api_mtls_endpoint = None
-
-            api_endpoint = (
-                client_options.api_endpoint
-                if client_options.api_endpoint
-                else self.DEFAULT_ENDPOINT
-            )
-
-            self._transport = self.DEFAULT_MTLS_TRANSPORT(
-                credentials=credentials,
-                host=api_endpoint,
-                api_mtls_endpoint=api_mtls_endpoint,
-                client_cert_source=client_options.client_cert_source,
-            )
+        {# NOTE(lidiz) Not using kwargs since we want the docstring and types. #}
+        self._client = {{ service.client_name }}(
+            credentials=credentials,
+            transport=transport,
+            client_options=client_options,
+        )
 
     {% for method in service.methods.values() -%}
     {% if not method.server_streaming %}async {% endif -%}def {{ method.name|snake_case }}(self,
@@ -296,7 +177,7 @@ class {{ service.async_client_name }}(metaclass={{ service.async_client_name }}M
         # Wrap the RPC method; this adds retry and timeout information,
         # and friendly error handling.
         rpc = gapic_v1.method_async.wrap_method(
-            self._transport.{{ method.name|snake_case }},
+            self._client._transport.{{ method.name|snake_case }},
             {%- if method.retry %}
             default_retry=retries.Retry(
                 {% if method.retry.initial_backoff %}initial={{ method.retry.initial_backoff }},{% endif %}
@@ -344,7 +225,7 @@ class {{ service.async_client_name }}(metaclass={{ service.async_client_name }}M
         # Wrap the response in an operation future.
         response = operation_async.from_gapic(
             response,
-            self._transport.operations_client,
+            self._client._transport.operations_client,
             {{ method.lro.response_type.ident }},
             metadata_type={{ method.lro.metadata_type.ident }},
         )
diff --git a/gapic/templates/%namespace/%name_%version/%sub/services/%service/client.py.j2 b/gapic/templates/%namespace/%name_%version/%sub/services/%service/client.py.j2
@@ -21,7 +21,8 @@ from google.oauth2 import service_account              # type: ignore
 {% endfor -%}
 {% endfilter %}
 from .transports.base import {{ service.name }}Transport
-from .transports.grpc import {{ service.name }}GrpcTransport
+from .transports.grpc import {{ service.grpc_transport_name }}
+from .transports.grpc_asyncio import {{ service.grpc_asyncio_transport_name }}
 
 
 class {{ service.client_name }}Meta(type):
@@ -32,16 +33,22 @@ class {{ service.client_name }}Meta(type):
     objects.
     """
     _transport_registry = OrderedDict()  # type: Dict[str, Type[{{ service.name }}Transport]]
-    _transport_registry['grpc'] = {{ service.name }}GrpcTransport
+    _transport_registry['grpc'] = {{ service.grpc_transport_name }}
+    _transport_registry['grpc_asyncio'] = {{ service.grpc_asyncio_transport_name }}
+
+    _mtls_support = {'grpc', 'grpc_asyncio'}
 
     def get_transport_class(cls,
             label: str = None,
+            enable_mtls: bool = False,
             ) -> Type[{{ service.name }}Transport]:
         """Return an appropriate transport class.
 
         Args:
             label: The name of the desired transport. If none is
                 provided, then the first transport in the registry is used.
+            enable_mtls: A bool indicates whether the transport needs to
+                support MTLS or not.
 
         Returns:
             The transport class to use.
@@ -91,7 +98,6 @@ class {{ service.client_name }}(metaclass={{ service.client_name }}Meta):
     DEFAULT_MTLS_ENDPOINT = _get_default_mtls_endpoint.__func__(  # type: ignore
         DEFAULT_ENDPOINT
     )
-    DEFAULT_MTLS_TRANSPORT = {{ service.grpc_transport_name }}
 
     @classmethod
     def from_service_account_file(cls, filename: str, *args, **kwargs):
@@ -192,7 +198,10 @@ class {{ service.client_name }}(metaclass={{ service.client_name }}Meta):
                 else self.DEFAULT_ENDPOINT
             )
 
-            self._transport = self.DEFAULT_MTLS_TRANSPORT(
+            self._transport = type(self).get_transport_class(
+                label=transport,
+                enable_mtls=True,
+            )(
                 credentials=credentials,
                 host=api_endpoint,
                 api_mtls_endpoint=api_mtls_endpoint,
diff --git a/gapic/templates/tests/unit/%name_%version/%sub/test_%service.py.j2 b/gapic/templates/tests/unit/%name_%version/%sub/test_%service.py.j2
@@ -71,7 +71,7 @@ def test_{{ service.client_name|snake_case }}_from_service_account_file(client_c
 ])
 def test_{{ service.client_name|snake_case }}_client_options(client_class, transport_class):
     # Check that if channel is provided we won't create a new one.
-    with mock.patch.object(client_class, 'get_transport_class') as gtc:
+    with mock.patch.object({{ service.client_name }}, 'get_transport_class') as gtc:
         transport = transport_class(
             credentials=credentials.AnonymousCredentials()
         )
@@ -80,10 +80,10 @@ def test_{{ service.client_name|snake_case }}_client_options(client_class, trans
 
     # Check mTLS is not triggered with empty client options.
     options = client_options.ClientOptions()
-    with mock.patch.object(client_class, 'get_transport_class') as gtc:
-        transport = gtc.return_value = mock.MagicMock()
+    with mock.patch.object(transport_class, '__init__') as patched:
+        patched.return_value = None
         client = client_class(client_options=options)
-        transport.assert_called_once_with(
+        patched.assert_called_once_with(
             credentials=None,
             host=client.DEFAULT_ENDPOINT,
         )
@@ -427,7 +427,7 @@ async def test_{{ method.name|snake_case }}_async_pager():
 
     # Mock the actual call within the gRPC stub, and fake the request.
     with mock.patch.object(
-            type(client._transport.{{ method.name|snake_case }}),
+            type(client._client._transport.{{ method.name|snake_case }}),
             '__call__', new_callable=mock.AsyncMock) as call:
         # Set the response to a series of pages.
         call.side_effect = (
@@ -475,7 +475,7 @@ async def test_{{ method.name|snake_case }}_async_pages():
 
     # Mock the actual call within the gRPC stub, and fake the request.
     with mock.patch.object(
-            type(client._transport.{{ method.name|snake_case }}),
+            type(client._client._transport.{{ method.name|snake_case }}),
             '__call__', new_callable=mock.AsyncMock) as call:
         # Set the response to a series of pages.
         call.side_effect = (
