Preload the count of associations similar to the familiar "preload"

preload_counts.rb

#
# 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
  #
  CounterCacheStore = Struct.new(:scope, :targets, :cache_data) do
    def initialize(scope, targets)
      self.scope      = scope
      self.targets    = targets
      self.cache_data = {}
    end

    def active?(target)
      targets && 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

  #
  # Linking the counter cache from the association
  #
  module Relation
    extend ActiveSupport::Concern

    def preload_counts(*association_names)
      @preload_counts ||= {}
      association_names.each do |association_name|
        @preload_counts[association_name] = true
      end
      # 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 do |records|
          records.map { |record| record.counter_cache = cache_store }
        end
      else
        to_a_without_counter_cache
      end
    end

    included do
      alias_method :to_a_without_counter_cache, :to_a
      alias_method :to_a, :to_a_with_counter_cache
    end
  end

  module Associations
    #
    # 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 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 counter_cache&.active?(@association.reflection.klass.table_name)

        counter_cache.cache_data[@association.reflection.name] ||= begin
          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
        counter_cache.cache_data[@association.reflection.name][@association.owner.id] || 0
      end

      included do
        alias_method :count_without_cache, :count
        alias_method :count, :count_with_cache
      end
    end
  end
end

This version works with Rails 5.