1 files changed, 55 insertions, 1 deletions
diff --git a/spec/support/helpers/database/database_helpers.rb b/spec/support/helpers/database/database_helpers.rb
index b8d7ea3662f..db093bcef85 100644
--- a/spec/support/helpers/database/database_helpers.rb
+++ b/spec/support/helpers/database/database_helpers.rb
@@ -5,11 +5,65 @@ module Database
     # In order to directly work with views using factories,
     # we can swapout the view for a table of identical structure.
     def swapout_view_for_table(view)
-      ActiveRecord::Base.connection.execute(<<~SQL)
+      ActiveRecord::Base.connection.execute(<<~SQL.squish)
         CREATE TABLE #{view}_copy (LIKE #{view});
         DROP VIEW #{view};
         ALTER TABLE #{view}_copy RENAME TO #{view};
       SQL
     end
+
+    # Set statement timeout temporarily.
+    # Useful when testing query timeouts.
+    #
+    # Note that this method cannot restore the timeout if a query
+    # was canceled due to e.g. a statement timeout.
+    # Refrain from using this transaction in these situations.
+    #
+    # @param timeout - Statement timeout in seconds
+    #
+    # Example:
+    #
+    #   with_statement_timeout(0.1) do
+    #     model.select('pg_sleep(0.11)')
+    #   end
+    def with_statement_timeout(timeout)
+      # Force a positive value and a minimum of 1ms for very small values.
+      timeout = (timeout * 1000).abs.ceil
+
+      raise ArgumentError, 'Using a timeout of `0` means to disable statement timeout.' if timeout == 0
+
+      previous_timeout = ActiveRecord::Base.connection
+        .exec_query('SHOW statement_timeout')[0].fetch('statement_timeout')
+
+      set_statement_timeout("#{timeout}ms")
+
+      yield
+    ensure
+      begin
+        set_statement_timeout(previous_timeout)
+      rescue ActiveRecord::StatementInvalid
+        # After a transaction was canceled/aborted due to e.g. a statement
+        # timeout commands are ignored and will raise in PG::InFailedSqlTransaction.
+        # We can safely ignore this error because the statement timeout was set
+        # for the currrent transaction which will be closed anyway.
+      end
+    end
+
+    # Set statement timeout for the current transaction.
+    #
+    # Note, that it does not restore the previous statement timeout.
+    # Use `with_statement_timeout` instead.
+    #
+    # @param timeout - Statement timeout in seconds
+    #
+    # Example:
+    #
+    #   set_statement_timeout(0.1)
+    #   model.select('pg_sleep(0.11)')
+    def set_statement_timeout(timeout)
+      ActiveRecord::Base.connection.execute(
+        format(%(SET LOCAL statement_timeout = '%s'), timeout)
+      )
+    end
   end
 end