diff options
Diffstat (limited to 'lib/gitlab/query_limiting/middleware.rb')
-rw-r--r-- | lib/gitlab/query_limiting/middleware.rb | 55 |
1 files changed, 55 insertions, 0 deletions
diff --git a/lib/gitlab/query_limiting/middleware.rb b/lib/gitlab/query_limiting/middleware.rb new file mode 100644 index 00000000000..949ae79a047 --- /dev/null +++ b/lib/gitlab/query_limiting/middleware.rb @@ -0,0 +1,55 @@ +# frozen_string_literal: true + +module Gitlab + module QueryLimiting + # Middleware for reporting (or raising) when a request performs more than a + # certain amount of database queries. + class Middleware + CONTROLLER_KEY = 'action_controller.instance'.freeze + ENDPOINT_KEY = 'api.endpoint'.freeze + + def initialize(app) + @app = app + end + + def call(env) + transaction, retval = Transaction.run do + @app.call(env) + end + + transaction.action = action_name(env) + transaction.act_upon_results + + retval + end + + def action_name(env) + if env[CONTROLLER_KEY] + action_for_rails(env) + elsif env[ENDPOINT_KEY] + action_for_grape(env) + end + end + + private + + def action_for_rails(env) + controller = env[CONTROLLER_KEY] + action = "#{controller.class.name}##{controller.action_name}" + + if controller.content_type == 'text/html' + action + else + "#{action} (#{controller.content_type})" + end + end + + def action_for_grape(env) + endpoint = env[ENDPOINT_KEY] + route = endpoint.route rescue nil + + "#{route.request_method} #{route.path}" if route + end + end + end +end |