From 52f3c2217322c2471473055f9b06f2c35bbbec8e Mon Sep 17 00:00:00 2001
From: evazion <noizave@gmail.com>
Date: Wed, 22 Sep 2021 05:34:18 -0500
Subject: [PATCH] rake: clean up rake tasks.

* Make danbooru:images:validate task run in parallel.
* Remove task for regenerating thumbnails.
* Move all rake tasks into same file.
* Add usage instructions.
---
 lib/tasks/danbooru.rake    | 116 ++++++++++++++++++++++++++++++++++++-
 lib/tasks/images.rake      |  87 ----------------------------
 lib/tasks/maintenance.rake |  23 --------
 3 files changed, 114 insertions(+), 112 deletions(-)
 delete mode 100644 lib/tasks/images.rake
 delete mode 100644 lib/tasks/maintenance.rake

diff --git a/lib/tasks/danbooru.rake b/lib/tasks/danbooru.rake
index b2daa7f0d..e50998835 100644
--- a/lib/tasks/danbooru.rake
+++ b/lib/tasks/danbooru.rake
@@ -1,9 +1,20 @@
+require "newrelic_rpm"
+require "tasks/newrelic"
+
+# This file contains commands for running various routine maintenance tasks on
+# a Danbooru instance. Run `bin/rails -T` to see a list of all available tasks.
+#
+# @see https://guides.rubyonrails.org/command_line.html#rake
+# @see https://guides.rubyonrails.org/command_line.html#custom-rake-tasks
 namespace :danbooru do
+  # Usage: bin/rails danbooru:cron
   desc "Run the cronjob scheduler"
   task cron: :environment do
     Clockwork::run
   end
 
+  # Usage: bin/rails danbooru:reindex_iqdb
+  #
   # Schedules all posts to be reindexed in IQDB. Requires the delayed jobs
   # worker (bin/delayed_job) to be running.
   desc "Reindex all posts in IQDB"
@@ -14,12 +25,113 @@ namespace :danbooru do
     end
   end
 
+  # Usage: TAGS="touhou" bin/rails danbooru:images:validate
+  #
+  # Check whether any images are missing, corrupt, or don't match the
+  # width/height/size/ext metadata in the database.
+  namespace :images do
+    task validate: :environment do
+      processes = ENV.fetch("PROCESSES", Etc.nprocessors).to_i
+      posts = Post.system_tag_match(ENV["TAGS"]).reorder(nil)
+
+      posts.parallel_each(in_processes: processes) do |post|
+        media_file = MediaFile.open(post.file(:original))
+
+        raise if post.md5 != media_file.md5
+        raise if post.image_width != media_file.width
+        raise if post.image_height != media_file.height
+        raise if post.file_size != media_file.file_size
+        raise if post.file_ext != media_file.file_ext.to_s
+
+        puts "post ##{post.id}"
+      rescue RuntimeError => e
+        hash = {
+          post: {
+            id: post.id,
+            md5: post.md5,
+            width: post.image_width,
+            height: post.image_height,
+            size: post.file_size,
+            ext: post.file_ext,
+          },
+          media_file: media_file.as_json.except("metadata"),
+          error: e.to_s,
+        }
+        STDERR.puts hash.to_json
+      rescue StandardError => e
+        hash = { id: post.id, md5: post.md5, error: e.to_s }
+        STDERR.puts hash.to_json
+      end
+    end
+
+    # Usage: bin/rails danbooru:images:populate_metadata
+    task populate_metadata: :environment do
+      sm = StorageManager::Local.new(base_url: "/", base_dir: ENV.fetch("DIR", Rails.root.join("public/data")))
+
+      MediaMetadata.joins(:media_asset).where(metadata: {}).find_each do |metadata|
+        asset = metadata.media_asset
+        file = sm.open(sm.file_path(asset.md5, asset.file_ext, :original))
+        media_file = MediaFile.open(file)
+
+        metadata.update!(metadata: media_file.metadata)
+        puts "metadata[id=#{metadata.id}, md5=#{asset.md5}]: #{media_file.metadata.count}"
+      rescue StandardError => e
+        puts "error[id=#{metadata.id}]: #{e}"
+      end
+    end
+
+    # Usage: TAGS="touhou" bin/rails danbooru:images:backup
+    desc "Backup images"
+    task backup: :environment do
+      CurrentUser.user = User.system
+      sm = Danbooru.config.backup_storage_manager
+      tags = ENV["TAGS"]
+      posts = Post.system_tag_match(tags)
+
+      posts.parallel_each do |post|
+        sm.store_file(post.file(:preview), post, :preview) if post.has_preview?
+        sm.store_file(post.file(:crop), post, :crop) if post.has_cropped?
+        sm.store_file(post.file(:sample), post, :sample) if post.has_large?
+        sm.store_file(post.file(:original), post, :original)
+      end
+    end
+  end
+
   namespace :docker do
-    # Note that uncommited changes won't be included in the image; commit
-    # changes first before building the image.
+    # Usage: bin/rails danbooru:docker:build
+    #
+    # Build a Docker image. Note that uncommited changes won't be included in
+    # the image; commit changes first before building the image.
     desc "Build a Docker image based on latest commit"
     task :build do
       system("git archive HEAD | docker buildx build - --build-arg SOURCE_COMMIT=$(git rev-parse HEAD) -t danbooru -f Dockerfile")
     end
   end
+
+  # @see app/logical/danbooru_maintenance.rb
+  namespace :maintenance do
+    # Usage: bin/rails danbooru:maintenance:hourly
+    desc "Run hourly maintenance jobs"
+    task hourly: :environment do
+      DanbooruMaintenance.hourly
+    end
+
+    # Usage: bin/rails danbooru:maintenance:daily
+    desc "Run daily maintenance jobs"
+    task daily: :environment do
+      DanbooruMaintenance.daily
+    end
+
+    # Usage: bin/rails danbooru:maintenance:weekly
+    desc "Run weekly maintenance jobs"
+    task weekly: :environment do
+      DanbooruMaintenance.weekly
+    end
+
+    # Usage: bin/rails danbooru:maintenance:monthly
+    desc "Run monthly maintenance jobs"
+    task monthly: :environment do
+      DanbooruMaintenance.monthly
+    end
+  end
 end
diff --git a/lib/tasks/images.rake b/lib/tasks/images.rake
deleted file mode 100644
index 1ffed087a..000000000
--- a/lib/tasks/images.rake
+++ /dev/null
@@ -1,87 +0,0 @@
-require "find"
-
-namespace :images do
-  task manifest: :environment do
-    root = ENV.fetch("DIR", Rails.root.join("public/data/original"))
-
-    paths = Find.find(root).lazy
-    paths = paths.reject { |path| File.directory?(path) }
-
-    paths.each do |path|
-      file = MediaFile.open(path)
-
-      hash = {
-        path: File.absolute_path(path),
-        name: File.basename(path, ".*"),
-        md5: file.md5,
-        size: file.file_size,
-        ext: file.file_ext,
-        w: file.width,
-        h: file.height,
-      }
-
-      puts hash.to_json
-    rescue StandardError => e
-      hash = {
-        path: File.absolute_path(path),
-        name: File.basename(path, ".*"),
-        md5: file&.md5,
-        size: file&.file_size,
-        error: e.message,
-      }
-
-      puts hash.to_json
-    end
-  end
-
-  task populate_media_metadata: :environment do
-    sm = StorageManager::Local.new(base_url: "/", base_dir: ENV.fetch("DIR", Rails.root.join("public/data")))
-
-    MediaMetadata.joins(:media_asset).where(metadata: {}).find_each do |metadata|
-      asset = metadata.media_asset
-      file = sm.open(sm.file_path(asset.md5, asset.file_ext, :original))
-      media_file = MediaFile.open(file)
-
-      metadata.update!(metadata: media_file.metadata)
-      puts "metadata[id=#{metadata.id}, md5=#{asset.md5}]: #{media_file.metadata.count}"
-    rescue StandardError => e
-      puts "error[id=#{metadata.id}]: #{e}"
-    end
-  end
-
-  desc "Backup images"
-  task :backup => :environment do
-    CurrentUser.user = User.system
-    sm = Danbooru.config.backup_storage_manager
-    tags = ENV["BACKUP_TAGS"]
-    posts = Post.system_tag_match(tags)
-
-    posts.find_each do |post|
-      sm.store_file(post.file(:preview), post, :preview) if post.has_preview?
-      sm.store_file(post.file(:crop), post, :crop) if post.has_cropped?
-      sm.store_file(post.file(:sample), post, :sample) if post.has_large?
-      sm.store_file(post.file(:original), post, :original)
-    end
-  end
-
-  desc "Regenerate thumbnails and samples"
-  task :regen => :environment do
-    tags = ENV["TAGS"]
-
-    Post.system_tag_match(tags).find_each do |post|
-      original_file = MediaFile.open(post.file(:original))
-      preview_file, crop_file, sample_file = UploadService::Utils.generate_resizes(original_file)
-
-      line = ""
-      line << "post ##{post.id}: "
-      line << "preview=#{preview_file.width}x#{preview_file.height}:#{preview_file.file_size.to_s(:human_size)} " if preview_file.present?
-      line << "crop=#{crop_file.width}x#{crop_file.height}:#{crop_file.file_size.to_s(:human_size)} " if crop_file.present?
-      line << "sample=#{sample_file.width}x#{sample_file.height}:#{sample_file.file_size.to_s(:human_size)} " if sample_file.present?
-      puts line
-
-      UploadService::Utils.distribute_files(preview_file, post, :preview) if preview_file.present?
-      UploadService::Utils.distribute_files(crop_file, post, :crop) if crop_file.present?
-      UploadService::Utils.distribute_files(sample_file, post, :large) if sample_file.present?
-    end
-  end
-end
diff --git a/lib/tasks/maintenance.rake b/lib/tasks/maintenance.rake
deleted file mode 100644
index 664c9d2a1..000000000
--- a/lib/tasks/maintenance.rake
+++ /dev/null
@@ -1,23 +0,0 @@
-require "tasks/newrelic" if defined?(NewRelic)
-
-namespace :maintenance do
-  desc "Run hourly maintenance jobs"
-  task hourly: :environment do
-    DanbooruMaintenance.hourly
-  end
-
-  desc "Run daily maintenance jobs"
-  task daily: :environment do
-    DanbooruMaintenance.daily
-  end
-
-  desc "Run weekly maintenance jobs"
-  task weekly: :environment do
-    DanbooruMaintenance.weekly
-  end
-
-  desc "Run monthly maintenance jobs"
-  task monthly: :environment do
-    DanbooruMaintenance.monthly
-  end
-end