fix #516: thumbnail_cleanup command for S3 and different source storages

easy_thumbnails/management/commands/thumbnail_cleanup.py

@@ -9,6 +9,7 @@
 
 from easy_thumbnails.conf import settings
 from easy_thumbnails.models import Source
+from easy_thumbnails.utils import get_storage_hash
 
 
 class ThumbnailCollectionCleaner:
@@ -25,8 +26,11 @@ def __init__(self, stdout, stderr):
         self.stdout = stdout
         self.stderr = stderr
 
-    def _get_absolute_path(self, path):
-        return os.path.join(settings.MEDIA_ROOT, path)
+    def _get_absolute_path(self, path, storage):
+        if hasattr(storage, 'location'):
+            return os.path.join(storage.location, path)
+        else:
+            return os.path.join(settings.MEDIA_ROOT, path)
 
     def _get_relative_path(self, path):
         return os.path.relpath(path, settings.MEDIA_ROOT)
@@ -42,7 +46,7 @@ def _delete_sources_by_id(self, ids):
         Source.objects.all().filter(id__in=ids).delete()
 
     def clean_up(self, dry_run=False, verbosity=1, last_n_days=0,
-                 cleanup_path=None, storage=None):
+                 cleanup_path=None, storage=None, source_storage_str=None):
         """
         Iterate through sources. Delete database references to sources
         not existing, including its corresponding thumbnails (files and
@@ -54,6 +58,10 @@ def clean_up(self, dry_run=False, verbosity=1, last_n_days=0,
         if not storage:
             storage = get_storage_class(settings.THUMBNAIL_DEFAULT_STORAGE)()
 
+        if not source_storage_str:
+            source_storage_str = settings.DEFAULT_FILE_STORAGE
+        source_storage = get_storage_class(source_storage_str)()
+
         sources_to_delete = []
         time_start = time.time()
 
@@ -66,18 +74,24 @@ def clean_up(self, dry_run=False, verbosity=1, last_n_days=0,
             query = query.filter(name__startswith=cleanup_path)
 
         for source in queryset_iterator(query):
+            source_storage_hash = source.storage_hash
+            if get_storage_hash(source_storage) != source_storage_hash:
+                print ("Source storage mismatch:", source)
+                print ("Try to set correct source hash through --source-storage management command parameter")
+                continue
+
             self.sources += 1
-            abs_source_path = self._get_absolute_path(source.name)
+            abs_source_path = self._get_absolute_path(source.name, source_storage)
 
-            if not self._check_if_exists(storage, abs_source_path):
+            if not self._check_if_exists(source_storage, abs_source_path):
                 if verbosity > 0:
                     self.stdout.write("Source not present: {}".format(abs_source_path))
                 self.source_refs_deleted += 1
                 sources_to_delete.append(source.id)
 
                 for thumb in source.thumbnails.all():
                     self.thumbnails_deleted += 1
-                    abs_thumbnail_path = self._get_absolute_path(thumb.name)
+                    abs_thumbnail_path = self._get_absolute_path(thumb.name, storage)
 
                     if self._check_if_exists(storage, abs_thumbnail_path):
                         if not dry_run:
@@ -146,12 +160,20 @@ def add_arguments(self, parser):
             dest='cleanup_path',
             type=str,
             help='Specify a path to clean up.')
+        parser.add_argument(
+            '--source-storage',
+            action='store',
+            dest='source_storage',
+            type=str,
+            help='Specify storage classpath for source files.')
 
     def handle(self, *args, **options):
         tcc = ThumbnailCollectionCleaner(self.stdout, self.stderr)
         tcc.clean_up(
             dry_run=options.get('dry_run', False),
             verbosity=int(options.get('verbosity', 1)),
             last_n_days=int(options.get('last_n_days', 0)),
-            cleanup_path=options.get('cleanup_path'))
+            cleanup_path=options.get('cleanup_path'),
+            source_storage_str=options.get('source_storage', None),
+        )
         tcc.print_stats()