fix: cancel upload when BlobWriter exits with exception

Author: ddelange

README.rst

@@ -72,6 +72,12 @@ setup.py file. Applications which do not import directly from
 `google-resumable-media` can safely disregard this dependency. This backwards
 compatibility feature will be removed in a future major version update.
 
+Miscellaneous
+~~~~~~~~~~~~~
+
+- The BlobWriter class now attempts to terminate an ongoing resumable upload if
+  the writer exits with an exception.
+
 Quick Start
 -----------
 

docs/storage/exceptions.rst

@@ -1,5 +1,5 @@
 Exceptions
-~~~~~~~~~
+~~~~~~~~~~
 
 .. automodule:: google.cloud.storage.exceptions
   :members:

google/cloud/storage/fileio.py

@@ -437,6 +437,19 @@ def close(self):
             self._upload_chunks_from_buffer(1)
         self._buffer.close()
 
+    def terminate(self):
+        """Cancel the ResumableUpload."""
+        if self._upload_and_transport:
+            upload, transport = self._upload_and_transport
+            transport.delete(upload.upload_url)
+        self._buffer.close()
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if exc_type is not None:
+            self.terminate()
+        else:
+            self.close()
+
     @property
     def closed(self):
         return self._buffer.closed

tests/system/test_fileio.py

@@ -14,6 +14,8 @@
 # limitations under the License.
 
 
+import pytest
+
 from .test_blob import _check_blob_hash
 
 
@@ -76,3 +78,41 @@ def test_blobwriter_and_blobreader_text_mode(
         assert text_data[:100] == reader.read(100)
         assert 0 == reader.seek(0)
         assert reader.read() == text_data
+
+
+
+def test_blobwriter_exit(
+    shared_bucket,
+    blobs_to_delete,
+    service_account,
+):
+    blob = shared_bucket.blob("NeverUploaded")
+
+    # no-op when nothing was uploaded yet
+    with pytest.raises(ValueError, match="SIGTERM received"):
+        with blob.open("wb") as writer:
+            writer.write(b"first chunk")  # not yet uploaded
+            raise ValueError("SIGTERM received")  # no upload to cancel in __exit__
+    # blob should not exist
+    assert not blob.exists()
+
+    # unhandled exceptions should cancel the upload
+    with pytest.raises(ValueError, match="SIGTERM received"):
+        with blob.open("wb") as writer:
+            writer.write(b"first chunk")  # not yet uploaded
+            writer.write(b"big chunk" * 1024 ** 8)  # uploaded
+            raise ValueError("SIGTERM received")  # upload is cancelled in __exit__
+    # blob should not exist
+    assert not blob.exists()
+
+    # handled exceptions should not cancel the upload
+    with blob.open("wb") as writer:
+        writer.write(b"first chunk")  # not yet uploaded
+        writer.write(b"big chunk" * 1024 ** 8)  # uploaded
+        try:
+            raise ValueError("This is fine")
+        except ValueError:
+            pass  # no exception context passed to __exit__
+    blobs_to_delete.append(blob)
+    # blob should have been uploaded
+    assert blob.exists()