implement external storage cleanup with subfolding

Author: dimitri-yatsenko

datajoint/erd.py

@@ -50,8 +50,8 @@ class ERD:
         """
 
         def __init__(self, *args, **kwargs):
-            warnings.warn('ERD functionality depends on matplotlib and pygraphviz. Please install both of these '
-                          'libraries to enable the ERD feature.')
+            warnings.warn('ERD functionality depends on matplotlib, networkx, and pygraphviz. '
+                          'Please install both of these libraries to enable the ERD feature.')
 else:
     class ERD(nx.DiGraph):
         """

datajoint/external.py

@@ -1,5 +1,5 @@
 import os
-from tqdm import tqdm
+import itertools
 from .settings import config
 from .errors import DataJointError
 from .hash import long_hash
@@ -165,22 +165,35 @@ def delete_garbage(self):
                 for ref in self.references) or "TRUE")
         print('Deleted %d items' % self.connection.query("SELECT ROW_COUNT()").fetchone()[0])
 
-    def clean_store(self, store, display_progress=True):
+    def clean_store(self, store, verbose=True):
         """
         Clean unused data in an external storage repository from unused blobs.
         This must be performed after delete_garbage during low-usage periods to reduce risks of data loss.
         """
         spec = config.get_store_spec(store)
-        progress = tqdm if display_progress else lambda x: x
-        in_use = set(self.fetch('hash'))
+        in_use = set(x for x in (self & '`hash` LIKE "%%{store}"'.format(store=store)).fetch('hash'))
         if spec['protocol'] == 'file':
-            for folder, _, files in progress(os.walk(os.path.join(spec['location'], self.database))):
-                for f in files:
-                    if f not in in_use:
+            count = itertools.count()
+            print('Deleting...')
+            deleted_folders = set()
+            for folder, dirs, files in os.walk(os.path.join(spec['location'], self.database), topdown=False):
+                if dirs and files:
+                    raise DataJointError('Invalid repository with files in non-terminal folder %s' % folder)
+                dirs = set(d for d in dirs if os.path.join(folder, d) not in deleted_folders)
+                if not dirs:
+                    files_not_in_use = [f for f in files if f not in in_use]
+                    for f in files_not_in_use:
                         filename = os.path.join(folder, f)
+                        next(count)
+                        if verbose:
+                            print(filename)
                         os.remove(filename)
+                    if len(files_not_in_use) == len(files):
+                        os.rmdir(folder)
+                        deleted_folders.add(folder)
+            print('Deleted %d objects' % next(count))
         elif spec['protocol'] == 's3':
             try:
-                s3.Folder(database=self.database, **spec).clean(in_use)
+                failed_deletes = s3.Folder(database=self.database, **spec).clean(in_use, verbose=verbose)
             except TypeError:
                 raise DataJointError('External store {store} configuration is incomplete.'.format(store=store))

datajoint/s3.py

@@ -30,13 +30,34 @@ def get(self, blob_hash):
         except minio.error.NoSuchKey:
             return None
 
-    def clean(self, exclude, max_count=None):
+    def clean(self, exclude, max_count=None, verbose=False):
         """
         Delete all objects except for those in the exclude
         :param exclude: a list of blob_hashes to skip.
         :param max_count: maximum number of object to delete
-        :return: generator of objects that failed to delete
+        :param verbose: If True, print deleted objects
+        :return: list of objects that failed to delete
         """
-        return self.client.remove_objects(self.bucket, itertools.islice(
-            (x.object_name for x in self.client.list_objects(self.bucket, self.remote_path + '/')
-             if x not in exclude), max_count))
+        count = itertools.count()
+        if verbose:
+            def out(name):
+                next(count)
+                print(name)
+                return name
+        else:
+            def out(name):
+                next(count)
+                return name
+
+        if verbose:
+            print('Deleting...')
+
+        names = (out(x.object_name)
+                 for x in self.client.list_objects(self.bucket, self.remote_path + '/', recursive=True)
+                 if x.object_name.split('/')[-1] not in exclude)
+
+        failed_deletes = list(
+            self.client.remove_objects(self.bucket, itertools.islice(names, max_count)))
+
+        print('Deleted: %i S3 objects' % next(count))
+        return failed_deletes

datajoint/schema.py

@@ -228,7 +228,7 @@ def jobs(self):
         return self._jobs
 
     @property
-    def external_table(self):
+    def external(self):
         """
         schema.external provides a view of the external hash table for the schema
         :return: external table
@@ -237,6 +237,8 @@ def external_table(self):
             self._external = ExternalTable(self.connection, self.database)
         return self._external
 
+    external_table = external  # for backward compatibility to pre-0.12.0
+
 
 def create_virtual_module(module_name, schema_name, create_schema=False, create_tables=False, connection=None):
     """

tests/test_erd.py

@@ -34,6 +34,7 @@ def test_dependencies():
 
     @staticmethod
     def test_erd():
+        assert_true(dj.erd.erd_active, 'Failed to import networkx and pydot')
         erd = dj.ERD(schema, context=namespace)
         graph = erd._make_graph()
         assert_true(set(cls.__name__ for cls in (A, B, D, E, L)).issubset(graph.nodes()))