Merge pull request rails#51336 from adrianna-chang-shopify/ac-retry-idempotent-queries-2

matthewd · web-flow · commit 530d77170ee3 · 2024-03-27T01:48:15.000+10:30
Retry known idempotent SELECT queries on connection-related exceptions
diff --git a/activerecord/CHANGELOG.md b/activerecord/CHANGELOG.md
@@ -1,3 +1,12 @@
+*   Retry known idempotent SELECT queries on connection-related exceptions
+
+    SELECT queries we construct by walking the Arel tree and / or with known model attributes
+    are idempotent and can safely be retried in the case of a connection error. Previously,
+    adapters such as `TrilogyAdapter` would raise `ActiveRecord::ConnectionFailed: Trilogy::EOFError`
+    when encountering a connection error mid-request.
+
+    *Adrianna Chang*
+
 *   Allow association's `foreign_key` to be composite.
 
     `query_constraints` option was the only way to configure a composite foreign key by passing an `Array`.
diff --git a/activerecord/lib/active_record/associations/join_dependency/join_association.rb b/activerecord/lib/active_record/associations/join_dependency/join_association.rb
@@ -93,7 +93,7 @@ def strict_loading?
           def append_constraints(connection, join, constraints)
             if join.is_a?(Arel::Nodes::StringJoin)
               join_string = Arel::Nodes::And.new(constraints.unshift join.left)
-              join.left = Arel.sql(connection.visitor.compile(join_string))
+              join.left = Arel.sql(connection.visitor.compile(join_string), retryable: true)
             else
               right = join.right
               right.expr = Arel::Nodes::And.new(constraints.unshift right.expr)
diff --git a/activerecord/lib/active_record/connection_adapters/abstract/database_statements.rb b/activerecord/lib/active_record/connection_adapters/abstract/database_statements.rb
@@ -14,7 +14,7 @@ def to_sql(arel_or_sql_string, binds = [])
         sql
       end
 
-      def to_sql_and_binds(arel_or_sql_string, binds = [], preparable = nil) # :nodoc:
+      def to_sql_and_binds(arel_or_sql_string, binds = [], preparable = nil, allow_retry = false) # :nodoc:
         # Arel::TreeManager -> Arel::Node
         if arel_or_sql_string.respond_to?(:ast)
           arel_or_sql_string = arel_or_sql_string.ast
@@ -27,6 +27,7 @@ def to_sql_and_binds(arel_or_sql_string, binds = [], preparable = nil) # :nodoc:
           end
 
           collector = collector()
+          collector.retryable = true
 
           if prepared_statements
             collector.preparable = true
@@ -41,10 +42,11 @@ def to_sql_and_binds(arel_or_sql_string, binds = [], preparable = nil) # :nodoc:
           else
             sql = visitor.compile(arel_or_sql_string, collector)
           end
-          [sql.freeze, binds, preparable]
+          allow_retry = collector.retryable
+          [sql.freeze, binds, preparable, allow_retry]
         else
           arel_or_sql_string = arel_or_sql_string.dup.freeze unless arel_or_sql_string.frozen?
-          [arel_or_sql_string, binds, preparable]
+          [arel_or_sql_string, binds, preparable, allow_retry]
         end
       end
       private :to_sql_and_binds
@@ -64,11 +66,15 @@ def cacheable_query(klass, arel) # :nodoc:
       end
 
       # Returns an ActiveRecord::Result instance.
-      def select_all(arel, name = nil, binds = [], preparable: nil, async: false)
+      def select_all(arel, name = nil, binds = [], preparable: nil, async: false, allow_retry: false)
         arel = arel_from_relation(arel)
-        sql, binds, preparable = to_sql_and_binds(arel, binds, preparable)
+        sql, binds, preparable, allow_retry = to_sql_and_binds(arel, binds, preparable, allow_retry)
 
-        select(sql, name, binds, prepare: prepared_statements && preparable, async: async && FutureResult::SelectAll)
+        select(sql, name, binds,
+          prepare: prepared_statements && preparable,
+          async: async && FutureResult::SelectAll,
+          allow_retry: allow_retry
+        )
       rescue ::RangeError
         ActiveRecord::Result.empty(async: async)
       end
@@ -495,7 +501,7 @@ def with_yaml_fallback(value) # :nodoc:
       end
 
       # This is a safe default, even if not high precision on all databases
-      HIGH_PRECISION_CURRENT_TIMESTAMP = Arel.sql("CURRENT_TIMESTAMP").freeze # :nodoc:
+      HIGH_PRECISION_CURRENT_TIMESTAMP = Arel.sql("CURRENT_TIMESTAMP", retryable: true).freeze # :nodoc:
       private_constant :HIGH_PRECISION_CURRENT_TIMESTAMP
 
       # Returns an Arel SQL literal for the CURRENT_TIMESTAMP for usage with
@@ -507,7 +513,7 @@ def high_precision_current_timestamp
         HIGH_PRECISION_CURRENT_TIMESTAMP
       end
 
-      def internal_exec_query(sql, name = "SQL", binds = [], prepare: false, async: false) # :nodoc:
+      def internal_exec_query(sql, name = "SQL", binds = [], prepare: false, async: false, allow_retry: false) # :nodoc:
         raise NotImplementedError
       end
 
@@ -606,7 +612,7 @@ def combine_multi_statements(total_sql)
         end
 
         # Returns an ActiveRecord::Result instance.
-        def select(sql, name = nil, binds = [], prepare: false, async: false)
+        def select(sql, name = nil, binds = [], prepare: false, async: false, allow_retry: false)
           if async && async_enabled?
             if current_transaction.joinable?
               raise AsynchronousQueryInsideTransactionError, "Asynchronous queries are not allowed inside transactions"
@@ -627,7 +633,7 @@ def select(sql, name = nil, binds = [], prepare: false, async: false)
             return future_result
           end
 
-          result = internal_exec_query(sql, name, binds, prepare: prepare)
+          result = internal_exec_query(sql, name, binds, prepare: prepare, allow_retry: allow_retry)
           if async
             FutureResult.wrap(result)
           else
diff --git a/activerecord/lib/active_record/connection_adapters/abstract/query_cache.rb b/activerecord/lib/active_record/connection_adapters/abstract/query_cache.rb
@@ -204,19 +204,19 @@ def clear_query_cache
         pool.clear_query_cache
       end
 
-      def select_all(arel, name = nil, binds = [], preparable: nil, async: false) # :nodoc:
+      def select_all(arel, name = nil, binds = [], preparable: nil, async: false, allow_retry: false) # :nodoc:
         arel = arel_from_relation(arel)
 
         # If arel is locked this is a SELECT ... FOR UPDATE or somesuch.
         # Such queries should not be cached.
         if @query_cache&.enabled? && !(arel.respond_to?(:locked) && arel.locked)
-          sql, binds, preparable = to_sql_and_binds(arel, binds, preparable)
+          sql, binds, preparable, allow_retry = to_sql_and_binds(arel, binds, preparable)
 
           if async
-            result = lookup_sql_cache(sql, name, binds) || super(sql, name, binds, preparable: preparable, async: async)
+            result = lookup_sql_cache(sql, name, binds) || super(sql, name, binds, preparable: preparable, async: async, allow_retry: allow_retry)
             FutureResult.wrap(result)
           else
-            cache_sql(sql, name, binds) { super(sql, name, binds, preparable: preparable, async: async) }
+            cache_sql(sql, name, binds) { super(sql, name, binds, preparable: preparable, async: async, allow_retry: allow_retry) }
           end
         else
           super
diff --git a/activerecord/lib/active_record/connection_adapters/abstract_mysql_adapter.rb b/activerecord/lib/active_record/connection_adapters/abstract_mysql_adapter.rb
@@ -229,12 +229,12 @@ def disable_referential_integrity # :nodoc:
       # Mysql2Adapter doesn't have to free a result after using it, but we use this method
       # to write stuff in an abstract way without concerning ourselves about whether it
       # needs to be explicitly freed or not.
-      def execute_and_free(sql, name = nil, async: false) # :nodoc:
+      def execute_and_free(sql, name = nil, async: false, allow_retry: false) # :nodoc:
         sql = transform_query(sql)
         check_if_write_query(sql)
 
         mark_transaction_written_if_write(sql)
-        yield raw_execute(sql, name, async: async)
+        yield raw_execute(sql, name, async: async, allow_retry: allow_retry)
       end
 
       def begin_db_transaction # :nodoc:
diff --git a/activerecord/lib/active_record/connection_adapters/mysql/database_statements.rb b/activerecord/lib/active_record/connection_adapters/mysql/database_statements.rb
@@ -11,7 +11,7 @@ module DatabaseStatements
 
         # https://dev.mysql.com/doc/refman/5.7/en/date-and-time-functions.html#function_current-timestamp
         # https://dev.mysql.com/doc/refman/5.7/en/date-and-time-type-syntax.html
-        HIGH_PRECISION_CURRENT_TIMESTAMP = Arel.sql("CURRENT_TIMESTAMP(6)").freeze # :nodoc:
+        HIGH_PRECISION_CURRENT_TIMESTAMP = Arel.sql("CURRENT_TIMESTAMP(6)", retryable: true).freeze # :nodoc:
         private_constant :HIGH_PRECISION_CURRENT_TIMESTAMP
 
         def write_query?(sql) # :nodoc:
diff --git a/activerecord/lib/active_record/connection_adapters/mysql2/database_statements.rb b/activerecord/lib/active_record/connection_adapters/mysql2/database_statements.rb
@@ -18,9 +18,9 @@ def select_all(*, **) # :nodoc:
           result
         end
 
-        def internal_exec_query(sql, name = "SQL", binds = [], prepare: false, async: false) # :nodoc:
+        def internal_exec_query(sql, name = "SQL", binds = [], prepare: false, async: false, allow_retry: false) # :nodoc:
           if without_prepared_statement?(binds)
-            execute_and_free(sql, name, async: async) do |result|
+            execute_and_free(sql, name, async: async, allow_retry: allow_retry) do |result|
               if result
                 build_result(columns: result.fields, rows: result.to_a)
               else
diff --git a/activerecord/lib/active_record/connection_adapters/postgresql/database_statements.rb b/activerecord/lib/active_record/connection_adapters/postgresql/database_statements.rb
@@ -124,7 +124,7 @@ def exec_restart_db_transaction # :nodoc:
         end
 
         # From https://www.postgresql.org/docs/current/functions-datetime.html#FUNCTIONS-DATETIME-CURRENT
-        HIGH_PRECISION_CURRENT_TIMESTAMP = Arel.sql("CURRENT_TIMESTAMP").freeze # :nodoc:
+        HIGH_PRECISION_CURRENT_TIMESTAMP = Arel.sql("CURRENT_TIMESTAMP", retryable: true).freeze # :nodoc:
         private_constant :HIGH_PRECISION_CURRENT_TIMESTAMP
 
         def high_precision_current_timestamp
diff --git a/activerecord/lib/active_record/connection_adapters/postgresql_adapter.rb b/activerecord/lib/active_record/connection_adapters/postgresql_adapter.rb
@@ -881,7 +881,7 @@ def exec_no_cache(sql, name, binds, async:, allow_retry:, materialize_transactio
 
           type_casted_binds = type_casted_binds(binds)
           log(sql, name, binds, type_casted_binds, async: async) do |notification_payload|
-            with_raw_connection(allow_retry: false, materialize_transactions: materialize_transactions) do |conn|
+            with_raw_connection(allow_retry: allow_retry, materialize_transactions: materialize_transactions) do |conn|
               result = conn.exec_params(sql, type_casted_binds)
               verified!
               notification_payload[:row_count] = result.count
@@ -895,7 +895,7 @@ def exec_cache(sql, name, binds, async:, allow_retry:, materialize_transactions:
 
           update_typemap_for_default_timezone
 
-          with_raw_connection(allow_retry: false, materialize_transactions: materialize_transactions) do |conn|
+          with_raw_connection(allow_retry: allow_retry, materialize_transactions: materialize_transactions) do |conn|
             stmt_key = prepare_statement(sql, binds, conn)
             type_casted_binds = type_casted_binds(binds)
 
diff --git a/activerecord/lib/active_record/connection_adapters/sqlite3/database_statements.rb b/activerecord/lib/active_record/connection_adapters/sqlite3/database_statements.rb
@@ -21,7 +21,7 @@ def explain(arel, binds = [], _options = [])
           SQLite3::ExplainPrettyPrinter.new.pp(result)
         end
 
-        def internal_exec_query(sql, name = nil, binds = [], prepare: false, async: false) # :nodoc:
+        def internal_exec_query(sql, name = nil, binds = [], prepare: false, async: false, allow_retry: false) # :nodoc:
           sql = transform_query(sql)
           check_if_write_query(sql)
 
@@ -106,7 +106,7 @@ def exec_rollback_db_transaction # :nodoc:
 
         # https://stackoverflow.com/questions/17574784
         # https://www.sqlite.org/lang_datefunc.html
-        HIGH_PRECISION_CURRENT_TIMESTAMP = Arel.sql("STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')").freeze # :nodoc:
+        HIGH_PRECISION_CURRENT_TIMESTAMP = Arel.sql("STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')", retryable: true).freeze # :nodoc:
         private_constant :HIGH_PRECISION_CURRENT_TIMESTAMP
 
         def high_precision_current_timestamp
diff --git a/activerecord/lib/active_record/connection_adapters/trilogy/database_statements.rb b/activerecord/lib/active_record/connection_adapters/trilogy/database_statements.rb
@@ -12,12 +12,12 @@ def select_all(*, **) # :nodoc:
           result
         end
 
-        def internal_exec_query(sql, name = "SQL", binds = [], prepare: false, async: false) # :nodoc:
+        def internal_exec_query(sql, name = "SQL", binds = [], prepare: false, async: false, allow_retry: false) # :nodoc:
           sql = transform_query(sql)
           check_if_write_query(sql)
           mark_transaction_written_if_write(sql)
 
-          result = raw_execute(sql, name, async: async)
+          result = raw_execute(sql, name, async: async, allow_retry: allow_retry)
           ActiveRecord::Result.new(result.fields, result.to_a)
         end
 
diff --git a/activerecord/lib/active_record/core.rb b/activerecord/lib/active_record/core.rb
@@ -431,7 +431,7 @@ def cached_find_by(keys, values)
           }
 
           begin
-            statement.execute(values.flatten, lease_connection).first
+            statement.execute(values.flatten, lease_connection, allow_retry: true).first
           rescue TypeError
             raise ActiveRecord::StatementInvalid
           end
diff --git a/activerecord/lib/active_record/internal_metadata.rb b/activerecord/lib/active_record/internal_metadata.rb
@@ -153,7 +153,7 @@ def update_entry(connection, key, new_value)
 
       def select_entry(connection, key)
         sm = Arel::SelectManager.new(arel_table)
-        sm.project(Arel::Nodes::SqlLiteral.new("*"))
+        sm.project(Arel::Nodes::SqlLiteral.new("*", retryable: true))
         sm.where(arel_table[primary_key].eq(Arel::Nodes::BindParam.new(key)))
         sm.order(arel_table[primary_key].asc)
         sm.limit = 1
diff --git a/activerecord/lib/active_record/querying.rb b/activerecord/lib/active_record/querying.rb
@@ -47,8 +47,8 @@ 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, &block)
-      _load_from_sql(_query_by_sql(sql, binds, preparable: preparable), &block)
+    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)
     end
 
     # Same as <tt>#find_by_sql</tt> but perform the query asynchronously and returns an ActiveRecord::Promise.
@@ -58,8 +58,8 @@ def async_find_by_sql(sql, binds = [], preparable: nil, &block)
       end
     end
 
-    def _query_by_sql(sql, binds = [], preparable: nil, async: false) # :nodoc:
-      lease_connection.select_all(sanitize_sql(sql), "#{name} Load", binds, preparable: preparable, async: async)
+    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)
     end
 
     def _load_from_sql(result_set, &block) # :nodoc:
diff --git a/activerecord/lib/active_record/relation/calculations.rb b/activerecord/lib/active_record/relation/calculations.rb
@@ -446,7 +446,7 @@ def aggregate_column(column_name)
         return column_name if Arel::Expressions === column_name
 
         arel_column(column_name.to_s) do |name|
-          Arel.sql(column_name == :all ? "*" : name)
+          column_name == :all ? Arel.sql("*", retryable: true) : Arel.sql(name)
         end
       end
 
@@ -643,7 +643,7 @@ def build_count_subquery(relation, column_name, distinct)
           relation.select_values = [ aggregate_column(column_name).as(column_alias) ]
         end
 
-        subquery_alias = Arel.sql("subquery_for_count")
+        subquery_alias = Arel.sql("subquery_for_count", retryable: true)
         select_value = operation_over_aggregate_column(column_alias, "count", false)
 
         relation.build_subquery(subquery_alias, select_value)
diff --git a/activerecord/lib/active_record/relation/predicate_builder.rb b/activerecord/lib/active_record/relation/predicate_builder.rb
@@ -28,9 +28,9 @@ def build_from_hash(attributes, &block)
     def self.references(attributes)
       attributes.each_with_object([]) do |(key, value), result|
         if value.is_a?(Hash)
-          result << Arel.sql(key)
+          result << Arel.sql(key, retryable: true)
         elsif (idx = key.rindex("."))
-          result << Arel.sql(key[0, idx])
+          result << Arel.sql(key[0, idx], retryable: true)
         end
       end
     end
diff --git a/activerecord/lib/active_record/relation/query_methods.rb b/activerecord/lib/active_record/relation/query_methods.rb
@@ -2013,7 +2013,7 @@ def order_column(field)
           if attr_name == "count" && !group_values.empty?
             table[attr_name]
           else
-            Arel.sql(adapter_class.quote_table_name(attr_name))
+            Arel.sql(adapter_class.quote_table_name(attr_name), retryable: true)
           end
         end
       end
diff --git a/activerecord/lib/active_record/statement_cache.rb b/activerecord/lib/active_record/statement_cache.rb
@@ -62,7 +62,7 @@ def sql_for(binds, connection)
     end
 
     class PartialQueryCollector
-      attr_accessor :preparable
+      attr_accessor :preparable, :retryable
 
       def initialize
         @parts = []
@@ -142,12 +142,12 @@ def initialize(query_builder, bind_map, klass)
       @klass = klass
     end
 
-    def execute(params, connection, &block)
+    def execute(params, connection, allow_retry: false, &block)
       bind_values = bind_map.bind params
 
       sql = query_builder.sql_for bind_values, connection
 
-      klass.find_by_sql(sql, bind_values, preparable: true, &block)
+      klass.find_by_sql(sql, bind_values, preparable: true, allow_retry: allow_retry, &block)
     rescue ::RangeError
       []
     end
diff --git a/activerecord/lib/arel.rb b/activerecord/lib/arel.rb
@@ -45,16 +45,20 @@ module Arel
   # that this behavior only applies when bind value parameters are
   # supplied in the call; without them, the placeholder tokens have no
   # special meaning, and will be passed through to the query as-is.
-  def self.sql(sql_string, *positional_binds, **named_binds)
+  #
+  # The +:retryable+ option can be used to mark the SQL as safe to retry.
+  # Use this option only if the SQL is idempotent, as it could be executed
+  # more than once.
+  def self.sql(sql_string, *positional_binds, retryable: false, **named_binds)
     if positional_binds.empty? && named_binds.empty?
-      Arel::Nodes::SqlLiteral.new sql_string
+      Arel::Nodes::SqlLiteral.new(sql_string, retryable: retryable)
     else
       Arel::Nodes::BoundSqlLiteral.new sql_string, positional_binds, named_binds
     end
   end
 
   def self.star # :nodoc:
-    sql "*"
+    sql("*", retryable: true)
   end
 
   def self.arel_node?(value) # :nodoc:
diff --git a/activerecord/lib/arel/alias_predication.rb b/activerecord/lib/arel/alias_predication.rb
@@ -3,7 +3,7 @@
 module Arel # :nodoc: all
   module AliasPredication
     def as(other)
-      Nodes::As.new self, Nodes::SqlLiteral.new(other)
+      Nodes::As.new self, Nodes::SqlLiteral.new(other, retryable: true)
     end
   end
 end
diff --git a/activerecord/lib/arel/collectors/bind.rb b/activerecord/lib/arel/collectors/bind.rb
@@ -3,6 +3,8 @@
 module Arel # :nodoc: all
   module Collectors
     class Bind
+      attr_accessor :retryable
+
       def initialize
         @binds = []
       end
diff --git a/activerecord/lib/arel/collectors/composite.rb b/activerecord/lib/arel/collectors/composite.rb
@@ -4,12 +4,19 @@ module Arel # :nodoc: all
   module Collectors
     class Composite
       attr_accessor :preparable
+      attr_reader :retryable
 
       def initialize(left, right)
         @left = left
         @right = right
       end
 
+      def retryable=(retryable)
+        left.retryable = retryable
+        right.retryable = retryable
+        @retryable = retryable
+      end
+
       def <<(str)
         left << str
         right << str
diff --git a/activerecord/lib/arel/collectors/sql_string.rb b/activerecord/lib/arel/collectors/sql_string.rb
diff --git a/activerecord/lib/arel/collectors/substitute_binds.rb b/activerecord/lib/arel/collectors/substitute_binds.rb
diff --git a/activerecord/lib/arel/nodes/sql_literal.rb b/activerecord/lib/arel/nodes/sql_literal.rb
diff --git a/activerecord/lib/arel/select_manager.rb b/activerecord/lib/arel/select_manager.rb
diff --git a/activerecord/lib/arel/visitors/mysql.rb b/activerecord/lib/arel/visitors/mysql.rb
diff --git a/activerecord/lib/arel/visitors/to_sql.rb b/activerecord/lib/arel/visitors/to_sql.rb
diff --git a/activerecord/test/cases/adapter_test.rb b/activerecord/test/cases/adapter_test.rb
diff --git a/activerecord/test/cases/arel/collectors/composite_test.rb b/activerecord/test/cases/arel/collectors/composite_test.rb
diff --git a/activerecord/test/cases/arel/visitors/to_sql_test.rb b/activerecord/test/cases/arel/visitors/to_sql_test.rb
