From 6eeb855aa0e6a7835d1d4f6e72951e43af22ab57 Mon Sep 17 00:00:00 2001 From: Peter Lamut Date: Tue, 21 Jul 2020 19:58:46 +0200 Subject: [PATCH] feat: add timeouts to Blob methods where missing (#185) * feat: add timeouts to Blob methods where missing * Require google-resumable-media version 0.6.0+ --- google/cloud/storage/blob.py | 154 ++++++++++++++- setup.py | 2 +- tests/unit/test_blob.py | 370 ++++++++++++++++++++++++++++++++--- 3 files changed, 487 insertions(+), 39 deletions(-) diff --git a/google/cloud/storage/blob.py b/google/cloud/storage/blob.py index d4b0956fe..efad9ae39 100644 --- a/google/cloud/storage/blob.py +++ b/google/cloud/storage/blob.py @@ -792,6 +792,7 @@ def _do_download( start=None, end=None, raw_download=False, + timeout=_DEFAULT_TIMEOUT, ): """Perform a download without any error handling. @@ -821,6 +822,14 @@ def _do_download( :type raw_download: bool :param raw_download: (Optional) If true, download the object without any expansion. + + :type timeout: float or tuple + :param timeout: + (Optional) The number of seconds the transport should wait for the + server response. Depending on the retry strategy, a request may be + repeated several times using the same timeout each time. + Can also be passed as a tuple (connect_timeout, read_timeout). + See :meth:`requests.Session.request` documentation for details. """ if self.chunk_size is None: if raw_download: @@ -831,7 +840,7 @@ def _do_download( download = klass( download_url, stream=file_obj, headers=headers, start=start, end=end ) - download.consume(transport) + download.consume(transport, timeout=timeout) else: @@ -850,7 +859,7 @@ def _do_download( ) while not download.finished: - download.consume_next_chunk(transport) + download.consume_next_chunk(transport, timeout=timeout) def download_to_file( self, @@ -863,6 +872,7 @@ def download_to_file( if_generation_not_match=None, if_metageneration_match=None, if_metageneration_not_match=None, + timeout=_DEFAULT_TIMEOUT, ): """Download the contents of this blob into a file-like object. @@ -931,6 +941,14 @@ def download_to_file( :param if_metageneration_not_match: (Optional) Make the operation conditional on whether the blob's current metageneration does not match the given value. + :type timeout: float or tuple + :param timeout: + (Optional) The number of seconds the transport should wait for the + server response. Depending on the retry strategy, a request may be + repeated several times using the same timeout each time. + Can also be passed as a tuple (connect_timeout, read_timeout). + See :meth:`requests.Session.request` documentation for details. + :raises: :class:`google.cloud.exceptions.NotFound` """ client = self._require_client(client) @@ -948,7 +966,14 @@ def download_to_file( transport = self._get_transport(client) try: self._do_download( - transport, file_obj, download_url, headers, start, end, raw_download + transport, + file_obj, + download_url, + headers, + start, + end, + raw_download, + timeout=timeout, ) except resumable_media.InvalidResponse as exc: _raise_from_invalid_response(exc) @@ -964,6 +989,7 @@ def download_to_filename( if_generation_not_match=None, if_metageneration_match=None, if_metageneration_not_match=None, + timeout=_DEFAULT_TIMEOUT, ): """Download the contents of this blob into a named file. @@ -1008,6 +1034,14 @@ def download_to_filename( :param if_metageneration_not_match: (Optional) Make the operation conditional on whether the blob's current metageneration does not match the given value. + :type timeout: float or tuple + :param timeout: + (Optional) The number of seconds the transport should wait for the + server response. Depending on the retry strategy, a request may be + repeated several times using the same timeout each time. + Can also be passed as a tuple (connect_timeout, read_timeout). + See :meth:`requests.Session.request` documentation for details. + :raises: :class:`google.cloud.exceptions.NotFound` """ try: @@ -1022,6 +1056,7 @@ def download_to_filename( if_generation_not_match=if_generation_not_match, if_metageneration_match=if_metageneration_match, if_metageneration_not_match=if_metageneration_not_match, + timeout=timeout, ) except resumable_media.DataCorruption: # Delete the corrupt downloaded file. @@ -1046,6 +1081,7 @@ def download_as_string( if_generation_not_match=None, if_metageneration_match=None, if_metageneration_not_match=None, + timeout=_DEFAULT_TIMEOUT, ): """Download the contents of this blob as a bytes object. @@ -1087,6 +1123,14 @@ def download_as_string( :param if_metageneration_not_match: (Optional) Make the operation conditional on whether the blob's current metageneration does not match the given value. + :type timeout: float or tuple + :param timeout: + (Optional) The number of seconds the transport should wait for the + server response. Depending on the retry strategy, a request may be + repeated several times using the same timeout each time. + Can also be passed as a tuple (connect_timeout, read_timeout). + See :meth:`requests.Session.request` documentation for details. + :rtype: bytes :returns: The data stored in this blob. @@ -1103,6 +1147,7 @@ def download_as_string( if_generation_not_match=if_generation_not_match, if_metageneration_match=if_metageneration_match, if_metageneration_not_match=if_metageneration_not_match, + timeout=timeout, ) return string_buffer.getvalue() @@ -1203,6 +1248,7 @@ def _do_multipart_upload( if_generation_not_match, if_metageneration_match, if_metageneration_not_match, + timeout=_DEFAULT_TIMEOUT, ): """Perform a multipart upload. @@ -1256,6 +1302,14 @@ def _do_multipart_upload( :param if_metageneration_not_match: (Optional) Make the operation conditional on whether the blob's current metageneration does not match the given value. + :type timeout: float or tuple + :param timeout: + (Optional) The number of seconds the transport should wait for the + server response. Depending on the retry strategy, a request may be + repeated several times using the same timeout each time. + Can also be passed as a tuple (connect_timeout, read_timeout). + See :meth:`requests.Session.request` documentation for details. + :rtype: :class:`~requests.Response` :returns: The "200 OK" response object returned after the multipart upload request. @@ -1318,7 +1372,9 @@ def _do_multipart_upload( max_retries=num_retries ) - response = upload.transmit(transport, data, object_metadata, content_type) + response = upload.transmit( + transport, data, object_metadata, content_type, timeout=timeout + ) return response @@ -1336,6 +1392,7 @@ def _initiate_resumable_upload( if_generation_not_match=None, if_metageneration_match=None, if_metageneration_not_match=None, + timeout=_DEFAULT_TIMEOUT, ): """Initiate a resumable upload. @@ -1402,6 +1459,14 @@ def _initiate_resumable_upload( :param if_metageneration_not_match: (Optional) Make the operation conditional on whether the blob's current metageneration does not match the given value. + :type timeout: float or tuple + :param timeout: + (Optional) The number of seconds the transport should wait for the + server response. Depending on the retry strategy, a request may be + repeated several times using the same timeout each time. + Can also be passed as a tuple (connect_timeout, read_timeout). + See :meth:`requests.Session.request` documentation for details. + :rtype: tuple :returns: Pair of @@ -1472,6 +1537,7 @@ def _initiate_resumable_upload( content_type, total_bytes=size, stream_final=False, + timeout=timeout, ) return upload, transport @@ -1488,6 +1554,7 @@ def _do_resumable_upload( if_generation_not_match, if_metageneration_match, if_metageneration_not_match, + timeout=_DEFAULT_TIMEOUT, ): """Perform a resumable upload. @@ -1544,6 +1611,14 @@ def _do_resumable_upload( :param if_metageneration_not_match: (Optional) Make the operation conditional on whether the blob's current metageneration does not match the given value. + :type timeout: float or tuple + :param timeout: + (Optional) The number of seconds the transport should wait for the + server response. Depending on the retry strategy, a request may be + repeated several times using the same timeout each time. + Can also be passed as a tuple (connect_timeout, read_timeout). + See :meth:`requests.Session.request` documentation for details. + :rtype: :class:`~requests.Response` :returns: The "200 OK" response object returned after the final chunk is uploaded. @@ -1559,10 +1634,11 @@ def _do_resumable_upload( if_generation_not_match=if_generation_not_match, if_metageneration_match=if_metageneration_match, if_metageneration_not_match=if_metageneration_not_match, + timeout=timeout, ) while not upload.finished: - response = upload.transmit_next_chunk(transport) + response = upload.transmit_next_chunk(transport, timeout=timeout) return response @@ -1578,6 +1654,7 @@ def _do_upload( if_generation_not_match, if_metageneration_match, if_metageneration_not_match, + timeout=_DEFAULT_TIMEOUT, ): """Determine an upload strategy and then perform the upload. @@ -1635,6 +1712,14 @@ def _do_upload( :param if_metageneration_not_match: (Optional) Make the operation conditional on whether the blob's current metageneration does not match the given value. + :type timeout: float or tuple + :param timeout: + (Optional) The number of seconds the transport should wait for the + server response. Depending on the retry strategy, a request may be + repeated several times using the same timeout each time. + Can also be passed as a tuple (connect_timeout, read_timeout). + See :meth:`requests.Session.request` documentation for details. + :rtype: dict :returns: The parsed JSON from the "200 OK" response. This will be the **only** response in the multipart case and it will be the @@ -1652,6 +1737,7 @@ def _do_upload( if_generation_not_match, if_metageneration_match, if_metageneration_not_match, + timeout=timeout, ) else: response = self._do_resumable_upload( @@ -1665,6 +1751,7 @@ def _do_upload( if_generation_not_match, if_metageneration_match, if_metageneration_not_match, + timeout=timeout, ) return response.json() @@ -1682,6 +1769,7 @@ def upload_from_file( if_generation_not_match=None, if_metageneration_match=None, if_metageneration_not_match=None, + timeout=_DEFAULT_TIMEOUT, ): """Upload the contents of this blob from a file-like object. @@ -1768,6 +1856,14 @@ def upload_from_file( :param if_metageneration_not_match: (Optional) Make the operation conditional on whether the blob's current metageneration does not match the given value. + :type timeout: float or tuple + :param timeout: + (Optional) The number of seconds the transport should wait for the + server response. Depending on the retry strategy, a request may be + repeated several times using the same timeout each time. + Can also be passed as a tuple (connect_timeout, read_timeout). + See :meth:`requests.Session.request` documentation for details. + :raises: :class:`~google.cloud.exceptions.GoogleCloudError` if the upload response returns an error status. @@ -1793,6 +1889,7 @@ def upload_from_file( if_generation_not_match, if_metageneration_match, if_metageneration_not_match, + timeout=timeout, ) self._set_properties(created_json) except resumable_media.InvalidResponse as exc: @@ -1808,6 +1905,7 @@ def upload_from_filename( if_generation_not_match=None, if_metageneration_match=None, if_metageneration_not_match=None, + timeout=_DEFAULT_TIMEOUT, ): """Upload this blob's contents from the content of a named file. @@ -1866,6 +1964,14 @@ def upload_from_filename( :type if_metageneration_not_match: long :param if_metageneration_not_match: (Optional) Make the operation conditional on whether the blob's current metageneration does not match the given value. + + :type timeout: float or tuple + :param timeout: + (Optional) The number of seconds the transport should wait for the + server response. Depending on the retry strategy, a request may be + repeated several times using the same timeout each time. + Can also be passed as a tuple (connect_timeout, read_timeout). + See :meth:`requests.Session.request` documentation for details. """ content_type = self._get_content_type(content_type, filename=filename) @@ -1881,6 +1987,7 @@ def upload_from_filename( if_generation_not_match=if_generation_not_match, if_metageneration_match=if_metageneration_match, if_metageneration_not_match=if_metageneration_not_match, + timeout=timeout, ) def upload_from_string( @@ -1893,6 +2000,7 @@ def upload_from_string( if_generation_not_match=None, if_metageneration_match=None, if_metageneration_not_match=None, + timeout=_DEFAULT_TIMEOUT, ): """Upload contents of this blob from the provided string. @@ -1946,6 +2054,14 @@ def upload_from_string( :type if_metageneration_not_match: long :param if_metageneration_not_match: (Optional) Make the operation conditional on whether the blob's current metageneration does not match the given value. + + :type timeout: float or tuple + :param timeout: + (Optional) The number of seconds the transport should wait for the + server response. Depending on the retry strategy, a request may be + repeated several times using the same timeout each time. + Can also be passed as a tuple (connect_timeout, read_timeout). + See :meth:`requests.Session.request` documentation for details. """ data = _to_bytes(data, encoding="utf-8") string_buffer = BytesIO(data) @@ -1959,10 +2075,16 @@ def upload_from_string( if_generation_not_match=if_generation_not_match, if_metageneration_match=if_metageneration_match, if_metageneration_not_match=if_metageneration_not_match, + timeout=timeout, ) def create_resumable_upload_session( - self, content_type=None, size=None, origin=None, client=None + self, + content_type=None, + size=None, + origin=None, + client=None, + timeout=_DEFAULT_TIMEOUT, ): """Create a resumable upload session. @@ -2020,6 +2142,14 @@ def create_resumable_upload_session( :param client: (Optional) The client to use. If not passed, falls back to the ``client`` stored on the blob's bucket. + :type timeout: float or tuple + :param timeout: + (Optional) The number of seconds the transport should wait for the + server response. Depending on the retry strategy, a request may be + repeated several times using the same timeout each time. + Can also be passed as a tuple (connect_timeout, read_timeout). + See :meth:`requests.Session.request` documentation for details. + :rtype: str :returns: The resumable upload session URL. The upload can be completed by making an HTTP PUT request with the @@ -2048,6 +2178,7 @@ def create_resumable_upload_session( predefined_acl=None, extra_headers=extra_headers, chunk_size=self._CHUNK_SIZE_MULTIPLE, + timeout=timeout, ) return upload.resumable_url @@ -2510,6 +2641,7 @@ def update_storage_class( if_source_generation_not_match=None, if_source_metageneration_match=None, if_source_metageneration_not_match=None, + timeout=_DEFAULT_TIMEOUT, ): """Update blob's storage class via a rewrite-in-place. This helper will wait for the rewrite to complete before returning, so it may take some @@ -2592,6 +2724,14 @@ def update_storage_class( conditional on whether the source object's current metageneration does not match the given value. + + :type timeout: float or tuple + :param timeout: + (Optional) The number of seconds the transport should wait for the + server response. Depending on the retry strategy, a request may be + repeated several times using the same timeout each time. + Can also be passed as a tuple (connect_timeout, read_timeout). + See :meth:`requests.Session.request` documentation for details. """ if new_class not in self.STORAGE_CLASSES: raise ValueError("Invalid storage class: %s" % (new_class,)) @@ -2610,6 +2750,7 @@ def update_storage_class( if_source_generation_not_match=if_source_generation_not_match, if_source_metageneration_match=if_source_metageneration_match, if_source_metageneration_not_match=if_source_metageneration_not_match, + timeout=timeout, ) while token is not None: token, _, _ = self.rewrite( @@ -2623,6 +2764,7 @@ def update_storage_class( if_source_generation_not_match=if_source_generation_not_match, if_source_metageneration_match=if_source_metageneration_match, if_source_metageneration_not_match=if_source_metageneration_not_match, + timeout=timeout, ) cache_control = _scalar_property("cacheControl") diff --git a/setup.py b/setup.py index 0c149d303..b2ded72b6 100644 --- a/setup.py +++ b/setup.py @@ -31,7 +31,7 @@ dependencies = [ "google-auth >= 1.11.0, < 2.0dev", "google-cloud-core >= 1.2.0, < 2.0dev", - "google-resumable-media >= 0.5.0, < 0.6dev", + "google-resumable-media >= 0.6.0, < 0.7dev", ] extras = {} diff --git a/tests/unit/test_blob.py b/tests/unit/test_blob.py index 017e86216..4635b050e 100644 --- a/tests/unit/test_blob.py +++ b/tests/unit/test_blob.py @@ -936,7 +936,7 @@ def _mock_requests_response(status_code, headers, content=b""): response.request = requests.Request("POST", "http://example.com").prepare() return response - def _do_download_helper_wo_chunks(self, w_range, raw_download): + def _do_download_helper_wo_chunks(self, w_range, raw_download, timeout=None): blob_name = "blob-name" client = mock.Mock() bucket = _Bucket(client) @@ -953,6 +953,13 @@ def _do_download_helper_wo_chunks(self, w_range, raw_download): else: patch = mock.patch("google.cloud.storage.blob.Download") + if timeout is None: + expected_timeout = self._get_default_timeout() + timeout_kwarg = {} + else: + expected_timeout = timeout + timeout_kwarg = {"timeout": timeout} + with patch as patched: if w_range: blob._do_download( @@ -963,6 +970,7 @@ def _do_download_helper_wo_chunks(self, w_range, raw_download): start=1, end=3, raw_download=raw_download, + **timeout_kwarg ) else: blob._do_download( @@ -971,6 +979,7 @@ def _do_download_helper_wo_chunks(self, w_range, raw_download): download_url, headers, raw_download=raw_download, + **timeout_kwarg ) if w_range: @@ -981,7 +990,10 @@ def _do_download_helper_wo_chunks(self, w_range, raw_download): patched.assert_called_once_with( download_url, stream=file_obj, headers=headers, start=None, end=None ) - patched.return_value.consume.assert_called_once_with(transport) + + patched.return_value.consume.assert_called_once_with( + transport, timeout=expected_timeout + ) def test__do_download_wo_chunks_wo_range_wo_raw(self): self._do_download_helper_wo_chunks(w_range=False, raw_download=False) @@ -995,7 +1007,12 @@ def test__do_download_wo_chunks_wo_range_w_raw(self): def test__do_download_wo_chunks_w_range_w_raw(self): self._do_download_helper_wo_chunks(w_range=True, raw_download=True) - def _do_download_helper_w_chunks(self, w_range, raw_download): + def test__do_download_wo_chunks_w_custom_timeout(self): + self._do_download_helper_wo_chunks( + w_range=False, raw_download=False, timeout=9.58 + ) + + def _do_download_helper_w_chunks(self, w_range, raw_download, timeout=None): blob_name = "blob-name" client = mock.Mock(_credentials=_make_credentials(), spec=["_credentials"]) bucket = _Bucket(client) @@ -1010,7 +1027,7 @@ def _do_download_helper_w_chunks(self, w_range, raw_download): download = mock.Mock(finished=False, spec=["finished", "consume_next_chunk"]) - def side_effect(_): + def side_effect(*args, **kwargs): download.finished = True download.consume_next_chunk.side_effect = side_effect @@ -1020,6 +1037,13 @@ def side_effect(_): else: patch = mock.patch("google.cloud.storage.blob.ChunkedDownload") + if timeout is None: + expected_timeout = self._get_default_timeout() + timeout_kwarg = {} + else: + expected_timeout = timeout + timeout_kwarg = {"timeout": timeout} + with patch as patched: patched.return_value = download if w_range: @@ -1031,6 +1055,7 @@ def side_effect(_): start=1, end=3, raw_download=raw_download, + **timeout_kwarg ) else: blob._do_download( @@ -1039,6 +1064,7 @@ def side_effect(_): download_url, headers, raw_download=raw_download, + **timeout_kwarg ) if w_range: @@ -1049,7 +1075,9 @@ def side_effect(_): patched.assert_called_once_with( download_url, chunk_size, file_obj, headers=headers, start=0, end=None ) - download.consume_next_chunk.assert_called_once_with(transport) + download.consume_next_chunk.assert_called_once_with( + transport, timeout=expected_timeout + ) def test__do_download_w_chunks_wo_range_wo_raw(self): self._do_download_helper_w_chunks(w_range=False, raw_download=False) @@ -1063,6 +1091,9 @@ def test__do_download_w_chunks_wo_range_w_raw(self): def test__do_download_w_chunks_w_range_w_raw(self): self._do_download_helper_w_chunks(w_range=True, raw_download=True) + def test__do_download_w_chunks_w_custom_timeout(self): + self._do_download_helper_w_chunks(w_range=True, raw_download=True, timeout=9.58) + def test_download_to_file_with_failure(self): import requests from google.resumable_media import InvalidResponse @@ -1091,7 +1122,14 @@ def test_download_to_file_with_failure(self): headers = {"accept-encoding": "gzip"} blob._do_download.assert_called_once_with( - client._http, file_obj, media_link, headers, None, None, False + client._http, + file_obj, + media_link, + headers, + None, + None, + False, + timeout=self._get_default_timeout(), ) def test_download_to_file_wo_media_link(self): @@ -1114,7 +1152,14 @@ def test_download_to_file_wo_media_link(self): ) headers = {"accept-encoding": "gzip"} blob._do_download.assert_called_once_with( - client._http, file_obj, expected_url, headers, None, None, False + client._http, + file_obj, + expected_url, + headers, + None, + None, + False, + timeout=self._get_default_timeout(), ) def test_download_to_file_w_generation_match(self): @@ -1136,10 +1181,17 @@ def test_download_to_file_w_generation_match(self): blob.download_to_file(file_obj, if_generation_not_match=GENERATION_NUMBER) blob._do_download.assert_called_once_with( - client._http, file_obj, EXPECTED_URL, HEADERS, None, None, False + client._http, + file_obj, + EXPECTED_URL, + HEADERS, + None, + None, + False, + timeout=self._get_default_timeout(), ) - def _download_to_file_helper(self, use_chunks, raw_download): + def _download_to_file_helper(self, use_chunks, raw_download, timeout=None): blob_name = "blob-name" client = mock.Mock(spec=[u"_http"]) bucket = _Bucket(client) @@ -1151,15 +1203,29 @@ def _download_to_file_helper(self, use_chunks, raw_download): blob.chunk_size = 3 blob._do_download = mock.Mock() + if timeout is None: + expected_timeout = self._get_default_timeout() + timeout_kwarg = {} + else: + expected_timeout = timeout + timeout_kwarg = {"timeout": timeout} + file_obj = io.BytesIO() if raw_download: - blob.download_to_file(file_obj, raw_download=True) + blob.download_to_file(file_obj, raw_download=True, **timeout_kwarg) else: - blob.download_to_file(file_obj) + blob.download_to_file(file_obj, **timeout_kwarg) headers = {"accept-encoding": "gzip"} blob._do_download.assert_called_once_with( - client._http, file_obj, media_link, headers, None, None, raw_download + client._http, + file_obj, + media_link, + headers, + None, + None, + raw_download, + timeout=expected_timeout, ) def test_download_to_file_wo_chunks_wo_raw(self): @@ -1174,7 +1240,12 @@ def test_download_to_file_wo_chunks_w_raw(self): def test_download_to_file_w_chunks_w_raw(self): self._download_to_file_helper(use_chunks=True, raw_download=True) - def _download_to_filename_helper(self, updated, raw_download): + def test_download_to_file_w_custom_timeout(self): + self._download_to_file_helper( + use_chunks=False, raw_download=False, timeout=9.58 + ) + + def _download_to_filename_helper(self, updated, raw_download, timeout=None): import os from google.cloud.storage._helpers import _convert_to_timestamp from google.cloud._testing import _NamedTemporaryFile @@ -1191,7 +1262,13 @@ def _download_to_filename_helper(self, updated, raw_download): blob._do_download = mock.Mock() with _NamedTemporaryFile() as temp: - blob.download_to_filename(temp.name, raw_download=raw_download) + if timeout is None: + blob.download_to_filename(temp.name, raw_download=raw_download) + else: + blob.download_to_filename( + temp.name, raw_download=raw_download, timeout=timeout, + ) + if updated is None: self.assertIsNone(blob.updated) else: @@ -1202,9 +1279,18 @@ def _download_to_filename_helper(self, updated, raw_download): updated_time = blob.updated.timestamp() self.assertEqual(mtime, updated_time) + expected_timeout = self._get_default_timeout() if timeout is None else timeout + headers = {"accept-encoding": "gzip"} blob._do_download.assert_called_once_with( - client._http, mock.ANY, media_link, headers, None, None, raw_download + client._http, + mock.ANY, + media_link, + headers, + None, + None, + raw_download, + timeout=expected_timeout, ) stream = blob._do_download.mock_calls[0].args[1] self.assertEqual(stream.name, temp.name) @@ -1228,7 +1314,14 @@ def test_download_to_filename_w_generation_match(self): blob.download_to_filename(temp.name, if_generation_match=GENERATION_NUMBER) blob._do_download.assert_called_once_with( - client._http, mock.ANY, EXPECTED_LINK, HEADERS, None, None, False + client._http, + mock.ANY, + EXPECTED_LINK, + HEADERS, + None, + None, + False, + timeout=self._get_default_timeout(), ) def test_download_to_filename_w_updated_wo_raw(self): @@ -1245,6 +1338,11 @@ def test_download_to_filename_w_updated_w_raw(self): def test_download_to_filename_wo_updated_w_raw(self): self._download_to_filename_helper(updated=None, raw_download=True) + def test_download_to_filename_w_custom_timeout(self): + self._download_to_filename_helper( + updated=None, raw_download=False, timeout=9.58 + ) + def test_download_to_filename_corrupted(self): from google.resumable_media import DataCorruption @@ -1273,7 +1371,14 @@ def test_download_to_filename_corrupted(self): headers = {"accept-encoding": "gzip"} blob._do_download.assert_called_once_with( - client._http, mock.ANY, media_link, headers, None, None, False + client._http, + mock.ANY, + media_link, + headers, + None, + None, + False, + timeout=self._get_default_timeout(), ) stream = blob._do_download.mock_calls[0].args[1] self.assertEqual(stream.name, filename) @@ -1300,12 +1405,19 @@ def test_download_to_filename_w_key(self): headers = {"accept-encoding": "gzip"} headers.update(_get_encryption_headers(key)) blob._do_download.assert_called_once_with( - client._http, mock.ANY, media_link, headers, None, None, False + client._http, + mock.ANY, + media_link, + headers, + None, + None, + False, + timeout=self._get_default_timeout(), ) stream = blob._do_download.mock_calls[0].args[1] self.assertEqual(stream.name, temp.name) - def _download_as_string_helper(self, raw_download): + def _download_as_string_helper(self, raw_download, timeout=None): blob_name = "blob-name" client = mock.Mock(spec=["_http"]) bucket = _Bucket(client) @@ -1314,12 +1426,27 @@ def _download_as_string_helper(self, raw_download): blob = self._make_one(blob_name, bucket=bucket, properties=properties) blob._do_download = mock.Mock() - fetched = blob.download_as_string(raw_download=raw_download) + if timeout is None: + expected_timeout = self._get_default_timeout() + fetched = blob.download_as_string(raw_download=raw_download) + else: + expected_timeout = timeout + fetched = blob.download_as_string( + raw_download=raw_download, timeout=timeout + ) + self.assertEqual(fetched, b"") headers = {"accept-encoding": "gzip"} blob._do_download.assert_called_once_with( - client._http, mock.ANY, media_link, headers, None, None, raw_download + client._http, + mock.ANY, + media_link, + headers, + None, + None, + raw_download, + timeout=expected_timeout, ) stream = blob._do_download.mock_calls[0].args[1] self.assertIsInstance(stream, io.BytesIO) @@ -1347,6 +1474,7 @@ def test_download_as_string_w_generation_match(self): if_generation_not_match=None, if_metageneration_match=None, if_metageneration_not_match=None, + timeout=self._get_default_timeout(), ) def test_download_as_string_wo_raw(self): @@ -1355,6 +1483,9 @@ def test_download_as_string_wo_raw(self): def test_download_as_string_w_raw(self): self._download_as_string_helper(raw_download=True) + def test_download_as_string_w_custom_timeout(self): + self._download_as_string_helper(raw_download=False, timeout=9.58) + def test__get_content_type_explicit(self): blob = self._make_one(u"blob-name", bucket=None) @@ -1471,6 +1602,7 @@ def _do_multipart_success( if_metageneration_match=None, if_metageneration_not_match=None, kms_key_name=None, + timeout=None, ): from six.moves.urllib.parse import urlencode @@ -1487,6 +1619,14 @@ def _do_multipart_success( data = b"data here hear hier" stream = io.BytesIO(data) content_type = u"application/xml" + + if timeout is None: + expected_timeout = self._get_default_timeout() + timeout_kwarg = {} + else: + expected_timeout = timeout + timeout_kwarg = {"timeout": timeout} + response = blob._do_multipart_upload( client, stream, @@ -1498,6 +1638,7 @@ def _do_multipart_success( if_generation_not_match, if_metageneration_match, if_metageneration_not_match, + **timeout_kwarg ) # Check the mocks and the returned value. @@ -1551,7 +1692,7 @@ def _do_multipart_success( ) headers = {"content-type": b'multipart/related; boundary="==0=="'} transport.request.assert_called_once_with( - "POST", upload_url, data=payload, headers=headers, timeout=mock.ANY + "POST", upload_url, data=payload, headers=headers, timeout=expected_timeout ) @mock.patch(u"google.resumable_media._upload.get_boundary", return_value=b"==0==") @@ -1598,6 +1739,10 @@ def test__do_multipart_upload_with_generation_match(self, mock_get_boundary): mock_get_boundary, if_generation_match=4, if_metageneration_match=4 ) + @mock.patch(u"google.resumable_media._upload.get_boundary", return_value=b"==0==") + def test__do_multipart_upload_with_custom_timeout(self, mock_get_boundary): + self._do_multipart_success(mock_get_boundary, timeout=9.58) + @mock.patch(u"google.resumable_media._upload.get_boundary", return_value=b"==0==") def test__do_multipart_upload_with_generation_not_match(self, mock_get_boundary): self._do_multipart_success( @@ -1635,6 +1780,7 @@ def _initiate_resumable_helper( if_metageneration_not_match=None, blob_chunk_size=786432, kms_key_name=None, + timeout=None, ): from six.moves.urllib.parse import urlencode from google.resumable_media.requests import ResumableUpload @@ -1665,6 +1811,14 @@ def _initiate_resumable_helper( data = b"hello hallo halo hi-low" stream = io.BytesIO(data) content_type = u"text/plain" + + if timeout is None: + expected_timeout = self._get_default_timeout() + timeout_kwarg = {} + else: + expected_timeout = timeout + timeout_kwarg = {"timeout": timeout} + upload, transport = blob._initiate_resumable_upload( client, stream, @@ -1678,6 +1832,7 @@ def _initiate_resumable_helper( if_generation_not_match=if_generation_not_match, if_metageneration_match=if_metageneration_match, if_metageneration_not_match=if_metageneration_not_match, + **timeout_kwarg ) # Check the returned values. @@ -1757,9 +1912,16 @@ def _initiate_resumable_helper( if extra_headers is not None: expected_headers.update(extra_headers) transport.request.assert_called_once_with( - "POST", upload_url, data=payload, headers=expected_headers, timeout=mock.ANY + "POST", + upload_url, + data=payload, + headers=expected_headers, + timeout=expected_timeout, ) + def test__initiate_resumable_upload_with_custom_timeout(self): + self._initiate_resumable_helper(timeout=9.58) + def test__initiate_resumable_upload_no_size(self): self._initiate_resumable_helper() @@ -1844,6 +2006,7 @@ def _do_resumable_upload_call0( if_generation_not_match=None, if_metageneration_match=None, if_metageneration_not_match=None, + timeout=None, ): # First mock transport.request() does initiates upload. upload_url = ( @@ -1861,7 +2024,7 @@ def _do_resumable_upload_call0( expected_headers["x-upload-content-length"] = str(size) payload = json.dumps({"name": blob.name}).encode("utf-8") return mock.call( - "POST", upload_url, data=payload, headers=expected_headers, timeout=mock.ANY + "POST", upload_url, data=payload, headers=expected_headers, timeout=timeout ) @staticmethod @@ -1876,6 +2039,7 @@ def _do_resumable_upload_call1( if_generation_not_match=None, if_metageneration_match=None, if_metageneration_not_match=None, + timeout=None, ): # Second mock transport.request() does sends first chunk. if size is None: @@ -1893,7 +2057,7 @@ def _do_resumable_upload_call1( resumable_url, data=payload, headers=expected_headers, - timeout=mock.ANY, + timeout=timeout, ) @staticmethod @@ -1908,6 +2072,7 @@ def _do_resumable_upload_call2( if_generation_not_match=None, if_metageneration_match=None, if_metageneration_not_match=None, + timeout=None, ): # Third mock transport.request() does sends last chunk. content_range = "bytes {:d}-{:d}/{:d}".format( @@ -1923,7 +2088,7 @@ def _do_resumable_upload_call2( resumable_url, data=payload, headers=expected_headers, - timeout=mock.ANY, + timeout=timeout, ) def _do_resumable_helper( @@ -1935,6 +2100,7 @@ def _do_resumable_helper( if_generation_not_match=None, if_metageneration_match=None, if_metageneration_not_match=None, + timeout=None, ): bucket = _Bucket(name="yesterday") blob = self._make_one(u"blob-name", bucket=bucket) @@ -1962,6 +2128,14 @@ def _do_resumable_helper( client._connection.API_BASE_URL = "https://storage.googleapis.com" stream = io.BytesIO(data) content_type = u"text/html" + + if timeout is None: + expected_timeout = self._get_default_timeout() + timeout_kwarg = {} + else: + expected_timeout = timeout + timeout_kwarg = {"timeout": timeout} + response = blob._do_resumable_upload( client, stream, @@ -1973,6 +2147,7 @@ def _do_resumable_helper( if_generation_not_match, if_metageneration_match, if_metageneration_not_match, + **timeout_kwarg ) # Check the returned values. @@ -1989,6 +2164,7 @@ def _do_resumable_helper( if_generation_not_match=if_generation_not_match, if_metageneration_match=if_metageneration_match, if_metageneration_not_match=if_metageneration_not_match, + timeout=expected_timeout, ) call1 = self._do_resumable_upload_call1( blob, @@ -2001,6 +2177,7 @@ def _do_resumable_helper( if_generation_not_match=if_generation_not_match, if_metageneration_match=if_metageneration_match, if_metageneration_not_match=if_metageneration_not_match, + timeout=expected_timeout, ) call2 = self._do_resumable_upload_call2( blob, @@ -2013,9 +2190,13 @@ def _do_resumable_helper( if_generation_not_match=if_generation_not_match, if_metageneration_match=if_metageneration_match, if_metageneration_not_match=if_metageneration_not_match, + timeout=expected_timeout, ) self.assertEqual(transport.request.mock_calls, [call0, call1, call2]) + def test__do_resumable_upload_with_custom_timeout(self): + self._do_resumable_helper(timeout=9.58) + def test__do_resumable_upload_no_size(self): self._do_resumable_helper() @@ -2038,6 +2219,7 @@ def _do_upload_helper( if_metageneration_match=None, if_metageneration_not_match=None, size=None, + timeout=None, ): from google.cloud.storage.blob import _MAX_MULTIPART_SIZE @@ -2061,6 +2243,14 @@ def _do_upload_helper( content_type = u"video/mp4" if size is None: size = 12345654321 + + if timeout is None: + expected_timeout = self._get_default_timeout() + timeout_kwarg = {} + else: + expected_timeout = timeout + timeout_kwarg = {"timeout": timeout} + # Make the request and check the mocks. created_json = blob._do_upload( client, @@ -2073,6 +2263,7 @@ def _do_upload_helper( if_generation_not_match, if_metageneration_match, if_metageneration_not_match, + **timeout_kwarg ) self.assertIs(created_json, mock.sentinel.json) response.json.assert_called_once_with() @@ -2088,6 +2279,7 @@ def _do_upload_helper( if_generation_not_match, if_metageneration_match, if_metageneration_not_match, + timeout=expected_timeout, ) blob._do_resumable_upload.assert_not_called() else: @@ -2103,6 +2295,7 @@ def _do_upload_helper( if_generation_not_match, if_metageneration_match, if_metageneration_not_match, + timeout=expected_timeout, ) def test__do_upload_uses_multipart(self): @@ -2110,12 +2303,25 @@ def test__do_upload_uses_multipart(self): self._do_upload_helper(size=_MAX_MULTIPART_SIZE) + def test__do_upload_uses_multipart_w_custom_timeout(self): + from google.cloud.storage.blob import _MAX_MULTIPART_SIZE + + self._do_upload_helper(size=_MAX_MULTIPART_SIZE, timeout=9.58) + def test__do_upload_uses_resumable(self): from google.cloud.storage.blob import _MAX_MULTIPART_SIZE chunk_size = 256 * 1024 # 256KB self._do_upload_helper(chunk_size=chunk_size, size=_MAX_MULTIPART_SIZE + 1) + def test__do_upload_uses_resumable_w_custom_timeout(self): + from google.cloud.storage.blob import _MAX_MULTIPART_SIZE + + chunk_size = 256 * 1024 # 256KB + self._do_upload_helper( + chunk_size=chunk_size, size=_MAX_MULTIPART_SIZE + 1, timeout=9.58 + ) + def test__do_upload_with_retry(self): self._do_upload_helper(num_retries=20) @@ -2150,6 +2356,8 @@ def _upload_from_file_helper(self, side_effect=None, **kwargs): new_updated = datetime.datetime(2017, 1, 1, 9, 9, 9, 81000, tzinfo=UTC) self.assertEqual(blob.updated, new_updated) + expected_timeout = kwargs.get("timeout", self._get_default_timeout()) + # Check the mock. num_retries = kwargs.get("num_retries") blob._do_upload.assert_called_once_with( @@ -2163,6 +2371,7 @@ def _upload_from_file_helper(self, side_effect=None, **kwargs): if_generation_not_match, if_metageneration_match, if_metageneration_not_match, + timeout=expected_timeout, ) return stream @@ -2183,6 +2392,9 @@ def test_upload_from_file_with_rewind(self): stream = self._upload_from_file_helper(rewind=True) assert stream.tell() == 0 + def test_upload_from_file_with_custom_timeout(self): + self._upload_from_file_helper(timeout=9.58) + def test_upload_from_file_failure(self): import requests @@ -2201,7 +2413,9 @@ def test_upload_from_file_failure(self): self.assertIn(message, exc_info.exception.message) self.assertEqual(exc_info.exception.errors, []) - def _do_upload_mock_call_helper(self, blob, client, content_type, size): + def _do_upload_mock_call_helper( + self, blob, client, content_type, size, timeout=None + ): self.assertEqual(blob._do_upload.call_count, 1) mock_call = blob._do_upload.mock_calls[0] call_name, pos_args, kwargs = mock_call @@ -2216,7 +2430,9 @@ def _do_upload_mock_call_helper(self, blob, client, content_type, size): self.assertIsNone(pos_args[7]) # if_generation_not_match self.assertIsNone(pos_args[8]) # if_metageneration_match self.assertIsNone(pos_args[9]) # if_metageneration_not_match - self.assertEqual(kwargs, {}) + + expected_timeout = self._get_default_timeout() if timeout is None else timeout + self.assertEqual(kwargs, {"timeout": expected_timeout}) return pos_args[1] @@ -2251,6 +2467,32 @@ def test_upload_from_filename(self): self.assertEqual(stream.mode, "rb") self.assertEqual(stream.name, temp.name) + def test_upload_from_filename_w_custom_timeout(self): + from google.cloud._testing import _NamedTemporaryFile + + blob = self._make_one("blob-name", bucket=None) + # Mock low-level upload helper on blob (it is tested elsewhere). + created_json = {"metadata": {"mint": "ice-cream"}} + blob._do_upload = mock.Mock(return_value=created_json, spec=[]) + # Make sure `metadata` is empty before the request. + self.assertIsNone(blob.metadata) + + data = b"soooo much data" + content_type = u"image/svg+xml" + client = mock.sentinel.client + with _NamedTemporaryFile() as temp: + with open(temp.name, "wb") as file_obj: + file_obj.write(data) + + blob.upload_from_filename( + temp.name, content_type=content_type, client=client, timeout=9.58 + ) + + # Check the mock. + self._do_upload_mock_call_helper( + blob, client, content_type, len(data), timeout=9.58 + ) + def _upload_from_string_helper(self, data, **kwargs): from google.cloud._helpers import _to_bytes @@ -2272,11 +2514,19 @@ def _upload_from_string_helper(self, data, **kwargs): # Check the mock. payload = _to_bytes(data, encoding="utf-8") stream = self._do_upload_mock_call_helper( - blob, client, "text/plain", len(payload) + blob, + client, + "text/plain", + len(payload), + kwargs.get("timeout", self._get_default_timeout()), ) self.assertIsInstance(stream, io.BytesIO) self.assertEqual(stream.getvalue(), payload) + def test_upload_from_string_w_custom_timeout(self): + data = b"XB]jb\xb8tad\xe0" + self._upload_from_string_helper(data, timeout=9.58) + def test_upload_from_string_w_bytes(self): data = b"XB]jb\xb8tad\xe0" self._upload_from_string_helper(data) @@ -2285,7 +2535,9 @@ def test_upload_from_string_w_text(self): data = u"\N{snowman} \N{sailboat}" self._upload_from_string_helper(data) - def _create_resumable_upload_session_helper(self, origin=None, side_effect=None): + def _create_resumable_upload_session_helper( + self, origin=None, side_effect=None, timeout=None + ): bucket = _Bucket(name="alex-trebek") blob = self._make_one("blob-name", bucket=bucket) chunk_size = 99 * blob._CHUNK_SIZE_MULTIPLE @@ -2303,8 +2555,20 @@ def _create_resumable_upload_session_helper(self, origin=None, side_effect=None) size = 10000 client = mock.Mock(_http=transport, _connection=_Connection, spec=[u"_http"]) client._connection.API_BASE_URL = "https://storage.googleapis.com" + + if timeout is None: + expected_timeout = self._get_default_timeout() + timeout_kwarg = {} + else: + expected_timeout = timeout + timeout_kwarg = {"timeout": timeout} + new_url = blob.create_resumable_upload_session( - content_type=content_type, size=size, origin=origin, client=client + content_type=content_type, + size=size, + origin=origin, + client=client, + **timeout_kwarg ) # Check the returned value and (lack of) side-effect. @@ -2326,12 +2590,19 @@ def _create_resumable_upload_session_helper(self, origin=None, side_effect=None) if origin is not None: expected_headers["Origin"] = origin transport.request.assert_called_once_with( - "POST", upload_url, data=payload, headers=expected_headers, timeout=mock.ANY + "POST", + upload_url, + data=payload, + headers=expected_headers, + timeout=expected_timeout, ) def test_create_resumable_upload_session(self): self._create_resumable_upload_session_helper() + def test_create_resumable_upload_session_with_custom_timeout(self): + self._create_resumable_upload_session_helper(timeout=9.58) + def test_create_resumable_upload_session_with_origin(self): self._create_resumable_upload_session_helper(origin="http://google.com") @@ -3252,6 +3523,41 @@ def test_update_storage_class_large_file(self): self.assertEqual(blob.storage_class, "NEARLINE") + def test_update_storage_class_with_custom_timeout(self): + BLOB_NAME = "blob-name" + STORAGE_CLASS = u"NEARLINE" + TOKEN = "TOKEN" + INCOMPLETE_RESPONSE = { + "totalBytesRewritten": 42, + "objectSize": 84, + "done": False, + "rewriteToken": TOKEN, + "resource": {"storageClass": STORAGE_CLASS}, + } + COMPLETE_RESPONSE = { + "totalBytesRewritten": 84, + "objectSize": 84, + "done": True, + "resource": {"storageClass": STORAGE_CLASS}, + } + response_1 = ({"status": http_client.OK}, INCOMPLETE_RESPONSE) + response_2 = ({"status": http_client.OK}, COMPLETE_RESPONSE) + connection = _Connection(response_1, response_2) + client = _Client(connection) + bucket = _Bucket(client=client) + blob = self._make_one(BLOB_NAME, bucket=bucket) + + blob.update_storage_class("NEARLINE", timeout=9.58) + + self.assertEqual(blob.storage_class, "NEARLINE") + + kw = connection._requested + self.assertEqual(len(kw), 2) + + for kw_item in kw: + self.assertIn("timeout", kw_item) + self.assertEqual(kw_item["timeout"], 9.58) + def test_update_storage_class_wo_encryption_key(self): BLOB_NAME = "blob-name" STORAGE_CLASS = u"NEARLINE"