Separate two groups of retryable DB exceptions

Connectivity issues need a reconnect to be retried, but there are also
in-server errors that can also be retried (lock timeouts and deadlocks).

To handle this better, we move translation into #with_raw_connection:
it's a much better representation of the boundary between the underlying
adapter library and Active Record's abstractions anyway. (It doesn't
know about the query being run, though, so we let #log enrich that into
the exception later.)

Author: matthewd

activerecord/lib/active_record/connection_adapters/abstract_adapter.rb

@@ -859,22 +859,39 @@ def with_raw_connection(allow_retry: false, uses_transaction: true)
               result = yield @raw_connection
               @verified = true
               result
-            rescue => ex
-              if retries_available > 0 && retryable_error?(ex) && reconnect_can_restore_state?
+            rescue => original_exception
+              translated_exception = translate_exception_class(original_exception, nil, nil)
+
+              if retries_available > 0
                 retries_available -= 1
-                reconnect!(restore_transactions: true)
-                retry
+
+                if retryable_query_error?(translated_exception)
+                  backoff(connection_retries - retries_available)
+                  retry
+                elsif retryable_connection_error?(translated_exception)
+                  reconnect!(restore_transactions: true)
+                  retry
+                end
               end
 
-              raise
+              raise translated_exception
             ensure
               dirty_current_transaction if uses_transaction
             end
           end
         end
 
-        def retryable_error?(exception)
-          false
+        def retryable_connection_error?(exception)
+          exception.is_a?(ConnectionNotEstablished)
+        end
+
+        def retryable_query_error?(exception)
+          exception.is_a?(Deadlocked) ||
+            exception.is_a?(LockWaitTimeout)
+        end
+
+        def backoff(counter)
+          sleep 0.1 * counter
         end
 
         # Returns a raw connection for internal use with methods that are known
@@ -923,7 +940,7 @@ def translate_exception_class(e, sql, binds)
           exception
         end
 
-        def log(sql, name = "SQL", binds = [], type_casted_binds = [], statement_name = nil, async: false) # :doc:
+        def log(sql, name = "SQL", binds = [], type_casted_binds = [], statement_name = nil, async: false, &block) # :doc:
           @instrumenter.instrument(
             "sql.active_record",
             sql:               sql,
@@ -932,11 +949,11 @@ def log(sql, name = "SQL", binds = [], type_casted_binds = [], statement_name =
             type_casted_binds: type_casted_binds,
             statement_name:    statement_name,
             async:             async,
-            connection:        self) do
-            yield
-          rescue => e
-            raise translate_exception_class(e, sql, binds)
-          end
+            connection:        self,
+            &block
+          )
+        rescue ActiveRecord::StatementInvalid => ex
+          raise ex.set_query(sql, binds)
         end
 
         def transform_query(sql)

activerecord/lib/active_record/connection_adapters/abstract_mysql_adapter.rb

@@ -723,11 +723,6 @@ def translate_exception(exception, message:, sql:, binds:)
           end
         end
 
-        def retryable_error?(exception)
-          error_number(exception).nil? &&
-            exception.message.match?(/MySQL client is not connected/i)
-        end
-
         def change_column_for_alter(table_name, column_name, type, **options)
           column = column_for(table_name, column_name)
           type ||= column.sql_type
@@ -865,7 +860,7 @@ def build_statement_pool
           StatementPool.new(self.class.type_cast_config_to_integer(@config[:statement_limit]))
         end
 
-        def mismatched_foreign_key(message, sql:, binds:)
+        def mismatched_foreign_key_details(message:, sql:)
           foreign_key_pat =
             /Referencing column '(\w+)' and referenced/i =~ message ? $1 : '\w+'
 
@@ -875,11 +870,7 @@ def mismatched_foreign_key(message, sql:, binds:)
             REFERENCES\s*(`?(?<target_table>\w+)`?)\s*\(`?(?<primary_key>\w+)`?\)
           /xmi.match(sql)
 
-          options = {
-            message: message,
-            sql: sql,
-            binds: binds,
-          }
+          options = {}
 
           if match
             options[:table] = match[:table]
@@ -889,6 +880,22 @@ def mismatched_foreign_key(message, sql:, binds:)
             options[:primary_key_column] = column_for(match[:target_table], match[:primary_key])
           end
 
+          options
+        end
+
+        def mismatched_foreign_key(message, sql:, binds:)
+          options = {
+            message: message,
+            sql: sql,
+            binds: binds,
+          }
+
+          if sql
+            options.update mismatched_foreign_key_details(message: message, sql: sql)
+          else
+            options[:query_parser] = ->(sql) { mismatched_foreign_key_details(message: message, sql: sql) }
+          end
+
           MismatchedForeignKey.new(**options)
         end
 

activerecord/lib/active_record/connection_adapters/postgresql_adapter.rb

@@ -693,6 +693,8 @@ def translate_exception(exception, message:, sql:, binds:)
           when nil
             if exception.message.match?(/connection is closed/i)
               ConnectionNotEstablished.new(exception)
+            elsif exception.is_a?(PG::ConnectionBad) && !exception.message.end_with?("\n")
+              ConnectionNotEstablished.new(exception)
             else
               super
             end
@@ -721,7 +723,10 @@ def translate_exception(exception, message:, sql:, binds:)
           end
         end
 
-        def retryable_error?(exception)
+        def retryable_query_error?(exception)
+        end
+
+        def retryable_connection_error?(exception)
           case exception
           when PG::ConnectionBad; !exception.message.end_with?("\n")
           end

activerecord/lib/active_record/connection_adapters/sqlite3_adapter.rb

@@ -577,10 +577,6 @@ def translate_exception(exception, message:, sql:, binds:)
           end
         end
 
-        def retryable_error?(exception)
-          exception.message.match?(/called on a closed database/i)
-        end
-
         COLLATE_REGEX = /.*"(\w+)".*collate\s+"(\w+)".*/i.freeze
 
         def table_structure_with_collation(table_name, basic_structure)

activerecord/lib/active_record/errors.rb

@@ -166,6 +166,15 @@ def initialize(message = nil, sql: nil, binds: nil)
     end
 
     attr_reader :sql, :binds
+
+    def set_query(sql, binds)
+      unless @sql
+        @sql = sql
+        @binds = binds
+      end
+
+      self
+    end
   end
 
   # Defunct wrapper class kept for compatibility.
@@ -192,8 +201,12 @@ def initialize(
       foreign_key: nil,
       target_table: nil,
       primary_key: nil,
-      primary_key_column: nil
+      primary_key_column: nil,
+      query_parser: nil
     )
+      @original_message = message
+      @query_parser = query_parser
+
       if table
         type = primary_key_column.bigint? ? :bigint : primary_key_column.type
         msg = <<~EOM.squish
@@ -211,8 +224,24 @@ def initialize(
       if message
         msg << "\nOriginal message: #{message}"
       end
+
       super(msg, sql: sql, binds: binds)
     end
+
+    def set_query(sql, binds)
+      if @query_parser && !@sql
+        self.class.new(
+          message: @original_message,
+          sql: sql,
+          binds: binds,
+          **@query_parser.call(sql)
+        ).tap do |exception|
+          exception.set_backtrace backtrace
+        end
+      else
+        super
+      end
+    end
   end
 
   # Raised when a record cannot be inserted or updated because it would violate a not null constraint.

activerecord/test/cases/statement_invalid_test.rb

@@ -21,7 +21,9 @@ def error_number
       sql = Book.where(author_id: 96, cover: "hard").to_sql
       error = assert_raises(ActiveRecord::StatementInvalid) do
         Book.connection.send(:log, sql, Book.name) do
-          raise MockDatabaseError
+          Book.connection.send(:with_raw_connection) do
+            raise MockDatabaseError
+          end
         end
       end
       assert_not error.message.include?("SELECT")
@@ -32,7 +34,9 @@ def error_number
       binds = [Minitest::Mock.new, Minitest::Mock.new]
       error = assert_raises(ActiveRecord::StatementInvalid) do
         Book.connection.send(:log, sql, Book.name, binds) do
-          raise MockDatabaseError
+          Book.connection.send(:with_raw_connection) do
+            raise MockDatabaseError
+          end
         end
       end
       assert_equal error.sql, sql