feat(orm): extend relation with preload and preload_association

This commit is contained in:
dualfroz committed 2026-09-08 17:43:16 +02:00
1 parent 9af9e0cd04
commit 84a04f212d
1 file changed
+172 -11
+172 -11
View File
@@ -194,29 +194,190 @@ module Prism
)
end
def preload(records)
return if @includes_value.empty? || records.empty?
@includes_value.each { |name| preload_association(name, records) }
end
def preload_association(name, records)
meta = @model.associations[name]
raise ArgumentError, "unknown association #{name.inspect} on #{@model}" unless meta
case meta[:type]
when :belongs_to then preload_belongs_to(name, meta, records)
when :has_one then preload_has_one(name, meta, records)
when :has_many then preload_has_many(name, meta, records)
when :has_many_through then preload_has_many_through(name, meta, records)
when :has_one_through then preload_has_one_through(name, meta, records)
else
raise ArgumentError, "includes does not support association type #{meta[:type].inspect} (#{name.inspect})"
end
end
def preload_belongs_to(name, meta, records)
return preload_belongs_to_polymorphic(name, meta, records) if meta[:polymorphic]
target = Object.const_get(meta[:class_name])
ids = records.map { |record| record[meta[:foreign_key]] }.compact.uniq
by_id = target.where(id: ids).each_with_object({}) { |record, hash| hash[record.id] = record }
records.each { |record| record.association_cache[name] = by_id[record[meta[:foreign_key]]] }
end
# Polymorphic belongs_to can point at rows in several different
# tables, so a single `where(id: ids)` on one target class would be
# wrong - group owners by their type column first and issue one
# batched query per distinct target type (still far fewer queries
# than one per owner for any realistic mix of types).
def preload_belongs_to_polymorphic(name, meta, records)
loaded = {}
records.group_by { |record| record[meta[:type_key]] }.each do |type_name, group|
next if type_name.nil?
ids = group.map { |record| record[meta[:foreign_key]] }.compact.uniq
Object.const_get(type_name).where(id: ids).each { |target| loaded[[type_name, target.id]] = target }
end
records.each do |record|
record.association_cache[name] = loaded[[record[meta[:type_key]], record[meta[:foreign_key]]]]
end
end
def preload_has_one(name, meta, records)
target = Object.const_get(meta[:class_name])
condition = { meta[:foreign_key] => records.map(&:id).compact.uniq }
condition[meta[:type_key]] = @model.name if meta[:type_key]
by_owner = target.where(condition).each_with_object({}) do |record, hash|
hash[record[meta[:foreign_key]]] = record
end
records.each { |record| record.association_cache[name] = by_owner[record.id] }
end
def preload_has_many(name, meta, records)
target = Object.const_get(meta[:class_name])
condition = { meta[:foreign_key] => records.map(&:id).compact.uniq }
condition[meta[:type_key]] = @model.name if meta[:type_key]
grouped = target.where(condition).to_a.group_by { |record| record[meta[:foreign_key]] }
records.each { |record| record.association_cache[name] = grouped[record.id] || [] }
end
# Batches the through-association's records for all owners in one
# query, then batches the source association on top of *those*
# records (recursing into another Relation's preloader), so an
# `includes(:commenters)` costs two queries total regardless of how
# many posts or comments are involved. The through step can itself be
# has_many/has_one (owner.id looked up by foreign key on the through
# table) or belongs_to (owner holds the foreign key, so several owners
# can share one through record) - each shape needs its own batching.
def preload_has_many_through(name, meta, records)
through_meta = @model.associations[meta[:through]]
unless through_meta
raise ArgumentError, "unknown through association #{meta[:through].inspect} on #{@model}"
end
case through_meta[:type]
when :has_many, :has_one
preload_has_many_through_by_foreign_key(name, meta, through_meta, records)
when :belongs_to
preload_has_many_through_by_belongs_to(name, meta, through_meta, records)
else
raise ArgumentError, "includes does not support #{name.inspect} (through a #{through_meta[:type]} " \
"association); through must be has_many, has_one, or belongs_to"
end
end
def preload_has_many_through_by_foreign_key(name, meta, through_meta, records)
through_target = Object.const_get(through_meta[:class_name])
owner_ids = records.map(&:id).compact.uniq
through_records = through_target.where(through_meta[:foreign_key] => owner_ids).to_a
through_by_owner = through_records.group_by { |record| record[through_meta[:foreign_key]] }
through_target.all.send(:preload_association, meta[:source], through_records) unless through_records.empty?
records.each do |record|
owned = through_by_owner[record.id] || []
record.association_cache[name] = Associable.collect_through(owned, meta[:source])
end
end
# has_many :through where the through step is a belongs_to, e.g.
# `Comment belongs_to :post` plus `Comment has_many :co_commenters,
# through: :post, source: :commenters` (everyone else who commented on
# the same post). Several owners can point at the same through record,
# so it is batched by id rather than by foreign key, then the source
# association is batched on top of the deduped through records - still
# two queries total no matter how many owners share a through record.
def preload_has_many_through_by_belongs_to(name, meta, through_meta, records)
through_target = Object.const_get(through_meta[:class_name])
through_ids = records.map { |record| record[through_meta[:foreign_key]] }.compact.uniq
through_by_id = through_target.where(id: through_ids).each_with_object({}) do |record, hash|
hash[record.id] = record
end
through_target.all.send(:preload_association, meta[:source], through_by_id.values) unless through_by_id.empty?
records.each do |record|
through_record = through_by_id[record[through_meta[:foreign_key]]]
record.association_cache[name] = Associable.collect_through(through_record, meta[:source])
end
end
# Single-record counterpart to preload_has_many_through: batches the
# through-association for all owners (has_one or belongs_to, so at
# most one record per owner), then batches the source association on
# top of those through-records, so an `includes(:account_history)`
# costs two queries total regardless of how many owners are involved.
def preload_has_one_through(name, meta, records)
through_meta = @model.associations[meta[:through]]
unless through_meta
raise ArgumentError, "unknown through association #{meta[:through].inspect} on #{@model}"
end
unless [:has_one, :belongs_to].include?(through_meta[:type])
raise ArgumentError, "includes does not support #{name.inspect} (through a #{through_meta[:type]} " \
"association); through must be has_one or belongs_to"
end
case through_meta[:type]
when :has_one then preload_has_one(meta[:through], through_meta, records)
when :belongs_to then preload_belongs_to(meta[:through], through_meta, records)
end
through_records = records.filter_map { |record| record.association_cache[meta[:through]] }.uniq
through_target = Object.const_get(through_meta[:class_name])
through_target.all.send(:preload_association, meta[:source], through_records) unless through_records.empty?
records.each do |record|
through_record = record.association_cache[meta[:through]]
record.association_cache[name] = through_record && through_record.association_cache[meta[:source]]
end
end
def rows
filtered = @model.store.all(@model.table_name).select { |row| matches?(row) }
candidates = indexed_candidates || @model.store.all(@model.table_name)
filtered = candidates.select { |row| @condition.matches?(row) }
filtered = apply_order(filtered)
filtered = filtered.uniq if @distinct_value
filtered = filtered.drop(@offset_value) if @offset_value
filtered = filtered.first(@limit_value) if @limit_value
filtered
end
def matches?(row)
@conditions.all? do |key, expected|
actual = row[key.to_s]
case expected
when Range then expected.cover?(actual)
when Array then expected.include?(actual)
else actual == expected
end
end
# nil means "no covering index, caller must scan the whole table";
# an Array (possibly empty) means an index narrowed it down already.
# The subsequent #select against the full @condition is still run
# unconditionally so this is always correct even when the equalities
# or range bounds extracted here are a strict subset of the real
# condition (e.g. an extra clause on a non-indexed column in the same
# #where). The actual equality-then-in_list-then-range-then-Or-union
# logic lives on the condition tree itself (ConditionNode#indexed_
# candidates) since OrNode needs to recurse into both branches, which
# a flat Hash of column=>value can't express.
def indexed_candidates
@condition.indexed_candidates(@model.store, @model.table_name)
end
def apply_order(rows)
return rows unless @order_by
column, direction = @order_by
sorted = rows.sort_by { |row| row[column] }
sorted = rows.sort_by { |row| [row[column].nil? ? 0 : 1, row[column]] }
direction == :desc ? sorted.reverse : sorted
end
end