diff --git a/storage/google/cloud/storage/blob.py b/storage/google/cloud/storage/blob.py index 15a7d83c129e..364e1c5f3293 100644 --- a/storage/google/cloud/storage/blob.py +++ b/storage/google/cloud/storage/blob.py @@ -338,6 +338,36 @@ def _get_download_url(self): else: return self.media_link + def _do_download(self, transport, file_obj, download_url, headers): + """Perform a download without any error handling. + + This is intended to be called by :meth:`download_to_file` so it can + be wrapped with error handling / remapping. + + :type transport: + :class:`~google.auth.transport.requests.AuthorizedSession` + :param transport: The transport (with credentials) that will + make authenticated requests. + + :type file_obj: file + :param file_obj: A file handle to which to write the blob's data. + + :type download_url: str + :param download_url: The URL where the media can be accessed. + + :type headers: dict + :param headers: Optional headers to be sent with the request(s). + """ + if self.chunk_size is None: + download = resumable_media.Download(download_url, headers=headers) + response = download.consume(transport) + file_obj.write(response.content) + else: + download = resumable_media.ChunkedDownload( + download_url, self.chunk_size, file_obj, headers=headers) + while not download.finished: + download.consume_next_chunk(transport) + def download_to_file(self, file_obj, client=None): """Download the contents of this blob into a file-like object. @@ -378,16 +408,13 @@ def download_to_file(self, file_obj, client=None): transport = google.auth.transport.requests.AuthorizedSession( client._credentials) - # Download the content. - if self.chunk_size is None: - download = resumable_media.Download(download_url, headers=headers) - response = download.consume(transport) - file_obj.write(response.content) - else: - download = resumable_media.ChunkedDownload( - download_url, self.chunk_size, file_obj, headers=headers) - while not download.finished: - download.consume_next_chunk(transport) + try: + self._do_download(transport, file_obj, download_url, headers) + except resumable_media.InvalidResponse as exc: + response = exc.response + faux_response = httplib2.Response({'status': response.status_code}) + raise make_exception(faux_response, response.content, + error_info=download_url, use_json=False) def download_to_filename(self, filename, client=None): """Download the contents of this blob into a named file. diff --git a/storage/tests/unit/test_blob.py b/storage/tests/unit/test_blob.py index a70248da1b76..517890663e0f 100644 --- a/storage/tests/unit/test_blob.py +++ b/storage/tests/unit/test_blob.py @@ -381,11 +381,102 @@ def _check_session_mocks(self, client, fake_session_factory, 'GET', expected_url, data=None, headers=headers) self.assertEqual(fake_transport.request.mock_calls, [call, call]) + def test__do_download_simple(self): + from io import BytesIO + from six.moves import http_client + + blob_name = 'blob-name' + # Create a fake client/bucket and use them in the Blob() constructor. + client = mock.Mock( + _credentials=_make_credentials(), spec=['_credentials']) + bucket = _Bucket(client) + blob = self._make_one(blob_name, bucket=bucket) + + # Make sure this will not be chunked. + self.assertIsNone(blob.chunk_size) + + transport = mock.Mock(spec=['request']) + transport.request.return_value = self._mock_requests_response( + http_client.OK, + {'content-length': '6', 'content-range': 'bytes 0-5/6'}, + content=b'abcdef') + file_obj = BytesIO() + download_url = 'http://test.invalid' + headers = {} + blob._do_download(transport, file_obj, download_url, headers) + # Make sure the download was as expected. + self.assertEqual(file_obj.getvalue(), b'abcdef') + + transport.request.assert_called_once_with( + 'GET', download_url, data=None, headers=headers) + + def test__do_download_chunked(self): + from io import BytesIO + + blob_name = 'blob-name' + # Create a fake client/bucket and use them in the Blob() constructor. + client = mock.Mock( + _credentials=_make_credentials(), spec=['_credentials']) + bucket = _Bucket(client) + blob = self._make_one(blob_name, bucket=bucket) + + # Modify the blob so there there will be 2 chunks of size 3. + blob._CHUNK_SIZE_MULTIPLE = 1 + blob.chunk_size = 3 + + transport = self._mock_transport() + file_obj = BytesIO() + download_url = 'http://test.invalid' + headers = {} + blob._do_download(transport, file_obj, download_url, headers) + # Make sure the download was as expected. + self.assertEqual(file_obj.getvalue(), b'abcdef') + + # Check that the transport was called exactly twice. + self.assertEqual(transport.request.call_count, 2) + # ``headers`` was modified (in place) once for each API call. + self.assertEqual(headers, {'range': 'bytes=3-5'}) + call = mock.call( + 'GET', download_url, data=None, headers=headers) + self.assertEqual(transport.request.mock_calls, [call, call]) + + @mock.patch('google.auth.transport.requests.AuthorizedSession') + def test_download_to_file_with_failure(self, fake_session_factory): + from io import BytesIO + from six.moves import http_client + from google.cloud import exceptions + + blob_name = 'blob-name' + transport = mock.Mock(spec=['request']) + bad_response_headers = { + 'Content-Length': '9', + 'Content-Type': 'text/html; charset=UTF-8', + } + transport.request.return_value = self._mock_requests_response( + http_client.NOT_FOUND, bad_response_headers, content=b'Not found') + fake_session_factory.return_value = transport + # Create a fake client/bucket and use them in the Blob() constructor. + client = mock.Mock( + _credentials=_make_credentials(), spec=['_credentials']) + bucket = _Bucket(client) + blob = self._make_one(blob_name, bucket=bucket) + # Set the media link on the blob + blob._properties['mediaLink'] = 'http://test.invalid' + + file_obj = BytesIO() + with self.assertRaises(exceptions.NotFound): + blob.download_to_file(file_obj) + + self.assertEqual(file_obj.tell(), 0) + # Check that exactly one transport was created. + fake_session_factory.assert_called_once_with(client._credentials) + # Check that the transport was called once. + transport.request.assert_called_once_with( + 'GET', blob.media_link, data=None, headers={}) + @mock.patch('google.auth.transport.requests.AuthorizedSession') def test_download_to_file_wo_media_link(self, fake_session_factory): from io import BytesIO - from six.moves.http_client import OK - from six.moves.http_client import PARTIAL_CONTENT blob_name = 'blob-name' fake_session_factory.return_value = self._mock_transport() @@ -413,7 +504,6 @@ def test_download_to_file_wo_media_link(self, fake_session_factory): def _download_to_file_helper(self, fake_session_factory, use_chunks=False): from io import BytesIO from six.moves.http_client import OK - from six.moves.http_client import PARTIAL_CONTENT blob_name = 'blob-name' fake_transport = self._mock_transport() @@ -459,8 +549,6 @@ def test_download_to_file_with_chunk_size(self): def test_download_to_filename(self, fake_session_factory): import os import time - from six.moves.http_client import OK - from six.moves.http_client import PARTIAL_CONTENT from google.cloud._testing import _NamedTemporaryFile blob_name = 'blob-name' @@ -493,8 +581,6 @@ def test_download_to_filename(self, fake_session_factory): def test_download_to_filename_w_key(self, fake_session_factory): import os import time - from six.moves.http_client import OK - from six.moves.http_client import PARTIAL_CONTENT from google.cloud._testing import _NamedTemporaryFile blob_name = 'blob-name' @@ -535,9 +621,6 @@ def test_download_to_filename_w_key(self, fake_session_factory): @mock.patch('google.auth.transport.requests.AuthorizedSession') def test_download_as_string(self, fake_session_factory): - from six.moves.http_client import OK - from six.moves.http_client import PARTIAL_CONTENT - blob_name = 'blob-name' fake_session_factory.return_value = self._mock_transport() # Create a fake client/bucket and use them in the Blob() constructor.