Add `RecordsUploads` module to record Upload records via callbacks

4 files changed, 46 insertions, 0 deletions

diff --git a/app/uploaders/attachment_uploader.rb b/app/uploaders/attachment_uploader.rb
index 6aa1f5a8c50..109eb2fea0b 100644
--- a/app/uploaders/attachment_uploader.rb
+++ b/app/uploaders/attachment_uploader.rb
@@ -1,4 +1,5 @@
 class AttachmentUploader < GitlabUploader
+  include RecordsUploads
   include UploaderHelper
 
   storage :file
diff --git a/app/uploaders/avatar_uploader.rb b/app/uploaders/avatar_uploader.rb
index b4c393c6f2c..66d3bcb998a 100644
--- a/app/uploaders/avatar_uploader.rb
+++ b/app/uploaders/avatar_uploader.rb
@@ -1,4 +1,5 @@
 class AvatarUploader < GitlabUploader
+  include RecordsUploads
   include UploaderHelper
 
   storage :file
diff --git a/app/uploaders/file_uploader.rb b/app/uploaders/file_uploader.rb
index 0d2edaeff3b..2cf97a1b6fd 100644
--- a/app/uploaders/file_uploader.rb
+++ b/app/uploaders/file_uploader.rb
@@ -1,5 +1,7 @@
 class FileUploader < GitlabUploader
+  include RecordsUploads
   include UploaderHelper
+
   MARKDOWN_PATTERN = %r{\!?\[.*?\]\(/uploads/(?<secret>[0-9a-f]{32})/(?<file>.*?)\)}
 
   storage :file
@@ -20,6 +22,10 @@ class FileUploader < GitlabUploader
     File.join(base_dir, 'tmp', @project.path_with_namespace, @secret)
   end
 
+  def model
+    project
+  end
+
   def to_markdown
     to_h[:markdown]
   end
diff --git a/app/uploaders/records_uploads.rb b/app/uploaders/records_uploads.rb
new file mode 100644
index 00000000000..7a0424b5adf
--- /dev/null
+++ b/app/uploaders/records_uploads.rb
@@ -0,0 +1,38 @@
+module RecordsUploads
+  extend ActiveSupport::Concern
+
+  included do
+    after :store,   :record_upload
+    before :remove, :destroy_upload
+  end
+
+  private
+
+  # After storing an attachment, create a corresponding Upload record
+  #
+  # NOTE: We're ignoring the argument passed to this callback because we want
+  # the `SanitizedFile` object from `CarrierWave::Uploader::Base#file`, not the
+  # `Tempfile` object the callback gets.
+  #
+  # Called `after :store`
+  def record_upload(_tempfile)
+    return unless file_storage?
+    return unless file.exists?
+
+    Upload.record(self)
+  end
+
+  # When removing an attachment, destroy any Upload records at the same path
+  #
+  # Note: this _will not work_ for Uploaders which relativize paths, such as
+  # `FileUploader`, but because that uploader uses different paths for every
+  # upload, that's an acceptable caveat.
+  #
+  # Called `before :remove`
+  def destroy_upload(*args)
+    return unless file_storage?
+    return unless file
+
+    Upload.remove_path(relative_path)
+  end
+end