1# Copyright 2017 Google Inc.
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
15"""Shared utilities used by both downloads and uploads.
16
17This utilities are explicitly catered to ``requests``-like transports.
18"""
19
20_SINGLE_GET_CHUNK_SIZE = 8192
21# The number of seconds to wait to establish a connection
22# (connect() call on socket). Avoid setting this to a multiple of 3 to not
23# Align with TCP Retransmission timing. (typically 2.5-3s)
24_DEFAULT_CONNECT_TIMEOUT = 61
25# The number of seconds to wait between bytes sent from the server.
26_DEFAULT_READ_TIMEOUT = 60
27
28
29class RequestsMixin(object):
30 """Mix-in class implementing ``requests``-specific behavior.
31
32 These are methods that are more general purpose, with implementations
33 specific to the types defined in ``requests``.
34 """
35
36 @staticmethod
37 def _get_status_code(response):
38 """Access the status code from an HTTP response.
39
40 Args:
41 response (~requests.Response): The HTTP response object.
42
43 Returns:
44 int: The status code.
45 """
46 return response.status_code
47
48 @staticmethod
49 def _get_headers(response):
50 """Access the headers from an HTTP response.
51
52 Args:
53 response (~requests.Response): The HTTP response object.
54
55 Returns:
56 ~requests.structures.CaseInsensitiveDict: The header mapping (keys
57 are case-insensitive).
58 """
59 return response.headers
60
61 @staticmethod
62 def _get_body(response):
63 """Access the response body from an HTTP response.
64
65 Args:
66 response (~requests.Response): The HTTP response object.
67
68 Returns:
69 bytes: The body of the ``response``.
70 """
71 return response.content
72
73
74class RawRequestsMixin(RequestsMixin):
75 @staticmethod
76 def _get_body(response):
77 """Access the response body from an HTTP response.
78
79 Args:
80 response (~requests.Response): The HTTP response object.
81
82 Returns:
83 bytes: The body of the ``response``.
84 """
85 if response._content is False:
86 response._content = b"".join(
87 response.raw.stream(_SINGLE_GET_CHUNK_SIZE, decode_content=False)
88 )
89 response._content_consumed = True
90 return response._content
91
92
93def wait_and_retry(func, retry_strategy):
94 """Attempts to retry a call to ``func`` until success.
95
96 Args:
97 func (Callable): A callable that takes no arguments and produces
98 an HTTP response which will be checked as retry-able.
99 retry_strategy (Optional[google.api_core.retry.Retry]): The
100 strategy to use if the request fails and must be retried.
101
102 Returns:
103 object: The return value of ``func``.
104 """
105 if retry_strategy:
106 func = retry_strategy(func)
107 return func()