Lease connections and renamed assert_queries to assert_queries_count

See rails/rails#50373

Author: aidanharan

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

@@ -8,23 +8,25 @@ module SQLServer
       module CoreExt
         module LoaderQuery
           def load_records_for_keys(keys, &block)
-            return super unless scope.connection.sqlserver?
+            scope.with_connection do |connection|
+              return super unless connection.sqlserver?
 
-            return [] if keys.empty?
+              return [] if keys.empty?
 
-            if association_key_name.is_a?(Array)
-              query_constraints = Hash.new { |hsh, key| hsh[key] = Set.new }
+              if association_key_name.is_a?(Array)
+                query_constraints = Hash.new { |hsh, key| hsh[key] = Set.new }
 
-              keys.each_with_object(query_constraints) do |values_set, constraints|
-                association_key_name.zip(values_set).each do |key_name, value|
-                  constraints[key_name] << value
+                keys.each_with_object(query_constraints) do |values_set, constraints|
+                  association_key_name.zip(values_set).each do |key_name, value|
+                    constraints[key_name] << value
+                  end
                 end
-              end
 
-              scope.where(query_constraints).load(&block)
-            else
-              keys.each_slice(in_clause_length).flat_map do |slice|
-                scope.where(association_key_name => slice).load(&block).records
+                scope.where(query_constraints).load(&block)
+              else
+                keys.each_slice(in_clause_length).flat_map do |slice|
+                  scope.where(association_key_name => slice).load(&block).records
+                end
               end
             end
           end

test/cases/adapter_test_sqlserver.rb

@@ -61,8 +61,8 @@ class AdapterTestSQLServer < ActiveRecord::TestCase
   end
 
   it "test table existence across database schemas" do
-    arunit_connection = Topic.connection
-    arunit2_connection = College.connection
+    arunit_connection = Topic.lease_connection
+    arunit2_connection = College.lease_connection
 
     arunit_database = arunit_connection.pool.db_config.database
     arunit2_database = arunit2_connection.pool.db_config.database

test/cases/coerced_tests.rb

@@ -56,7 +56,7 @@ def test_partial_index_coerced
 
     t = Topic.create!(title: "abc")
     t.author_name = "John"
-    assert_queries(1) do
+    assert_queries_count(1) do
       t.valid?
     end
   end
@@ -234,7 +234,7 @@ def test_update_date_time_attributes_with_default_timezone_local
   coerce_tests! %r{an empty transaction does not raise if preventing writes}
   test "an empty transaction does not raise if preventing writes coerced" do
     ActiveRecord::Base.while_preventing_writes do
-      assert_queries(1, ignore_none: true) do
+      assert_queries_count(1, ignore_none: true) do
         Bird.transaction do
           ActiveRecord::Base.lease_connection.materialize_transactions
         end
@@ -490,8 +490,8 @@ def test_limit_with_offset_is_kept_coerced
   coerce_tests! :test_distinct_count_all_with_custom_select_and_order
   def test_distinct_count_all_with_custom_select_and_order_coerced
     accounts = Account.distinct.select("credit_limit % 10 AS the_limit").order(Arel.sql("credit_limit % 10"))
-    assert_queries(1) { assert_equal 3, accounts.count(:all) }
-    assert_queries(1) { assert_equal 3, accounts.load.size }
+    assert_queries_count(1) { assert_equal 3, accounts.count(:all) }
+    assert_queries_count(1) { assert_equal 3, accounts.load.size }
   end
 
   # Leave it up to users to format selects/functions so HAVING works correctly.
@@ -1018,7 +1018,7 @@ def test_implicit_order_column_is_configurable_coerced
     assert_equal topics(:fifth), Topic.first
     assert_equal topics(:third), Topic.last
 
-    c = Topic.connection
+    c = Topic.lease_connection
     assert_sql(/ORDER BY #{Regexp.escape(c.quote_table_name("topics.title"))} DESC, #{Regexp.escape(c.quote_table_name("topics.id"))} DESC OFFSET 0 ROWS FETCH NEXT @0 ROWS ONLY.*@0 = 1/i) {
       Topic.last
     }
@@ -1032,7 +1032,7 @@ def test_implicit_order_set_to_primary_key_coerced
     old_implicit_order_column = Topic.implicit_order_column
     Topic.implicit_order_column = "id"
 
-    c = Topic.connection
+    c = Topic.lease_connection
     assert_sql(/ORDER BY #{Regexp.escape(c.quote_table_name("topics.id"))} DESC OFFSET 0 ROWS FETCH NEXT @0 ROWS ONLY.*@0 = 1/i) {
       Topic.last
     }
@@ -1046,7 +1046,7 @@ def test_implicit_order_for_model_without_primary_key_coerced
     old_implicit_order_column = NonPrimaryKey.implicit_order_column
     NonPrimaryKey.implicit_order_column = "created_at"
 
-    c = NonPrimaryKey.connection
+    c = NonPrimaryKey.lease_connection
 
     assert_sql(/ORDER BY #{Regexp.escape(c.quote_table_name("non_primary_keys.created_at"))} DESC OFFSET 0 ROWS FETCH NEXT @0 ROWS ONLY.*@0 = 1/i) {
       NonPrimaryKey.last
@@ -1067,7 +1067,7 @@ def test_member_on_unloaded_relation_with_composite_primary_key_coerced
   # Check for `FETCH NEXT x ROWS` rather then `LIMIT`.
   coerce_tests! :test_implicit_order_column_prepends_query_constraints
   def test_implicit_order_column_prepends_query_constraints_coerced
-    c = ClothingItem.connection
+    c = ClothingItem.lease_connection
     ClothingItem.implicit_order_column = "description"
     quoted_type = Regexp.escape(c.quote_table_name("clothing_items.clothing_type"))
     quoted_color = Regexp.escape(c.quote_table_name("clothing_items.color"))
@@ -1083,7 +1083,7 @@ def test_implicit_order_column_prepends_query_constraints_coerced
   # Check for `FETCH NEXT x ROWS` rather then `LIMIT`.
   coerce_tests! %r{#last for a model with composite query constraints}
   test "#last for a model with composite query constraints coerced" do
-    c = ClothingItem.connection
+    c = ClothingItem.lease_connection
     quoted_type = Regexp.escape(c.quote_table_name("clothing_items.clothing_type"))
     quoted_color = Regexp.escape(c.quote_table_name("clothing_items.color"))
 
@@ -1095,7 +1095,7 @@ def test_implicit_order_column_prepends_query_constraints_coerced
   # Check for `FETCH NEXT x ROWS` rather then `LIMIT`.
   coerce_tests! %r{#first for a model with composite query constraints}
   test "#first for a model with composite query constraints coerced" do
-    c = ClothingItem.connection
+    c = ClothingItem.lease_connection
     quoted_type = Regexp.escape(c.quote_table_name("clothing_items.clothing_type"))
     quoted_color = Regexp.escape(c.quote_table_name("clothing_items.color"))
 
@@ -1107,7 +1107,7 @@ def test_implicit_order_column_prepends_query_constraints_coerced
   # Check for `FETCH NEXT x ROWS` rather then `LIMIT`.
   coerce_tests! :test_implicit_order_column_reorders_query_constraints
   def test_implicit_order_column_reorders_query_constraints_coerced
-    c = ClothingItem.connection
+    c = ClothingItem.lease_connection
     ClothingItem.implicit_order_column = "color"
     quoted_type = Regexp.escape(c.quote_table_name("clothing_items.clothing_type"))
     quoted_color = Regexp.escape(c.quote_table_name("clothing_items.color"))
@@ -1131,7 +1131,7 @@ def test_include_on_unloaded_relation_with_composite_primary_key_coerced
   # Check for `FETCH NEXT x ROWS` rather then `LIMIT`.
   coerce_tests! :test_nth_to_last_with_order_uses_limit
   def test_nth_to_last_with_order_uses_limit_coerced
-    c = Topic.connection
+    c = Topic.lease_connection
     assert_sql(/ORDER BY #{Regexp.escape(c.quote_table_name("topics.id"))} DESC OFFSET @(\d) ROWS FETCH NEXT @(\d) ROWS ONLY.*@\1 = 1.*@\2 = 1/i) do
       Topic.second_to_last
     end
@@ -1324,7 +1324,7 @@ def test_create_without_primary_key_no_extra_query_coerced
       self.table_name = "dashboards"
     end
     klass.create! # warmup schema cache
-    assert_queries(2, ignore_none: true) { klass.create! }
+    assert_queries_count(2, ignore_none: true) { klass.create! }
   end
 end
 
@@ -1354,7 +1354,7 @@ def test_query_cached_even_when_types_are_reset_coerced
       Task.initialize_find_by_cache
       Task.define_attribute_methods
 
-      assert_queries(1, ignore_none: true) do
+      assert_queries_count(1, ignore_none: true) do
         Task.find(1)
       end
 
@@ -1454,11 +1454,11 @@ def test_relations_dont_load_all_records_in_pretty_print_coerced
   def test_empty_complex_chained_relations_coerced
     posts = Post.select("comments_count").where("id is not null").group("author_id", "id").where("legacy_comments_count > 0")
 
-    assert_queries(1) { assert_equal false, posts.empty? }
+    assert_queries_count(1) { assert_equal false, posts.empty? }
     assert_not_predicate posts, :loaded?
 
     no_posts = posts.where(title: "")
-    assert_queries(1) { assert_equal true, no_posts.empty? }
+    assert_queries_count(1) { assert_equal true, no_posts.empty? }
     assert_not_predicate no_posts, :loaded?
   end
 
@@ -2338,7 +2338,7 @@ def test_preloads_has_many_on_model_with_a_composite_primary_key_through_id_attr
     assert_equal 2, sql.size
     preload_sql = sql.last
 
-    c = Cpk::OrderAgreement.connection
+    c = Cpk::OrderAgreement.lease_connection
     order_id_column = Regexp.escape(c.quote_table_name("cpk_order_agreements.order_id"))
     order_id_constraint = /#{order_id_column} = @0.*@0 = \d+$/
     expectation = /SELECT.*WHERE.* #{order_id_constraint}/
@@ -2362,7 +2362,7 @@ def test_preloads_belongs_to_a_composite_primary_key_model_through_id_attribute_
     assert_equal 2, sql.size
     preload_sql = sql.last
 
-    c = Cpk::Order.connection
+    c = Cpk::Order.lease_connection
     order_id = Regexp.escape(c.quote_table_name("cpk_orders.id"))
     order_constraint = /#{order_id} = @0.*@0 = \d+$/
     expectation = /SELECT.*WHERE.* #{order_constraint}/
@@ -2377,7 +2377,7 @@ class BasePreventWritesTest < ActiveRecord::TestCase
   coerce_tests! %r{an empty transaction does not raise if preventing writes}
   test "an empty transaction does not raise if preventing writes coerced" do
     ActiveRecord::Base.while_preventing_writes do
-      assert_queries(1, ignore_none: true) do
+      assert_queries_count(1, ignore_none: true) do
         Bird.transaction do
           ActiveRecord::Base.lease_connection.materialize_transactions
         end

test/cases/eager_load_too_many_ids_test_sqlserver.rb

@@ -11,7 +11,7 @@ def test_batch_preloading_too_many_ids
     # 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
+    assert_queries_count(expected_query_count) do
       Citation.preload(:reference_of).to_a.size
     end
   end