Change `tootctl media remove-orphans` to work for all classes (#13316)

Change `tootctl media lookup` to not use an interactive prompt

2 files changed, 96 insertions, 28 deletions

diff --git a/lib/mastodon/media_cli.rb b/lib/mastodon/media_cli.rb
index d842b986f..b4ad78fe5 100644
--- a/lib/mastodon/media_cli.rb
+++ b/lib/mastodon/media_cli.rb
@@ -45,6 +45,7 @@ module Mastodon
     end
 
     option :start_after
+    option :prefix
     option :dry_run, type: :boolean, default: false
     desc 'remove-orphans', 'Scan storage and check for files that do not belong to existing media attachments'
     long_desc <<~LONG_DESC
@@ -58,6 +59,7 @@ module Mastodon
       reclaimed_bytes = 0
       removed         = 0
       dry_run         = options[:dry_run] ? ' (DRY RUN)' : ''
+      prefix          = options[:prefix]
 
       case Paperclip::Attachment.default_options[:storage]
       when :s3
@@ -69,7 +71,7 @@ module Mastodon
         loop do
           objects = begin
             begin
-              bucket.objects(start_after: last_key, prefix: 'media_attachments/files/').limit(1000).map { |x| x }
+              bucket.objects(start_after: last_key, prefix: prefix).limit(1000).map { |x| x }
             rescue => e
               progress.log(pastel.red("Error fetching list of files: #{e}"))
               progress.log("If you want to continue from this point, add --start-after=#{last_key} to your command") if last_key
@@ -79,16 +81,21 @@ module Mastodon
 
           break if objects.empty?
 
-          last_key        = objects.last.key
-          attachments_map = MediaAttachment.where(id: objects.map { |object| object.key.split('/')[2..-2].join.to_i }).each_with_object({}) { |attachment, map| map[attachment.id] = attachment }
+          last_key   = objects.last.key
+          record_map = preload_records_from_mixed_objects(objects)
 
           objects.each do |object|
-            attachment_id = object.key.split('/')[2..-2].join.to_i
-            filename      = object.key.split('/').last
+            path_segments   = object.key.split('/')
+            model_name      = path_segments.first.classify
+            attachment_name = path_segments[1].singularize
+            record_id       = path_segments[2..-2].join.to_i
+            file_name       = path_segments.last
+            record          = record_map.dig(model_name, record_id)
+            attachment      = record&.public_send(attachment_name)
 
             progress.increment
 
-            next unless attachments_map[attachment_id].nil? || !attachments_map[attachment_id].variant?(filename)
+            next unless attachment.blank? || !attachment.variant?(file_name)
 
             begin
               object.delete unless options[:dry_run]
@@ -110,17 +117,24 @@ module Mastodon
 
         root_path = ENV.fetch('RAILS_ROOT_PATH', File.join(':rails_root', 'public', 'system')).gsub(':rails_root', Rails.root.to_s)
 
-        Find.find(File.join(root_path, 'media_attachments', 'files')) do |path|
+        Find.find(File.join(*[root_path, prefix].compact)) do |path|
           next if File.directory?(path)
 
-          key           = path.gsub("#{root_path}#{File::SEPARATOR}", '')
-          attachment_id = key.split(File::SEPARATOR)[2..-2].join.to_i
-          filename      = key.split(File::SEPARATOR).last
-          attachment    = MediaAttachment.find_by(id: attachment_id)
+          key             = path.gsub("#{root_path}#{File::SEPARATOR}", '')
+          path_segments   = key.split(File::SEPARATOR)
+          model_name      = path_segments.first.classify
+          record_id       = path_segments[2..-2].join.to_i
+          attachment_name = path_segments[1].singularize
+          file_name       = path_segments.last
+
+          next unless PRELOAD_MODEL_WHITELIST.include?(model_name)
+
+          record     = model_name.constantize.find_by(id: record_id)
+          attachment = record&.public_send(attachment_name)
 
           progress.increment
 
-          next unless attachment.nil? || !attachment.variant?(filename)
+          next unless attachment.blank? || !attachment.variant?(file_name)
 
           begin
             size = File.size(path)
@@ -213,25 +227,66 @@ module Mastodon
       say("Settings:\t#{number_to_human_size(SiteUpload.sum(:file_file_size))}")
     end
 
-    desc 'lookup', 'Lookup where media is displayed by passing a media URL'
-    def lookup
-      prompt = TTY::Prompt.new
+    desc 'lookup URL', 'Lookup where media is displayed by passing a media URL'
+    def lookup(url)
+      path          = Addressable::URI.parse(url).path
+      path_segments = path.split('/')[2..-1]
+      model_name    = path_segments.first.classify
+      record_id     = path_segments[2..-2].join.to_i
 
-      url = prompt.ask('Please enter a URL to the media to lookup:', required: true)
+      unless PRELOAD_MODEL_WHITELIST.include?(model_name)
+        say("Cannot find corresponding model: #{model_name}", :red)
+        exit(1)
+      end
 
-      attachment_id = url
-                      .split('/')[0..-2]
-                      .grep(/\A\d+\z/)
-                      .join('')
+      record = model_name.constantize.find_by(id: record_id)
+      record = record.status if record.respond_to?(:status)
 
-      if url.split('/')[0..-2].include? 'media_attachments'
-        model = MediaAttachment.find(attachment_id).status
-        prompt.say(ActivityPub::TagManager.instance.url_for(model))
-      elsif url.split('/')[0..-2].include? 'accounts'
-        model = Account.find(attachment_id)
-        prompt.say(ActivityPub::TagManager.instance.url_for(model))
-      else
-        prompt.say('Not found')
+      unless record
+        say('Cannot find corresponding record', :red)
+        exit(1)
+      end
+
+      display_url = ActivityPub::TagManager.instance.url_for(record)
+
+      if display_url.blank?
+        say('No public URL for this type of record', :red)
+        exit(1)
+      end
+
+      say(display_url, :blue)
+    rescue Addressable::URI::InvalidURIError
+      say('Invalid URL', :red)
+      exit(1)
+    end
+
+    private
+
+    PRELOAD_MODEL_WHITELIST = %w(
+      Account
+      Backup
+      CustomEmoji
+      Import
+      MediaAttachment
+      PreviewCard
+      SiteUpload
+    ).freeze
+
+    def preload_records_from_mixed_objects(objects)
+      preload_map = Hash.new { |hash, key| hash[key] = [] }
+
+      objects.map do |object|
+        segments   = object.key.split('/').first
+        model_name = segments.first.classify
+        record_id  = segments[2..-2].join.to_i
+
+        next unless PRELOAD_MODEL_WHITELIST.include?(model_name)
+
+        preload_map[model_name] << record_id
+      end
+
+      preload_map.each_with_object({}) do |(model_name, record_ids), model_map|
+        model_map[model_name] = model_name.constantize.where(id: record_ids).each_with_object({}) { |record, record_map| record_map[record.id] = record }
       end
     end
   end
diff --git a/lib/paperclip/attachment_extensions.rb b/lib/paperclip/attachment_extensions.rb
index 3b308af5f..d9ec0159a 100644
--- a/lib/paperclip/attachment_extensions.rb
+++ b/lib/paperclip/attachment_extensions.rb
@@ -24,6 +24,19 @@ module Paperclip
         flush_deletes
       end
     end
+
+    def variant?(other_filename)
+      return true  if original_filename == other_filename
+      return false if original_filename.nil?
+
+      formats = styles.values.map(&:format).compact
+
+      return false if formats.empty?
+
+      other_extension = File.extname(other_filename)
+
+      formats.include?(other_extension.delete('.')) && File.basename(other_filename, other_extension) == File.basename(original_filename, File.extname(original_filename))
+    end
   end
 end