Skip to content

Commit

Permalink
feat: download_to_filename deletes the empty file on a 404 (#1394)
Browse files Browse the repository at this point in the history
  • Loading branch information
andrewsg authored Dec 11, 2024
1 parent 6aef9b7 commit 066be2d
Show file tree
Hide file tree
Showing 3 changed files with 47 additions and 2 deletions.
2 changes: 2 additions & 0 deletions README.rst
Original file line number Diff line number Diff line change
Expand Up @@ -99,6 +99,8 @@ Miscellaneous
- Retry behavior is now identical between media operations (uploads and
downloads) and other operations, and custom predicates are now supported for
media operations as well.
- Blob.download_as_filename() will now delete the empty file if it results in a
google.cloud.exceptions.NotFound exception (HTTP 404).

Quick Start
-----------
Expand Down
4 changes: 2 additions & 2 deletions google/cloud/storage/blob.py
Original file line number Diff line number Diff line change
Expand Up @@ -1252,8 +1252,8 @@ def _handle_filename_and_download(self, filename, *args, **kwargs):
**kwargs,
)

except DataCorruption:
# Delete the corrupt downloaded file.
except (DataCorruption, NotFound):
# Delete the corrupt or empty downloaded file.
os.remove(filename)
raise

Expand Down
43 changes: 43 additions & 0 deletions tests/unit/test_blob.py
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@
import mock
import pytest

from google.cloud.exceptions import NotFound
from google.cloud.storage import _helpers
from google.cloud.storage._helpers import _get_default_headers
from google.cloud.storage._helpers import _get_default_storage_base_url
Expand Down Expand Up @@ -1848,6 +1849,48 @@ def test_download_to_filename_corrupted(self):
stream = blob._prep_and_do_download.mock_calls[0].args[0]
self.assertEqual(stream.name, filename)

def test_download_to_filename_notfound(self):
blob_name = "blob-name"
client = self._make_client()
bucket = _Bucket(client)
blob = self._make_one(blob_name, bucket=bucket)

with mock.patch.object(blob, "_prep_and_do_download"):
blob._prep_and_do_download.side_effect = NotFound("testing")

# Try to download into a temporary file (don't use
# `_NamedTemporaryFile` it will try to remove after the file is
# already removed)
filehandle, filename = tempfile.mkstemp()
os.close(filehandle)
self.assertTrue(os.path.exists(filename))

with self.assertRaises(NotFound):
blob.download_to_filename(filename)

# Make sure the file was cleaned up.
self.assertFalse(os.path.exists(filename))

expected_timeout = self._get_default_timeout()
blob._prep_and_do_download.assert_called_once_with(
mock.ANY,
client=None,
start=None,
end=None,
if_etag_match=None,
if_etag_not_match=None,
if_generation_match=None,
if_generation_not_match=None,
if_metageneration_match=None,
if_metageneration_not_match=None,
raw_download=False,
timeout=expected_timeout,
checksum="auto",
retry=DEFAULT_RETRY,
)
stream = blob._prep_and_do_download.mock_calls[0].args[0]
self.assertEqual(stream.name, filename)

def _download_as_bytes_helper(self, raw_download, timeout=None, **extra_kwargs):
blob_name = "blob-name"
client = self._make_client()
Expand Down

0 comments on commit 066be2d

Please sign in to comment.