Support nested eager loading on Sequel models that have already been loaded into memory.

model-eager-loading.rb

module Sequel
  module Plugins
    # This plugin adds support for eager loading associations _after_ a record, or
    # multiple records, have been fetched from the database. It is useful for ensuring
    # nested associations are efficiently loaded when it is difficult to modify the
    # dataset used to initially fetch the data.
    #
    # Example usage:
    #
    #   artist = Artist.first
    #
    #   # SELECT * FROM albums WHERE artist_id = ?
    #   # SELECT * FROM tracks WHERE album_id IN (?, ?, ?, ...)
    #   Artist.eager_load(artist, albums: :tracks)
    #
    #   # No queries issued
    #   artist.albums[2].tracks
    #
    # Additionally, an array of models may be passed in to eager load associations for all
    # of them:
    #
    #   albums = Album.dataset.limit(10).all
    #
    #   # SELECT * FROM tracks WHERE album_id IN (?, ?, ?, ...)
    #   Album.eager_load(albums, :tracks)
    #
    # The plugin will traverse through associations already loaded to eager load nested
    # associations:
    #
    #   artist = Artist.first
    #   artist.albums
    #
    #   # SELECT * FROM tracks WHERE album_id IN (?, ?, ?, ...)
    #   Artist.eager_load(artist, albums: :tracks)
    #
    # This plugin supports the same capabilities as the regular eager method, namely
    # using procs as callbacks for filtering association dataset. If an association has
    # already been loaded, however, the proc will not apply to that dataset.
    module ModelEagerLoading
      module ClassMethods
        # This method borrows a lot from the #eager_load Sequel::Dataset method defined in
        # Sequels lib/sequel/model/associations.rb.
        def eager_load(model_or_models, *associations)
          models = model_or_models.is_a?(Array) ? model_or_models : [model_or_models]

          # Cache to avoid building id maps multiple times.
          key_hash = {}

          normalize_associations(associations).each do |association, nested_associations|
            eager_block = nil
            if nested_associations&.count == 1 && nested_associations.keys[0].is_a?(Proc)
              eager_block = nested_associations.keys[0]
              nested_associations = nested_associations.values[0]
            end

            reflection = get_reflection(association)

            # If all of the models have already loaded the association, we'll just
            # recursively call ::eager_load to load nested associations.
            if models.all? {|m| m.associations.key?(association)}
              if nested_associations
                reflection.associated_class.eager_load(
                  models.map {|m| m.associations[association]}.flatten,
                  nested_associations,
                )
              end
            else
              key = reflection.eager_loader_key
              id_map = nil

              if key && !key_hash[key]
                id_map = Hash.new {|h, k| h[k] = []}

                models.each do |model|
                  if key.is_a?(Array)
                    model_id = key.map {|c| model.get_column_value(c)}
                    id_map[model_id] << model if model_id.all?
                  else
                    model_id = model.get_column_value(key)
                    id_map[model_id] << model if model_id
                  end
                end
              end

              loader = reflection[:eager_loader]
              loader.call(
                key_hash: key_hash,
                rows: models,
                associations: nested_associations,
                self: self,
                eager_block: eager_block,
                id_map: id_map,
              )

              if reflection[:after_load]
                models.each do |model|
                  model.send(
                    :run_association_callbacks,
                    reflection,
                    :after_load,
                    model.associations[association],
                  )
                end
              end
            end
          end
        end

        private

        # Convert nested associations into a normalized hash form. Arrays are flattened
        # and single associations are converted to hashes with nil as the keys, and
        # all hashes are merged together.
        #
        # For example,
        #   [:a, {b: :c}, {d: <proc>}, [:e, {f: :g}], h: [:h1, :h2, {h3: <proc>}]]
        # becomes
        #   {
        #     a: nil,
        #     b: :c,
        #     d: {<proc> => nil},
        #     e: nil,
        #     f: :g,
        #     h: {
        #       h1: nil,
        #       h2: nil,
        #       h3: {<proc> => nil},
        #     }
        #   }
        #
        # Note that this will allow later declarations to overwrite previous ones. This
        # is Sequel's current behavior. The following query will not execute the proc
        # callback when eager loading the albums association.
        #
        # Artist.dataset.eager({albums: <proc>}, :albums).all
        #
        # This is somewhat similar to Sequel's eager_options_for_associations, but
        # recursive and without the association checks.
        def normalize_associations(*associations)
          association_hash = {}

          associations.flatten.each do |association|
            case association
            when Symbol, Proc
              association_hash[association] = nil
            when Hash
              association.each do |association_name, nested_associations|
                association_hash[association_name] =
                  normalize_associations(nested_associations)
              end
            end
          end

          association_hash
        end

        # Get the related AssociationReflection for a model, ensuring that it's valid to
        # eager load.
        #
        # Copied from Sequel's check_association in lib/sequel/model/associations.rb.
        def get_reflection(association)
          reflection = association_reflection(association)

          if !reflection
            fail(
              Sequel::UndefinedAssociation,
              "Invalid association #{association} for #{name}",
            )
          end

          if reflection[:allow_eager] == false
            fail(
              Sequel::Error,
              "Eager loading is not allowed for #{model.name} association #{association}",
            )
          end

          reflection
        end
      end
    end
  end
end