Simplify logics under gapic_v1

lidizheng · Apr 9, 2020 · Apr 9, 2020 · Apr 10, 2020 · Apr 10, 2020 · Apr 10, 2020
commit 84703177d51afa27a05249cb7a784fcd53993663
diff --git a/google/api_core/gapic_v1/config.py b/google/api_core/gapic_v1/config.py
@@ -45,7 +45,7 @@ def _exception_class_for_grpc_status_name(name):
     return exceptions.exception_class_for_grpc_status(getattr(grpc.StatusCode, name))
 
 
-def _retry_from_retry_config(retry_params, retry_codes):
+def _retry_from_retry_config(retry_params, retry_codes, retry_impl=retry.Retry):
     """Creates a Retry object given a gapic retry configuration.
 
     Args:
@@ -70,7 +70,7 @@ def _retry_from_retry_config(retry_params, retry_codes):
     exception_classes = [
         _exception_class_for_grpc_status_name(code) for code in retry_codes
     ]
-    return retry.Retry(
+    return retry_impl(
         retry.if_exception_type(*exception_classes),
         initial=(retry_params["initial_retry_delay_millis"] / _MILLIS_PER_SECOND),
         maximum=(retry_params["max_retry_delay_millis"] / _MILLIS_PER_SECOND),
@@ -110,7 +110,7 @@ def _timeout_from_retry_config(retry_params):
 MethodConfig = collections.namedtuple("MethodConfig", ["retry", "timeout"])
 
 
-def parse_method_configs(interface_config):
+def parse_method_configs(interface_config, retry_impl=retry.Retry):
     """Creates default retry and timeout objects for each method in a gapic
     interface config.
 
@@ -120,6 +120,8 @@ def parse_method_configs(interface_config):
             an interface named ``google.example.v1.ExampleService`` you would
             pass in just that interface's configuration, for example
             ``gapic_config['interfaces']['google.example.v1.ExampleService']``.
+        retry_impl (Callable): The constructor that creates a retry decorator
+            that will be applied to the method based on method configs.
 
     Returns:
         Mapping[str, MethodConfig]: A mapping of RPC method names to their
@@ -151,7 +153,7 @@ def parse_method_configs(interface_config):
         if retry_params_name is not None:
             retry_params = retry_params_map[retry_params_name]
             retry_ = _retry_from_retry_config(
-                retry_params, retry_codes_map[method_params["retry_codes_name"]]
+                retry_params, retry_codes_map[method_params["retry_codes_name"]], retry_impl
             )
             timeout_ = _timeout_from_retry_config(retry_params)
 

diff --git a/google/api_core/gapic_v1/config_async.py b/google/api_core/gapic_v1/config_async.py
@@ -19,45 +19,13 @@
 
 from google.api_core import retry_async
 from google.api_core import timeout
-from google.api_core.gapic_v1.config import _timeout_from_retry_config, MethodConfig, _exception_class_for_grpc_status_name, _MILLIS_PER_SECOND
-
-
-def _retry_from_retry_config(retry_params, retry_codes):
-    """Creates an AsyncRetry object given a gapic retry configuration.
-
-    Args:
-        retry_params (dict): The retry parameter values, for example::
-
-            {
-                "initial_retry_delay_millis": 1000,
-                "retry_delay_multiplier": 2.5,
-                "max_retry_delay_millis": 120000,
-                "initial_rpc_timeout_millis": 120000,
-                "rpc_timeout_multiplier": 1.0,
-                "max_rpc_timeout_millis": 120000,
-                "total_timeout_millis": 600000
-            }
-
-        retry_codes (sequence[str]): The list of retryable gRPC error code
-            names.
-
-    Returns:
-        google.api_core.retry_async.AsyncRetry: The default retry object for the method.
-    """
-    exception_classes = [
-        _exception_class_for_grpc_status_name(code) for code in retry_codes
-    ]
-    return retry_async.AsyncRetry(
-        retry_async.if_exception_type(*exception_classes),
-        initial=(retry_params["initial_retry_delay_millis"] / _MILLIS_PER_SECOND),
-        maximum=(retry_params["max_retry_delay_millis"] / _MILLIS_PER_SECOND),
-        multiplier=retry_params["retry_delay_multiplier"],
-        deadline=retry_params["total_timeout_millis"] / _MILLIS_PER_SECOND)
+from google.api_core.gapic_v1.config import MethodConfig
+from google.api_core.gapic_v1 import config
 
 
 def parse_method_configs(interface_config):
     """Creates default retry and timeout objects for each method in a gapic
-    interface config.
+    interface config with AsyncIO semantics.
 
     Args:
         interface_config (Mapping): The interface config section of the full
@@ -70,42 +38,6 @@ def parse_method_configs(interface_config):
         Mapping[str, MethodConfig]: A mapping of RPC method names to their
             configuration.
     """
-    # Grab all the retry codes
-    retry_codes_map = {
-        name: retry_codes
-        for name, retry_codes in interface_config.get("retry_codes",
-                                                      {}).items()
-    }
-
-    # Grab all of the retry params
-    retry_params_map = {
-        name: retry_params
-        for name, retry_params in interface_config.get("retry_params",
-                                                       {}).items()
-    }
-
-    # Iterate through all the API methods and create a flat MethodConfig
-    # instance for each one.
-    method_configs = {}
-
-    for method_name, method_params in interface_config.get("methods",
-                                                           {}).items():
-        retry_params_name = method_params.get("retry_params_name")
-
-        if retry_params_name is not None:
-            retry_params = retry_params_map[retry_params_name]
-            retry_ = _retry_from_retry_config(
-                retry_params,
-                retry_codes_map[method_params["retry_codes_name"]])
-            timeout_ = _timeout_from_retry_config(retry_params)
-
-        # No retry config, so this is a non-retryable method.
-        else:
-            retry_ = None
-            timeout_ = timeout.ConstantTimeout(
-                method_params["timeout_millis"] / _MILLIS_PER_SECOND)
-
-        method_configs[method_name] = MethodConfig(retry=retry_,
-                                                   timeout=timeout_)
-
-    return method_configs
+    return config.parse_method_configs(
+        interface_config,
+        retry_impl=retry_async.AsyncRetry)
diff --git a/google/api_core/gapic_v1/method_async.py b/google/api_core/gapic_v1/method_async.py
@@ -18,42 +18,35 @@
 pagination, and long-running operations to gRPC methods.
 """
 
-import sys
-
-from google.api_core.gapic_v1 import method
-from google.api_core import general_helpers
-from google.api_core import grpc_helpers_async
-from google.api_core.gapic_v1 import client_info
-from google.api_core.gapic_v1.method import GapicCallable, DEFAULT, USE_DEFAULT_METADATA
-
-if sys.version_info[0] >= 3 and sys.version_info[1] >= 6:
-    from grpc.experimental import aio
-
-    def wrap_method(
-        func,
-        default_retry=None,
-        default_timeout=None,
-        client_info=client_info.DEFAULT_CLIENT_INFO,
-    ):
-        """Wrap an async RPC method with common behavior.
-
-        Returns:
-            Callable: A new callable that takes optional ``retry`` and ``timeout``
-                arguments and applies the common error mapping, retry, timeout,
-                and metadata behavior to the low-level RPC method.
-        """
-        func = grpc_helpers_async.wrap_errors(func)
-
-        if client_info is not None:
-            user_agent_metadata = [client_info.to_grpc_metadata()]
-        else:
-            user_agent_metadata = None
-
-        return general_helpers.wraps(func)(
-            GapicCallable(
-                func, default_retry, default_timeout, metadata=user_agent_metadata
-            )
+from google.api_core import general_helpers, grpc_helpers_async
+from google.api_core.gapic_v1 import client_info, method
+from google.api_core.gapic_v1.method import (DEFAULT, USE_DEFAULT_METADATA,
+                                             GapicCallable)
+from grpc.experimental import aio
+
+
+def wrap_method(
+    func,
+    default_retry=None,
+    default_timeout=None,
+    client_info=client_info.DEFAULT_CLIENT_INFO,
+):
+    """Wrap an async RPC method with common behavior.
+
+    Returns:
+        Callable: A new callable that takes optional ``retry`` and ``timeout``
+            arguments and applies the common error mapping, retry, timeout,
+            and metadata behavior to the low-level RPC method.
+    """
+    func = grpc_helpers_async.wrap_errors(func)
+
+    if client_info is not None:
+        user_agent_metadata = [client_info.to_grpc_metadata()]
+    else:
+        user_agent_metadata = None
+
+    return general_helpers.wraps(func)(
+        GapicCallable(
+            func, default_retry, default_timeout, metadata=user_agent_metadata
         )
-
-else:
-    raise RuntimeError('gRPC AsyncIO supports Python 3.6+, please upgrade Python to use it.')
+    )