From 5eaf214594811e5a8a597fc372ac36e5e14f410c Mon Sep 17 00:00:00 2001
From: Bianca Nenciu <nbianca@users.noreply.github.com>
Date: Wed, 16 Feb 2022 09:00:30 +0200
Subject: [PATCH] FEATURE: New plugin API to check if upload is used (#15545)

This commit introduces two new APIs for handling unused uploads, one
can be used to exclude uploads in bulk when the data model allow and
the other one excludes uploads one by one.
---
 app/jobs/scheduled/clean_up_uploads.rb | 11 ++---
 app/models/upload.rb                   | 24 +++++++++
 lib/plugin/instance.rb                 |  8 +++
 spec/jobs/clean_up_uploads_spec.rb     | 68 ++++++++++++++++++++++++++
 4 files changed, 104 insertions(+), 7 deletions(-)

diff --git a/app/jobs/scheduled/clean_up_uploads.rb b/app/jobs/scheduled/clean_up_uploads.rb
index 561691a7972..df55eeb1f10 100644
--- a/app/jobs/scheduled/clean_up_uploads.rb
+++ b/app/jobs/scheduled/clean_up_uploads.rb
@@ -26,6 +26,8 @@ module Jobs
       s3_cdn_hostname = URI.parse(SiteSetting.Upload.s3_cdn_url || "").hostname
 
       result = Upload.by_users
+      Upload.unused_callbacks&.each { |handler| result = handler.call(result) }
+      result = result
         .where("uploads.retain_hours IS NULL OR uploads.created_at < current_timestamp - interval '1 hour' * uploads.retain_hours")
         .where("uploads.created_at < ?", grace_period.hour.ago)
         .where("uploads.access_control_post_id IS NULL")
@@ -39,14 +41,9 @@ module Jobs
           next if ReviewableQueuedPost.pending.where("payload->>'raw' LIKE '%#{upload.sha1}%' OR payload->>'raw' LIKE '%#{encoded_sha}%'").exists?
           next if Draft.where("data LIKE '%#{upload.sha1}%' OR data LIKE '%#{encoded_sha}%'").exists?
           next if UserProfile.where("bio_raw LIKE '%#{upload.sha1}%' OR bio_raw LIKE '%#{encoded_sha}%'").exists?
-          if defined?(ChatMessage)
-            # TODO after May 2022 - remove this. No longer needed as chat uploads are in a table
-            next if ChatMessage.where("message LIKE ? OR message LIKE ?", "%#{upload.sha1}%", "%#{encoded_sha}%").exists?
-          end
 
-          if defined?(ChatUpload)
-            next if ChatUpload.where(upload: upload).exists?
-          end
+          next if Upload.in_use_callbacks&.any? { |callback| callback.call(upload) }
+
           upload.destroy
         else
           upload.delete
diff --git a/app/models/upload.rb b/app/models/upload.rb
index d0611ba84af..061cc266af3 100644
--- a/app/models/upload.rb
+++ b/app/models/upload.rb
@@ -64,6 +64,30 @@ class Upload < ActiveRecord::Base
     )
   end
 
+  def self.add_unused_callback(&block)
+    (@unused_callbacks ||= []) << block
+  end
+
+  def self.unused_callbacks
+    @unused_callbacks
+  end
+
+  def self.reset_unused_callbacks
+    @unused_callbacks = []
+  end
+
+  def self.add_in_use_callback(&block)
+    (@in_use_callbacks ||= []) << block
+  end
+
+  def self.in_use_callbacks
+    @in_use_callbacks
+  end
+
+  def self.reset_in_use_callbacks
+    @in_use_callbacks = []
+  end
+
   def self.with_no_non_post_relations
     scope = self
       .joins(<<~SQL)
diff --git a/lib/plugin/instance.rb b/lib/plugin/instance.rb
index 50f4ae04e10..8482ff02109 100644
--- a/lib/plugin/instance.rb
+++ b/lib/plugin/instance.rb
@@ -252,6 +252,14 @@ class Plugin::Instance
     Site.add_categories_callbacks(&block)
   end
 
+  def register_upload_unused(&block)
+    Upload.add_unused_callback(&block)
+  end
+
+  def register_upload_in_use(&block)
+    Upload.add_in_use_callback(&block)
+  end
+
   def custom_avatar_column(column)
     reloadable_patch do |plugin|
       UserLookup.lookup_columns << column
diff --git a/spec/jobs/clean_up_uploads_spec.rb b/spec/jobs/clean_up_uploads_spec.rb
index 093a760cc3d..bb0c2ecfe99 100644
--- a/spec/jobs/clean_up_uploads_spec.rb
+++ b/spec/jobs/clean_up_uploads_spec.rb
@@ -51,6 +51,74 @@ describe Jobs::CleanUpUploads do
     expect(Upload.exists?(id: expired_upload.id)).to eq(false)
   end
 
+  describe 'unused callbacks' do
+    before do
+      Upload.add_unused_callback do |uploads|
+        uploads.where.not(id: expired_upload.id)
+      end
+    end
+
+    after do
+      Upload.reset_unused_callbacks
+    end
+
+    it 'does not delete uploads skipped by an unused callback' do
+      expect do
+        Jobs::CleanUpUploads.new.execute(nil)
+      end.to change { Upload.count }.by(0)
+
+      expect(Upload.exists?(id: expired_upload.id)).to eq(true)
+    end
+
+    it 'deletes other uploads not skipped by an unused callback' do
+      expired_upload2 = fabricate_upload
+      upload = fabricate_upload
+      PostUpload.create(post: Fabricate(:post), upload: upload)
+
+      expect do
+        Jobs::CleanUpUploads.new.execute(nil)
+      end.to change { Upload.count }.by(-1)
+
+      expect(Upload.exists?(id: expired_upload.id)).to eq(true)
+      expect(Upload.exists?(id: expired_upload2.id)).to eq(false)
+      expect(Upload.exists?(id: upload.id)).to eq(true)
+    end
+  end
+
+  describe 'in use callbacks' do
+    before do
+      Upload.add_in_use_callback do |upload|
+        expired_upload.id == upload.id
+      end
+    end
+
+    after do
+      Upload.reset_in_use_callbacks
+    end
+
+    it 'does not delete uploads that are in use by callback' do
+      expect do
+        Jobs::CleanUpUploads.new.execute(nil)
+      end.to change { Upload.count }.by(0)
+
+      expect(Upload.exists?(id: expired_upload.id)).to eq(true)
+    end
+
+    it 'deletes other uploads that are not in use by callback' do
+      expired_upload2 = fabricate_upload
+      upload = fabricate_upload
+      PostUpload.create(post: Fabricate(:post), upload: upload)
+
+      expect do
+        Jobs::CleanUpUploads.new.execute(nil)
+      end.to change { Upload.count }.by(-1)
+
+      expect(Upload.exists?(id: expired_upload.id)).to eq(true)
+      expect(Upload.exists?(id: expired_upload2.id)).to eq(false)
+      expect(Upload.exists?(id: upload.id)).to eq(true)
+    end
+  end
+
   describe 'when clean_up_uploads is disabled' do
     before do
       SiteSetting.clean_up_uploads = false