Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Adding error remapping for Blob.download_to_file(). #3338

Merged
merged 1 commit into from
Apr 27, 2017
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
47 changes: 37 additions & 10 deletions storage/google/cloud/storage/blob.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Expand Down Expand Up @@ -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.
Expand Down
103 changes: 93 additions & 10 deletions storage/tests/unit/test_blob.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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'
Expand Down Expand Up @@ -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'
Expand Down Expand Up @@ -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.
Expand Down