Merge pull request rails#54436 from skipkayhil/hm-retryable-statement-cached-queries

Enable statement-cached queries to be retryable

Author: byroot

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]