Enable statement-cached queries to be retryable

The StatementCache uses a different code path for generating SQL and
executing queries, so it doesn't currently benefit from the ability to
mark queries as retryable. This mainly affects association SELECTs,
because the other user of StatementCache is `find_by` and it explicitly
passes `allow_retry: true` when executing a cached query.

This commit enables queries that would be retryable without using the
StatementCache to also be retryable when using it. Just like in
`to_sql_and_binds`, `cacheable_query` now sets `retryable: true` on the
`collector` that constructs the query and additionally copies the
`collector`'s final `retryable` value to the query object. Then, instead
of accepting an `allow_retry` parameter when executing the cached query,
`execute` now uses the query object's `retryable` attribute.

The test for this case uses Author/Books instead of Post/Comments
because Comment has a `default_scope`, so `post.comments` does not use
the statement cache.

Author: skipkayhil

activerecord/CHANGELOG.md

@@ -1,3 +1,8 @@
+*   Enable automatically retrying idempotent association queries on connection
+    errors.
+
+    *Hartley McGuire*
+
 *   Add `allow_retry` to `sql.active_record` instrumentation.
 
     This enables identifying queries which are and are not automatically

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

@@ -55,12 +55,15 @@ def to_sql_and_binds(arel_or_sql_string, binds = [], preparable = nil, allow_ret
       # can be used to query the database repeatedly.
       def cacheable_query(klass, arel) # :nodoc:
         if prepared_statements
+          collector = collector()
+          collector.retryable = true
           sql, binds = visitor.compile(arel.ast, collector)
-          query = klass.query(sql)
+          query = klass.query(sql, retryable: collector.retryable)
         else
           collector = klass.partial_query_collector
+          collector.retryable = true
           parts, binds = visitor.compile(arel.ast, collector)
-          query = klass.partial_query(parts)
+          query = klass.partial_query(parts, retryable: collector.retryable)
         end
         [query, binds]
       end

activerecord/lib/active_record/core.rb

@@ -450,7 +450,7 @@ def cached_find_by(keys, values)
               where(wheres).limit(1)
             }
 
-            statement.execute(values.flatten, connection, allow_retry: true).then do |r|
+            statement.execute(values.flatten, connection).then do |r|
               r.first
             rescue TypeError
               raise ActiveRecord::StatementInvalid

activerecord/lib/active_record/statement_cache.rb

@@ -31,8 +31,11 @@ class StatementCache # :nodoc:
     class Substitute; end # :nodoc:
 
     class Query # :nodoc:
-      def initialize(sql)
+      attr_reader :retryable
+
+      def initialize(sql, retryable:)
         @sql = sql
+        @retryable = retryable
       end
 
       def sql_for(binds, connection)
@@ -41,11 +44,12 @@ def sql_for(binds, connection)
     end
 
     class PartialQuery < Query # :nodoc:
-      def initialize(values)
+      def initialize(values, retryable:)
         @values = values
         @indexes = values.each_with_index.find_all { |thing, i|
           Substitute === thing
         }.map(&:last)
+        @retryable = retryable
       end
 
       def sql_for(binds, connection)
@@ -94,12 +98,12 @@ def value
       end
     end
 
-    def self.query(sql)
-      Query.new(sql)
+    def self.query(...)
+      Query.new(...)
     end
 
-    def self.partial_query(values)
-      PartialQuery.new(values)
+    def self.partial_query(...)
+      PartialQuery.new(...)
     end
 
     def self.partial_query_collector
@@ -142,14 +146,14 @@ def initialize(query_builder, bind_map, model)
       @model = model
     end
 
-    def execute(params, connection, allow_retry: false, async: false, &block)
+    def execute(params, connection, async: false, &block)
       bind_values = @bind_map.bind params
       sql = @query_builder.sql_for bind_values, connection
 
       if async
-        @model.async_find_by_sql(sql, bind_values, preparable: true, allow_retry: allow_retry, &block)
+        @model.async_find_by_sql(sql, bind_values, preparable: true, allow_retry: @query_builder.retryable, &block)
       else
-        @model.find_by_sql(sql, bind_values, preparable: true, allow_retry: allow_retry, &block)
+        @model.find_by_sql(sql, bind_values, preparable: true, allow_retry: @query_builder.retryable, &block)
       end
     rescue ::RangeError
       async ? Promise.wrap([]) : []

activerecord/test/cases/adapter_test.rb

@@ -702,14 +702,15 @@ def teardown
 
       test "idempotent SELECT queries allow retries" do
         notifications = capture_notifications("sql.active_record") do
-          assert Post.first
+          assert (a = Author.first)
           assert Post.where(id: [1, 2]).first
           assert Post.find(1)
           assert Post.find_by(title: "Welcome to the weblog")
           assert_predicate Post, :exists?
+          a.books.to_a
         end.select { |n| n.payload[:name] != "SCHEMA" }
 
-        assert_equal 5, notifications.length
+        assert_equal 6, notifications.length
 
         notifications.each do |n|
           assert n.payload[:allow_retry]