pipeloader 0.0.1 → 0.0.2

data/lib/pipeloader/batch/batch_proxy.rb

@@ -0,0 +1,204 @@
+# frozen_string_literal: true
+
+module Pipeloader
+  module Batch
+    # A lazy, chainable stand-in for a has_many association whose LOAD is batched
+    # across all live siblings of the owner. Query builders (where/order/limit/
+    # select/...) accumulate an ActiveRecord::Relation; materializing (each/to_a/
+    # map/...) runs ONE query for every sibling, applies the accumulated scope,
+    # partitions rows by foreign key, and slices limit/offset PER GROUP.
+    #
+    # It duck-types the common Relation + Enumerable surface (~80% of a real
+    # CollectionProxy). Anything not handled here falls through to the real AR
+    # association, so writers (<<, build, create, ...) still work.
+    class BatchProxy
+      include Enumerable
+
+      # Pure query refinements that map 1:1 onto AR::Relation and return a new proxy.
+      # `where` is defined explicitly (below) so `where.not(...)` chains too.
+      QUERY_METHODS = %i[
+        rewhere order reorder reselect distinct group having
+        joins left_outer_joins includes preload eager_load references unscope readonly
+      ].freeze
+
+      # The ONLY methods delegated to the real AR association. These are writes, so
+      # they can't escape the (batched) read path. Everything not handled by this
+      # proxy raises NoMethodError instead of silently falling through to a
+      # per-record query.
+      WRITE_METHODS = %i[
+        << concat push create create! build new
+        delete destroy delete_all destroy_all clear replace
+      ].freeze
+
+      def initialize(owner, reflection, relation: nil, limit_value: nil, offset_value: nil)
+        @owner = owner
+        @reflection = reflection
+        @relation = relation
+        @limit_value = limit_value
+        @offset_value = offset_value
+      end
+
+      attr_reader :owner, :limit_value, :offset_value
+
+      def name
+        @reflection.name
+      end
+
+      def foreign_key
+        @reflection.foreign_key.to_s
+      end
+
+      def owner_key
+        @reflection.active_record_primary_key.to_sym
+      end
+
+      def relation
+        @relation ||= apply_reflection_scope(@reflection.klass.all)
+      end
+
+      QUERY_METHODS.each do |meth|
+        define_method(meth) do |*args, **kwargs, &block|
+          spawn(relation: relation.public_send(meth, *args, **kwargs, &block))
+        end
+      end
+
+      WRITE_METHODS.each do |meth|
+        define_method(meth) do |*args, **kwargs, &block|
+          real_association.public_send(meth, *args, **kwargs, &block)
+        end
+      end
+
+      # limit/offset are honored PER GROUP (top-N per owner), applied after the
+      # single batched query — not as a global SQL LIMIT that would collapse every
+      # owner into one window.
+      def limit(value)
+        spawn(limit_value: value)
+      end
+
+      def offset(value)
+        spawn(offset_value: value)
+      end
+
+      # where(conditions) refines the batched query; bare `where` returns a chain so
+      # `where.not(...)` works like a relation's.
+      def where(*args, **kwargs, &block)
+        return WhereChain.new(self) if args.empty? && kwargs.empty? && block.nil?
+
+        spawn(relation: relation.where(*args, **kwargs, &block))
+      end
+
+      # select(:cols) refines the query; select { block } filters loaded records
+      # (matching ActiveRecord::Relation#select's dual behavior).
+      def select(*columns, &block)
+        return records.select(&block) if block
+
+        spawn(relation: relation.select(*columns))
+      end
+
+      def each(&block)
+        records.each(&block)
+      end
+
+      def records
+        return @records if defined?(@records)
+
+        @records = Pipeloader::Batch::BatchLoader.load(self)
+      end
+      alias_method :to_a, :records
+      alias_method :to_ary, :records
+      alias_method :load, :records
+
+      def size
+        records.size
+      end
+      alias_method :length, :size
+
+      def empty?
+        records.empty?
+      end
+
+      def last(*args)
+        records.last(*args)
+      end
+
+      def [](*args)
+        records[*args]
+      end
+
+      def ids
+        records.map(&:id)
+      end
+
+      def pluck(*columns)
+        records.map { |record| columns.one? ? record[columns.first] : columns.map { |column| record[column] } }
+      end
+
+      # Batched, under our control: both reuse the batched scope rather than
+      # issuing a per-owner query.
+      def find_by(*args, **kwargs)
+        where(*args, **kwargs).first
+      end
+
+      def exists?(*args, **kwargs)
+        return !empty? if args.empty? && kwargs.empty?
+
+        where(*args, **kwargs).any?
+      end
+
+      # Signature of the accumulated scope (without the per-owner FK filter), so
+      # siblings asking for the same scope share one batch.
+      def cache_signature
+        [relation.to_sql, @limit_value, @offset_value]
+      end
+
+      def inspect
+        records.inspect
+      end
+
+      private
+
+      def spawn(relation: nil, limit_value: @limit_value, offset_value: @offset_value)
+        self.class.new(
+          @owner, @reflection,
+          relation: relation || @relation, limit_value: limit_value, offset_value: offset_value
+        )
+      end
+
+      # Start from the association's own scope, so `batch_has_many :open, -> { where(state: "open") }`
+      # actually narrows. An owner-dependent scope (arity > 0) can't be applied uniformly
+      # across owners, so refuse it rather than silently drop or misapply it.
+      def apply_reflection_scope(base)
+        scope = @reflection.scope
+        return base unless scope
+        unless scope.arity.zero?
+          raise Pipeloader::Batch::Error,
+                "#{@reflection.name}'s scope takes the owner as an argument, which can't be batched; " \
+                "filter at the call site (owner.#{@reflection.name}.where(...)) instead"
+        end
+
+        base.instance_exec(&scope)
+      end
+
+      # The real AR association reader — used ONLY to delegate the WRITE_METHODS
+      # above. The read path never touches it, so reads cannot silently escape
+      # batching.
+      def real_association
+        @real_association ||= @owner.association(@reflection.name).reader
+      end
+
+      # Backs bare `where` so `where.not(...)` / `.missing` / `.associated` chain:
+      # forward to the relation's own WhereChain, then re-wrap the refined relation.
+      class WhereChain
+        def initialize(proxy)
+          @proxy = proxy
+        end
+
+        %i[not missing associated].each do |meth|
+          define_method(meth) do |*args, **kwargs|
+            @proxy.__send__(:spawn, relation: @proxy.relation.where.public_send(meth, *args, **kwargs))
+          end
+        end
+      end
+    end
+  end
+end

data/lib/pipeloader/batch/context.rb

@@ -0,0 +1,43 @@
+# frozen_string_literal: true
+
+module Pipeloader
+  module Batch
+    # A sibling group: the records loaded together by one query. It lives ON the
+    # records (each carries a reference, see Model#_pipeloader_batch_context), not in
+    # any thread/fiber/request storage. That is what makes batching correct under
+    # GraphQL::Dataloader fibers, fiber-per-request servers, and plain threads alike:
+    # the context travels with the data instead of being looked up ambiently.
+    class Context
+      # Re-entrancy guard for the singular interceptor: while a Preloader fills an
+      # association, its own load_target calls must not re-enter. Fiber-local, since a
+      # preload runs synchronously in one fiber and sibling fibers preload on their own.
+      PRELOADING_KEY = :pipeloader_batch_preloading
+
+      class << self
+        def preloading?
+          Thread.current[PRELOADING_KEY] || false
+        end
+
+        def while_preloading
+          Thread.current[PRELOADING_KEY] = true
+          yield
+        ensure
+          Thread.current[PRELOADING_KEY] = false
+        end
+      end
+
+      def initialize
+        @objs = Hash.new { |hash, key| hash[key] = [] }
+      end
+
+      def add(instance)
+        @objs[instance.class] << instance
+      end
+
+      # The records of class `cls` in this group (an empty array if none).
+      def all(cls)
+        @objs[cls]
+      end
+    end
+  end
+end

data/lib/pipeloader/batch/fetcher.rb

@@ -0,0 +1,30 @@
+# frozen_string_literal: true
+
+module Pipeloader
+  module Batch
+    # A class-level relationship definition. `fetch` gathers `owner`'s sibling group,
+    # runs the loader ONCE for all their keys, then distributes each slice into the
+    # siblings' FetcherStates — the configured `default` when a key has no value.
+    class Fetcher
+      attr_reader :target_class, :name, :getter, :loader, :default
+
+      def initialize(target_class, name, getter, loader, default: nil)
+        @target_class = target_class
+        @name = name
+        @getter = getter
+        @loader = loader
+        @default = default
+      end
+
+      def fetch(owner)
+        siblings = owner._pipeloader_batch_context.all(target_class)
+        ids = siblings.map { |sibling| sibling.send(getter) }
+        results = loader.call(ids, self)
+        siblings.each_with_index do |sibling, i|
+          value = results[ids[i]]
+          sibling._fetcher_state(name).results = value.nil? ? @default : value
+        end
+      end
+    end
+  end
+end

data/lib/pipeloader/batch/fetcher_state.rb

@@ -0,0 +1,27 @@
+# frozen_string_literal: true
+
+module Pipeloader
+  module Batch
+    # Per-instance, per-relationship cache. The first access triggers the
+    # (batched) fetch; subsequent accesses return the cached slice.
+    class FetcherState
+      attr_reader :results, :fetched
+
+      def initialize(fetcher)
+        @fetcher = fetcher
+        @fetched = false
+        @results = nil
+      end
+
+      def fetch(owner)
+        @fetcher.fetch(owner) unless @fetched
+        @results
+      end
+
+      def results=(results)
+        @results = results
+        @fetched = true
+      end
+    end
+  end
+end

data/lib/pipeloader/batch/load_grouping.rb

@@ -0,0 +1,28 @@
+# frozen_string_literal: true
+
+module Pipeloader
+  module Batch
+    # Prepended onto ActiveRecord::Relation: every set of records loaded by ONE query
+    # is stamped with a shared Context (its sibling group), carried on the records.
+    # One hook covers every path that matters — a root collection, the IN query a
+    # batched has_many runs, and the queries AR's Preloader runs for a batched
+    # belongs_to / has_one all materialize through here. A record loaded on its own
+    # (find / find_by, one row) gets no group and falls back to a solo one
+    # (Model#_pipeloader_batch_context), which is correct: it has nothing to batch with.
+    module LoadGrouping
+      def records
+        recs = super
+        if recs.length > 1 && klass.include?(Pipeloader::Batch::Model)
+          group = Context.new
+          recs.each do |record|
+            next if record.instance_variable_defined?(:@_pipeloader_batch_context)
+
+            record.instance_variable_set(:@_pipeloader_batch_context, group)
+            group.add(record)
+          end
+        end
+        recs
+      end
+    end
+  end
+end

data/lib/pipeloader/batch/load_interceptor.rb

@@ -0,0 +1,44 @@
+# frozen_string_literal: true
+
+module Pipeloader
+  module Batch
+    # Prepended onto AR's SingularAssociation (belongs_to / has_one). When a
+    # batch-declared association's target is first loaded for one record, it is
+    # loaded for ALL live siblings in the current Context at once (via AR's own
+    # Preloader). Non-batch associations, and models that don't include
+    # Pipeloader::Batch::Model, fall straight through to normal AR.
+    module LoadInterceptor
+      def load_target
+        Pipeloader::Batch.batch_fill(self) unless loaded?
+        super
+      end
+    end
+
+    # The singular interceptor's hook: only fire for an association this owner has
+    # batch-declared, then defer to fill_association.
+    def self.batch_fill(association)
+      klass = association.owner.class
+      return unless klass.respond_to?(:pipeloader_batched_association?)
+      return unless klass.pipeloader_batched_association?(association.reflection.name)
+
+      fill_association(association.owner, association.reflection.name)
+    end
+
+    # Preload `name` across every live sibling of `owner` in one shot, so each
+    # sibling's target is set without a per-record query. Uses AR's own Preloader,
+    # which walks plain, polymorphic, and :through associations alike. After this,
+    # the caller's `load_target` finds the target loaded and returns it.
+    def self.fill_association(owner, name)
+      return if Context.preloading?
+
+      siblings = owner._pipeloader_batch_context.all(owner.class)
+      records = siblings.select { |record| record.persisted? && !record.association(name).loaded? }
+      records << owner if owner.persisted? && records.none? { |record| record.equal?(owner) }
+      return if records.empty?
+
+      Context.while_preloading do
+        ::ActiveRecord::Associations::Preloader.new(records: records, associations: name).call
+      end
+    end
+  end
+end

data/lib/pipeloader/batch/model.rb

@@ -0,0 +1,170 @@
+# frozen_string_literal: true
+
+require "set"
+require "active_support/concern"
+require_relative "relationship"
+
+module Pipeloader
+  module Batch
+    # Include into an ActiveRecord model to enable batch-loaded relationships.
+    #
+    #   class Author < ApplicationRecord
+    #     include Pipeloader::Batch::Model
+    #     batch_has_many   :books            # chainable, batched (where/order/limit)
+    #     batch_has_one    :profile
+    #     batch_belongs_to :publisher
+    #     batch_count      :books_count
+    #     batch_aggregate  :total_pages, of: :books, function: :sum, column: :pages
+    #   end
+    #
+    # `batch_has_many` returns a BatchProxy: a lazy, chainable relation-like object
+    # whose load is batched across every live instance of the class in the current
+    # Context — and whose .where/.order/.limit are applied INSIDE that one batched
+    # query (limit/offset per group). belongs_to/has_one return a single native
+    # record, batched via the load interceptor. Aggregates return batched scalars.
+    module Model
+      extend ActiveSupport::Concern
+
+      included do
+        class_attribute :_pipeloader_batch_associations, instance_accessor: false, default: Set.new
+        class_attribute :_pipeloader_batch_fetchers, instance_accessor: false, default: {}
+      end
+
+      module ClassMethods
+        # Sentinel marking "no explicit default given" so we can pick a sensible
+        # per-function default (0 for count/sum, nil for min/max/average).
+        AGGREGATE_DEFAULTS = { count: 0, sum: 0 }.freeze
+
+        # --- Collection: a real has_many, but the reader returns a BatchProxy ---
+
+        def batch_has_many(name, scope = nil, **options, &extension)
+          has_many(name, scope, **options, &extension)
+          reflection = reflect_on_association(name)
+          if reflection.through_reflection
+            # A :through collection's target has no direct FK to the owner, so the
+            # flat-FK BatchProxy can't build its query. Batch it through AR's
+            # Preloader (which walks the join), filling every live sibling at once;
+            # the reader returns a plain loaded array (no chainable proxy).
+            define_method(name) do
+              assoc = association(name)
+              Pipeloader::Batch.fill_association(self, name) unless assoc.loaded?
+              assoc.load_target
+            end
+          else
+            define_method(name) { Pipeloader::Batch::BatchProxy.new(self, reflection) }
+          end
+        end
+
+        # --- Singular: a real association whose load is batched (single record) ---
+
+        def batch_has_one(name, scope = nil, **options, &extension)
+          has_one(name, scope, **options, &extension)
+          _pipeloader_batch_association(name)
+        end
+
+        def batch_belongs_to(name, scope = nil, **options, &extension)
+          belongs_to(name, scope, **options, &extension)
+          _pipeloader_batch_association(name)
+        end
+
+        def pipeloader_batched_association?(name)
+          _pipeloader_batch_associations.include?(name)
+        end
+
+        # --- Scalar aggregates: COUNT / SUM / AVERAGE / MINIMUM / MAXIMUM ---
+
+        def batch_count(name, of: nil, class_name: nil, foreign_key: nil, primary_key: nil, default: 0)
+          _pipeloader_batch_register_aggregate(name, Relationship.aggregate(
+            self, name,
+            of: of, function: :count, column: nil,
+            class_name: class_name, foreign_key: foreign_key, primary_key: primary_key,
+            default: default
+          ))
+        end
+
+        # function: :sum / :average / :minimum / :maximum (or :count). Empty groups
+        # default to 0 for count/sum and nil for the rest, override with :default.
+        # For anything a plain GROUP BY can't express, use `batch` with a loader.
+        def batch_aggregate(name, function:, of: nil, column: nil, class_name: nil, foreign_key: nil,
+                            primary_key: nil, default: AGGREGATE_DEFAULTS)
+          if function != :count && column.nil?
+            raise Pipeloader::Batch::Error, "batch_aggregate #{name.inspect} (#{function}) requires a :column"
+          end
+
+          resolved_default = default.equal?(AGGREGATE_DEFAULTS) ? AGGREGATE_DEFAULTS[function] : default
+          _pipeloader_batch_register_aggregate(name, Relationship.aggregate(
+            self, name,
+            of: of, function: function, column: column,
+            class_name: class_name, foreign_key: foreign_key, primary_key: primary_key,
+            default: resolved_default
+          ))
+        end
+
+        # --- General: any batched value, computed once per Context ---
+
+        # Define `name` as a value loaded ONCE across every live sibling. The
+        # loader gets the array of owner keys (the `key:` column, default the
+        # primary key) and returns a `{ key => value }` Hash; each instance reads
+        # its own value, or `default` when its key is absent. This is the escape
+        # hatch behind batch_count / batch_aggregate, surfaced for everything that
+        # isn't a plain association — existence checks, viewer-scoped flags,
+        # lookups by a non-PK column, derived rows:
+        #
+        #   batch :viewer_has_starred, default: false do |repo_ids|
+        #     Star.where(user_id: Current.user.id, repo_id: repo_ids)
+        #         .pluck(:repo_id).index_with(true)         # missing -> false
+        #   end
+        def batch(name, key: nil, default: nil, &loader)
+          raise Pipeloader::Batch::Error, "batch #{name.inspect} requires a loader block" unless loader
+
+          getter = (key || primary_key).to_sym
+          _pipeloader_batch_register_fetcher(name, getter, ->(keys, _fetcher) { loader.call(keys) }, default)
+        end
+
+        def _pipeloader_batch_association(name)
+          # merge (not mutate) so subclasses get their own copy — STI-safe.
+          self._pipeloader_batch_associations = _pipeloader_batch_associations + [name]
+        end
+
+        def _pipeloader_batch_register_aggregate(name, relationship)
+          _pipeloader_batch_register_fetcher(name, relationship.getter, relationship.loader, relationship.default)
+        end
+
+        def _pipeloader_batch_register_fetcher(name, getter, loader, default)
+          self._pipeloader_batch_fetchers = _pipeloader_batch_fetchers.merge(
+            name => Pipeloader::Batch::Fetcher.new(self, name, getter, loader, default: default)
+          )
+          define_method(name) { batch_load(name) }
+        end
+      end
+
+      # This record's sibling group: the records loaded with it (stamped at load time
+      # by Pipeloader::Batch::LoadGrouping), or a solo group when it was loaded or
+      # built on its own. Every batch method reads its siblings from here.
+      def _pipeloader_batch_context
+        @_pipeloader_batch_context ||= Pipeloader::Batch::Context.new.tap { |c| c.add(self) }
+      end
+
+      # Aggregate readers route through here; record associations use BatchProxy
+      # (has_many) or AR's own reader (belongs_to/has_one).
+      def batch_load(name)
+        _fetcher_state(name).fetch(self)
+      end
+
+      def _fetcher_state(name)
+        ivar = :"@_pipeloader_batch_state_#{name}"
+        return instance_variable_get(ivar) if instance_variable_defined?(ivar)
+
+        fetcher = self.class._pipeloader_batch_fetchers[name]
+        raise Pipeloader::Batch::Error, "no batch relationship named #{name.inspect}" unless fetcher
+
+        instance_variable_set(ivar, Pipeloader::Batch::FetcherState.new(fetcher))
+      end
+
+      # Per-instance cache for batched collection scopes, keyed by [name, scope].
+      def _pipeloader_batch_scope_cache
+        @_pipeloader_batch_scope_cache ||= {}
+      end
+    end
+  end
+end

data/lib/pipeloader/batch/relationship.rb

@@ -0,0 +1,68 @@
+# frozen_string_literal: true
+
+require "active_support/core_ext/string/inflections"
+
+module Pipeloader
+  module Batch
+    # Builds the pieces an aggregate Fetcher needs: a `getter` (the owner key read
+    # off each sibling), a `loader` lambda (one GROUP BY query returning a
+    # { key => scalar } hash), and a `default` for keys with no rows.
+    #
+    # The child class and keys are resolved with this precedence:
+    #   explicit option > AR reflection (if the source names a real association) > convention.
+    # The target class constant is resolved lazily so autoload order doesn't matter.
+    class Relationship
+      attr_reader :getter, :loader, :default
+
+      def self.aggregate(owner, name, of:, function:, column:, class_name:, foreign_key:, primary_key:, default:)
+        source = of || default_source(name, function)
+        resolve, fk, getter = child_keys(owner, source, class_name, foreign_key, primary_key)
+
+        derived = lambda do |ids, _fetcher|
+          scope = resolve.call.where(fk => ids).group(fk)
+          function == :count ? scope.count : scope.public_send(function, column)
+        end
+        new(getter, derived, default)
+      end
+
+      # Resolve [class-resolver, foreign-key, owner-getter] for a has_many-style
+      # source rooted at `assoc_name` on `owner`.
+      def self.child_keys(owner, assoc_name, class_name, foreign_key, primary_key)
+        reflection = owner.reflect_on_association(assoc_name)
+        fk = (foreign_key || reflection&.foreign_key || owner.model_name.element.foreign_key).to_s
+        getter = (primary_key || reflection&.active_record_primary_key || owner.primary_key).to_sym
+        resolve = class_resolver(class_name, reflection) { assoc_name.to_s.singularize.camelize.constantize }
+        [resolve, fk, getter]
+      end
+
+      # For `batch_count :books_count`, derive the source association ("books") by
+      # stripping the function suffix when present.
+      def self.default_source(name, function)
+        suffix = "_#{function}"
+        name.to_s.end_with?(suffix) ? name.to_s.delete_suffix(suffix) : name.to_s
+      end
+
+      # Memoized lambda producing the target class constant. The convention block
+      # is only used when neither class_name nor a reflection is available.
+      def self.class_resolver(class_name, reflection, &convention)
+        resolved = nil
+        lambda do
+          resolved ||=
+            if class_name
+              class_name.to_s.constantize
+            elsif reflection
+              reflection.klass
+            else
+              convention.call
+            end
+        end
+      end
+
+      def initialize(getter, loader, default)
+        @getter = getter
+        @loader = loader
+        @default = default
+      end
+    end
+  end
+end

data/lib/pipeloader/batch.rb

@@ -0,0 +1,44 @@
+# frozen_string_literal: true
+
+require "active_record"
+
+require_relative "batch/context"
+require_relative "batch/fetcher"
+require_relative "batch/fetcher_state"
+require_relative "batch/relationship"
+require_relative "batch/batch_proxy"
+require_relative "batch/batch_loader"
+require_relative "batch/load_interceptor"
+require_relative "batch/load_grouping"
+require_relative "batch/model"
+
+module Pipeloader
+  # Automatic N+1 elimination for plain ActiveRecord — no GraphQL, no `includes`,
+  # no DataLoader keys. Declare a relationship with a batch macro, then traverse
+  # records one at a time; the first access loads the relationship for every record
+  # loaded alongside it (its sibling group) in a single query.
+  #
+  #   class Author < ApplicationRecord
+  #     include Pipeloader::Batch::Model
+  #     batch_has_many :books            # chainable, batched (where/order/limit)
+  #     batch_has_one  :profile
+  #     batch_count    :books_count
+  #   end
+  #
+  #   Author.all.to_a.each { |a| a.books.to_a }   # ONE query for everyone's books
+  #
+  # Siblings are the records loaded by the same query: the group is stamped onto them
+  # as they load (Pipeloader::Batch::LoadGrouping) and carried on the records, so it
+  # needs no surrounding block and is correct under threads, fibers, and
+  # fiber-per-request servers alike.
+  module Batch
+    class Error < StandardError; end
+  end
+end
+
+# belongs_to / has_one (SingularAssociation) batch their load via the interceptor +
+# AR's Preloader. has_many returns a Pipeloader::Batch::BatchProxy instead (so
+# where/order/limit stay batched), so CollectionAssociation needs no hook.
+ActiveRecord::Associations::SingularAssociation.prepend(Pipeloader::Batch::LoadInterceptor)
+# Stamp co-loaded records with their sibling group as they materialize.
+ActiveRecord::Relation.prepend(Pipeloader::Batch::LoadGrouping)