Merge pull request #63 from clio/run_preload_query_base_on_dive_limit

Update has_many_aggregate preload function to slice query into multiple ones base on dive limit

Author: wendy-clio

lib/jit_preloader.rb

@@ -26,6 +26,16 @@ def self.globally_enabled=(value)
     @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

lib/jit_preloader/active_record/base.rb

@@ -90,7 +90,7 @@ def self.prepended(base)
     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 @@ def has_many_aggregate(assoc, name, aggregate, field, table_alias_name: nil, def
           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,27 +122,29 @@ def has_many_aggregate(assoc, name, aggregate, field, table_alias_name: nil, def
             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?
               conditions[reflection.type] = self.class.name
             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 ||= {}

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 @@
         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

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, "*"