Merge pull request #46273 from adrianna-chang-shopify/ac-execute-takes-allow-retry

Allow adapter `#execute` methods  to take `allow_retry` option

activerecord/CHANGELOG.md

@@ -1,3 +1,8 @@
+*   Adapter `#execute` methods now accept an `allow_retry` option. When set to `true`, the SQL statement will be
+    retried, up to the database's configured `connection_retries` value, upon encountering connection-related errors.
+
+    *Adrianna Chang*
+
 *   Only trigger `after_commit :destroy` callbacks when a database row is deleted.
 
     This prevents `after_commit :destroy` callbacks from being triggered again

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

@@ -110,10 +110,15 @@ def write_query?(sql)
 
       # Executes the SQL statement in the context of this connection and returns
       # the raw result from the connection adapter.
+      #
+      # Setting +allow_retry+ to true causes the db to reconnect and retry
+      # executing the SQL statement in case of a connection-related exception.
+      # This option should only be enabled for known idempotent queries.
+      #
       # Note: depending on your database connector, the result returned by this
       # method may be manually memory managed. Consider using the exec_query
       # wrapper instead.
-      def execute(sql, name = nil)
+      def execute(sql, name = nil, allow_retry: false)
         raise NotImplementedError
       end
 

activerecord/lib/active_record/connection_adapters/abstract_mysql_adapter.rb

@@ -221,11 +221,15 @@ def disable_referential_integrity # :nodoc:
       #++
 
       # Executes the SQL statement in the context of this connection.
-      def execute(sql, name = nil, async: false)
+      #
+      # Setting +allow_retry+ to true causes the db to reconnect and retry
+      # executing the SQL statement in case of a connection-related exception.
+      # This option should only be enabled for known idempotent queries.
+      def execute(sql, name = nil, async: false, allow_retry: false)
         sql = transform_query(sql)
         check_if_write_query(sql)
 
-        raw_execute(sql, name, async: async)
+        raw_execute(sql, name, async: async, allow_retry: allow_retry)
       end
 
       # Mysql2Adapter doesn't have to free a result after using it, but we use this method

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

@@ -33,15 +33,20 @@ def write_query?(sql) # :nodoc:
 
         # Executes an SQL statement, returning a PG::Result object on success
         # or raising a PG::Error exception otherwise.
+        #
+        # Setting +allow_retry+ to true causes the db to reconnect and retry
+        # executing the SQL statement in case of a connection-related exception.
+        # This option should only be enabled for known idempotent queries.
+        #
         # Note: the PG::Result object is manually memory managed; if you don't
         # need it specifically, you may want consider the <tt>exec_query</tt> wrapper.
-        def execute(sql, name = nil)
+        def execute(sql, name = nil, allow_retry: false)
           sql = transform_query(sql)
           check_if_write_query(sql)
 
           mark_transaction_written_if_write(sql)
 
-          with_raw_connection do |conn|
+          with_raw_connection(allow_retry: allow_retry) do |conn|
             log(sql, name) do
               conn.async_exec(sql)
             end

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

@@ -20,14 +20,14 @@ def explain(arel, binds = [])
           SQLite3::ExplainPrettyPrinter.new.pp(exec_query(sql, "EXPLAIN", []))
         end
 
-        def execute(sql, name = nil) # :nodoc:
+        def execute(sql, name = nil, allow_retry: false) # :nodoc:
           sql = transform_query(sql)
           check_if_write_query(sql)
 
           mark_transaction_written_if_write(sql)
 
           log(sql, name) do
-            with_raw_connection do |conn|
+            with_raw_connection(allow_retry: allow_retry) do |conn|
               conn.execute(sql)
             end
           end

activerecord/test/cases/adapter_test.rb

@@ -690,6 +690,21 @@ def teardown
         end
       end
 
+      unless current_adapter?(:SQLite3Adapter)
+        test "#execute is retryable" do
+          conn_id = case @connection.class::ADAPTER_NAME
+                    when "Mysql2"
+                      @connection.execute("SELECT CONNECTION_ID()").to_a[0][0]
+                    when "PostgreSQL"
+                      @connection.execute("SELECT pg_backend_pid()").to_a[0]["pg_backend_pid"]
+          end
+
+          kill_connection_from_server(conn_id)
+
+          @connection.execute("SELECT 1", allow_retry: true)
+        end
+      end
+
       private
         def raw_transaction_open?(connection)
           case connection.class::ADAPTER_NAME
@@ -731,6 +746,18 @@ def remote_disconnect(connection)
             skip
           end
         end
+
+        def kill_connection_from_server(connection_id)
+          conn = @connection.pool.checkout
+          case conn.class::ADAPTER_NAME
+          when "Mysql2"
+            conn.execute("KILL #{connection_id}")
+          when "PostgreSQL"
+            conn.execute("SELECT pg_cancel_backend(#{connection_id})")
+          end
+
+          conn.close
+        end
     end
   end
 end