Don't upload a partial file when the writer fails

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

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")