diff --git a/fs_s3fs/_s3fs.py b/fs_s3fs/_s3fs.py index d7f98e2..81b444a 100644 --- a/fs_s3fs/_s3fs.py +++ b/fs_s3fs/_s3fs.py @@ -82,7 +82,10 @@ def __enter__(self): return self def __exit__(self, exc_type, exc_value, traceback): - self.close() + if exc_type is None: + self.close() + else: + self.discard() @property def raw(self): @@ -92,6 +95,17 @@ def close(self): if self._on_close is not None: self._on_close(self) + def discard(self): + """Close the file without uploading, abandoning anything written to it. + + Writes are buffered in a temporary file and uploaded in a single call when + the file is closed, so a write that fails part way through has not reached + S3 yet. Uploading at that point publishes a truncated object that no reader + can tell apart from a complete one, so the buffer is thrown away instead. + """ + self._on_close = None + self._f.close() + @property def closed(self): return self._f.closed diff --git a/fs_s3fs/tests/test_s3fs.py b/fs_s3fs/tests/test_s3fs.py index ae113ae..e1f4f26 100644 --- a/fs_s3fs/tests/test_s3fs.py +++ b/fs_s3fs/tests/test_s3fs.py @@ -5,7 +5,10 @@ from nose.plugins.attrib import attr from fs.test import FSTestCases +from fs.mode import Mode + from fs_s3fs import S3FS +from fs_s3fs._s3fs import S3File import boto3 @@ -48,6 +51,46 @@ def _delete_bucket_contents(self): self.client.delete_object(Bucket=self.bucket_name, Key=obj["Key"]) +class TestS3FileClose(unittest.TestCase): + """A failed write must not publish a partial object. + + S3File buffers writes in a temporary file and uploads once, when the file is + closed. If the writer raises part way through, the buffer holds a truncated copy + that nothing downstream can distinguish from a complete file, so it is discarded + rather than uploaded. + """ + + def _make_file(self, uploaded): + return S3File.factory( + "test.bin", Mode("wb"), on_close=lambda s3file: uploaded.append(s3file) + ) + + def test_uploads_when_the_block_completes(self): + uploaded = [] + with self._make_file(uploaded) as s3file: + s3file.write(b"complete") + self.assertEqual(len(uploaded), 1) + + def test_does_not_upload_when_the_block_raises(self): + uploaded = [] + with self.assertRaises(ValueError): + with self._make_file(uploaded) as s3file: + s3file.write(b"partial") + raise ValueError("writer failed part way through") + self.assertEqual(uploaded, []) + + def test_discarding_leaves_nothing_to_upload_later(self): + # io.IOBase.__del__ calls close(), so discard() has to make that a no-op or + # the abandoned buffer is uploaded anyway once the object is collected. + uploaded = [] + s3file = self._make_file(uploaded) + s3file.write(b"partial") + s3file.discard() + s3file.close() + self.assertEqual(uploaded, []) + self.assertTrue(s3file.closed) + + class TestS3FSHelpers(unittest.TestCase): def test_path_to_key(self): s3 = S3FS("foo")