1# Copyright 2020 Google LLC
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#     http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14
15import requests
16import requests.exceptions as requests_exceptions
17
18from google.api_core import exceptions as api_exceptions
19from google.api_core import retry
20from google.auth import exceptions as auth_exceptions
21
22
23# ConnectionError is a built-in exception only in Python3 and not in Python2.
24try:
25    _RETRYABLE_STDLIB_TYPES = (ConnectionError,)
26except NameError:
27    _RETRYABLE_STDLIB_TYPES = ()
28
29
30_RETRYABLE_TYPES = _RETRYABLE_STDLIB_TYPES + (
31    api_exceptions.TooManyRequests,  # 429
32    api_exceptions.InternalServerError,  # 500
33    api_exceptions.BadGateway,  # 502
34    api_exceptions.ServiceUnavailable,  # 503
35    api_exceptions.GatewayTimeout,  # 504
36    requests.ConnectionError,
37    requests_exceptions.ChunkedEncodingError,
38)
39
40
41# Some retriable errors don't have their own custom exception in api_core.
42_ADDITIONAL_RETRYABLE_STATUS_CODES = (408,)
43
44
45def _should_retry(exc):
46    """Predicate for determining when to retry."""
47    if isinstance(exc, _RETRYABLE_TYPES):
48        return True
49    elif isinstance(exc, api_exceptions.GoogleAPICallError):
50        return exc.code in _ADDITIONAL_RETRYABLE_STATUS_CODES
51    elif isinstance(exc, auth_exceptions.TransportError):
52        return _should_retry(exc.args[0])
53    else:
54        return False
55
56
57DEFAULT_RETRY = retry.Retry(predicate=_should_retry)
58"""The default retry object.
59
60This retry setting will retry all _RETRYABLE_TYPES and any status codes from
61_ADDITIONAL_RETRYABLE_STATUS_CODES.
62
63To modify the default retry behavior, create a new retry object modeled after
64this one by calling it a ``with_XXX`` method. For example, to create a copy of
65DEFAULT_RETRY with a deadline of 30 seconds, pass
66``retry=DEFAULT_RETRY.with_deadline(30)``. See google-api-core reference
67(https://googleapis.dev/python/google-api-core/latest/retry.html) for details.
68"""
69
70
71class ConditionalRetryPolicy(object):
72    """A class for use when an API call is only conditionally safe to retry.
73
74    This class is intended for use in inspecting the API call parameters of an
75    API call to verify that any flags necessary to make the API call idempotent
76    (such as specifying an ``if_generation_match`` or related flag) are present.
77
78    It can be used in place of a ``retry.Retry`` object, in which case
79    ``_http.Connection.api_request`` will pass the requested api call keyword
80    arguments into the ``conditional_predicate`` and return the ``retry_policy``
81    if the conditions are met.
82
83    :type retry_policy: class:`google.api_core.retry.Retry`
84    :param retry_policy: A retry object defining timeouts, persistence and which
85        exceptions to retry.
86
87    :type conditional_predicate: callable
88    :param conditional_predicate: A callable that accepts exactly the number of
89        arguments in ``required_kwargs``, in order, and returns True if the
90        arguments have sufficient data to determine that the call is safe to
91        retry (idempotent).
92
93    :type required_kwargs: list(str)
94    :param required_kwargs:
95        A list of keyword argument keys that will be extracted from the API call
96        and passed into the ``conditional predicate`` in order.
97    """
98
99    def __init__(self, retry_policy, conditional_predicate, required_kwargs):
100        self.retry_policy = retry_policy
101        self.conditional_predicate = conditional_predicate
102        self.required_kwargs = required_kwargs
103
104    def get_retry_policy_if_conditions_met(self, **kwargs):
105        if self.conditional_predicate(*[kwargs[key] for key in self.required_kwargs]):
106            return self.retry_policy
107        return None
108
109
110def is_generation_specified(query_params):
111    """Return True if generation or if_generation_match is specified."""
112    generation = query_params.get("generation") is not None
113    if_generation_match = query_params.get("ifGenerationMatch") is not None
114    return generation or if_generation_match
115
116
117def is_metageneration_specified(query_params):
118    """Return True if if_metageneration_match is specified."""
119    if_metageneration_match = query_params.get("ifMetagenerationMatch") is not None
120    return if_metageneration_match
121
122
123def is_etag_in_data(data):
124    """Return True if an etag is contained in the request body.
125
126    :type data: dict or None
127    :param data: A dict representing the request JSON body. If not passed, returns False.
128    """
129    return data is not None and "etag" in data
130
131
132def is_etag_in_json(data):
133    """
134    ``is_etag_in_json`` is supported for backwards-compatibility reasons only;
135    please use ``is_etag_in_data`` instead.
136    """
137    return is_etag_in_data(data)
138
139
140DEFAULT_RETRY_IF_GENERATION_SPECIFIED = ConditionalRetryPolicy(
141    DEFAULT_RETRY, is_generation_specified, ["query_params"]
142)
143"""Conditional wrapper for the default retry object.
144
145This retry setting will retry all _RETRYABLE_TYPES and any status codes from
146_ADDITIONAL_RETRYABLE_STATUS_CODES, but only if the request included an
147``ifGenerationMatch`` header.
148"""
149
150DEFAULT_RETRY_IF_METAGENERATION_SPECIFIED = ConditionalRetryPolicy(
151    DEFAULT_RETRY, is_metageneration_specified, ["query_params"]
152)
153"""Conditional wrapper for the default retry object.
154
155This retry setting will retry all _RETRYABLE_TYPES and any status codes from
156_ADDITIONAL_RETRYABLE_STATUS_CODES, but only if the request included an
157``ifMetagenerationMatch`` header.
158"""
159
160DEFAULT_RETRY_IF_ETAG_IN_JSON = ConditionalRetryPolicy(
161    DEFAULT_RETRY, is_etag_in_json, ["data"]
162)
163"""Conditional wrapper for the default retry object.
164
165This retry setting will retry all _RETRYABLE_TYPES and any status codes from
166_ADDITIONAL_RETRYABLE_STATUS_CODES, but only if the request included an
167``ETAG`` entry in its payload.
168"""
169