Eliminate missed lease_connection calls

Followup: rails#51353
Fix: rails#51722

Not too sure what happened, but I somehow missed quite a number
of `lease_connection` calls inside Active Record.

Author: byroot

activerecord/lib/active_record/associations/alias_tracker.rb

@@ -6,21 +6,23 @@ module ActiveRecord
   module Associations
     # Keeps track of table aliases for ActiveRecord::Associations::JoinDependency
     class AliasTracker # :nodoc:
-      def self.create(connection, initial_table, joins, aliases = nil)
-        if joins.empty?
-          aliases ||= Hash.new(0)
-        elsif aliases
-          default_proc = aliases.default_proc || proc { 0 }
-          aliases.default_proc = proc { |h, k|
-            h[k] = initial_count_for(connection, k, joins) + default_proc.call(h, k)
-          }
-        else
-          aliases = Hash.new { |h, k|
-            h[k] = initial_count_for(connection, k, joins)
-          }
+      def self.create(pool, initial_table, joins, aliases = nil)
+        pool.with_connection do |connection|
+          if joins.empty?
+            aliases ||= Hash.new(0)
+          elsif aliases
+            default_proc = aliases.default_proc || proc { 0 }
+            aliases.default_proc = proc { |h, k|
+              h[k] = initial_count_for(connection, k, joins) + default_proc.call(h, k)
+            }
+          else
+            aliases = Hash.new { |h, k|
+              h[k] = initial_count_for(connection, k, joins)
+            }
+          end
+          aliases[initial_table] = 1
+          new(connection.table_alias_length, aliases)
         end
-        aliases[initial_table] = 1
-        new(connection, aliases)
       end
 
       def self.initial_count_for(connection, name, table_joins)
@@ -46,9 +48,9 @@ def self.initial_count_for(connection, name, table_joins)
       end
 
       # table_joins is an array of arel joins which might conflict with the aliases we assign here
-      def initialize(connection, aliases)
-        @aliases    = aliases
-        @connection = connection
+      def initialize(table_alias_length, aliases)
+        @aliases = aliases
+        @table_alias_length = table_alias_length
       end
 
       def aliased_table_for(arel_table, table_name = nil)
@@ -60,7 +62,7 @@ def aliased_table_for(arel_table, table_name = nil)
           arel_table = arel_table.alias(table_name) if arel_table.name != table_name
         else
           # Otherwise, we need to use an alias
-          aliased_name = @connection.table_alias_for(yield)
+          aliased_name = table_alias_for(yield)
 
           # Update the count
           count = aliases[aliased_name] += 1
@@ -76,8 +78,12 @@ def aliased_table_for(arel_table, table_name = nil)
       attr_reader :aliases
 
       private
+        def table_alias_for(table_name)
+          table_name[0...@table_alias_length].tr(".", "_")
+        end
+
         def truncate(name)
-          name.slice(0, @connection.table_alias_length - 2)
+          name.slice(0, @table_alias_length - 2)
         end
     end
   end

activerecord/lib/active_record/attributes.rb

@@ -239,8 +239,10 @@ def define_attribute(
 
       def _default_attributes # :nodoc:
         @default_attributes ||= begin
-          attributes_hash = columns_hash.transform_values do |column|
-            ActiveModel::Attribute.from_database(column.name, column.default, type_for_column(column))
+          attributes_hash = with_connection do |connection|
+            columns_hash.transform_values do |column|
+              ActiveModel::Attribute.from_database(column.name, column.default, type_for_column(connection, column))
+            end
           end
 
           attribute_set = ActiveModel::AttributeSet.new(attributes_hash)
@@ -295,7 +297,7 @@ def resolve_type_name(name, **options)
           Type.lookup(name, **options, adapter: Type.adapter_name_from(self))
         end
 
-        def type_for_column(column)
+        def type_for_column(connection, column)
           hook_attribute_type(column.name, super)
         end
     end

activerecord/lib/active_record/connection_adapters/abstract/query_cache.rb

@@ -93,8 +93,8 @@ def initialize(...)
             end
         end
 
-        def lease_connection(**)
-          connection = super
+        def checkout_and_verify(connection)
+          super
           connection.query_cache ||= query_cache
           connection
         end
@@ -147,19 +147,19 @@ def clear_query_cache
           end
         end
 
+        def query_cache
+          @thread_query_caches.compute_if_absent(ActiveSupport::IsolatedExecutionState.context) do
+            Store.new(@query_cache_max_size)
+          end
+        end
+
         private
           def prune_thread_cache
             dead_threads = @thread_query_caches.keys.reject(&:alive?)
             dead_threads.each do |dead_thread|
               @thread_query_caches.delete(dead_thread)
             end
           end
-
-          def query_cache
-            @thread_query_caches.compute_if_absent(ActiveSupport::IsolatedExecutionState.context) do
-              Store.new(@query_cache_max_size)
-            end
-          end
       end
 
       attr_accessor :query_cache

activerecord/lib/active_record/core.rb

@@ -376,9 +376,9 @@ def type_caster # :nodoc:
         TypeCaster::Map.new(self)
       end
 
-      def cached_find_by_statement(key, &block) # :nodoc:
-        cache = @find_by_statement_cache[lease_connection.prepared_statements]
-        cache.compute_if_absent(key) { StatementCache.create(lease_connection, &block) }
+      def cached_find_by_statement(connection, key, &block) # :nodoc:
+        cache = @find_by_statement_cache[connection.prepared_statements]
+        cache.compute_if_absent(key) { StatementCache.create(connection, &block) }
       end
 
       private
@@ -419,21 +419,23 @@ def table_metadata
         end
 
         def cached_find_by(keys, values)
-          statement = cached_find_by_statement(keys) { |params|
-            wheres = keys.index_with do |key|
-              if key.is_a?(Array)
-                [key.map { params.bind }]
-              else
-                params.bind
+          with_connection do |connection|
+            statement = cached_find_by_statement(connection, keys) { |params|
+              wheres = keys.index_with do |key|
+                if key.is_a?(Array)
+                  [key.map { params.bind }]
+                else
+                  params.bind
+                end
               end
-            end
-            where(wheres).limit(1)
-          }
+              where(wheres).limit(1)
+            }
 
-          begin
-            statement.execute(values.flatten, lease_connection, allow_retry: true).first
-          rescue TypeError
-            raise ActiveRecord::StatementInvalid
+            begin
+              statement.execute(values.flatten, lease_connection, allow_retry: true).first
+            rescue TypeError
+              raise ActiveRecord::StatementInvalid
+            end
           end
         end
     end

activerecord/lib/active_record/model_schema.rb

@@ -379,7 +379,7 @@ def sequence_name
 
       def reset_sequence_name # :nodoc:
         @explicit_sequence_name = false
-        @sequence_name          = lease_connection.default_sequence_name(table_name, primary_key)
+        @sequence_name          = with_connection { |c| c.default_sequence_name(table_name, primary_key) }
       end
 
       # Sets the name of the sequence to use when generating ids to the given
@@ -404,13 +404,13 @@ def sequence_name=(value)
       # Determines if the primary key values should be selected from their
       # corresponding sequence before the insert statement.
       def prefetch_primary_key?
-        lease_connection.prefetch_primary_key?(table_name)
+        with_connection { |c| c.prefetch_primary_key?(table_name) }
       end
 
       # Returns the next value that will be used as the primary key on
       # an insert statement.
       def next_sequence_value
-        lease_connection.next_sequence_value(sequence_name)
+        with_connection { |c| c.next_sequence_value(sequence_name) }
       end
 
       # Indicates whether the table associated with this class exists
@@ -435,10 +435,10 @@ def columns
         @columns ||= columns_hash.values.freeze
       end
 
-      def _returning_columns_for_insert # :nodoc:
+      def _returning_columns_for_insert(connection) # :nodoc:
         @_returning_columns_for_insert ||= begin
           auto_populated_columns = columns.filter_map do |c|
-            c.name if lease_connection.return_value_after_insert?(c)
+            c.name if connection.return_value_after_insert?(c)
           end
 
           auto_populated_columns.empty? ? Array(primary_key) : auto_populated_columns
@@ -523,7 +523,7 @@ def content_columns
       #    end
       #  end
       def reset_column_information
-        lease_connection.clear_cache!
+        connection_pool.active_connection&.clear_cache!
         ([self] + descendants).each(&:undefine_attribute_methods)
         schema_cache.clear_data_source_cache!(table_name)
 
@@ -617,8 +617,8 @@ def compute_table_name
           end
         end
 
-        def type_for_column(column)
-          type = lease_connection.lookup_cast_type_from_column(column)
+        def type_for_column(connection, column)
+          type = connection.lookup_cast_type_from_column(column)
 
           if immutable_strings_by_default && type.respond_to?(:to_immutable_string)
             type = type.to_immutable_string

activerecord/lib/active_record/persistence.rb

@@ -568,7 +568,7 @@ def delete(id_or_array)
         delete_by(primary_key => id_or_array)
       end
 
-      def _insert_record(values, returning) # :nodoc:
+      def _insert_record(connection, values, returning) # :nodoc:
         primary_key = self.primary_key
         primary_key_value = nil
 
@@ -583,12 +583,12 @@ def _insert_record(values, returning) # :nodoc:
 
         with_connection do |c|
           if values.empty?
-            im.insert(c.empty_insert_statement_value(primary_key))
+            im.insert(connection.empty_insert_statement_value(primary_key))
           else
             im.insert(values.transform_keys { |name| arel_table[name] })
           end
 
-          c.insert(
+          connection.insert(
             im, "#{self} Create", primary_key || false, primary_key_value,
             returning: returning
           )
@@ -1255,16 +1255,19 @@ def _update_record(attribute_names = self.attribute_names)
     def _create_record(attribute_names = self.attribute_names)
       attribute_names = attributes_for_create(attribute_names)
 
-      returning_columns = self.class._returning_columns_for_insert
+      self.class.with_connection do |connection|
+        returning_columns = self.class._returning_columns_for_insert(connection)
 
-      returning_values = self.class._insert_record(
-        attributes_with_values(attribute_names),
-        returning_columns
-      )
+        returning_values = self.class._insert_record(
+          connection,
+          attributes_with_values(attribute_names),
+          returning_columns
+        )
 
-      returning_columns.zip(returning_values).each do |column, value|
-        _write_attribute(column, value) if !_read_attribute(column)
-      end if returning_values
+        returning_columns.zip(returning_values).each do |column, value|
+          _write_attribute(column, value) if !_read_attribute(column)
+        end if returning_values
+      end
 
       @new_record = false
       @previously_new_record = true

activerecord/lib/active_record/querying.rb

@@ -48,18 +48,25 @@ module Querying
     # Note that building your own SQL query string from user input may expose your application to
     # injection attacks (https://guides.rubyonrails.org/security.html#sql-injection).
     def find_by_sql(sql, binds = [], preparable: nil, allow_retry: false, &block)
-      _load_from_sql(_query_by_sql(sql, binds, preparable: preparable, allow_retry: allow_retry), &block)
+      result = with_connection do |c|
+        _query_by_sql(c, sql, binds, preparable: preparable, allow_retry: allow_retry)
+      end
+      _load_from_sql(result, &block)
     end
 
     # Same as <tt>#find_by_sql</tt> but perform the query asynchronously and returns an ActiveRecord::Promise.
     def async_find_by_sql(sql, binds = [], preparable: nil, &block)
-      _query_by_sql(sql, binds, preparable: preparable, async: true).then do |result|
+      result = with_connection do |c|
+        _query_by_sql(c, sql, binds, preparable: preparable, async: true)
+      end
+
+      result.then do |result|
         _load_from_sql(result, &block)
       end
     end
 
-    def _query_by_sql(sql, binds = [], preparable: nil, async: false, allow_retry: false) # :nodoc:
-      lease_connection.select_all(sanitize_sql(sql), "#{name} Load", binds, preparable: preparable, async: async, allow_retry: allow_retry)
+    def _query_by_sql(connection, sql, binds = [], preparable: nil, async: false, allow_retry: false) # :nodoc:
+      connection.select_all(sanitize_sql(sql), "#{name} Load", binds, preparable: preparable, async: async, allow_retry: allow_retry)
     end
 
     def _load_from_sql(result_set, &block) # :nodoc:
@@ -99,12 +106,16 @@ def _load_from_sql(result_set, &block) # :nodoc:
     #
     # * +sql+ - An SQL statement which should return a count query from the database, see the example above.
     def count_by_sql(sql)
-      lease_connection.select_value(sanitize_sql(sql), "#{name} Count").to_i
+      with_connection do |c|
+        c.select_value(sanitize_sql(sql), "#{name} Count").to_i
+      end
     end
 
     # Same as <tt>#count_by_sql</tt> but perform the query asynchronously and returns an ActiveRecord::Promise.
     def async_count_by_sql(sql)
-      lease_connection.select_value(sanitize_sql(sql), "#{name} Count", async: true).then(&:to_i)
+      with_connection do |c|
+        c.select_value(sanitize_sql(sql), "#{name} Count", async: true).then(&:to_i)
+      end
     end
   end
 end

activerecord/lib/active_record/reflection.rb

@@ -529,7 +529,9 @@ def association_scope_cache(klass, owner, &block)
         if polymorphic?
           key = [key, owner._read_attribute(@foreign_type)]
         end
-        klass.cached_find_by_statement(key, &block)
+        klass.with_connection do |connection|
+          klass.cached_find_by_statement(connection, key, &block)
+        end
       end
 
       def join_table