Explorar el Código

[aws] Avoid bulk delete on empty directory

Romain Rigaux hace 4 años
padre
commit
8884e7971f
Se han modificado 2 ficheros con 127 adiciones y 21 borrados
  1. 23 21
      desktop/libs/aws/src/aws/s3/s3fs.py
  2. 104 0
      desktop/libs/aws/src/aws/s3/s3fs_test.py

+ 23 - 21
desktop/libs/aws/src/aws/s3/s3fs.py

@@ -345,27 +345,29 @@ class S3FileSystem(object):
       key = self._get_key(path, validate=False)
 
       if key.exists():
-        to_delete = iter([key])
-      else:
-        to_delete = iter([])
-
-      if self.isdir(path):
-        # add `/` to prevent removing of `s3://b/a_new` trying to remove `s3://b/a`
-        prefix = self._append_separator(key.name)
-        keys = key.bucket.list(prefix=prefix)
-        to_delete = itertools.chain(keys, to_delete)
-        result = key.bucket.delete_keys(to_delete)
-        if result.errors:
-          msg = "%d errors occurred while attempting to delete the following S3 paths:\n%s" % (
-            len(result.errors), '\n'.join(['%s: %s' % (error.key, error.message) for error in result.errors])
-          )
-          LOG.error(msg)
-          raise S3FileSystemException(msg)
-      else:
-        # Avoid Raz issues
-        deleted_key = key.bucket.delete_key(next(to_delete))
-        if deleted_key.exists():
-          raise S3FileSystemException('Could not delete key %s' % deleted_key)
+        to_delete = [key]
+        dir_keys = []
+
+        if self.isdir(path):
+          # add `/` to prevent removing of `s3://b/a_new` trying to remove `s3://b/a`
+          prefix = self._append_separator(key.name)
+          dir_keys = key.bucket.list(prefix=prefix)
+          to_delete = itertools.chain(dir_keys, to_delete)
+
+        if not dir_keys:
+          # Avoid Raz bulk delete issue
+          deleted_key = key.delete()
+          if deleted_key.exists():
+            raise S3FileSystemException('Could not delete key %s' % deleted_key)
+        else:
+          result = key.bucket.delete_keys(to_delete)
+          if result.errors:
+            msg = "%d errors occurred while attempting to delete the following S3 paths:\n%s" % (
+              len(result.errors), '\n'.join(['%s: %s' % (error.key, error.message) for error in result.errors])
+            )
+            LOG.error(msg)
+            raise S3FileSystemException(msg)
+
 
   @translate_s3_error
   @auth_error_handler

+ 104 - 0
desktop/libs/aws/src/aws/s3/s3fs_test.py

@@ -19,6 +19,7 @@ import json
 import os
 import tempfile
 import string
+import sys
 
 from nose.tools import assert_equal, assert_false, assert_true, assert_raises, eq_
 
@@ -32,6 +33,109 @@ from aws.s3.s3test_utils import S3TestBase, generate_id
 from aws.s3.upload import DEFAULT_WRITE_SIZE
 
 
+if sys.version_info[0] > 2:
+  from unittest.mock import patch, Mock
+else:
+  from mock import patch, Mock
+
+
+class TestS3FileSystem():
+
+  def test_rmtree_bucket(self):
+    with patch('aws.s3.s3fs.S3FileSystem._delete_bucket') as _delete_bucket:
+
+      fs = S3FileSystem(s3_connection=Mock())
+
+      fs.rmtree(path='s3a://gethue')
+
+      _delete_bucket.assert_called()
+
+  def test_rmtree_key(self):
+    with patch('aws.s3.s3fs.S3FileSystem._get_key') as _get_key:
+      with patch('aws.s3.s3fs.S3FileSystem.isdir') as isdir:
+
+        key = Mock(
+          name='data',
+          exists=Mock(return_value=True),
+          bucket=Mock(
+            list=Mock(return_value=[]),
+            delete_key=Mock()
+          ),
+          delete=Mock(
+            return_value=Mock(
+              exists=Mock(return_value=False)
+            )
+          )
+        )
+        _get_key.return_value = key
+        isdir.return_value = False
+
+        fs = S3FileSystem(s3_connection=Mock())
+
+        fs.rmtree(path='s3a://gethue/data')
+
+        key.delete.assert_called()
+        key.bucket.delete_keys.assert_not_called()
+
+  def test_rmtree_empty_dir(self):
+    with patch('aws.s3.s3fs.S3FileSystem._get_key') as _get_key:
+      with patch('aws.s3.s3fs.S3FileSystem.isdir') as isdir:
+
+        key = Mock(
+          name='data',
+          exists=Mock(return_value=True),
+          bucket=Mock(
+            list=Mock(return_value=[]),
+            delete_key=Mock()
+          ),
+          delete=Mock(
+            return_value=Mock(
+              exists=Mock(return_value=False)
+            )
+          )
+        )
+        _get_key.return_value = key
+        isdir.return_value = True
+
+        fs = S3FileSystem(s3_connection=Mock())
+
+        fs.rmtree(path='s3a://gethue/data')
+
+        key.delete.assert_called()
+        key.bucket.delete_keys.assert_not_called()
+
+  def test_rmtree_non_empty_dir(self):
+    with patch('aws.s3.s3fs.S3FileSystem._get_key') as _get_key:
+      with patch('aws.s3.s3fs.S3FileSystem.isdir') as isdir:
+
+        key = Mock(
+          name='data',
+          exists=Mock(return_value=True),
+          bucket=Mock(
+            list=Mock(return_value=['s3a://gethue/data/1', 's3a://gethue/data/2']),
+            delete_keys=Mock(
+              return_value=Mock(
+                errors=[]
+              )
+            )
+          ),
+          delete=Mock(
+            return_value=Mock(
+              exists=Mock(return_value=False)
+            )
+          )
+        )
+        _get_key.return_value = key
+        isdir.return_value = True
+
+        fs = S3FileSystem(s3_connection=Mock())
+
+        fs.rmtree(path='s3a://gethue/data')
+
+        key.delete.assert_not_called()
+        key.bucket.delete_keys.assert_called()
+
+
 class S3FSTest(S3TestBase):
 
   @classmethod