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

@@ -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("wt") 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("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__
+    # 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()