Last active
July 7, 2022 00:50
-
-
Save apauly/38f3e88d8f35b6bcf323 to your computer and use it in GitHub Desktop.
Preload the count of associations similar to the familiar "preload"
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
## | |
# PreloadCounts | |
# | |
# Usage: | |
# collection = Post.limit(10).preload_counts(:users) | |
# collection.first.users.count # fires one query to fetch all counts | |
# collection[1].users.count # uses the cached value | |
# collection.last.users.count # uses the cached value | |
# | |
# Call `::PreloadCounts.enable!` inside of an initializer to enable preloading of association counts | |
# | |
module PreloadCounts | |
def self.enable! | |
::ActiveRecord::Base.send(:include, ::PreloadCounts::Base) | |
::ActiveRecord::Relation.send(:include, ::PreloadCounts::Relation) | |
::ActiveRecord::Associations::CollectionProxy.send(:include, ::PreloadCounts::Associations::CollectionProxy) | |
end | |
# | |
# Wrapper class that will hold the original scope, the desired targets and our cache date | |
# | |
class CounterCacheStore < Struct.new(:scope, :targets, :cache_data) | |
def initialize(scope, targets) | |
self.scope = scope | |
self.targets = targets | |
self.cache_data = {} | |
end | |
def active?(target) | |
self.targets && self.targets[target.to_sym] | |
end | |
end | |
# | |
# Every record of our collection will need access to the counter_cache | |
# | |
module Base | |
extend ActiveSupport::Concern | |
included do | |
attr_accessor :counter_cache | |
end | |
end | |
module Relation | |
extend ActiveSupport::Concern | |
def preload_counts(*association_names) | |
@preload_counts ||= {} | |
association_names.each{|association_name| | |
@preload_counts[association_name] = true | |
} | |
# return self for easy chaining | |
self | |
end | |
# | |
# Let every record remember the original scope. It will | |
# be the starting point for our count query. | |
# This functionality is injected into `to_a` to be as lazy as possible. | |
# | |
def to_a_with_counter_cache | |
if @preload_counts | |
cache_store = ::PreloadCounts::CounterCacheStore.new(self, @preload_counts) | |
to_a_without_counter_cache.tap{|records| | |
records.map{|record| record.counter_cache = cache_store } | |
} | |
else | |
to_a_without_counter_cache | |
end | |
end | |
included do | |
alias_method_chain :to_a, :counter_cache | |
end | |
end | |
# | |
# We only modify `CollectionProxy` - not `Relation`. That means that calling | |
# `User.limit(5).first.posts.count` will use our cache but | |
# `User.limit(5).first.posts.where(released: true).count` will not. | |
# | |
module Associations | |
module CollectionProxy | |
extend ActiveSupport::Concern | |
# | |
# Ask for specific count. This fires the grouped count query once. | |
# Runs a normal count query if the cache was not enabled. | |
# | |
def count_with_cache(*args) | |
cached_count || count_without_cache(*args) | |
end | |
# | |
# just a shorthand | |
# | |
def counter_cache | |
@association.owner.counter_cache | |
end | |
# | |
# if the cache is enabled, we run a single query to fetch the count for all records | |
# | |
def cached_count | |
# return unless the cache was enabled | |
return nil unless self.counter_cache && self.counter_cache.active?(@association.reflection.klass.table_name) | |
self.counter_cache.cache_data[@association.reflection.name] ||= begin | |
self.counter_cache.scope. | |
limit(nil). # reset the limit | |
joins(@association.reflection.name). # join the desired association | |
distinct(false). # reset distinct | |
group("#{@association.reflection.active_record.table_name}.id"). # group by the owners id to get grouped counts | |
select("DISTINCT #{@association.reflection.klass.table_name}.id"). # distinct by the targets primary key | |
count | |
end | |
self.counter_cache.cache_data[@association.reflection.name][@association.owner.id] || 0 | |
end | |
included do | |
alias_method_chain :count, :cache | |
end | |
end | |
end | |
end |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
If you need a version that works with Rails 5, please take a look at my fork: https://gist.github.com/cmoran92/620ec601e1b5e8732221d638e75f39cc