jit_preloader 1.0.4 → 2.1.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 515aa37a3edcd60ff85c7490bf7a96906a612c7b5645d3d81542c5e7997a9026
-  data.tar.gz: 4cab1990cd0346da95d0d0785d28b18d5db95fecf66dbfcf14be26d0b020af5d
+  metadata.gz: bca887aa87a283d55ee26b1fe606948733ecfecc7ce40d4cedc8b89d90efb834
+  data.tar.gz: 9a7c463ccacf8daeebc23d2eaa375d8c4f908b333c492baaa4c4c931bbc7d1b3
 SHA512:
-  metadata.gz: bf81adb00759b36e3bea5dfcaf9431c3259cac06412b57008e799585c7094d97073f836300a9631bd24fa0de7820874646861e7ea68e27cb773662d2e3826dec
-  data.tar.gz: b96b45f91474ca01a5235d5ba071c750268e7c029fcfa371c308747433a5cb9f54bfb2560a99051c335352c8a536a01887ae01bb840337ee2e9ab3a808f5ec84
+  metadata.gz: 30920baa0426de8bc5188fc80ccfacf6e040f6bf1f7c934214bb75600080a090756ca957da9a59327a63b30550d7a4efdf5f4e66c873783e0c5aa627fe558e9f
+  data.tar.gz: 9adfa72af9bfcc684c3c22f6e82b353af9b36081224cae92e248e5c52db7f028a9c6f49dc3760d059f476980a48f4852b6f94c5e5faf82dbef28087e0b8dd3cd

data/.github/workflows/ci.yml

@@ -0,0 +1,32 @@
+name: CI
+
+on:
+  push:
+    branches: [master]
+  pull_request:
+    branches: [master]
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        gemfile:
+          - Gemfile
+          - Gemfile.5.2
+          - Gemfile.6.0
+          - Gemfile.6.1
+    env:
+      BUNDLE_GEMFILE: ${{ matrix.gemfile }}
+    steps:
+      - uses: actions/checkout@v4
+      - name: Set up Ruby ${{ matrix.ruby-version }}
+        uses: ruby/setup-ruby@v1
+        with:
+          ruby-version: 2.7
+      - name: Install dependencies
+        run: bundle install
+      - name: Run tests
+        run:
+          bundle exec rspec

data/.github/workflows/gem-push.yml

@@ -1,7 +1,6 @@
 name: Ruby Gem
 
 on:
-  workflow_dispatch:
   release:
     types: [ published ]
 
@@ -13,11 +12,11 @@ jobs:
       contents: read
 
     steps:
-    - uses: actions/checkout@v2
+    - uses: actions/checkout@v4
     - name: Set up Ruby 2.7
-      uses: actions/setup-ruby@v1
+      uses: ruby/setup-ruby@v1
       with:
-        ruby-version: 2.7.x
+        ruby-version: 2.7
 
     - name: Publish to RubyGems
       env:

data/.gitignore

@@ -10,6 +10,7 @@
 /test/tmp/
 /test/version_tmp/
 /tmp/
+.idea/*
 
 # Used by dotenv library to load environment variables.
 # .env

data/Gemfile

@@ -1,4 +1,5 @@
 source 'https://rubygems.org'
 
+gem "activerecord", ">=7"
 # Specify your gem's dependencies in jit_preloader.gemspec
 gemspec

data/README.md

@@ -173,6 +173,20 @@ end
 
 ```
 
+Furthermore, there is an argument `max_ids_per_query` setting max ids per query. This helps prevent running a single query with too large list of ids which may be less efficient than splitting into multiple queries.
+```ruby
+class Contact < ActiveRecord::Base
+  has_many :addresses
+  has_many_aggregate :addresses, :count_all, :count, "*", max_ids_per_query: 10
+end
+
+Contact.jit_preload.each do |contact|
+  contact.addresses_count_all
+end
+# SELECT contact_id, COUNT(*) FROM addresses WHERE contact_id IN (1, 2, 3, ... ,10) GROUP BY contact_id
+# SELECT contact_id, COUNT(*) FROM addresses WHERE contact_id IN (11, 12, 13) GROUP BY contact_id
+```
+
 ### Preloading a subset of an association
 
 There are often times when you want to preload a subset of an association, or change how the SQL statement is generated. For example, if a `Contact` model has
@@ -213,6 +227,7 @@ end
 ### Jit preloading globally across your application
 
 The JitPreloader can be globally enabled, in which case most N+1 queries in your app should just disappear. It is off by default.
+The `max_ids_per_query` argument on loading aggregate methods can also apply on a global level.
 
 ```ruby
 # Can be true or false
@@ -223,12 +238,24 @@ JitPreloader.globally_enabled = true
 # so that you can turn it on or off dynamically.
 JitPreloader.globally_enabled = ->{ $redis.get('always_jit_preload') == 'on' }
 
+# Setting global max ids constraint on all aggregation methods.
+JitPreloader.max_ids_per_query = 10
+
+class Contact < ActiveRecord::Base
+  has_many :emails
+  has_many_aggregate :emails, :count_all, :count, "*"
+end
+
 # When enabled globally, this would not generate an N+1 query.
 Contact.all.each do |contact|
   contact.emails.each do |email|
     # do something
   end
+  # When max_ides_per_query is set globally, the aggregate method will split query base on the limit.
+  contact.emails_count_all
 end
+
+
 ```
 
 ## What it doesn't solve

data/jit_preloader.gemspec

@@ -18,7 +18,7 @@ Gem::Specification.new do |spec|
   spec.test_files    = spec.files.grep(%r{^(test|spec|features)/})
   spec.require_paths = ["lib"]
 
-  spec.add_dependency "activerecord", ">= 5.2", "< 7"
+  spec.add_dependency "activerecord", "< 8"
   spec.add_dependency "activesupport"
 
   spec.add_development_dependency "bundler"

data/lib/jit_preloader/active_record/associations/preloader/ar7_association.rb

@@ -0,0 +1,63 @@
+module JitPreloader
+  module PreloaderAssociation
+
+    # A monkey patch to ActiveRecord. The old method looked like the snippet
+    # below. Our changes here are that we remove records that are already
+    # part of the target, then attach all of the records to a new jit preloader.
+    #
+    # def run
+    #   records = records_by_owner
+
+    #   owners.each do |owner|
+    #     associate_records_to_owner(owner, records[owner] || [])
+    #   end if @associate
+
+    #   self
+    # end
+
+    def run
+      return unless (reflection.scope.nil? || reflection.scope.arity == 0) && klass.ancestors.include?(ActiveRecord::Base)
+
+      super.tap do
+        if preloaded_records.any? && preloaded_records.none?(&:jit_preloader)
+          JitPreloader::Preloader.attach(preloaded_records) if owners.any?(&:jit_preloader) || JitPreloader.globally_enabled?
+        end
+      end
+    end
+
+    # Original method:
+    # def associate_records_to_owner(owner, records)
+    #   return if loaded?(owner)
+    #
+    #   association = owner.association(reflection.name)
+    #
+    #   if reflection.collection?
+    #     association.target = records
+    #   else
+    #     association.target = records.first
+    #   end
+    # end
+    def associate_records_to_owner(owner, records)
+      return if loaded?(owner)
+
+      association = owner.association(reflection.name)
+
+      if reflection.collection?
+        new_records = association.target.any? ? records - association.target : records
+        association.target.concat(new_records)
+        association.loaded!
+      else
+        association.target = records.first
+      end
+    end
+
+    def build_scope
+      super.tap do |scope|
+        scope.jit_preload! if owners.any?(&:jit_preloader) || JitPreloader.globally_enabled?
+      end
+    end
+  end
+end
+
+ActiveRecord::Associations::Preloader::Association.prepend(JitPreloader::PreloaderAssociation)
+ActiveRecord::Associations::Preloader::ThroughAssociation.prepend(JitPreloader::PreloaderAssociation)

data/lib/jit_preloader/active_record/associations/preloader/ar7_branch.rb

@@ -0,0 +1,22 @@
+module JitPreloader
+  module PreloaderBranch
+    """
+    ActiveRecord version >= 7.x.x introduced an improvement for preloading associations in batches:
+    https://github.com/rails/rails/blob/main/activerecord/lib/active_record/associations/preloader.rb#L121
+
+    Our existing monkey-patches will ignore associations whose classes are not descendants of
+    ActiveRecord::Base (example: https://github.com/clio/jit_preloader/blob/master/lib/jit_preloader/active_record/associations/preloader/ar6_association.rb#L19).
+    But this change breaks that behaviour because now Batch is calling `klass.base_class` (a method defined by ActiveRecord::Base)
+    before we have a chance to filter out the non-AR classes.
+    This patch for AR 7.x makes the Branch class ignore any association loaders that aren't for ActiveRecord::Base subclasses.
+    """
+
+    def loaders
+      @loaders = super.find_all do |loader|
+        loader.klass < ::ActiveRecord::Base
+      end
+    end
+  end
+end
+
+ActiveRecord::Associations::Preloader::Branch.prepend(JitPreloader::PreloaderBranch)

data/lib/jit_preloader/active_record/base.rb

@@ -18,7 +18,7 @@ module JitPreloadExtension
     end
   end
 
-  if Gem::Version.new(ActiveRecord::VERSION::STRING) >= Gem::Version.new("6.0.0")
+  if Gem::Version.new(ActiveRecord::VERSION::STRING) >= Gem::Version.new("7.0.0")
     def preload_scoped_relation(name:, base_association:, preload_scope: nil)
       return jit_preload_scoped_relations[name] if jit_preload_scoped_relations&.key?(name)
 
@@ -33,11 +33,11 @@ module JitPreloadExtension
         end
       end
 
-      preloader_association = ActiveRecord::Associations::Preloader.new.preload(
-        records,
-        base_association,
-        preload_scope
-      ).first
+      preloader_association = ActiveRecord::Associations::Preloader.new(
+        records: records,
+        associations: base_association,
+        scope: preload_scope
+      ).call.first
 
       records.each do |record|
         record.jit_preload_scoped_relations ||= {}
@@ -90,7 +90,7 @@ module JitPreloadExtension
     class << base
       delegate :jit_preload, to: :all
 
-      def has_many_aggregate(assoc, name, aggregate, field, table_alias_name: nil, default: 0)
+      def has_many_aggregate(assoc, name, aggregate, field, table_alias_name: nil, default: 0, max_ids_per_query: nil)
         method_name = "#{assoc}_#{name}"
 
         define_method(method_name) do |conditions={}|
@@ -101,6 +101,13 @@ module JitPreloadExtension
           if jit_preloader
             reflection = association(assoc).reflection
             primary_ids = jit_preloader.records.collect{|r| r[reflection.active_record_primary_key] }
+            max_ids_per_query = max_ids_per_query || JitPreloader.max_ids_per_query
+            if max_ids_per_query
+              slices = primary_ids.each_slice(max_ids_per_query)
+            else
+              slices = [primary_ids]
+            end
+
             klass = reflection.klass
 
             aggregate_association = reflection
@@ -115,15 +122,13 @@ module JitPreloadExtension
             table_reference = table_alias_name
             table_reference ||= association_scope.references_values.first || aggregate_association.table_name
 
-            conditions[table_reference] = { aggregate_association.foreign_key => primary_ids }
-
             # If the association is a STI child model, specify its type in the condition so that it
             # doesn't include results from other child models
             parent_is_base_class = aggregate_association.klass.superclass.abstract_class? || aggregate_association.klass.superclass == ActiveRecord::Base
             has_type_column = aggregate_association.klass.column_names.include?(aggregate_association.klass.inheritance_column)
             is_child_sti_model = !parent_is_base_class && has_type_column
             if is_child_sti_model
-              conditions[table_reference].merge!({ aggregate_association.klass.inheritance_column => aggregate_association.klass.sti_name })
+              conditions[table_reference] = { aggregate_association.klass.inheritance_column => aggregate_association.klass.sti_name }
             end
 
             if reflection.type.present?
@@ -131,11 +136,15 @@ module JitPreloadExtension
             end
             group_by = "#{table_reference}.#{aggregate_association.foreign_key}"
 
-            preloaded_data = Hash[association_scope
-              .where(conditions)
-              .group(group_by)
-              .send(aggregate, field)
-            ]
+            preloaded_data = {}
+            slices.each do |slice|
+              data = Hash[association_scope
+                            .where(conditions.deep_merge(table_reference => { aggregate_association.foreign_key => slice }))
+                            .group(group_by)
+                            .send(aggregate, field)
+              ]
+              preloaded_data.merge!(data)
+            end
 
             jit_preloader.records.each do |record|
               record.jit_preload_aggregates ||= {}

data/lib/jit_preloader/preloader.rb

@@ -3,22 +3,47 @@ module JitPreloader
 
     attr_accessor :records
 
-    def self.attach(records)
-      new.tap do |loader|
-        loader.records = records.dup
-        records.each do |record|
-          record.jit_preloader = loader
+    if Gem::Version.new(ActiveRecord::VERSION::STRING) >= Gem::Version.new("7.0.0")
+      def self.attach(records)
+        new(records: records.dup, associations: nil).tap do |loader|
+          records.each do |record|
+            record.jit_preloader = loader
+          end
         end
       end
-    end
 
-    def jit_preload(association)
-      # It is possible that the records array has multiple different classes (think single table inheritance).
-      # Thus, it is possible that some of the records don't have an association.
-      records_with_association = records.reject{|r| r.class.reflect_on_association(association).nil? }
-      preload records_with_association, association
+      def jit_preload(associations)
+        # It is possible that the records array has multiple different classes (think single table inheritance).
+        # Thus, it is possible that some of the records don't have an association.
+        records_with_association = records.reject{|r| r.class.reflect_on_association(associations).nil? }
+
+        # Some of the records may already have the association loaded and we should not load them again
+        records_requiring_loading = records_with_association.select{|r| !r.association(associations).loaded? }
+
+        self.class.new(records: records_requiring_loading, associations: associations).call
+      end
+    else
+      def self.attach(records)
+        new.tap do |loader|
+          loader.records = records.dup
+          records.each do |record|
+            record.jit_preloader = loader
+          end
+        end
+      end
+
+      def jit_preload(associations)
+        # It is possible that the records array has multiple different classes (think single table inheritance).
+        # Thus, it is possible that some of the records don't have an association.
+        records_with_association = records.reject{ |record| record.class.reflect_on_association(associations).nil? }
+
+        # Some of the records may already have the association loaded and we should not load them again
+        records_requiring_loading = records_with_association.select{ |record| !record.association(associations).loaded? }
+        preload records_with_association, associations
+      end
     end
 
+
     # We do not want the jit_preloader to be dumpable
     # If you dump a ActiveRecord::Base object that has a jit_preloader instance variable
     # you will also end up dumping all of the records the preloader has reference to.

data/lib/jit_preloader/version.rb

@@ -1,3 +1,3 @@
 module JitPreloader
-  VERSION = "1.0.4"
+  VERSION = "2.1.0"
 end

data/lib/jit_preloader.rb

@@ -8,7 +8,10 @@ require 'jit_preloader/active_record/base'
 require 'jit_preloader/active_record/relation'
 require 'jit_preloader/active_record/associations/collection_association'
 require 'jit_preloader/active_record/associations/singular_association'
-if Gem::Version.new(ActiveRecord::VERSION::STRING) >= Gem::Version.new("6.0.0")
+if Gem::Version.new(ActiveRecord::VERSION::STRING) >= Gem::Version.new("7.0.0")
+  require 'jit_preloader/active_record/associations/preloader/ar7_association'
+  require 'jit_preloader/active_record/associations/preloader/ar7_branch'
+elsif Gem::Version.new(ActiveRecord::VERSION::STRING) >= Gem::Version.new("6.0.0")
   require 'jit_preloader/active_record/associations/preloader/ar6_association'
 elsif Gem::Version.new(ActiveRecord::VERSION::STRING) >= Gem::Version.new("5.2.2")
   require 'jit_preloader/active_record/associations/preloader/ar5_association'
@@ -23,6 +26,16 @@ module JitPreloader
     @enabled = value
   end
 
+  def self.max_ids_per_query=(max_ids)
+    if max_ids && max_ids >= 1
+      @max_ids_per_query = max_ids
+    end
+  end
+
+  def self.max_ids_per_query
+    @max_ids_per_query
+  end
+
   def self.globally_enabled?
     if @enabled && @enabled.respond_to?(:call)
       @enabled.call

data/spec/lib/jit_preloader/preloader_spec.rb

@@ -1,4 +1,5 @@
 require "spec_helper"
+require "db-query-matchers"
 
 RSpec.describe JitPreloader::Preloader do
   let!(:contact1) do
@@ -479,6 +480,69 @@ RSpec.describe JitPreloader::Preloader do
         end
       end
     end
+
+    context "with dive limit set" do
+      let!(:contact_book_1) { ContactBook.create(name: "The Yellow Pages") }
+      let!(:contact_book_2) { ContactBook.create(name: "The Yellow Pages") }
+      let!(:contact_book_3) { ContactBook.create(name: "The Yellow Pages") }
+      let!(:company1) { Company.create(name: "Company1", contact_book: contact_book_1) }
+      let!(:company2) { Company.create(name: "Company2", contact_book: contact_book_1) }
+      let!(:company3) { Company.create(name: "Company2", contact_book: contact_book_2) }
+      let!(:company4) { Company.create(name: "Company4", contact_book: contact_book_3) }
+      let!(:company5) { Company.create(name: "Company5", contact_book: contact_book_3) }
+
+      context "from the global value" do
+        before do
+          JitPreloader.max_ids_per_query = 2
+        end
+
+        after do
+          JitPreloader.max_ids_per_query = nil
+        end
+
+        it "can handle queries" do
+          contact_books = ContactBook.jit_preload.to_a
+
+          expect(contact_books.first.companies_count).to eq 2
+          expect(contact_books.second.companies_count).to eq 1
+          expect(contact_books.last.companies_count).to eq 2
+        end
+
+        it "makes the right number of queries based on dive limit" do
+          contact_books = ContactBook.jit_preload.to_a
+          expect do
+            contact_books.first.companies_count
+          end.to make_database_queries(count: 2)
+
+          expect do
+            contact_books.second.companies_count
+            contact_books.last.companies_count
+          end.to_not make_database_queries
+        end
+      end
+
+      context "from aggregate argument" do
+        it "can handle queries" do
+          contact_books = ContactBook.jit_preload.to_a
+
+          expect(contact_books.first.companies_count_with_max_ids_set).to eq 2
+          expect(contact_books.second.companies_count_with_max_ids_set).to eq 1
+          expect(contact_books.last.companies_count_with_max_ids_set).to eq 2
+        end
+
+        it "makes the right number of queries based on dive limit" do
+          contact_books = ContactBook.jit_preload.to_a
+          expect do
+            contact_books.first.companies_count_with_max_ids_set
+          end.to make_database_queries(count: 2)
+
+          expect do
+            contact_books.second.companies_count_with_max_ids_set
+            contact_books.last.companies_count_with_max_ids_set
+          end.to_not make_database_queries
+        end
+      end
+    end
   end
 
 end

data/spec/support/models.rb

@@ -10,6 +10,7 @@ class ContactBook < ActiveRecord::Base
   has_many :children, through: :parents
 
   has_many_aggregate :companies, :count, :count, "*"
+  has_many_aggregate :companies, :count_with_max_ids_set, :count, "*", max_ids_per_query: 2
   has_many_aggregate :employees, :count, :count, "*"
   has_many_aggregate :company_employees, :count, :count, "*"
   has_many_aggregate :children, :count, :count, "*"

metadata

@@ -1,35 +1,29 @@
 --- !ruby/object:Gem::Specification
 name: jit_preloader
 version: !ruby/object:Gem::Version
-  version: 1.0.4
+  version: 2.1.0
 platform: ruby
 authors:
 - Kyle d'Oliveira
 autorequire: 
 bindir: bin
 cert_chain: []
-date: 2021-10-14 00:00:00.000000000 Z
+date: 2024-04-29 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: activerecord
   requirement: !ruby/object:Gem::Requirement
     requirements:
-    - - ">="
-      - !ruby/object:Gem::Version
-        version: '5.2'
     - - "<"
       - !ruby/object:Gem::Version
-        version: '7'
+        version: '8'
   type: :runtime
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
-    - - ">="
-      - !ruby/object:Gem::Version
-        version: '5.2'
     - - "<"
       - !ruby/object:Gem::Version
-        version: '7'
+        version: '8'
 - !ruby/object:Gem::Dependency
   name: activesupport
   requirement: !ruby/object:Gem::Requirement
@@ -151,7 +145,8 @@ executables: []
 extensions: []
 extra_rdoc_files: []
 files:
-- ".github/workflows/CODEOWNERS"
+- ".github/CODEOWNERS"
+- ".github/workflows/ci.yml"
 - ".github/workflows/gem-push.yml"
 - ".gitignore"
 - ".rspec"
@@ -170,6 +165,8 @@ files:
 - lib/jit_preloader/active_record/associations/collection_association.rb
 - lib/jit_preloader/active_record/associations/preloader/ar5_association.rb
 - lib/jit_preloader/active_record/associations/preloader/ar6_association.rb
+- lib/jit_preloader/active_record/associations/preloader/ar7_association.rb
+- lib/jit_preloader/active_record/associations/preloader/ar7_branch.rb
 - lib/jit_preloader/active_record/associations/preloader/collection_association.rb
 - lib/jit_preloader/active_record/associations/preloader/singular_association.rb
 - lib/jit_preloader/active_record/associations/singular_association.rb