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