From 1d034a322348b3e1509c20594973636a5cab1654 Mon Sep 17 00:00:00 2001
From: evazion <noizave@gmail.com>
Date: Sun, 17 Oct 2021 21:59:09 -0500
Subject: [PATCH] media assets: move more file-handling logic into MediaAsset.

Move more of the file-handling logic from UploadService and
StorageManager into MediaAsset. This is part of refactoring posts and
uploads to allow multiple images per post.
---
 app/logical/storage_manager.rb        |  12 +-
 app/logical/storage_manager/local.rb  |   8 +-
 app/logical/storage_manager/rclone.rb |   2 +-
 app/logical/storage_manager/sftp.rb   |   5 +-
 app/logical/upload_service/utils.rb   |  46 +-----
 app/models/media_asset.rb             | 203 ++++++++++++++++++++++++--
 app/models/post.rb                    |   2 +-
 app/models/upload.rb                  |   7 +-
 test/unit/storage_manager_test.rb     |   6 +-
 test/unit/upload_service_test.rb      |  51 ++++---
 10 files changed, 238 insertions(+), 104 deletions(-)

diff --git a/app/logical/storage_manager.rb b/app/logical/storage_manager.rb
index 905475099..2d46e6514 100644
--- a/app/logical/storage_manager.rb
+++ b/app/logical/storage_manager.rb
@@ -109,13 +109,13 @@ class StorageManager
 
     case type
     when :preview
-      "#{base_dir}/preview/#{subdir}#{file}"
+      "/preview/#{subdir}#{file}"
     when :crop
-      "#{base_dir}/crop/#{subdir}#{file}"
+      "/crop/#{subdir}#{file}"
     when :large
-      "#{base_dir}/sample/#{subdir}#{file}"
+      "/sample/#{subdir}#{file}"
     when :original
-      "#{base_dir}/original/#{subdir}#{file}"
+      "/original/#{subdir}#{file}"
     end
   end
 
@@ -145,4 +145,8 @@ class StorageManager
     tags = post.presenter.humanized_essential_tag_string.gsub(/[^a-z0-9]+/, "_").gsub(/(?:^_+)|(?:_+$)/, "").gsub(/_{2,}/, "_")
     "__#{tags}__"
   end
+
+  def full_path(path)
+    File.join(base_dir, path)
+  end
 end
diff --git a/app/logical/storage_manager/local.rb b/app/logical/storage_manager/local.rb
index f6c404819..bb68b4fde 100644
--- a/app/logical/storage_manager/local.rb
+++ b/app/logical/storage_manager/local.rb
@@ -3,7 +3,7 @@ class StorageManager::Local < StorageManager
   DEFAULT_PERMISSIONS = 0o644
 
   def store(io, dest_path)
-    temp_path = dest_path + "-" + SecureRandom.uuid + ".tmp"
+    temp_path = full_path(dest_path) + "-" + SecureRandom.uuid + ".tmp"
 
     FileUtils.mkdir_p(File.dirname(temp_path))
     io.rewind
@@ -11,17 +11,17 @@ class StorageManager::Local < StorageManager
     raise Error, "store failed: #{bytes_copied}/#{io.size} bytes copied" if bytes_copied != io.size
 
     FileUtils.chmod(DEFAULT_PERMISSIONS, temp_path)
-    File.rename(temp_path, dest_path)
+    File.rename(temp_path, full_path(dest_path))
   rescue StandardError => e
     FileUtils.rm_f(temp_path)
     raise Error, e
   end
 
   def delete(path)
-    FileUtils.rm_f(path)
+    FileUtils.rm_f(full_path(path))
   end
 
   def open(path)
-    File.open(path, "r", binmode: true)
+    File.open(full_path(path), "r", binmode: true)
   end
 end
diff --git a/app/logical/storage_manager/rclone.rb b/app/logical/storage_manager/rclone.rb
index ab17d94b2..9fa9d843b 100644
--- a/app/logical/storage_manager/rclone.rb
+++ b/app/logical/storage_manager/rclone.rb
@@ -35,6 +35,6 @@ class StorageManager::Rclone < StorageManager
   end
 
   def key(path)
-    ":#{remote}:#{bucket}#{path}"
+    ":#{remote}:#{bucket}#{full_path(path)}"
   end
 end
diff --git a/app/logical/storage_manager/sftp.rb b/app/logical/storage_manager/sftp.rb
index ed7a9c7b4..2984be307 100644
--- a/app/logical/storage_manager/sftp.rb
+++ b/app/logical/storage_manager/sftp.rb
@@ -19,6 +19,7 @@ class StorageManager::SFTP < StorageManager
   end
 
   def store(file, dest_path)
+    dest_path = full_path(dest_path)
     temp_upload_path = dest_path + "-" + SecureRandom.uuid + ".tmp"
     dest_backup_path = dest_path + "-" + SecureRandom.uuid + ".bak"
 
@@ -42,7 +43,7 @@ class StorageManager::SFTP < StorageManager
 
   def delete(dest_path)
     each_host do |_host, sftp|
-      force { sftp.remove!(dest_path) }
+      force { sftp.remove!(full_path(dest_path)) }
     end
   end
 
@@ -50,7 +51,7 @@ class StorageManager::SFTP < StorageManager
     file = Tempfile.new(binmode: true)
 
     Net::SFTP.start(hosts.first, nil, ssh_options) do |sftp|
-      sftp.download!(dest_path, file.path)
+      sftp.download!(full_path(dest_path), file.path)
     end
 
     file
diff --git a/app/logical/upload_service/utils.rb b/app/logical/upload_service/utils.rb
index 45bc3f811..ed23c8fc9 100644
--- a/app/logical/upload_service/utils.rb
+++ b/app/logical/upload_service/utils.rb
@@ -2,52 +2,10 @@ class UploadService
   module Utils
     module_function
 
-    def distribute_files(file, record, type, original_post_id: nil)
-      # need to do this for hybrid storage manager
-      post = Post.new
-      post.id = original_post_id if original_post_id.present?
-      post.md5 = record.md5
-      post.file_ext = record.file_ext
-      [Danbooru.config.storage_manager, Danbooru.config.backup_storage_manager].each do |sm|
-        sm.store_file(file, post, type)
-      end
-    end
-
     def is_downloadable?(source)
       source =~ %r{\Ahttps?://}
     end
 
-    def generate_resizes(media_file)
-      preview_file = media_file.preview(Danbooru.config.small_image_width, Danbooru.config.small_image_width)
-      crop_file = media_file.crop(Danbooru.config.small_image_width, Danbooru.config.small_image_width)
-
-      if media_file.is_ugoira?
-        sample_file = media_file.convert
-      elsif media_file.is_image? && media_file.width > Danbooru.config.large_image_width
-        sample_file = media_file.preview(Danbooru.config.large_image_width, media_file.height)
-      else
-        sample_file = nil
-      end
-
-      [preview_file, crop_file, sample_file]
-    end
-
-    def process_resizes(upload, file, original_post_id, media_file: nil)
-      media_file ||= upload.media_file
-      preview_file, crop_file, sample_file = Utils.generate_resizes(media_file)
-
-      begin
-        Utils.distribute_files(file, upload, :original, original_post_id: original_post_id) if file.present?
-        Utils.distribute_files(sample_file, upload, :large, original_post_id: original_post_id) if sample_file.present?
-        Utils.distribute_files(preview_file, upload, :preview, original_post_id: original_post_id) if preview_file.present?
-        Utils.distribute_files(crop_file, upload, :crop, original_post_id: original_post_id) if crop_file.present?
-      ensure
-        preview_file.try(:close!)
-        crop_file.try(:close!)
-        sample_file.try(:close!)
-      end
-    end
-
     def process_file(upload, file, original_post_id: nil)
       upload.file = file
       media_file = upload.media_file
@@ -61,9 +19,7 @@ class UploadService
       upload.validate!(:file)
       upload.tag_string = "#{upload.tag_string} #{Utils.automatic_tags(media_file)}"
 
-      process_resizes(upload, file, original_post_id)
-
-      MediaAsset.create!(file: media_file)
+      MediaAsset.upload!(media_file)
     end
 
     def automatic_tags(media_file)
diff --git a/app/models/media_asset.rb b/app/models/media_asset.rb
index 19b6a4076..dff00ebf1 100644
--- a/app/models/media_asset.rb
+++ b/app/models/media_asset.rb
@@ -1,7 +1,7 @@
 class MediaAsset < ApplicationRecord
   has_one :media_metadata, dependent: :destroy
   delegate :metadata, to: :media_metadata
-  delegate :is_animated?, :is_animated_gif?, :is_animated_png?, :is_non_repeating_animation?, :is_greyscale?, :is_rotated?, to: :metadata
+  delegate :is_non_repeating_animation?, :is_greyscale?, :is_rotated?, to: :metadata
 
   enum status: {
     processing: 100,
@@ -10,25 +10,198 @@ class MediaAsset < ApplicationRecord
     expunged: 400,
   }
 
-  def self.search(params)
-    q = search_attributes(params, :id, :created_at, :updated_at, :md5, :file_ext, :file_size, :image_width, :image_height, :media_metadata)
+  class Variant
+    attr_reader :media_asset, :variant
+    delegate :md5, :storage_service, :backup_storage_service, to: :media_asset
 
-    if params[:metadata].present?
-      q = q.joins(:media_metadata).merge(MediaMetadata.search(metadata: params[:metadata]))
+    def initialize(media_asset, variant)
+      @media_asset = media_asset
+      @variant = variant
+
+      raise ArgumentError, "asset doesn't have #{variant} variant" unless Variant.exists?(media_asset, variant)
     end
 
-    q.apply_default_order(params)
+    def store_file!(original_file)
+      file = convert_file(original_file)
+      storage_service.store(file, file_path)
+      backup_storage_service.store(file, file_path)
+    end
+
+    def delete_file!
+      storage_service.delete(file_path)
+      backup_storage_service.delete(file_path)
+    end
+
+    def open_file
+      storage_service.open(file_path)
+    end
+
+    def convert_file(media_file)
+      case variant
+      in :preview
+        media_file.preview(Danbooru.config.small_image_width, Danbooru.config.small_image_width)
+      in :crop
+        media_file.crop(Danbooru.config.small_image_width, Danbooru.config.small_image_width)
+      in :sample if media_asset.is_ugoira?
+        media_file.convert
+      in :sample if media_asset.is_static_image? && media_asset.image_width > Danbooru.config.large_image_width
+        media_file.preview(Danbooru.config.large_image_width, media_asset.image_height)
+      in :original
+        media_file
+      end
+    end
+
+    def file_url(slug = "")
+      storage_service.file_url(file_path(slug))
+    end
+
+    def file_path(slug = "")
+      if variant.in?(%i[preview crop]) && media_asset.is_flash?
+        "/images/download-preview.png"
+      else
+        slug = "__#{slug}__" if slug.present?
+        "/#{variant}/#{md5[0..1]}/#{md5[2..3]}/#{slug}#{file_name}"
+      end
+    end
+
+    # The file name of this variant.
+    def file_name
+      case variant
+      when :sample
+        "sample-#{md5}.#{file_ext}"
+      else
+        "#{md5}.#{file_ext}"
+      end
+    end
+
+    # The file extension of this variant.
+    def file_ext
+      case variant
+      when :preview
+        "jpg"
+      when :crop
+        "jpg"
+      when :sample
+        media_asset.is_ugoira? ? "webm" : "jpg"
+      when :original
+        media_asset.file_ext
+      end
+    end
+
+    def self.exists?(media_asset, variant)
+      case variant
+      when :preview
+        true
+      when :crop
+        true
+      when :sample
+        media_asset.is_ugoira? || (media_asset.is_static_image? && media_asset.image_width > Danbooru.config.large_image_width)
+      when :original
+        true
+      end
+    end
   end
 
-  def file=(file_or_path)
-    media_file = file_or_path.is_a?(MediaFile) ? file_or_path : MediaFile.open(file_or_path)
+  concerning :SearchMethods do
+    class_methods do
+      def search(params)
+        q = search_attributes(params, :id, :created_at, :updated_at, :md5, :file_ext, :file_size, :image_width, :image_height)
 
-    self.md5 = media_file.md5
-    self.file_ext = media_file.file_ext
-    self.file_size = media_file.file_size
-    self.image_width = media_file.width
-    self.image_height = media_file.height
-    self.duration = media_file.duration
-    self.media_metadata = MediaMetadata.new(file: media_file)
+        if params[:metadata].present?
+          q = q.joins(:media_metadata).merge(MediaMetadata.search(metadata: params[:metadata]))
+        end
+
+        q.apply_default_order(params)
+      end
+    end
+  end
+
+  concerning :FileMethods do
+    class_methods do
+      def upload!(media_file)
+        media_asset = create!(file: media_file, status: :processing)
+        media_asset.distribute_files!(media_file)
+        media_asset.update!(status: :active)
+        media_asset
+      end
+    end
+
+    def file=(file_or_path)
+      media_file = file_or_path.is_a?(MediaFile) ? file_or_path : MediaFile.open(file_or_path)
+
+      self.md5 = media_file.md5
+      self.file_ext = media_file.file_ext
+      self.file_size = media_file.file_size
+      self.image_width = media_file.width
+      self.image_height = media_file.height
+      self.duration = media_file.duration
+      self.media_metadata = MediaMetadata.new(file: media_file)
+    end
+
+    def delete_files!
+      variants.each(&:delete_file!)
+    end
+
+    def distribute_files!(media_file)
+      variants.each do |variant|
+        variant.store_file!(media_file)
+      end
+    end
+
+    def storage_service
+      Danbooru.config.storage_manager
+    end
+
+    def backup_storage_service
+      Danbooru.config.backup_storage_manager
+    end
+  end
+
+  concerning :VariantMethods do
+    def variant(type)
+      Variant.new(self, type)
+    end
+
+    def has_variant?(variant)
+      Variant.exists?(self, variant)
+    end
+
+    def variants
+      %i[preview crop sample original].select { |v| has_variant?(v) }.map { |v| variant(v) }
+    end
+  end
+
+  concerning :FileTypeMethods do
+    def is_image?
+      file_ext.in?(%w[jpg png gif])
+    end
+
+    def is_static_image?
+      is_image? && !is_animated?
+    end
+
+    def is_video?
+      file_ext.in?(%w[webm mp4])
+    end
+
+    def is_ugoira?
+      file_ext == "zip"
+    end
+
+    def is_flash?
+      file_ext == "swf"
+    end
+
+    def is_animated?
+      duration.present?
+    end
+
+    def is_animated_gif?
+      is_animated? && file_ext == "gif"
+    end
+
+    def is_animated_png?
+      is_animated? && file_ext == "png"
+    end
   end
 end
diff --git a/app/models/post.rb b/app/models/post.rb
index a074169da..d1e11e3f5 100644
--- a/app/models/post.rb
+++ b/app/models/post.rb
@@ -1194,7 +1194,7 @@ class Post < ApplicationRecord
         ModAction.log("<@#{user.name}> regenerated IQDB for post ##{id}", :post_regenerate_iqdb, user)
       else
         media_file = MediaFile.open(file, frame_data: pixiv_ugoira_frame_data&.data.to_a)
-        UploadService::Utils.process_resizes(self, nil, id, media_file: media_file)
+        media_asset.distribute_files!(media_file)
 
         update!(
           image_width: media_file.width,
diff --git a/app/models/upload.rb b/app/models/upload.rb
index 8a142787a..665dacbd5 100644
--- a/app/models/upload.rb
+++ b/app/models/upload.rb
@@ -116,13 +116,8 @@ class Upload < ApplicationRecord
       end
 
       media_asset&.destroy!
+      media_asset&.delete_files!
       DanbooruLogger.info("Uploads: Deleting files for upload md5=#{md5}")
-      Danbooru.config.storage_manager.delete_file(nil, md5, file_ext, :original)
-      Danbooru.config.storage_manager.delete_file(nil, md5, file_ext, :large)
-      Danbooru.config.storage_manager.delete_file(nil, md5, file_ext, :preview)
-      Danbooru.config.backup_storage_manager.delete_file(nil, md5, file_ext, :original)
-      Danbooru.config.backup_storage_manager.delete_file(nil, md5, file_ext, :large)
-      Danbooru.config.backup_storage_manager.delete_file(nil, md5, file_ext, :preview)
     end
   end
 
diff --git a/test/unit/storage_manager_test.rb b/test/unit/storage_manager_test.rb
index 778a711bb..41437a8e6 100644
--- a/test/unit/storage_manager_test.rb
+++ b/test/unit/storage_manager_test.rb
@@ -12,14 +12,14 @@ class StorageManagerTest < ActiveSupport::TestCase
 
     context "#store method" do
       should "store the file" do
-        @storage_manager.store(StringIO.new("data"), "#{@temp_dir}/test.txt")
+        @storage_manager.store(StringIO.new("data"), "test.txt")
 
         assert("data", File.read("#{@temp_dir}/test.txt"))
       end
 
       should "overwrite the file if it already exists" do
-        @storage_manager.store(StringIO.new("foo"), "#{@temp_dir}/test.txt")
-        @storage_manager.store(StringIO.new("bar"), "#{@temp_dir}/test.txt")
+        @storage_manager.store(StringIO.new("foo"), "test.txt")
+        @storage_manager.store(StringIO.new("bar"), "test.txt")
 
         assert("bar", File.read("#{@temp_dir}/test.txt"))
       end
diff --git a/test/unit/upload_service_test.rb b/test/unit/upload_service_test.rb
index b16fe466a..56eb59339 100644
--- a/test/unit/upload_service_test.rb
+++ b/test/unit/upload_service_test.rb
@@ -14,6 +14,14 @@ class UploadServiceTest < ActiveSupport::TestCase
     }
   }
 
+  def assert_file_exists(upload, variant)
+    assert_nothing_raised { upload.media_asset.variant(variant).open_file }
+  end
+
+  def assert_file_does_not_exist(upload, variant)
+    assert_raise { upload.media_asset.variant(variant).open_file }
+  end
+
   context "::Utils" do
     context "#get_file_for_upload" do
       context "for a non-source site" do
@@ -75,13 +83,11 @@ class UploadServiceTest < ActiveSupport::TestCase
 
       context "with an original_post_id" do
         should "run" do
-          UploadService::Utils.expects(:distribute_files).times(3)
           UploadService::Utils.process_file(@upload, @upload.file.tempfile, original_post_id: 12345)
         end
       end
 
       should "run" do
-        UploadService::Utils.expects(:distribute_files).times(3)
         UploadService::Utils.process_file(@upload, @upload.file.tempfile)
         assert_equal("jpg", @upload.file_ext)
         assert_equal(28086, @upload.file_size)
@@ -91,7 +97,6 @@ class UploadServiceTest < ActiveSupport::TestCase
       end
 
       should "create a media asset" do
-        UploadService::Utils.expects(:distribute_files).times(3)
         UploadService::Utils.process_file(@upload, @upload.file.tempfile)
 
         @media_asset = @upload.media_asset
@@ -138,8 +143,8 @@ class UploadServiceTest < ActiveSupport::TestCase
           assert_equal(9800, @upload.file_size)
           assert_equal("png", @upload.file_ext)
           assert_equal("f5fe24f3a3a13885285f6627e04feec9", @upload.md5)
-          assert(File.exist?(Danbooru.config.storage_manager.file_path(@upload.md5, "png", :original)))
-          assert(File.exist?(Danbooru.config.storage_manager.file_path(@upload.md5, "png", :preview)))
+          assert_file_exists(@upload, :preview)
+          assert_file_exists(@upload, :original)
         end
       end
 
@@ -158,8 +163,8 @@ class UploadServiceTest < ActiveSupport::TestCase
           assert_equal(317733, @upload.file_size)
           assert_equal("jpg", @upload.file_ext)
           assert_equal("4c71da5638b897aa6da1150e742e2982", @upload.md5)
-          assert(File.exist?(Danbooru.config.storage_manager.file_path(@upload.md5, "jpg", :original)))
-          assert(File.exist?(Danbooru.config.storage_manager.file_path(@upload.md5, "jpg", :preview)))
+          assert_file_exists(@upload, :preview)
+          assert_file_exists(@upload, :original)
         end
       end
 
@@ -179,8 +184,8 @@ class UploadServiceTest < ActiveSupport::TestCase
           assert_equal(2804, @upload.file_size)
           assert_equal("zip", @upload.file_ext)
           assert_equal("cad1da177ef309bf40a117c17b8eecf5", @upload.md5)
-          assert(File.exist?(Danbooru.config.storage_manager.file_path(@upload.md5, "zip", :original)))
-          assert(File.exist?(Danbooru.config.storage_manager.file_path(@upload.md5, "zip", :large)))
+          assert_file_exists(@upload, :sample)
+          assert_file_exists(@upload, :original)
         end
       end
 
@@ -197,9 +202,9 @@ class UploadServiceTest < ActiveSupport::TestCase
           assert_equal(181309, @upload.file_size)
           assert_equal("jpg", @upload.file_ext)
           assert_equal("93f4dd66ef1eb11a89e56d31f9adc8d0", @upload.md5)
-          assert(File.exist?(Danbooru.config.storage_manager.file_path(@upload.md5, "jpg", :original)))
-          assert(File.exist?(Danbooru.config.storage_manager.file_path(@upload.md5, "jpg", :large)))
-          assert(File.exist?(Danbooru.config.storage_manager.file_path(@upload.md5, "jpg", :preview)))
+          assert_file_exists(@upload, :preview)
+          assert_file_exists(@upload, :sample)
+          assert_file_exists(@upload, :original)
         end
       end
 
@@ -216,8 +221,8 @@ class UploadServiceTest < ActiveSupport::TestCase
           assert_equal("mp4", @upload.file_ext)
           assert_operator(@upload.file_size, :>, 0)
           assert_not_nil(@upload.source)
-          assert(File.exist?(Danbooru.config.storage_manager.file_path(@upload.md5, "mp4", :original)))
-          assert(File.exist?(Danbooru.config.storage_manager.file_path(@upload.md5, "mp4", :preview)))
+          assert_file_exists(@upload, :preview)
+          assert_file_exists(@upload, :original)
         end
       end
 
@@ -982,31 +987,31 @@ class UploadServiceTest < ActiveSupport::TestCase
 
     should "delete unused files after deleting the upload" do
       @upload = as(@user) { UploadService::Preprocessor.new(file: upload_file("test/files/test.jpg")).start! }
-      assert(File.exist?(Danbooru.config.storage_manager.file_path(@upload.md5, "jpg", :original)))
+      assert_file_exists(@upload, :original)
 
       @upload.destroy!
-      refute(File.exist?(Danbooru.config.storage_manager.file_path(@upload.md5, "jpg", :original)))
+      assert_file_does_not_exist(@upload, :original)
     end
 
     should "not delete files that are still in use by a post" do
       @upload = as(@user) { UploadService.new(file: upload_file("test/files/test.jpg")).start! }
-      assert(File.exist?(Danbooru.config.storage_manager.file_path(@upload.md5, "jpg", :original)))
+      assert_file_exists(@upload, :original)
 
       @upload.destroy!
-      assert(File.exist?(Danbooru.config.storage_manager.file_path(@upload.md5, "jpg", :original)))
+      assert_file_exists(@upload, :original)
     end
 
     should "not delete files if they're still in use by another upload" do
       @upload1 = as(@user) { UploadService::Preprocessor.new(file: upload_file("test/files/test.jpg")).start! }
       @upload2 = as(@user) { UploadService::Preprocessor.new(file: upload_file("test/files/test.jpg")).start! }
       assert_equal(@upload1.md5, @upload2.md5)
-      assert(File.exist?(Danbooru.config.storage_manager.file_path(@upload1.md5, "jpg", :original)))
+      assert_file_exists(@upload1, :original)
 
       @upload1.destroy!
-      assert(File.exist?(Danbooru.config.storage_manager.file_path(@upload1.md5, "jpg", :original)))
+      assert_file_exists(@upload1, :original)
 
       @upload2.destroy!
-      refute(File.exist?(Danbooru.config.storage_manager.file_path(@upload2.md5, "jpg", :original)))
+      assert_file_does_not_exist(@upload2, :original)
     end
 
     should "not delete files that were replaced after upload and are still pending deletion" do
@@ -1019,11 +1024,11 @@ class UploadServiceTest < ActiveSupport::TestCase
       # after replacement the uploaded file is no longer in use, but it shouldn't be
       # deleted yet. it should only be deleted by the replacer after the grace period.
       @upload.destroy!
-      assert(File.exist?(Danbooru.config.storage_manager.file_path(@upload.md5, "jpg", :original)))
+      assert_file_exists(@upload, :original)
 
       travel (PostReplacement::DELETION_GRACE_PERIOD + 1).days
       perform_enqueued_jobs
-      refute(File.exist?(Danbooru.config.storage_manager.file_path(@upload.md5, "jpg", :original)))
+      assert_file_does_not_exist(@upload, :original)
     end
 
     should "work on uploads without a file" do