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

Author: dhermes

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.

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.