Skip to content

Commit

Permalink
fix: cancel upload when BlobWriter exits with exception
Browse files Browse the repository at this point in the history
  • Loading branch information
ddelange committed Mar 22, 2024
1 parent e1bae03 commit 7324818
Show file tree
Hide file tree
Showing 2 changed files with 46 additions and 0 deletions.
12 changes: 12 additions & 0 deletions google/cloud/storage/fileio.py
Original file line number Diff line number Diff line change
Expand Up @@ -437,6 +437,18 @@ 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)

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
Expand Down
34 changes: 34 additions & 0 deletions tests/system/test_fileio.py
Original file line number Diff line number Diff line change
Expand Up @@ -76,3 +76,37 @@ 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")

# 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()

0 comments on commit 7324818

Please sign in to comment.