fix: cancel upload when BlobWriter exits with exception

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

@@ -76,3 +76,48 @@ 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
+    try:
+        with blob.open("wt") as writer:
+            writer.write(b'first chunk')  # not yet uploaded
+            raise ValueError('SIGTERM received')  # no upload to cancel in __exit__
+    except ValueError:
+        pass
+
+    # blob should not exist
+    assert not blob.exists()
+
+    # unhandled exceptions should cancel the upload
+    try:
+        with blob.open("wt") 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__
+    except ValueError:
+        pass
+
+    # blob should not exist
+    assert not blob.exists()
+
+    # handled exceptions should not cancel the upload
+    with blob.open("wt") 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()