Pass optional record in blob finder methods

Allow record to be optionally passed to blob finders to make sharding
easier.

activestorage/CHANGELOG.md

@@ -1,3 +1,8 @@
+*   Allow record to be optionally passed to blob finders to make sharding
+    easier.
+
+    *Gannon McGibbon*
+
 *   Switch from `azure-storage` gem to `azure-storage-blob` gem for Azure service.
 
     *Peter Zhu*

activestorage/app/models/active_storage/blob.rb

@@ -45,19 +45,19 @@ class << self
     # that was created ahead of the upload itself on form submission.
     #
     # The signed ID is also used to create stable URLs for the blob through the BlobsController.
-    def find_signed(id)
+    def find_signed(id, record: nil)
       find ActiveStorage.verifier.verify(id, purpose: :blob_id)
     end
 
     # Returns a new, unsaved blob instance after the +io+ has been uploaded to the service.
     # When providing a content type, pass <tt>identify: false</tt> to bypass automatic content type inference.
-    def build_after_upload(io:, filename:, content_type: nil, metadata: nil, identify: true)
+    def build_after_upload(io:, filename:, content_type: nil, metadata: nil, identify: true, record: nil)
       new(filename: filename, content_type: content_type, metadata: metadata).tap do |blob|
         blob.upload(io, identify: identify)
       end
     end
 
-    def build_after_unfurling(io:, filename:, content_type: nil, metadata: nil, identify: true) #:nodoc:
+    def build_after_unfurling(io:, filename:, content_type: nil, metadata: nil, identify: true, record: nil) #:nodoc:
       new(filename: filename, content_type: content_type, metadata: metadata).tap do |blob|
         blob.unfurl(io, identify: identify)
       end
@@ -67,7 +67,7 @@ def build_after_unfurling(io:, filename:, content_type: nil, metadata: nil, iden
     # then the +io+ is uploaded, then the blob is saved. This is done this way to avoid uploading (which may take
     # time), while having an open database transaction.
     # When providing a content type, pass <tt>identify: false</tt> to bypass automatic content type inference.
-    def create_after_upload!(io:, filename:, content_type: nil, metadata: nil, identify: true)
+    def create_after_upload!(io:, filename:, content_type: nil, metadata: nil, identify: true, record: nil)
       build_after_upload(io: io, filename: filename, content_type: content_type, metadata: metadata, identify: identify).tap(&:save!)
     end
 
@@ -76,7 +76,7 @@ def create_after_upload!(io:, filename:, content_type: nil, metadata: nil, ident
     # in order to produce the signed URL for uploading. This signed URL points to the key generated by the blob.
     # Once the form using the direct upload is submitted, the blob can be associated with the right record using
     # the signed ID.
-    def create_before_direct_upload!(filename:, byte_size:, checksum:, content_type: nil, metadata: nil)
+    def create_before_direct_upload!(filename:, byte_size:, checksum:, content_type: nil, metadata: nil, record: nil)
       create! filename: filename, byte_size: byte_size, checksum: checksum, content_type: content_type, metadata: metadata
     end
 

activestorage/lib/active_storage/attached/changes/create_one.rb

@@ -53,14 +53,16 @@ def find_or_build_blob
         when ActiveStorage::Blob
           attachable
         when ActionDispatch::Http::UploadedFile, Rack::Test::UploadedFile
-          ActiveStorage::Blob.build_after_unfurling \
+          ActiveStorage::Blob.build_after_unfurling(
             io: attachable.open,
             filename: attachable.original_filename,
-            content_type: attachable.content_type
+            content_type: attachable.content_type,
+            record: record,
+          )
         when Hash
-          ActiveStorage::Blob.build_after_unfurling(attachable)
+          ActiveStorage::Blob.build_after_unfurling(attachable.merge(record: record))
         when String
-          ActiveStorage::Blob.find_signed(attachable)
+          ActiveStorage::Blob.find_signed(attachable, record: record)
         else
           raise ArgumentError, "Could not find or build blob: expected attachable, got #{attachable.inspect}"
         end

activestorage/test/models/attached/one_test.rb

@@ -2,9 +2,11 @@
 
 require "test_helper"
 require "database/setup"
+require "active_support/testing/method_call_assertions"
 
 class ActiveStorage::OneAttachedTest < ActiveSupport::TestCase
   include ActiveJob::TestHelper
+  include ActiveSupport::Testing::MethodCallAssertions
 
   setup do
     @user = User.create!(name: "Josh")
@@ -25,16 +27,45 @@ class ActiveStorage::OneAttachedTest < ActiveSupport::TestCase
     assert_equal "funky.jpg", @user.avatar.filename.to_s
   end
 
+  test "attaching an existing blob from a signed ID passes record" do
+    blob = create_blob(filename: "funky.jpg")
+    arguments = [blob.signed_id, record: @user]
+    assert_called_with(ActiveStorage::Blob, :find_signed, arguments, returns: blob) do
+      @user.avatar.attach blob.signed_id
+    end
+  end
+
   test "attaching a new blob from a Hash to an existing record" do
     @user.avatar.attach io: StringIO.new("STUFF"), filename: "town.jpg", content_type: "image/jpg"
     assert_equal "town.jpg", @user.avatar.filename.to_s
   end
 
+  test "attaching a new blob from a Hash to an existing record passes record" do
+    hash = { io: StringIO.new("STUFF"), filename: "town.jpg", content_type: "image/jpg" }
+    blob = ActiveStorage::Blob.build_after_unfurling(hash)
+    arguments = [hash.merge(record: @user)]
+    assert_called_with(ActiveStorage::Blob, :build_after_unfurling, arguments, returns: blob) do
+      @user.avatar.attach hash
+    end
+  end
+
   test "attaching a new blob from an uploaded file to an existing record" do
     @user.avatar.attach fixture_file_upload("racecar.jpg")
     assert_equal "racecar.jpg", @user.avatar.filename.to_s
   end
 
+  test "attaching a new blob from an uploaded file to an existing record passes record" do
+    upload = fixture_file_upload("racecar.jpg")
+    def upload.open
+      @io ||= StringIO.new("")
+    end
+    arguments = [io: upload.open, filename: upload.original_filename, content_type: upload.content_type, record: @user]
+    blob = ActiveStorage::Blob.build_after_unfurling(*arguments)
+    assert_called_with(ActiveStorage::Blob, :build_after_unfurling, arguments, returns: blob) do
+      @user.avatar.attach upload
+    end
+  end
+
   test "attaching an existing blob to an existing, changed record" do
     @user.name = "Tina"
     assert @user.changed?

activestorage/test/models/blob_test.rb

@@ -51,6 +51,12 @@ class ActiveStorage::BlobTest < ActiveSupport::TestCase
     assert_match(/^[a-z0-9]{28}$/, create_blob.key)
   end
 
+  test "create after upload accepts a record for overrides" do
+    assert_nothing_raised do
+      create_blob(record: User.new)
+    end
+  end
+
   test "image?" do
     blob = create_file_blob filename: "racecar.jpg"
     assert_predicate blob, :image?

activestorage/test/test_helper.rb

@@ -51,24 +51,24 @@ class ActiveSupport::TestCase
   end
 
   private
-    def create_blob(data: "Hello world!", filename: "hello.txt", content_type: "text/plain", identify: true)
-      ActiveStorage::Blob.create_after_upload! io: StringIO.new(data), filename: filename, content_type: content_type, identify: identify
+    def create_blob(data: "Hello world!", filename: "hello.txt", content_type: "text/plain", identify: true, record: nil)
+      ActiveStorage::Blob.create_after_upload! io: StringIO.new(data), filename: filename, content_type: content_type, identify: identify, record: record
     end
 
-    def create_file_blob(filename: "racecar.jpg", content_type: "image/jpeg", metadata: nil)
-      ActiveStorage::Blob.create_after_upload! io: file_fixture(filename).open, filename: filename, content_type: content_type, metadata: metadata
+    def create_file_blob(filename: "racecar.jpg", content_type: "image/jpeg", metadata: nil, record: nil)
+      ActiveStorage::Blob.create_after_upload! io: file_fixture(filename).open, filename: filename, content_type: content_type, metadata: metadata, record: record
     end
 
-    def create_blob_before_direct_upload(filename: "hello.txt", byte_size:, checksum:, content_type: "text/plain")
-      ActiveStorage::Blob.create_before_direct_upload! filename: filename, byte_size: byte_size, checksum: checksum, content_type: content_type
+    def create_blob_before_direct_upload(filename: "hello.txt", byte_size:, checksum:, content_type: "text/plain", record: nil)
+      ActiveStorage::Blob.create_before_direct_upload! filename: filename, byte_size: byte_size, checksum: checksum, content_type: content_type, record: record
     end
 
-    def directly_upload_file_blob(filename: "racecar.jpg", content_type: "image/jpeg")
+    def directly_upload_file_blob(filename: "racecar.jpg", content_type: "image/jpeg", record: nil)
       file = file_fixture(filename)
       byte_size = file.size
       checksum = Digest::MD5.file(file).base64digest
 
-      create_blob_before_direct_upload(filename: filename, byte_size: byte_size, checksum: checksum, content_type: content_type).tap do |blob|
+      create_blob_before_direct_upload(filename: filename, byte_size: byte_size, checksum: checksum, content_type: content_type, record: record).tap do |blob|
         service = ActiveStorage::Blob.service.try(:primary) || ActiveStorage::Blob.service
         service.upload(blob.key, file.open)
       end