debian-mirror-gitlab/lib/gitlab/query_limiting/transaction.rb

108 lines
2.9 KiB
Ruby
Raw Normal View History

2019-02-15 15:39:39 +05:30
# frozen_string_literal: true
2018-03-17 18:26:18 +05:30
module Gitlab
module QueryLimiting
class Transaction
THREAD_KEY = :__gitlab_query_counts_transaction
2021-04-29 21:17:54 +05:30
attr_accessor :count
2018-03-17 18:26:18 +05:30
# The name of the action (e.g. `UsersController#show`) that is being
# executed.
attr_accessor :action
# The maximum number of SQL queries that can be executed in a request. For
# the sake of keeping things simple we hardcode this value here, it's not
# supposed to be changed very often anyway.
THRESHOLD = 100
2021-04-17 20:07:23 +05:30
LOG_THRESHOLD = THRESHOLD * 1.5
2018-03-17 18:26:18 +05:30
# Error that is raised whenever exceeding the maximum number of queries.
ThresholdExceededError = Class.new(StandardError)
def self.current
Thread.current[THREAD_KEY]
end
# Starts a new transaction and returns it and the blocks' return value.
#
# Example:
#
# transaction, retval = Transaction.run do
# 10
# end
#
# retval # => 10
def self.run
transaction = new
Thread.current[THREAD_KEY] = transaction
[transaction, yield]
ensure
Thread.current[THREAD_KEY] = nil
end
def initialize
@action = nil
@count = 0
2021-04-17 20:07:23 +05:30
@sql_executed = []
2018-03-17 18:26:18 +05:30
end
# Sends a notification based on the number of executed SQL queries.
def act_upon_results
return unless threshold_exceeded?
error = ThresholdExceededError.new(error_message)
raise(error) if raise_error?
end
2022-07-16 23:28:13 +05:30
def increment(sql = nil)
@count += 1 if enabled? && !ignorable?(sql)
end
GEO_NODES_LOAD = 'SELECT 1 AS one FROM "geo_nodes" LIMIT 1'
LICENSES_LOAD = 'SELECT "licenses".* FROM "licenses" ORDER BY "licenses"."id"'
ATTR_INTROSPECTION = %r/SELECT .*\ba.attname\b.* (FROM|JOIN) pg_attribute a/m.freeze
# queries can be safely ignored if they are amoritized in regular usage
# (i.e. only requested occasionally and otherwise cached).
def ignorable?(sql)
return true if sql&.include?(GEO_NODES_LOAD)
return true if sql&.include?(LICENSES_LOAD)
return true if ATTR_INTROSPECTION =~ sql
false
2018-03-17 18:26:18 +05:30
end
2021-04-17 20:07:23 +05:30
def executed_sql(sql)
2022-07-16 23:28:13 +05:30
return if @count > LOG_THRESHOLD || ignorable?(sql)
@sql_executed << sql
2021-04-17 20:07:23 +05:30
end
2018-03-17 18:26:18 +05:30
def raise_error?
Rails.env.test?
end
def threshold_exceeded?
count > THRESHOLD
end
def error_message
header = 'Too many SQL queries were executed'
2019-02-15 15:39:39 +05:30
header = "#{header} in #{action}" if action
2021-04-17 20:07:23 +05:30
msg = "a maximum of #{THRESHOLD} is allowed but #{count} SQL queries were executed"
log = @sql_executed.each_with_index.map { |sql, i| "#{i}: #{sql}" }.join("\n").presence
ellipsis = '...' if @count > LOG_THRESHOLD
2018-03-17 18:26:18 +05:30
2021-04-17 20:07:23 +05:30
["#{header}: #{msg}", log, ellipsis].compact.join("\n")
2018-03-17 18:26:18 +05:30
end
2021-04-29 21:17:54 +05:30
def enabled?
::Gitlab::QueryLimiting.enabled?
end
2018-03-17 18:26:18 +05:30
end
end
end