activerecord-precompiled 0.1.0

data/lib/active_record_precompiled/keys.rb

@@ -0,0 +1,41 @@
+# frozen_string_literal: true
+
+module ActiveRecordPrecompiled
+  # The columns an association joins on. Rails 8.0 and 8.1 have
+  # join_primary_key/join_foreign_key on reflections and
+  # LoaderQuery#association_key_name; Rails main dropped the first two (for
+  # *_key_definition and association routes) and renamed the third
+  # destination_key_name. These answer as the old methods did, on either.
+  module Keys
+    JOIN_KEYS = ActiveRecord::Reflection::AbstractReflection.method_defined?(:join_foreign_key)
+
+    LOADER_KEY = if ActiveRecord::Associations::Preloader::Association::LoaderQuery.method_defined?(:association_key_name)
+      :association_key_name
+    else
+      :destination_key_name
+    end
+
+    class << self
+      # The key on the associated table: what its rows are matched on.
+      def join_primary_key(reflection, klass = nil)
+        return klass ? reflection.join_primary_key(klass) : reflection.join_primary_key if JOIN_KEYS
+
+        reflection = reflection.source_reflection while reflection.through_reflection?
+        return reflection.foreign_key unless reflection.belongs_to?
+
+        reflection.polymorphic? ? reflection.association_primary_key(klass) : reflection.association_primary_key
+      end
+
+      # The key on the owner's table: the value the associated rows match.
+      def join_foreign_key(reflection)
+        return reflection.join_foreign_key if JOIN_KEYS
+
+        reflection = reflection.source_reflection while reflection.through_reflection?
+        reflection.belongs_to? ? reflection.foreign_key : reflection.active_record_primary_key
+      end
+
+      # The column a Preloader query looks its keys up in.
+      def loader_key(loader_query) = loader_query.public_send(LOADER_KEY)
+    end
+  end
+end

data/lib/active_record_precompiled/preload.rb

@@ -0,0 +1,619 @@
+# frozen_string_literal: true
+
+module ActiveRecordPrecompiled
+  # Precompiled association queries for a Site's includes/preload.
+  #
+  # ActiveRecord's Preloader runs one query per association, always through
+  # LoaderQuery#load_records_for_keys: `scope.where(key => keys).load`. The
+  # build captures that query's SQL for one key (`key = ?`) and for two
+  # (`key IN (?, ?)`), so any number of keys gets ActiveRecord's exact SQL.
+  # Everything else (grouping by owner, inverses, nested and through
+  # associations) stays ActiveRecord's.
+  #
+  # At runtime a plan is used only when the LoaderQuery is `eql?` to one built
+  # from the same reflection: the identity ActiveRecord itself uses to batch
+  # identical association queries.
+  module Preload
+    # Above this, building records dominates and the gain is small; it also
+    # keeps well clear of SQLite's bind parameter limit.
+    MAX_KEYS = 100
+    THREAD_KEY = :__ar_precompiled_preload
+    # Placeholder keys for the build-time capture; recognised among the binds.
+    PROBES = { integer: [-7_000_001, -7_000_002], string: %w[__ar_precompiled_key_0__ __ar_precompiled_key_1__] }.freeze
+    IN_TWO = " IN (?, ?)"
+
+    # --- build time ----------------------------------------------------------
+
+    # Plans for every association in a preload tree that has a fixed query.
+    # The others are left to the Preloader as is.
+    def self.compile(model, values)
+      plans = []
+      walk(model, values) do |owner, reflection|
+        if reflection.through_reflection?
+          plans.concat(compile_through(owner, reflection))
+        elsif (plan = compile_one(owner, reflection))
+          plans << plan
+        end
+      end
+      plans.uniq
+    end
+
+    def self.walk(model, value, &block)
+      case value
+      when Array then value.each { walk(model, it, &block) }
+      when Hash then value.each { |k, v| walk(model, k, &block) && walk(model._reflect_on_association(k.to_sym).klass, v, &block) }
+      when Symbol, String
+        # As the Preloader looks it up: has_and_belongs_to_many is a through association here.
+        reflection = model._reflect_on_association(value.to_sym) or return false
+        yield model, reflection
+        !reflection.polymorphic?
+      end
+    end
+
+    # A has_many/has_one :through runs two preloads of its own (owners to
+    # through records, through records to targets), with scopes it builds
+    # (ThroughAssociation#through_scope and #scope). Plans for those, found
+    # again at runtime by the same path (`via`). Nested through associations
+    # are left to ActiveRecord.
+    def self.compile_through(owner, reflection)
+      through = reflection.through_reflection
+      source = reflection.source_reflection
+      return [] if through.through_reflection? || source.through_reflection? || source.polymorphic?
+      return [] if reflection.scope && !static_scope?(reflection)
+
+      parts = through_parts(owner, reflection) or return []
+      via = [owner.name, reflection.name]
+      [
+        compile_one(owner, through, preload_scope: parts[:through], via: via + ["through"]),
+        compile_one(through.klass, source, preload_scope: parts[:source], via: via + ["source"]),
+      ].compact
+    end
+
+    # The scopes ThroughAssociation hands its two inner Preloaders.
+    def self.through_parts(owner, reflection)
+      association = ActiveRecord::Associations::Preloader::ThroughAssociation.new(
+        reflection.klass, [owner.allocate], reflection, nil, nil, false,
+      )
+      { through: association.send(:through_scope), source: association.send(:scope) }
+    rescue StandardError
+      nil
+    end
+
+    def self.compile_one(owner, reflection, preload_scope: nil, via: nil)
+      return if reflection.through_reflection? || reflection.polymorphic?
+      return if reflection.scope && !static_scope?(reflection)
+
+      klass = reflection.klass
+      return if klass.default_scopes.any? || !klass.table_exists?
+
+      loader = loader_for(owner, reflection, preload_scope)
+      # references only matters for eager_load, which a preload query isn't.
+      return unless loader && loader.scope.values.keys.all? { %i[where order references].include?(it) }
+
+      key = Keys.loader_key(loader)
+      return if key.is_a?(Array)
+
+      probes = PROBES[klass.type_for_attribute(key).type] or return
+      one, two = [probes.first(1), probes].map do |keys|
+        queries = klass.with_connection do |conn|
+          return unless conn.is_a?(ActiveRecord::ConnectionAdapters::SQLite3Adapter)
+
+          prepared(conn) { Capture.queries { loader.load_records_for_keys(keys) } }
+        end
+        return unless queries.size == 1
+
+        queries.first
+      end
+      # The key's own list: the scope may have IN lists of its own.
+      column = klass.with_connection { "#{it.quote_table_name(klass.table_name)}.#{it.quote_column_name(key)}" }
+      marker = "#{column}#{IN_TWO}"
+      return unless two.sql.scan(marker).size == 1
+
+      prefix, suffix = two.sql.split(marker, 2)
+      prefix += column
+      many = bind_specs(klass, two, probes)
+      at = many.index(:key)
+      return unless many.count(:key) == 2 && many[at + 1] == :key && bind_specs(klass, one, probes).count(:key) == 1
+
+      { owner: owner.name, name: reflection.name, schema: Site.schema_of(klass),
+        one: [one.sql, bind_specs(klass, one, probes)], many: [prefix, suffix, many[0, at], many[(at + 2)..]],
+        via: via }.compact
+    end
+
+    # Association scopes such as `-> { order(position: :asc) }`; scopes that
+    # read the time, the owner, or anything else when they run are not.
+    def self.static_scope?(reflection)
+      scope = original_scope(reflection.scope)
+      return false unless scope.arity.zero?
+
+      ScopeAnalyzer.new.analyze_proc(reflection.klass, scope)
+      true
+    rescue ScopeAnalyzer::NotStatic
+      false
+    end
+
+    # Rails wraps a scope without arguments as `proc { instance_exec(&scope) }`
+    # (Associations::Builder::Association.build_scope); analyze the original.
+    def self.original_scope(scope)
+      wrapped = scope.source_location&.first.to_s.end_with?("active_record/associations/builder/association.rb")
+      wrapped && scope.binding.local_variable_defined?(:scope) ? scope.binding.local_variable_get(:scope) : scope
+    end
+
+    def self.bind_specs(klass, query, probes)
+      query.binds.map do |bind|
+        next :key if probes.include?(bind.value_before_type_cast)
+
+        value = klass.with_connection { Site.db_constant(it, bind.value_for_database) }
+        return [] if value.equal?(Site::FALLBACK) # no plan: count(:key) won't match
+
+        [:val, value, bind.name]
+      end
+    end
+
+    def self.prepared(conn)
+      previous = conn.instance_variable_get(:@prepared_statements)
+      conn.instance_variable_set(:@prepared_statements, true)
+      yield
+    ensure
+      conn.instance_variable_set(:@prepared_statements, previous)
+    end
+
+    # The LoaderQuery the Preloader would build for this reflection. The
+    # owner is only asked for its class, so it needn't be initialized.
+    def self.loader_for(owner, reflection, preload_scope = nil)
+      association = ActiveRecord::Associations::Preloader::Association.new(
+        reflection.klass, [owner.allocate], reflection, preload_scope, nil, false,
+      )
+      association.send(:loader_query)
+    end
+
+    # --- runtime ---------------------------------------------------------------
+
+    class Plan
+      PLANS = {}
+
+      # Sites that preload the same association share one plan.
+      def self.for(definition)
+        definition = definition.transform_values { ActiveRecordPrecompiled.intern(it) }.freeze
+        PLANS[definition] ||= new(**definition)
+      end
+
+      # `via` is [owner, through association, "through" or "source"] for the
+      # inner preloads of a through association.
+      def initialize(owner:, name:, schema:, one:, many:, via: nil)
+        @owner_name = owner
+        @name = name
+        @via = via
+        @schema = schema
+        @one_sql, @one_binds = one
+        @prefix, @suffix, @before, @after = many
+      end
+
+      attr_reader :owner_name, :name, :via
+
+      # The verified reference, computed once; false if this process would
+      # build different SQL.
+      def loader_query
+        return @loader_query if defined?(@loader_query)
+
+        @loader_query = reference || false
+      end
+
+      # The LoaderQuery this plan answers for, or nil if this process would
+      # build different SQL than the build captured (model or schema changed):
+      # then the plan is never used. Checked once per Site.
+      def reference
+        owner, scope = @via ? via_owner_and_scope : [constant(@owner_name), nil]
+        return unless owner
+
+        reflection = owner._reflect_on_association(@name)
+        return unless reflection && Site.schema_of(reflection.klass) == @schema
+
+        loader = Preload.loader_for(owner, reflection, scope)
+        loader if Preload.sql_for_one_key(loader) == @one_sql
+      end
+
+      # The model this inner preload of a through association runs on (the
+      # join model of a has_and_belongs_to_many is anonymous), and the preload
+      # scope ThroughAssociation passes it.
+      private def via_owner_and_scope
+        through_owner = constant(@via[0]) or return nil
+        through = through_owner._reflect_on_association(@via[1])
+        return nil unless through&.through_reflection?
+
+        parts = Preload.through_parts(through_owner, through) or return nil
+        if @via[2] == "through"
+          [through_owner, parts[:through]]
+        else
+          [through.through_reflection.klass, parts[:source]]
+        end
+      end
+
+      private def constant(name)
+        Object.const_get(name)
+      rescue NameError
+        nil
+      end
+
+      # Records for `keys`, as `scope.where(key => keys).load(&block)` would
+      # return them; Site::FALLBACK when a key isn't a plain value.
+      def load(loader_query, keys, &block)
+        load_keys(loader_query.scope.klass, Keys.loader_key(loader_query), keys, &block)
+      end
+
+      def load_keys(klass, key_name, keys, &block)
+        type = klass.type_for_attribute(key_name)
+        klass.with_connection do |conn|
+          return Site::FALLBACK unless conn.is_a?(ActiveRecord::ConnectionAdapters::SQLite3Adapter)
+
+          values = keys.map do |key|
+            value = type.serialize(key)
+            return Site::FALLBACK if value.nil? || !Site::SCALARS.any? { it === value }
+
+            Site.db_value(conn, value)
+          end
+          sql, binds = statement_for(values)
+          columns, rows = Site.execute(conn, sql, binds, "#{klass.name} Load", self, [klass, key_name, values.size])
+          records = klass._load_from_sql(ActiveRecord::Result.new(columns, rows), &block)
+          Site.lazy_preload(records)
+          records
+        end
+      end
+
+      private
+
+      def statement_for(values)
+        if values.size == 1
+          [@one_sql, @one_binds.map { it == :key ? values[0] : it[1] }]
+        else
+          sql = (@sql_by_count ||= {})[values.size] ||= "#{@prefix} IN (#{Array.new(values.size, "?").join(", ")})#{@suffix}"
+          [sql, @before.map { it[1] } + values + @after.map { it[1] }]
+        end
+      end
+
+      public
+
+      # [name, type] per bind, for the `sql.active_record` payload.
+      def bind_names((klass, key_name, count))
+        key = [key_name, klass.type_for_attribute(key_name)]
+        constant = ->(spec) { [spec[2], ActiveModel::Type.default_value] }
+        if count == 1
+          @one_binds.map { it == :key ? key : constant.(it) }
+        else
+          @before.map(&constant) + Array.new(count, key) + @after.map(&constant)
+        end
+      end
+    end
+
+    # LoaderQuery => Plan for a Site's verified plans, built on first use.
+    def self.registry(plans)
+      # Equal LoaderQueries (the same table, key and scope) share one plan:
+      # they run the same SQL.
+      plans.each_with_object({}) do |plan, registry|
+        (query = plan.loader_query) && (registry[query] ||= plan)
+      end
+    end
+
+    # ActiveRecord's SQL for one key, as the preloader would send it with
+    # prepared statements on.
+    def self.sql_for_one_key(loader)
+      klass = loader.scope.klass
+      key = Keys.loader_key(loader)
+      probe = PROBES[klass.type_for_attribute(key).type]&.first or return
+      klass.with_connection do |conn|
+        prepared(conn) { conn.__send__(:to_sql_and_binds, loader.scope.where(key => [probe]).arel, [], nil).first }
+      end
+    end
+
+    # A preload tree run without the Preloader's objects, for trees where
+    # every association has a verified plan. It is Preloader::Batch in
+    # miniature, and does what Batch, Branch, Association and
+    # ThroughAssociation do, in the same order: rounds of runnable loaders,
+    # loaders deferred while another branch will still load their table,
+    # loaders with equal LoaderQueries loaded by one query (so records are
+    # shared between associations as ActiveRecord shares them), owners
+    # grouped by key, inverses, already-loaded targets.
+    class Tree
+      # A through association has no plan of its own: `through` and `source`
+      # are the two inner preloads ThroughAssociation runs.
+      Node = Struct.new(:reflection, :plan, :owner_key, :key, :convert, :children, :through, :source, :future)
+
+      # nil when the tree doesn't qualify; the Site then uses the Preloader.
+      # `plans`: [owner, association] => Plan; `inner`: [owner, association,
+      # "through" | "source"] => Plan for through associations.
+      def self.build(model, value, plans, inner = {})
+        nodes = nodes_for(model, value, plans, inner)
+        nodes ? new(nodes) : nil
+      end
+
+      def self.nodes_for(model, value, plans, inner)
+        return nil unless model.descendants.empty? # STI subclasses may redefine associations
+
+        case value
+        when Array
+          nodes = value.map { nodes_for(model, it, plans, inner) }
+          nodes.all? ? nodes.flatten : nil
+        when Hash
+          value.map do |name, children|
+            node = node_for(model, name, plans, inner) or return nil
+            node.children = nodes_for(node.reflection.klass, children, plans, inner) or return nil
+            node
+          end
+        when Symbol, String
+          (node = node_for(model, value, plans, inner)) ? [node] : nil
+        end
+      end
+
+      def self.node_for(model, name, plans, inner)
+        reflection = model._reflect_on_association(name) or return nil
+        return nil if reflection.respond_to?(:deprecated?) && reflection.deprecated?
+        return through_node(model, reflection, inner) if reflection.through_reflection?
+
+        plan = plans[[model.name, reflection.name]] or return nil
+        plain_node(model, reflection, plan)
+      end
+
+      def self.plain_node(model, reflection, plan)
+        klass = reflection.klass
+        key = Keys.join_primary_key(reflection, klass).to_s
+        owner_key = Keys.join_foreign_key(reflection).to_s
+        convert = klass.type_for_attribute(key).type != model.type_for_attribute(owner_key).type
+        Node.new(reflection, plan, owner_key, key, convert, [])
+      end
+
+      # Only when both inner preloads assign what they load (their scopes are
+      # empty: no conditions, order or distinct on the through association),
+      # so an owner's targets are the source records of its through records.
+      def self.through_node(model, reflection, inner)
+        through = reflection.through_reflection
+        source = reflection.source_reflection
+        return nil if through.through_reflection? || source.through_reflection? || source.polymorphic?
+        return nil unless through.klass.descendants.empty?
+
+        parts = Preload.through_parts(model, reflection) or return nil
+        return nil unless parts[:through].empty_scope? && parts[:source].empty_scope?
+
+        through_plan = inner[[model.name, reflection.name, "through"]] or return nil
+        source_plan = inner[[model.name, reflection.name, "source"]] or return nil
+        Node.new(reflection, nil, nil, nil, nil, [], plain_node(model, through, through_plan), plain_node(through.klass, source, source_plan))
+      end
+
+      def initialize(nodes)
+        @nodes = nodes
+        nodes.each { compute_future(it) }
+      end
+
+      # What a branch that hasn't started yet will load: Branch#future_classes
+      # through #likely_reflections (each reflection's chain), and its children's.
+      def compute_future(node)
+        node.children.each { compute_future(it) }
+        node.future = (node.reflection.chain.map(&:klass) + node.children.flat_map(&:future)).uniq
+      end
+
+      # Preloader::Batch#call.
+      def call(records, registry)
+        branches = @nodes.map { Branch.new(it, records) }
+        until branches.empty?
+          loaders = branches.flat_map(&:runnable_loaders)
+          if loaders.any?
+            future_tables = branches.flat_map { |branch| branch.future_classes - branch.runnable_loaders.map(&:klass) }.map(&:table_name).uniq
+            targets = loaders.reject { future_tables.include?(it.klass.table_name) }
+            targets = loaders if targets.empty?
+
+            targets.grep(Loader).group_by { [it.loader_query, it.klass] }.each_value { load_batch(it) }
+            targets.each(&:run)
+          end
+          finished, in_progress = branches.partition(&:done?)
+          branches = in_progress + finished.flat_map(&:children)
+        end
+      end
+
+      private
+
+      # LoaderQuery#load_records_in_batch with LoaderRecords: one query for
+      # the keys of every loader in the group, plus their loaded targets.
+      def load_batch(loaders)
+        keys = {} # an ordered set
+        already = {}
+        loaders.each do |loader|
+          loader.owners_by_key.each do |key, owners|
+            if (loaded_owner = owners.find { loader.loaded?(it) })
+              already[key] = loader.target_for(loaded_owner)
+            else
+              keys[key] = true
+            end
+          end
+        end
+        keys = keys.keys - already.keys
+
+        set_inverse = ->(record) { loaders.each { it.set_inverse(record) } }
+        first = loaders.first
+        fetched = keys.empty? ? [] : first.node.plan.load_keys(first.klass, first.node.key, keys, &set_inverse) if keys.size <= MAX_KEYS
+        # Too many keys, or one that isn't a plain value: ActiveRecord's own query.
+        fetched = first.loader_query.load_records_for_keys(keys, &set_inverse).to_a if fetched.nil? || fetched.equal?(Site::FALLBACK)
+
+        raw = fetched + already.values.flatten
+        loaders.each do |loader|
+          loader.load_records(raw)
+          loader.run
+        end
+      end
+
+      # Preloader::Branch for one association of the tree.
+      class Branch
+        def initialize(node, owners)
+          @node = node
+          @loader = node.through ? ThroughLoader.new(node, owners) : Loader.new(node, owners)
+        end
+
+        def runnable_loaders = @loader.runnable_loaders.reject(&:run?)
+        def done? = @loader.run?
+        def future_classes = (@loader.future_classes + @node.children.flat_map(&:future)).uniq
+        def children = @node.children.map { Branch.new(it, @loader.preloaded_records) }
+      end
+
+      # Preloader::Association.
+      class Loader
+        attr_reader :node, :klass, :owners
+
+        def initialize(node, owners)
+          @node = node
+          @klass = node.reflection.klass
+          @owners = owners.uniq(&:__id__)
+          @name = node.reflection.name
+          @run = false
+        end
+
+        def loader_query = @node.plan.loader_query
+        def run? = @run
+        def runnable_loaders = [self]
+        def future_classes = run? ? [] : [klass]
+        def loaded?(owner) = owner.association(@name).loaded?
+        def target_for(owner) = Array.wrap(owner.association(@name).target)
+
+        def owners_by_key
+          @owners_by_key ||= owners.each_with_object({}) do |owner, result|
+            key = key_of(owner, @node.owner_key)
+            (result[key] ||= []) << owner if key
+          end
+        end
+
+        def set_inverse(record)
+          if (owners = owners_by_key[key_of(record, @node.key)])
+            owners.first.association(@name).set_inverse_instance(record)
+          end
+        end
+
+        def load_records(raw)
+          @records_by_owner = {}.compare_by_identity
+          collection = @node.reflection.collection?
+          @preloaded_records = raw.select do |record|
+            assigned = false
+            owners_by_key[key_of(record, @node.key)]&.each do |owner|
+              entries = (@records_by_owner[owner] ||= [])
+              if collection || entries.empty?
+                entries << record
+                assigned = true
+              end
+            end
+            assigned
+          end
+        end
+
+        def records_by_owner = @records_by_owner || {}
+        def preloaded_records = @preloaded_records || []
+
+        def run
+          return self if run?
+
+          @run = true
+          owners.each { Tree.assign(it.association(@name), @node.reflection, records_by_owner[it] || []) }
+          self
+        end
+
+        private
+
+        def key_of(record, attribute)
+          key = record._read_attribute(attribute)
+          @node.convert ? key.to_s : key
+        end
+      end
+
+      # Preloader::ThroughAssociation, whose two inner preloads assign.
+      class ThroughLoader
+        attr_reader :klass, :owners
+
+        def initialize(node, owners)
+          @node = node
+          @klass = node.reflection.klass
+          @owners = owners.uniq(&:__id__)
+          @name = node.reflection.name
+          @through = Loader.new(node.through, @owners)
+          @run = false
+        end
+
+        def run? = @run
+        def loaded?(owner) = owner.association(@name).loaded?
+
+        # Inner Preloader for the source: none when there are no through records.
+        def source
+          return @source if defined?(@source)
+
+          middle = @through.records_by_owner.values.flatten
+          @source = middle.empty? ? nil : Loader.new(@node.source, middle)
+        end
+
+        def data_available?
+          owners.all? { loaded?(it) } || (@through.run? && (source.nil? || source.run?))
+        end
+
+        def runnable_loaders
+          if data_available? then [self]
+          elsif @through.run? then source.runnable_loaders
+          else @through.runnable_loaders
+          end
+        end
+
+        def future_classes
+          if run? then []
+          elsif @through.run? then source ? source.future_classes : []
+          else (@through.future_classes + @node.reflection.source_reflection.chain.map(&:klass)).uniq
+          end
+        end
+
+        def preloaded_records = source ? source.preloaded_records : []
+
+        def run
+          return self if run?
+
+          @run = true
+          source_by_owner = source ? source.records_by_owner : {}
+          owners.each do |owner|
+            association = owner.association(@name)
+            next if association.loaded?
+
+            records = (@through.records_by_owner[owner] || []).flat_map { source_by_owner[it] || [] }
+            Tree.assign(association, @node.reflection, records)
+          end
+          self
+        end
+      end
+
+      # Association#associate_records_to_owner.
+      def self.assign(association, reflection, records)
+        return if association.loaded?
+
+        if reflection.collection?
+          association.target = records + association.target.reject(&:persisted?)
+        else
+          association.target = records.first
+        end
+      end
+    end
+
+    # Makes plans available to load_records_for_keys while the block runs.
+    def self.with(registry)
+      previous = Thread.current[THREAD_KEY]
+      Thread.current[THREAD_KEY] = registry
+      yield
+    ensure
+      Thread.current[THREAD_KEY] = previous
+    end
+
+    module Hook
+      def load_records_for_keys(keys, &block)
+        registry = Thread.current[Preload::THREAD_KEY]
+        return super unless registry && !keys.empty? && keys.size <= Preload::MAX_KEYS
+        return super unless (plan = registry[self])
+
+        records = plan.load(self, keys.to_a, &block)
+        records.equal?(Site::FALLBACK) ? super : records
+      end
+    end
+
+    def self.install!
+      loader = ActiveRecord::Associations::Preloader::Association::LoaderQuery
+      loader.prepend(Hook) unless loader <= Hook
+    end
+  end
+end