Merge pull request rails#40842 from ghiculescu/activestorage-variant-nplusone

Add support for eager loading Active Storage variants

Author: pixeltrix

activestorage/CHANGELOG.md

@@ -1,3 +1,8 @@
+*   Add `with_all_variant_records` method to eager load all variant records on an attachment at once.
+    `with_attached_image` scope now eager loads variant records if using variant tracking.
+
+    *Alex Ghiculescu*
+
 *   Add metadata value for presence of audio channel in video blobs
 
     The `metadata` attribute of video blobs has a new boolean key named `audio` that is set to

activestorage/app/models/active_storage/attachment.rb

@@ -19,6 +19,8 @@ class ActiveStorage::Attachment < ActiveStorage::Record
   after_create_commit :mirror_blob_later, :analyze_blob_later
   after_destroy_commit :purge_dependent_blob_later
 
+  scope :with_all_variant_records, -> { includes(blob: :variant_records) }
+
   # Synchronously deletes the attachment and {purges the blob}[rdoc-ref:ActiveStorage::Blob#purge].
   def purge
     transaction do

activestorage/app/models/active_storage/variant_with_record.rb

@@ -49,6 +49,10 @@ def create_or_find_record(image:)
     end
 
     def record
-      @record ||= blob.variant_records.find_by(variation_digest: variation.digest)
+      @record ||= if blob.variant_records.loaded?
+        blob.variant_records.find { |v| v.variation_digest == variation.digest }
+      else
+        blob.variant_records.find_by(variation_digest: variation.digest)
+      end
     end
 end

activestorage/lib/active_storage/attached/model.rb

@@ -166,7 +166,13 @@ def purge_later
         end
         has_many :"#{name}_blobs", through: :"#{name}_attachments", class_name: "ActiveStorage::Blob", source: :blob, strict_loading: strict_loading
 
-        scope :"with_attached_#{name}", -> { includes("#{name}_attachments": :blob) }
+        scope :"with_attached_#{name}", -> {
+          if ActiveStorage.track_variants
+            includes("#{name}_attachments": { blob: :variant_records })
+          else
+            includes("#{name}_attachments": :blob)
+          end
+        }
 
         after_save { attachment_changes[name.to_s]&.save }
 

activestorage/test/models/variant_with_record_test.rb

@@ -56,4 +56,82 @@ class ActiveStorage::VariantWithRecordTest < ActiveSupport::TestCase
 
     assert_equal "local_public", variant.image.blob.service_name
   end
+
+  test "eager loading is supported" do
+    user = User.create!(name: "Josh")
+
+    blob1 = directly_upload_file_blob(filename: "racecar.jpg")
+    assert_difference -> { ActiveStorage::VariantRecord.count }, +1 do
+      blob1.representation(resize: "100x100").process
+    end
+
+    blob2 = directly_upload_file_blob(filename: "racecar_rotated.jpg")
+    assert_difference -> { ActiveStorage::VariantRecord.count }, +1 do
+      blob2.representation(resize: "100x100").process
+    end
+
+    assert_no_difference -> { ActiveStorage::VariantRecord.count } do
+      user.vlogs.attach(blob1)
+      user.vlogs.attach(blob2)
+    end
+
+    user.reload
+
+    assert_no_difference -> { ActiveStorage::VariantRecord.count } do
+      assert_queries(5) do
+        # 5 queries:
+        # attachments (vlogs) x 1
+        # blob x 2
+        # variant record x 2
+        user.vlogs.map do |vlog|
+          vlog.representation(resize: "100x100").processed
+        end
+      end
+    end
+
+    user.reload
+
+    assert_no_difference -> { ActiveStorage::VariantRecord.count } do
+      assert_queries(3) do
+        # 3 queries:
+        # attachments (vlogs) x 1
+        # blob x 1
+        # variant record x 1
+        user.vlogs.includes(blob: :variant_records).map do |vlog|
+          vlog.representation(resize: "100x100").processed
+        end
+      end
+    end
+
+    user.reload
+
+    assert_no_difference -> { ActiveStorage::VariantRecord.count } do
+      assert_queries(3) do
+        # 3 queries:
+        # attachments (vlogs) x 1
+        # blob x 1
+        # variant record x 1
+        user.vlogs.with_all_variant_records.map do |vlog|
+          vlog.representation(resize: "100x100").processed
+        end
+      end
+    end
+
+    user.reload
+
+    assert_no_difference -> { ActiveStorage::VariantRecord.count } do
+      assert_queries(4) do
+        # 4 queries:
+        # user x 1
+        # attachments (vlogs) x 1
+        # blob x 1
+        # variant record x 1
+        User.where(id: user.id).with_attached_vlogs.map do |u|
+          u.vlogs.map do |vlog|
+            vlog.representation(resize: "100x100").processed
+          end
+        end
+      end
+    end
+  end
 end

activestorage/test/test_helper.rb

@@ -60,6 +60,23 @@ class ActiveSupport::TestCase
     ActiveStorage::Current.reset
   end
 
+  def assert_queries(expected_count)
+    ActiveRecord::Base.connection.materialize_transactions
+
+    queries = []
+    ActiveSupport::Notifications.subscribe("sql.active_record") do |*, payload|
+      queries << payload[:sql] unless %w[ SCHEMA TRANSACTION ].include?(payload[:name])
+    end
+
+    yield.tap do
+      assert_equal expected_count, queries.size, "#{queries.size} instead of #{expected_count} queries were executed. #{queries.inspect}"
+    end
+  end
+
+  def assert_no_queries(&block)
+    assert_queries(0, &block)
+  end
+
   private
     def create_blob(key: nil, data: "Hello world!", filename: "hello.txt", content_type: "text/plain", identify: true, service_name: nil, record: nil)
       ActiveStorage::Blob.create_and_upload! key: key, io: StringIO.new(data), filename: filename, content_type: content_type, identify: identify, service_name: service_name, record: record