Merge pull request #205 from mariusvniekerk/autocommit-missing

Missing autocommit for fs.open()

Author: martindurant

s3fs/core.py

@@ -253,7 +253,7 @@ def get_delegated_s3pars(self, exp=3600):
                 'token': cred['SessionToken'], 'anon': False}
 
     def _open(self, path, mode='rb', block_size=None, acl='', version_id=None,
-              fill_cache=None, cache_type='bytes', **kwargs):
+              fill_cache=None, cache_type='bytes', autocommit=True, **kwargs):
         """ Open a file for reading or writing
 
         Parameters
@@ -299,7 +299,8 @@ def _open(self, path, mode='rb', block_size=None, acl='', version_id=None,
 
         return S3File(self, path, mode, block_size=block_size, acl=acl,
                       version_id=version_id, fill_cache=fill_cache,
-                      s3_additional_kwargs=kw, cache_type=cache_type)
+                      s3_additional_kwargs=kw, cache_type=cache_type,
+                      autocommit=autocommit)
 
     def _lsdir(self, path, refresh=False, max_items=None):
         if path.startswith('s3://'):
@@ -899,8 +900,11 @@ class S3File(AbstractBufferedFile):
     def __init__(self, s3, path, mode='rb', block_size=5 * 2 ** 20, acl="",
                  version_id=None, fill_cache=True, s3_additional_kwargs=None,
                  autocommit=True, cache_type='bytes'):
-        if not split_path(path)[1]:
+        bucket, key = split_path(path)
+        if not key:
             raise ValueError('Attempt to open non key-like path: %s' % path)
+        self.bucket = bucket
+        self.key = key
         self.version_id = version_id
         self.acl = acl
         self.mpu = None
@@ -935,7 +939,6 @@ def _call_s3(self, method, *kwarglist, **kwargs):
                                 **kwargs)
 
     def _initiate_upload(self):
-        bucket, key = split_path(self.path)
         if self.acl and self.acl not in key_acls:
             raise ValueError('ACL not in %s', key_acls)
         self.parts = []
@@ -945,7 +948,7 @@ def _initiate_upload(self):
         try:
             self.mpu = self._call_s3(
                 self.fs.s3.create_multipart_upload,
-                Bucket=bucket, Key=key, ACL=self.acl)
+                Bucket=self.bucket, Key=self.key, ACL=self.acl)
         except ClientError as e:
             raise translate_boto_error(e)
         except ParamValidationError as e:
@@ -958,8 +961,9 @@ def _initiate_upload(self):
                 out = self.fs._call_s3(
                     self.fs.s3.upload_part_copy,
                     self.s3_additional_kwargs,
-                    Bucket=bucket,
-                    Key=key, PartNumber=1,
+                    Bucket=self.bucket,
+                    Key=self.key, 
+                    PartNumber=1,
                     UploadId=self.mpu['UploadId'],
                     CopySource=self.path)
                 self.parts.append({'PartNumber': 1,
@@ -1003,8 +1007,7 @@ def url(self, **kwargs):
         return self.fs.url(self.path, **kwargs)
 
     def _fetch_range(self, start, end):
-        bucket, key = self.path.split('/', 1)
-        return _fetch_range(self.fs.s3, bucket, key, self.version_id, start, end)
+        return _fetch_range(self.fs.s3, self.bucket, self.key, self.version_id, start, end)
 
     def _upload_chunk(self, final=False):
         bucket, key = split_path(self.path)
@@ -1050,12 +1053,11 @@ def _upload_chunk(self, final=False):
 
     def commit(self):
         logger.debug("COMMIT")
-        bucket, key = self.path.split('/', 1)
         part_info = {'Parts': self.parts}
         write_result = self._call_s3(
             self.fs.s3.complete_multipart_upload,
-            Bucket=bucket,
-            Key=key,
+            Bucket=self.bucket,
+            Key=self.key,
             UploadId=self.mpu['UploadId'],
             MultipartUpload=part_info)
         if self.fs.version_aware:
@@ -1072,6 +1074,16 @@ def commit(self):
                 self.fs.invalidate_cache(path)
             path = path + '/' + p
 
+    def discard(self):
+        if self.autocommit:
+            raise ValueError("Cannot discard when autocommit is enabled")
+        self._call_s3(
+            self.fs.s3.abort_multipart_upload,
+            Bucket=self.bucket,
+            Key=self.key,
+            UploadId=self.mpu['UploadId'],
+        )
+
 
 def _fetch_range(client, bucket, key, version_id, start, end, max_attempts=10,
                  req_kw=None):

s3fs/tests/test_s3fs.py

@@ -1228,3 +1228,35 @@ def test_cache_after_copy(s3):
     assert 'test/afile' in s3.ls('s3://test', False)
     s3.cp('test/afile', 'test/bfile')
     assert 'test/bfile' in s3.ls('s3://test', False)
+
+
+def test_autocommit(s3):
+    auto_file = test_bucket_name + '/auto_file'
+    committed_file = test_bucket_name + '/commit_file'
+    aborted_file = test_bucket_name + '/aborted_file'
+    s3 = S3FileSystem(anon=False, version_aware=True)
+    
+    def write_and_flush(path, autocommit):
+        with s3.open(path, 'wb', autocommit=autocommit) as fo:
+            fo.write(b'1')
+        return fo
+
+    # regular behavior
+    fo = write_and_flush(auto_file, autocommit=True)
+    assert fo.autocommit
+    assert s3.exists(auto_file)
+    
+    fo = write_and_flush(committed_file, autocommit=False)
+    assert not fo.autocommit
+    assert not s3.exists(committed_file)
+    fo.commit()
+    assert s3.exists(committed_file)
+
+    fo = write_and_flush(aborted_file,autocommit=False)
+    assert not s3.exists(aborted_file)
+    fo.discard()
+    assert not s3.exists(aborted_file)
+    # Cannot commit a file that was discarded
+    with pytest.raises(Exception):
+        fo.commit()
+