Merge pull request #986 from yellowspot/reimplement-preloader-monkey-patch

[Rails 7] Update preloader monkey patch

Author: aidanharan

lib/active_record/connection_adapters/sqlserver/core_ext/preloader.rb

@@ -6,20 +6,12 @@ module ActiveRecord
   module ConnectionAdapters
     module SQLServer
       module CoreExt
-        module Preloader
-          private
+        module LoaderQuery
+          def load_records_for_keys(keys, &block)
+            return super unless scope.connection.adapter_name == "SQLServer"
 
-          def records_for(ids)
-            return super unless klass.connection.adapter_name == "SQLServer"
-
-            ids.each_slice(in_clause_length).flat_map do |slice|
-              scope.where(association_key_name => slice).load do |record|
-                # Processing only the first owner
-                # because the record is modified but not an owner
-                owner = owners_by_key[convert_key(record[association_key_name])].first
-                association = owner.association(reflection.name)
-                association.set_inverse_instance(record)
-              end.records
+            keys.each_slice(in_clause_length).flat_map do |slice|
+              scope.where(association_key_name => slice).load(&block).records
             end
           end
 
@@ -33,6 +25,6 @@ def in_clause_length
 end
 
 ActiveSupport.on_load(:active_record) do
-  mod = ActiveRecord::ConnectionAdapters::SQLServer::CoreExt::Preloader
-  ActiveRecord::Associations::Preloader::Association.prepend(mod)
+  mod = ActiveRecord::ConnectionAdapters::SQLServer::CoreExt::LoaderQuery
+  ActiveRecord::Associations::Preloader::Association::LoaderQuery.prepend(mod)
 end

test/cases/eager_load_too_many_ids_test_sqlserver.rb

@@ -0,0 +1,18 @@
+require "cases/helper_sqlserver"
+require "models/citation"
+require "models/book"
+
+class EagerLoadingTooManyIdsTest < ActiveRecord::TestCase
+  fixtures :citations
+
+  def test_batch_preloading_too_many_ids
+    in_clause_length = 10_000
+
+    # We Monkey patch Preloader to work with batches of 10_000 records.
+    # Expect: N Books queries + Citation query
+    expected_query_count = (Citation.count / in_clause_length.to_f).ceil + 1
+    assert_queries(expected_query_count) do
+      Citation.preload(:reference_of).to_a.size
+    end
+  end
+end