Add new error class `StatementTimeout` which will be raised when statement timeout exceeded (#31129)

We are sometimes using The MAX_EXECUTION_TIME hint for MySQL depending
on the situation. It will prevent catastrophic performance down by wrong
performing queries.

The new error class `StatementTimeout` will make to be easier to handle
that case.

https://dev.mysql.com/doc/refman/5.7/en/optimizer-hints.html#optimizer-hints-execution-time

6 files changed, 73 insertions, 2 deletions

diff --git a/activerecord/CHANGELOG.md b/activerecord/CHANGELOG.md
index 34ef5c79e1..c95e80755d 100644
--- a/activerecord/CHANGELOG.md
+++ b/activerecord/CHANGELOG.md
@@ -1,3 +1,8 @@
+*   Add new error class `StatementTimeout` which will be raised
+    when statement timeout exceeded.
+
+    *Ryuta Kamizono*
+
 *   Fix `bin/rails db:migrate` with specified `VERSION`.
     `bin/rails db:migrate` with empty VERSION behaves as without `VERSION`.
     Check a format of `VERSION`: Allow a migration version number
diff --git a/activerecord/lib/active_record/connection_adapters/abstract_mysql_adapter.rb b/activerecord/lib/active_record/connection_adapters/abstract_mysql_adapter.rb
index bfec6fb784..ca651ef390 100644
--- a/activerecord/lib/active_record/connection_adapters/abstract_mysql_adapter.rb
+++ b/activerecord/lib/active_record/connection_adapters/abstract_mysql_adapter.rb
@@ -635,6 +635,7 @@ module ActiveRecord
         ER_CANNOT_ADD_FOREIGN   = 1215
         ER_CANNOT_CREATE_TABLE  = 1005
         ER_LOCK_WAIT_TIMEOUT    = 1205
+        ER_QUERY_TIMEOUT        = 3024
 
         def translate_exception(exception, message)
           case error_number(exception)
@@ -660,6 +661,8 @@ module ActiveRecord
             Deadlocked.new(message)
           when ER_LOCK_WAIT_TIMEOUT
             TransactionTimeout.new(message)
+          when ER_QUERY_TIMEOUT
+            StatementTimeout.new(message)
           else
             super
           end
diff --git a/activerecord/lib/active_record/connection_adapters/postgresql_adapter.rb b/activerecord/lib/active_record/connection_adapters/postgresql_adapter.rb
index 46863c41ab..5ce6765dd8 100644
--- a/activerecord/lib/active_record/connection_adapters/postgresql_adapter.rb
+++ b/activerecord/lib/active_record/connection_adapters/postgresql_adapter.rb
@@ -392,6 +392,7 @@ module ActiveRecord
         SERIALIZATION_FAILURE = "40001"
         DEADLOCK_DETECTED     = "40P01"
         LOCK_NOT_AVAILABLE    = "55P03"
+        QUERY_CANCELED        = "57014"
 
         def translate_exception(exception, message)
           return exception unless exception.respond_to?(:result)
@@ -413,6 +414,8 @@ module ActiveRecord
             Deadlocked.new(message)
           when LOCK_NOT_AVAILABLE
             TransactionTimeout.new(message)
+          when QUERY_CANCELED
+            StatementTimeout.new(message)
           else
             super
           end
diff --git a/activerecord/lib/active_record/errors.rb b/activerecord/lib/active_record/errors.rb
index 9ef3316393..f77cd23e22 100644
--- a/activerecord/lib/active_record/errors.rb
+++ b/activerecord/lib/active_record/errors.rb
@@ -335,8 +335,11 @@ module ActiveRecord
   class IrreversibleOrderError < ActiveRecordError
   end
 
-  # TransactionTimeout will be raised when lock wait timeout expires.
-  # Wait time value is set by innodb_lock_wait_timeout.
+  # TransactionTimeout will be raised when lock wait timeout exceeded.
   class TransactionTimeout < StatementInvalid
   end
+
+  # StatementTimeout will be raised when statement timeout exceeded.
+  class StatementTimeout < StatementInvalid
+  end
 end
diff --git a/activerecord/test/cases/adapters/mysql2/transaction_test.rb b/activerecord/test/cases/adapters/mysql2/transaction_test.rb
index ac9a8d9dfb..4a3a4503de 100644
--- a/activerecord/test/cases/adapters/mysql2/transaction_test.rb
+++ b/activerecord/test/cases/adapters/mysql2/transaction_test.rb
@@ -87,5 +87,34 @@ module ActiveRecord
         end
       end
     end
+
+    test "raises StatementTimeout when statement timeout exceeded" do
+      skip unless ActiveRecord::Base.connection.show_variable("max_execution_time")
+      assert_raises(ActiveRecord::StatementTimeout) do
+        s = Sample.create!(value: 1)
+        latch1 = Concurrent::CountDownLatch.new
+        latch2 = Concurrent::CountDownLatch.new
+
+        thread = Thread.new do
+          Sample.transaction do
+            Sample.lock.find(s.id)
+            latch1.count_down
+            latch2.wait
+          end
+        end
+
+        begin
+          Sample.transaction do
+            latch1.wait
+            Sample.connection.execute("SET max_execution_time = 1")
+            Sample.lock.find(s.id)
+          end
+        ensure
+          Sample.connection.execute("SET max_execution_time = DEFAULT")
+          latch2.count_down
+          thread.join
+        end
+      end
+    end
   end
 end
diff --git a/activerecord/test/cases/adapters/postgresql/transaction_test.rb b/activerecord/test/cases/adapters/postgresql/transaction_test.rb
index b6aec8e993..4d63bbce59 100644
--- a/activerecord/test/cases/adapters/postgresql/transaction_test.rb
+++ b/activerecord/test/cases/adapters/postgresql/transaction_test.rb
@@ -120,6 +120,34 @@ module ActiveRecord
       end
     end
 
+    test "raises StatementTimeout when statement timeout exceeded" do
+      assert_raises(ActiveRecord::StatementTimeout) do
+        s = Sample.create!(value: 1)
+        latch1 = Concurrent::CountDownLatch.new
+        latch2 = Concurrent::CountDownLatch.new
+
+        thread = Thread.new do
+          Sample.transaction do
+            Sample.lock.find(s.id)
+            latch1.count_down
+            latch2.wait
+          end
+        end
+
+        begin
+          Sample.transaction do
+            latch1.wait
+            Sample.connection.execute("SET statement_timeout = 1")
+            Sample.lock.find(s.id)
+          end
+        ensure
+          Sample.connection.execute("SET statement_timeout = DEFAULT")
+          latch2.count_down
+          thread.join
+        end
+      end
+    end
+
     private
 
       def with_warning_suppression