feat: download_to_filename deletes the empty file on a 404 (#1394)

andrewsg · web-flow · commit 066be2db789c · 2024-12-11T10:11:09.000-08:00
diff --git a/README.rst b/README.rst
@@ -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
 -----------
diff --git a/google/cloud/storage/blob.py b/google/cloud/storage/blob.py
@@ -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
 
diff --git a/tests/unit/test_blob.py b/tests/unit/test_blob.py
@@ -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
@@ -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()
