From df04844e4991cff7c455a70543450003d565da52 Mon Sep 17 00:00:00 2001 From: Danny Hermes Date: Tue, 1 Aug 2017 13:13:25 -0700 Subject: [PATCH] Streaming directly into file on storage downloads. --- storage/google/cloud/storage/blob.py | 5 ++--- storage/setup.py | 2 +- storage/tests/unit/test_blob.py | 25 ++++++++++++++++++------- 3 files changed, 21 insertions(+), 11 deletions(-) diff --git a/storage/google/cloud/storage/blob.py b/storage/google/cloud/storage/blob.py index 836cfb645f42..dd76def82ba7 100644 --- a/storage/google/cloud/storage/blob.py +++ b/storage/google/cloud/storage/blob.py @@ -414,9 +414,8 @@ def _do_download(self, transport, file_obj, download_url, headers): :param headers: Optional headers to be sent with the request(s). """ if self.chunk_size is None: - download = Download(download_url, headers=headers) - response = download.consume(transport) - file_obj.write(response.content) + download = Download(download_url, stream=file_obj, headers=headers) + download.consume(transport) else: download = ChunkedDownload( download_url, self.chunk_size, file_obj, headers=headers) diff --git a/storage/setup.py b/storage/setup.py index 8d11055fac77..0cf3de9cab4a 100644 --- a/storage/setup.py +++ b/storage/setup.py @@ -53,7 +53,7 @@ REQUIREMENTS = [ 'google-cloud-core >= 0.25.0, < 0.26dev', 'google-auth >= 1.0.0', - 'google-resumable-media >= 0.2.1', + 'google-resumable-media >= 0.2.2', 'requests >= 2.0.0', ] diff --git a/storage/tests/unit/test_blob.py b/storage/tests/unit/test_blob.py index 7cc0dadc2691..e0a41ee793d2 100644 --- a/storage/tests/unit/test_blob.py +++ b/storage/tests/unit/test_blob.py @@ -375,13 +375,20 @@ def test__get_download_url_on_the_fly_with_generation(self): self.assertEqual(download_url, expected_url) @staticmethod - def _mock_requests_response(status_code, headers, content=b''): + def _mock_requests_response( + status_code, headers, content=b'', stream=False): import requests response = requests.Response() response.status_code = status_code response.headers.update(headers) - response._content = content + if stream: + response.raw = io.BytesIO(content) + response._content = False + else: + response.raw = None + response._content = content + response.request = requests.Request( 'POST', 'http://example.com').prepare() return response @@ -429,7 +436,9 @@ def test__do_download_simple(self): transport.request.return_value = self._mock_requests_response( http_client.OK, {'content-length': '6', 'content-range': 'bytes 0-5/6'}, - content=b'abcdef') + content=b'abcdef', + stream=True, + ) file_obj = io.BytesIO() download_url = 'http://test.invalid' headers = {} @@ -438,7 +447,7 @@ def test__do_download_simple(self): self.assertEqual(file_obj.getvalue(), b'abcdef') transport.request.assert_called_once_with( - 'GET', download_url, data=None, headers=headers) + 'GET', download_url, data=None, headers=headers, stream=True) def test__do_download_chunked(self): blob_name = 'blob-name' @@ -493,7 +502,7 @@ def test_download_to_file_with_failure(self): self.assertEqual(file_obj.tell(), 0) # Check that the transport was called once. transport.request.assert_called_once_with( - 'GET', blob.media_link, data=None, headers={}) + 'GET', blob.media_link, data=None, headers={}, stream=True) def test_download_to_file_wo_media_link(self): blob_name = 'blob-name' @@ -535,7 +544,9 @@ def _download_to_file_helper(self, use_chunks=False): single_chunk_response = self._mock_requests_response( http_client.OK, {'content-length': '6', 'content-range': 'bytes 0-5/6'}, - content=b'abcdef') + content=b'abcdef', + stream=True, + ) transport.request.side_effect = [single_chunk_response] file_obj = io.BytesIO() @@ -546,7 +557,7 @@ def _download_to_file_helper(self, use_chunks=False): self._check_session_mocks(client, transport, media_link) else: transport.request.assert_called_once_with( - 'GET', media_link, data=None, headers={}) + 'GET', media_link, data=None, headers={}, stream=True) def test_download_to_file_default(self): self._download_to_file_helper()