activerecord 7.1.3.4 → 7.2.0.beta1

data/lib/active_record/normalization.rb

@@ -86,10 +86,8 @@ module ActiveRecord # :nodoc:
       #
       #   User.normalize_value_for(:phone, "+1 (555) 867-5309") # => "5558675309"
       def normalizes(*names, with:, apply_to_nil: false)
-        names.each do |name|
-          attribute(name) do |cast_type|
-            NormalizedValueType.new(cast_type: cast_type, normalizer: with, normalize_nil: apply_to_nil)
-          end
+        decorate_attributes(names) do |name, cast_type|
+          NormalizedValueType.new(cast_type: cast_type, normalizer: with, normalize_nil: apply_to_nil)
         end
 
         self.normalized_attributes += names.map(&:to_sym)
@@ -154,9 +152,7 @@ module ActiveRecord # :nodoc:
           [self.class, cast_type, normalizer, normalize_nil?].hash
         end
 
-        def inspect
-          Kernel.instance_method(:inspect).bind_call(self)
-        end
+        define_method(:inspect, Kernel.instance_method(:inspect))
 
         private
           def normalize(value)

data/lib/active_record/persistence.rb

@@ -87,282 +87,6 @@ module ActiveRecord
         end
       end
 
-      # Inserts a single record into the database in a single SQL INSERT
-      # statement. It does not instantiate any models nor does it trigger
-      # Active Record callbacks or validations. Though passed values
-      # go through Active Record's type casting and serialization.
-      #
-      # See #insert_all for documentation.
-      def insert(attributes, returning: nil, unique_by: nil, record_timestamps: nil)
-        insert_all([ attributes ], returning: returning, unique_by: unique_by, record_timestamps: record_timestamps)
-      end
-
-      # Inserts multiple records into the database in a single SQL INSERT
-      # statement. It does not instantiate any models nor does it trigger
-      # Active Record callbacks or validations. Though passed values
-      # go through Active Record's type casting and serialization.
-      #
-      # The +attributes+ parameter is an Array of Hashes. Every Hash determines
-      # the attributes for a single row and must have the same keys.
-      #
-      # Rows are considered to be unique by every unique index on the table. Any
-      # duplicate rows are skipped.
-      # Override with <tt>:unique_by</tt> (see below).
-      #
-      # Returns an ActiveRecord::Result with its contents based on
-      # <tt>:returning</tt> (see below).
-      #
-      # ==== Options
-      #
-      # [:returning]
-      #   (PostgreSQL and SQLite3 only) An array of attributes to return for all successfully
-      #   inserted records, which by default is the primary key.
-      #   Pass <tt>returning: %w[ id name ]</tt> for both id and name
-      #   or <tt>returning: false</tt> to omit the underlying <tt>RETURNING</tt> SQL
-      #   clause entirely.
-      #
-      #   You can also pass an SQL string if you need more control on the return values
-      #   (for example, <tt>returning: Arel.sql("id, name as new_name")</tt>).
-      #
-      # [:unique_by]
-      #   (PostgreSQL and SQLite only) By default rows are considered to be unique
-      #   by every unique index on the table. Any duplicate rows are skipped.
-      #
-      #   To skip rows according to just one unique index pass <tt>:unique_by</tt>.
-      #
-      #   Consider a Book model where no duplicate ISBNs make sense, but if any
-      #   row has an existing id, or is not unique by another unique index,
-      #   ActiveRecord::RecordNotUnique is raised.
-      #
-      #   Unique indexes can be identified by columns or name:
-      #
-      #     unique_by: :isbn
-      #     unique_by: %i[ author_id name ]
-      #     unique_by: :index_books_on_isbn
-      #
-      # [:record_timestamps]
-      #   By default, automatic setting of timestamp columns is controlled by
-      #   the model's <tt>record_timestamps</tt> config, matching typical
-      #   behavior.
-      #
-      #   To override this and force automatic setting of timestamp columns one
-      #   way or the other, pass <tt>:record_timestamps</tt>:
-      #
-      #     record_timestamps: true  # Always set timestamps automatically
-      #     record_timestamps: false # Never set timestamps automatically
-      #
-      # Because it relies on the index information from the database
-      # <tt>:unique_by</tt> is recommended to be paired with
-      # Active Record's schema_cache.
-      #
-      # ==== Example
-      #
-      #   # Insert records and skip inserting any duplicates.
-      #   # Here "Eloquent Ruby" is skipped because its id is not unique.
-      #
-      #   Book.insert_all([
-      #     { id: 1, title: "Rework", author: "David" },
-      #     { id: 1, title: "Eloquent Ruby", author: "Russ" }
-      #   ])
-      #
-      #   # insert_all works on chained scopes, and you can use create_with
-      #   # to set default attributes for all inserted records.
-      #
-      #   author.books.create_with(created_at: Time.now).insert_all([
-      #     { id: 1, title: "Rework" },
-      #     { id: 2, title: "Eloquent Ruby" }
-      #   ])
-      def insert_all(attributes, returning: nil, unique_by: nil, record_timestamps: nil)
-        InsertAll.new(self, attributes, on_duplicate: :skip, returning: returning, unique_by: unique_by, record_timestamps: record_timestamps).execute
-      end
-
-      # Inserts a single record into the database in a single SQL INSERT
-      # statement. It does not instantiate any models nor does it trigger
-      # Active Record callbacks or validations. Though passed values
-      # go through Active Record's type casting and serialization.
-      #
-      # See #insert_all! for more.
-      def insert!(attributes, returning: nil, record_timestamps: nil)
-        insert_all!([ attributes ], returning: returning, record_timestamps: record_timestamps)
-      end
-
-      # Inserts multiple records into the database in a single SQL INSERT
-      # statement. It does not instantiate any models nor does it trigger
-      # Active Record callbacks or validations. Though passed values
-      # go through Active Record's type casting and serialization.
-      #
-      # The +attributes+ parameter is an Array of Hashes. Every Hash determines
-      # the attributes for a single row and must have the same keys.
-      #
-      # Raises ActiveRecord::RecordNotUnique if any rows violate a
-      # unique index on the table. In that case, no rows are inserted.
-      #
-      # To skip duplicate rows, see #insert_all. To replace them, see #upsert_all.
-      #
-      # Returns an ActiveRecord::Result with its contents based on
-      # <tt>:returning</tt> (see below).
-      #
-      # ==== Options
-      #
-      # [:returning]
-      #   (PostgreSQL and SQLite3 only) An array of attributes to return for all successfully
-      #   inserted records, which by default is the primary key.
-      #   Pass <tt>returning: %w[ id name ]</tt> for both id and name
-      #   or <tt>returning: false</tt> to omit the underlying <tt>RETURNING</tt> SQL
-      #   clause entirely.
-      #
-      #   You can also pass an SQL string if you need more control on the return values
-      #   (for example, <tt>returning: Arel.sql("id, name as new_name")</tt>).
-      #
-      # [:record_timestamps]
-      #   By default, automatic setting of timestamp columns is controlled by
-      #   the model's <tt>record_timestamps</tt> config, matching typical
-      #   behavior.
-      #
-      #   To override this and force automatic setting of timestamp columns one
-      #   way or the other, pass <tt>:record_timestamps</tt>:
-      #
-      #     record_timestamps: true  # Always set timestamps automatically
-      #     record_timestamps: false # Never set timestamps automatically
-      #
-      # ==== Examples
-      #
-      #   # Insert multiple records
-      #   Book.insert_all!([
-      #     { title: "Rework", author: "David" },
-      #     { title: "Eloquent Ruby", author: "Russ" }
-      #   ])
-      #
-      #   # Raises ActiveRecord::RecordNotUnique because "Eloquent Ruby"
-      #   # does not have a unique id.
-      #   Book.insert_all!([
-      #     { id: 1, title: "Rework", author: "David" },
-      #     { id: 1, title: "Eloquent Ruby", author: "Russ" }
-      #   ])
-      def insert_all!(attributes, returning: nil, record_timestamps: nil)
-        InsertAll.new(self, attributes, on_duplicate: :raise, returning: returning, record_timestamps: record_timestamps).execute
-      end
-
-      # Updates or inserts (upserts) a single record into the database in a
-      # single SQL INSERT statement. It does not instantiate any models nor does
-      # it trigger Active Record callbacks or validations. Though passed values
-      # go through Active Record's type casting and serialization.
-      #
-      # See #upsert_all for documentation.
-      def upsert(attributes, **kwargs)
-        upsert_all([ attributes ], **kwargs)
-      end
-
-      # Updates or inserts (upserts) multiple records into the database in a
-      # single SQL INSERT statement. It does not instantiate any models nor does
-      # it trigger Active Record callbacks or validations. Though passed values
-      # go through Active Record's type casting and serialization.
-      #
-      # The +attributes+ parameter is an Array of Hashes. Every Hash determines
-      # the attributes for a single row and must have the same keys.
-      #
-      # Returns an ActiveRecord::Result with its contents based on
-      # <tt>:returning</tt> (see below).
-      #
-      # By default, +upsert_all+ will update all the columns that can be updated when
-      # there is a conflict. These are all the columns except primary keys, read-only
-      # columns, and columns covered by the optional +unique_by+.
-      #
-      # ==== Options
-      #
-      # [:returning]
-      #   (PostgreSQL and SQLite3 only) An array of attributes to return for all successfully
-      #   inserted records, which by default is the primary key.
-      #   Pass <tt>returning: %w[ id name ]</tt> for both id and name
-      #   or <tt>returning: false</tt> to omit the underlying <tt>RETURNING</tt> SQL
-      #   clause entirely.
-      #
-      #   You can also pass an SQL string if you need more control on the return values
-      #   (for example, <tt>returning: Arel.sql("id, name as new_name")</tt>).
-      #
-      # [:unique_by]
-      #   (PostgreSQL and SQLite only) By default rows are considered to be unique
-      #   by every unique index on the table. Any duplicate rows are skipped.
-      #
-      #   To skip rows according to just one unique index pass <tt>:unique_by</tt>.
-      #
-      #   Consider a Book model where no duplicate ISBNs make sense, but if any
-      #   row has an existing id, or is not unique by another unique index,
-      #   ActiveRecord::RecordNotUnique is raised.
-      #
-      #   Unique indexes can be identified by columns or name:
-      #
-      #     unique_by: :isbn
-      #     unique_by: %i[ author_id name ]
-      #     unique_by: :index_books_on_isbn
-      #
-      # Because it relies on the index information from the database
-      # <tt>:unique_by</tt> is recommended to be paired with
-      # Active Record's schema_cache.
-      #
-      # [:on_duplicate]
-      #   Configure the SQL update sentence that will be used in case of conflict.
-      #
-      #   NOTE: If you use this option you must provide all the columns you want to update
-      #   by yourself.
-      #
-      #   Example:
-      #
-      #     Commodity.upsert_all(
-      #       [
-      #         { id: 2, name: "Copper", price: 4.84 },
-      #         { id: 4, name: "Gold", price: 1380.87 },
-      #         { id: 6, name: "Aluminium", price: 0.35 }
-      #       ],
-      #       on_duplicate: Arel.sql("price = GREATEST(commodities.price, EXCLUDED.price)")
-      #     )
-      #
-      #   See the related +:update_only+ option. Both options can't be used at the same time.
-      #
-      # [:update_only]
-      #   Provide a list of column names that will be updated in case of conflict. If not provided,
-      #   +upsert_all+ will update all the columns that can be updated. These are all the columns
-      #   except primary keys, read-only columns, and columns covered by the optional +unique_by+
-      #
-      #   Example:
-      #
-      #     Commodity.upsert_all(
-      #       [
-      #         { id: 2, name: "Copper", price: 4.84 },
-      #         { id: 4, name: "Gold", price: 1380.87 },
-      #         { id: 6, name: "Aluminium", price: 0.35 }
-      #       ],
-      #       update_only: [:price] # Only prices will be updated
-      #     )
-      #
-      #   See the related +:on_duplicate+ option. Both options can't be used at the same time.
-      #
-      # [:record_timestamps]
-      #   By default, automatic setting of timestamp columns is controlled by
-      #   the model's <tt>record_timestamps</tt> config, matching typical
-      #   behavior.
-      #
-      #   To override this and force automatic setting of timestamp columns one
-      #   way or the other, pass <tt>:record_timestamps</tt>:
-      #
-      #     record_timestamps: true  # Always set timestamps automatically
-      #     record_timestamps: false # Never set timestamps automatically
-      #
-      # ==== Examples
-      #
-      #   # Inserts multiple records, performing an upsert when records have duplicate ISBNs.
-      #   # Here "Eloquent Ruby" overwrites "Rework" because its ISBN is duplicate.
-      #
-      #   Book.upsert_all([
-      #     { title: "Rework", author: "David", isbn: "1" },
-      #     { title: "Eloquent Ruby", author: "Russ", isbn: "1" }
-      #   ], unique_by: :isbn)
-      #
-      #   Book.find_by(isbn: "1").title # => "Eloquent Ruby"
-      def upsert_all(attributes, on_duplicate: :update, update_only: nil, returning: nil, unique_by: nil, record_timestamps: nil)
-        InsertAll.new(self, attributes, on_duplicate: on_duplicate, update_only: update_only, returning: returning, unique_by: unique_by, record_timestamps: record_timestamps).execute
-      end
-
       # Given an attributes hash, +instantiate+ returns a new instance of
       # the appropriate class. Accepts only keys as strings.
       #
@@ -456,7 +180,7 @@ module ActiveRecord
       end
 
       # Accepts a list of attribute names to be used in the WHERE clause
-      # of SELECT / UPDATE / DELETE queries and in the ORDER BY clause for `#first` and `#last` finder methods.
+      # of SELECT / UPDATE / DELETE queries and in the ORDER BY clause for +#first+ and +#last+ finder methods.
       #
       #   class Developer < ActiveRecord::Base
       #     query_constraints :company_id, :id
@@ -469,7 +193,7 @@ module ActiveRecord
       #   developer.update!(name: "Nikita")
       #   # UPDATE "developers" SET "name" = 'Nikita' WHERE "developers"."company_id" = 1 AND "developers"."id" = 1
       #
-      #   It is possible to update attribute used in the query_by clause:
+      #   # It is possible to update an attribute used in the query_constraints clause:
       #   developer.update!(company_id: 2)
       #   # UPDATE "developers" SET "company_id" = 2 WHERE "developers"."company_id" = 1 AND "developers"."id" = 1
       #
@@ -511,62 +235,7 @@ module ActiveRecord
         @composite_query_constraints_list ||= query_constraints_list || Array(primary_key)
       end
 
-      # Destroy an object (or multiple objects) that has the given id. The object is instantiated first,
-      # therefore all callbacks and filters are fired off before the object is deleted. This method is
-      # less efficient than #delete but allows cleanup methods and other actions to be run.
-      #
-      # This essentially finds the object (or multiple objects) with the given id, creates a new object
-      # from the attributes, and then calls destroy on it.
-      #
-      # ==== Parameters
-      #
-      # * +id+ - This should be the id or an array of ids to be destroyed.
-      #
-      # ==== Examples
-      #
-      #   # Destroy a single object
-      #   Todo.destroy(1)
-      #
-      #   # Destroy multiple objects
-      #   todos = [1,2,3]
-      #   Todo.destroy(todos)
-      def destroy(id)
-        multiple_ids = if composite_primary_key?
-          id.first.is_a?(Array)
-        else
-          id.is_a?(Array)
-        end
-
-        if multiple_ids
-          find(id).each(&:destroy)
-        else
-          find(id).destroy
-        end
-      end
-
-      # Deletes the row with a primary key matching the +id+ argument, using an
-      # SQL +DELETE+ statement, and returns the number of rows deleted. Active
-      # Record objects are not instantiated, so the object's callbacks are not
-      # executed, including any <tt>:dependent</tt> association options.
-      #
-      # You can delete multiple rows at once by passing an Array of <tt>id</tt>s.
-      #
-      # Note: Although it is often much faster than the alternative, #destroy,
-      # skipping callbacks might bypass business logic in your application
-      # that ensures referential integrity or performs other essential jobs.
-      #
-      # ==== Examples
-      #
-      #   # Delete a single row
-      #   Todo.delete(1)
-      #
-      #   # Delete multiple rows
-      #   Todo.delete([2,3,4])
-      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
 
@@ -579,16 +248,18 @@ module ActiveRecord
 
         im = Arel::InsertManager.new(arel_table)
 
-        if values.empty?
-          im.insert(connection.empty_insert_statement_value(primary_key))
-        else
-          im.insert(values.transform_keys { |name| arel_table[name] })
-        end
+        with_connection do |c|
+          if values.empty?
+            im.insert(connection.empty_insert_statement_value(primary_key))
+          else
+            im.insert(values.transform_keys { |name| arel_table[name] })
+          end
 
-        connection.insert(
-          im, "#{self} Create", primary_key || false, primary_key_value,
-          returning: returning
-        )
+          connection.insert(
+            im, "#{self} Create", primary_key || false, primary_key_value,
+            returning: returning
+          )
+        end
       end
 
       def _update_record(values, constraints) # :nodoc:
@@ -605,7 +276,9 @@ module ActiveRecord
         um.set(values.transform_keys { |name| arel_table[name] })
         um.wheres = constraints
 
-        connection.update(um, "#{self} Update")
+        with_connection do |c|
+          c.update(um, "#{self} Update")
+        end
       end
 
       def _delete_record(constraints) # :nodoc:
@@ -621,7 +294,9 @@ module ActiveRecord
         dm = Arel::DeleteManager.new(arel_table)
         dm.wheres = constraints
 
-        connection.delete(dm, "#{self} Destroy")
+        with_connection do |c|
+          c.delete(dm, "#{self} Destroy")
+        end
       end
 
       private
@@ -1067,7 +742,7 @@ module ActiveRecord
     #   end
     #
     def reload(options = nil)
-      self.class.connection.clear_query_cache
+      self.class.connection_pool.clear_query_cache
 
       fresh_object = if apply_scoping?(options)
         _find_record((options || {}).merge(all_queries: true))
@@ -1247,16 +922,19 @@ module ActiveRecord
     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

data/lib/active_record/query_cache.rb

@@ -8,7 +8,13 @@ module ActiveRecord
       # If it's not, it will execute the given block.
       def cache(&block)
         if connected? || !configurations.empty?
-          connection.cache(&block)
+          pool = connection_pool
+          was_enabled = pool.query_cache_enabled
+          begin
+            pool.enable_query_cache(&block)
+          ensure
+            pool.clear_query_cache unless was_enabled
+          end
         else
           yield
         end
@@ -16,9 +22,12 @@ module ActiveRecord
 
       # Disable the query cache within the block if Active Record is configured.
       # If it's not, it will execute the given block.
-      def uncached(&block)
+      #
+      # Set <tt>dirties: false</tt> to prevent query caches on all connections from being cleared by write operations.
+      # (By default, write operations dirty all connections' query caches in case they are replicas whose cache would now be outdated.)
+      def uncached(dirties: true, &block)
         if connected? || !configurations.empty?
-          connection.uncached(&block)
+          connection_pool.disable_query_cache(dirties: dirties, &block)
         else
           yield
         end
@@ -26,14 +35,17 @@ module ActiveRecord
     end
 
     def self.run
-      ActiveRecord::Base.connection_handler.each_connection_pool.reject { |p| p.query_cache_enabled }.each { |p| p.enable_query_cache! }
+      ActiveRecord::Base.connection_handler.each_connection_pool.reject(&:query_cache_enabled).each(&:enable_query_cache!)
     end
 
     def self.complete(pools)
-      pools.each { |pool| pool.disable_query_cache! }
+      pools.each do |pool|
+        pool.disable_query_cache!
+        pool.clear_query_cache
+      end
 
       ActiveRecord::Base.connection_handler.each_connection_pool do |pool|
-        pool.release_connection if pool.active_connection? && !pool.connection.transaction_open?
+        pool.release_connection if pool.active_connection? && !pool.lease_connection.transaction_open?
       end
     end
 

data/lib/active_record/query_logs.rb

@@ -26,6 +26,7 @@ module ActiveRecord
   # * +socket+
   # * +db_host+
   # * +database+
+  # * +source_location+
   #
   # Action Controller adds default tags when loaded:
   #
@@ -108,6 +109,20 @@ module ActiveRecord
           end
       end
 
+      if Thread.respond_to?(:each_caller_location)
+        def query_source_location # :nodoc:
+          Thread.each_caller_location do |location|
+            frame = LogSubscriber.backtrace_cleaner.clean_frame(location.path)
+            return frame if frame
+          end
+          nil
+        end
+      else
+        def query_source_location # :nodoc:
+          LogSubscriber.backtrace_cleaner.clean(caller_locations(1).each).first
+        end
+      end
+
       ActiveSupport::ExecutionContext.after_change { ActiveRecord::QueryLogs.clear_cache }
 
       private

data/lib/active_record/query_logs_formatter.rb

@@ -28,7 +28,7 @@ module ActiveRecord
       end
 
       def format(pairs)
-        pairs.sort_by!(&:first)
+        pairs.sort_by! { |pair| pair.first.to_s }
         super
       end
 

data/lib/active_record/querying.rb

@@ -10,15 +10,16 @@ module ActiveRecord
       :first_or_create, :first_or_create!, :first_or_initialize,
       :find_or_create_by, :find_or_create_by!, :find_or_initialize_by,
       :create_or_find_by, :create_or_find_by!,
-      :destroy_all, :delete_all, :update_all, :touch_all, :destroy_by, :delete_by,
+      :destroy, :destroy_all, :delete, :delete_all, :update_all, :touch_all, :destroy_by, :delete_by,
       :find_each, :find_in_batches, :in_batches,
       :select, :reselect, :order, :regroup, :in_order_of, :reorder, :group, :limit, :offset, :joins, :left_joins, :left_outer_joins,
       :where, :rewhere, :invert_where, :preload, :extract_associated, :eager_load, :includes, :from, :lock, :readonly,
       :and, :or, :annotate, :optimizer_hints, :extending,
       :having, :create_with, :distinct, :references, :none, :unscope, :merge, :except, :only,
       :count, :average, :minimum, :maximum, :sum, :calculate,
-      :pluck, :pick, :ids, :async_ids, :strict_loading, :excluding, :without, :with,
+      :pluck, :pick, :ids, :async_ids, :strict_loading, :excluding, :without, :with, :with_recursive,
       :async_count, :async_average, :async_minimum, :async_maximum, :async_sum, :async_pluck, :async_pick,
+      :insert, :insert_all, :insert!, :insert_all!, :upsert, :upsert_all
     ].freeze # :nodoc:
     delegate(*QUERYING_METHODS, to: :all)
 
@@ -47,19 +48,26 @@ module ActiveRecord
     #
     # 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, &block)
-      _load_from_sql(_query_by_sql(sql, binds, preparable: preparable), &block)
+    def find_by_sql(sql, binds = [], preparable: nil, allow_retry: false, &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) # :nodoc:
-      connection.select_all(sanitize_sql(sql), "#{name} Load", binds, preparable: preparable, async: async)
+    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 +107,16 @@ module ActiveRecord
     #
     # * +sql+ - An SQL statement which should return a count query from the database, see the example above.
     def count_by_sql(sql)
-      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)
-      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