Move with_timeout / without_timeout to ApplicationRecord.

app/logical/bulk_revert.rb

@@ -10,7 +10,7 @@ class BulkRevert
     ModAction.log("Processed bulk revert for #{constraints.inspect} by #{creator.name}")
 
     CurrentUser.scoped(creator) do
-      ActiveRecord::Base.without_timeout do
+      ApplicationRecord.without_timeout do
         find_post_versions.order("updated_at, id").each do |version|
           version.undo!
         end

app/logical/moderator/dashboard/report.rb

@@ -9,37 +9,37 @@ module Moderator
       end
 
       def artists
-        ActiveRecord::Base.without_timeout do
+        ApplicationRecord.without_timeout do
           Queries::Artist.all(min_date, max_level)
         end
       end
 
       def comments
-        ActiveRecord::Base.without_timeout do
+        ApplicationRecord.without_timeout do
           Queries::Comment.all(min_date, max_level)
         end
       end
 
       def mod_actions
-        ActiveRecord::Base.without_timeout do
+        ApplicationRecord.without_timeout do
           Queries::ModAction.all
         end
       end
 
       def notes
-        ActiveRecord::Base.without_timeout do
+        ApplicationRecord.without_timeout do
           Queries::Note.all(min_date, max_level)
         end
       end
 
       def appeals
-        ActiveRecord::Base.without_timeout do
+        ApplicationRecord.without_timeout do
           Queries::PostAppeal.all(min_date)
         end
       end
 
       def flags
-        ActiveRecord::Base.without_timeout do
+        ApplicationRecord.without_timeout do
           Queries::PostFlag.all(min_date)
         end
       end
@@ -49,19 +49,19 @@ module Moderator
       end
 
       def posts
-        ActiveRecord::Base.without_timeout do
+        ApplicationRecord.without_timeout do
           Queries::Upload.all(min_date, max_level)
         end
       end
 
       def user_feedbacks
-        ActiveRecord::Base.without_timeout do
+        ApplicationRecord.without_timeout do
           Queries::UserFeedback.all
         end
       end
 
       def wiki_pages
-        ActiveRecord::Base.without_timeout do
+        ApplicationRecord.without_timeout do
           Queries::WikiPage.all(min_date, max_level)
         end
       end

app/models/application_record.rb

@@ -43,6 +43,29 @@ class ApplicationRecord < ActiveRecord::Base
     end
   end
 
+  concerning :ActiveRecordExtensions do
+    class_methods do
+      def without_timeout
+        connection.execute("SET STATEMENT_TIMEOUT = 0") unless Rails.env == "test"
+        yield
+      ensure
+        connection.execute("SET STATEMENT_TIMEOUT = #{CurrentUser.user.try(:statement_timeout) || 3_000}") unless Rails.env == "test"
+      end
+
+      def with_timeout(n, default_value = nil, new_relic_params = {})
+        connection.execute("SET STATEMENT_TIMEOUT = #{n}") unless Rails.env == "test"
+        yield
+      rescue ::ActiveRecord::StatementInvalid => x
+        if Rails.env.production?
+          NewRelic::Agent.notice_error(x, :custom_params => new_relic_params.merge(:user_id => CurrentUser.id, :user_ip_addr => CurrentUser.ip_addr))
+        end
+        return default_value
+      ensure
+        connection.execute("SET STATEMENT_TIMEOUT = #{CurrentUser.user.try(:statement_timeout) || 3_000}") unless Rails.env == "test"
+      end
+    end
+  end
+
   concerning :PostgresExtensions do
     class_methods do
       def columns(*params)