Improve ActiveStorageLoader example for STI models

examples/active_storage_loader.rb

@@ -42,34 +42,51 @@
 #     end
 #   end
 # end
-module Loaders
-  class ActiveStorageLoader < GraphQL::Batch::Loader
-    attr_reader :record_type, :attachment_name, :association_type # should be has_one_attached or has_many_attached
 
-    def initialize(record_type, attachment_name, association_type: :has_one_attached)
-      super()
-      @record_type = record_type
-      @attachment_name = attachment_name
-      @association_type = association_type
-    end
+class ActiveStorageLoader < GraphQL::Batch::Loader
+  attr_reader :record_type, :attachment_name, :association_type
 
-    def perform(record_ids)
-      # find records and fulfill promises
-      attachments = ActiveStorage::Attachment.includes(:blob).where(
-        record_type: record_type, record_id: record_ids, name: attachment_name
-      )
+  def initialize(record_type, attachment_name, association_type: :has_one_attached)
+    super()
+    @record_type = record_type
+    @attachment_name = attachment_name
+    @association_type = association_type
+  end
 
-      if @association_type == :has_one_attached
-        attachments.each do |attachment|
-          fulfill(attachment.record_id, attachment)
-        end
+  def perform(record_ids)
+    # find records and fulfill promises
+    attachments = ActiveStorage::Attachment.includes(:blob, :record).where(
+      record_type: ancestors_record_types, record_id: record_ids, name: attachment_name
+    )
 
-        record_ids.each { |id| fulfill(id, nil) unless fulfilled?(id) }
-      else
-        record_ids.each do |record_id|
-          fulfill(record_id, attachments.select { |attachment| attachment.record_id == record_id })
-        end
+    if @association_type == :has_one_attached
+      attachments.each do |attachment|
+        fulfill(attachment.record_id, attachment)
+      end
+
+      record_ids.each { |id| fulfill(id, nil) unless fulfilled?(id) }
+    else
+      record_ids.each do |record_id|
+        fulfill(record_id, attachments.select { |attachment| attachment.record_id == record_id })
       end
     end
   end
+
+  private
+
+  def ancestors_record_types
+    # Get all ancestor classes of record_type that are descendants of ActiveRecord::Base
+    # This is necessary because in a Single Table Inheritance (STI) setup,
+    # the `has_one_attached` or `has_many_attached`
+    # could be defined on any of the ancestor classes of the model, not just the model itself,
+    # which determines whether the `record_type` string is stored as the model's class name
+    # or the ancestor's class name.
+    # So we for any of the ancestor classes to ensure we don't miss the attachment
+    # we are looking for:
+
+    @record_type.to_s.constantize.ancestors.select { |ancestor|
+      ancestor < ActiveRecord::Base
+    }.map(&:to_s)
+  end
+end
 end