47 lines
1.6 KiB
Ruby
47 lines
1.6 KiB
Ruby
module Danbooru
|
|
module Extensions
|
|
module ActiveRecord
|
|
extend ActiveSupport::Concern
|
|
|
|
module ClassMethods
|
|
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.user.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
|
|
|
|
%w(execute select_value select_values select_all).each do |method_name|
|
|
define_method("#{method_name}_sql") do |sql, *params|
|
|
self.class.connection.__send__(method_name, self.class.sanitize_sql_array([sql, *params]))
|
|
end
|
|
|
|
self.class.__send__(:define_method, "#{method_name}_sql") do |sql, *params|
|
|
connection.__send__(method_name, sanitize_sql_array([sql, *params]))
|
|
end
|
|
end
|
|
end
|
|
end
|
|
end
|
|
|
|
class ActiveRecord::Base
|
|
class << self
|
|
public :sanitize_sql_array
|
|
end
|
|
|
|
include Danbooru::Extensions::ActiveRecord
|
|
end
|