Eliminate uses of lease_connection (#1157)

Author: aidanharan

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

@@ -10,11 +10,13 @@ module AttributeMethods
           private
 
           def attributes_for_update(attribute_names)
-            return super unless self.class.lease_connection.adapter_name == "SQLServer"
+            self.class.with_connection do |connection|
+              return super(attribute_names) unless connection.sqlserver?
 
-            super.reject do |name|
-              column = self.class.columns_hash[name]
-              column && column.respond_to?(:is_identity?) && column.is_identity?
+              super(attribute_names).reject do |name|
+                column = self.class.columns_hash[name]
+                column && column.respond_to?(:is_identity?) && column.is_identity?
+              end
             end
           end
         end

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

@@ -9,10 +9,12 @@ module SQLServer
       module CoreExt
         module Calculations
           def calculate(operation, column_name)
-            if klass.lease_connection.sqlserver?
-              _calculate(operation, column_name)
-            else
-              super
+            klass.with_connection do |connection|
+              if connection.sqlserver?
+                _calculate(operation, column_name)
+              else
+                super
+              end
             end
           end
 
@@ -54,9 +56,10 @@ def _calculate(operation, column_name)
           end
 
           def build_count_subquery(relation, column_name, distinct)
-            return super unless klass.lease_connection.adapter_name == "SQLServer"
-
-            super(relation.unscope(:order), column_name, distinct)
+            klass.with_connection do |connection|
+              relation = relation.unscope(:order) if connection.sqlserver?
+              super(relation, column_name, distinct)
+            end
           end
         end
       end

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

@@ -9,12 +9,15 @@ module Explain
           SQLSERVER_STATEMENT_REGEXP = /N'(.+)', N'(.+)', (.+)/
 
           def exec_explain(queries, options = [])
-            return super unless lease_connection.adapter_name == "SQLServer"
+            with_connection do |connection|
+              return super(queries, options) unless connection.sqlserver?
 
-            unprepared_queries = queries.map do |(sql, binds)|
-              [unprepare_sqlserver_statement(sql, binds), binds]
+              unprepared_queries = queries.map do |(sql, binds)|
+                [unprepare_sqlserver_statement(sql, binds), binds]
+              end
+
+              super(unprepared_queries, options)
             end
-            super(unprepared_queries, options)
           end
 
           private

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

@@ -11,10 +11,12 @@ module FinderMethods
           private
 
           def construct_relation_for_exists(conditions)
-            if klass.lease_connection.sqlserver?
-              _construct_relation_for_exists(conditions)
-            else
-              super
+            klass.with_connection do |connection|
+              if connection.sqlserver?
+                _construct_relation_for_exists(conditions)
+              else
+                super
+              end
             end
           end
 

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

lib/active_record/connection_adapters/sqlserver/database_statements.rb

@@ -30,7 +30,7 @@ def raw_execute(sql, name, async: false, allow_retry: false, materialize_transac
           result
         end
 
-        def internal_exec_query(sql, name = "SQL", binds = [], prepare: false, async: false)
+        def internal_exec_query(sql, name = "SQL", binds = [], prepare: false, async: false, allow_retry: false)
           result = nil
           sql = transform_query(sql)
 

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

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