Track and act upon the number of executed queries

app/controllers/admin/services_controller.rb

 class Admin::ServicesController < Admin::ApplicationController
   include ServiceParams
 
+  before_action :whitelist_query_limiting, only: [:index]
   before_action :service, only: [:edit, :update]
 
   def index
@@ -37,4 +38,8 @@ class Admin::ServicesController < Admin::ApplicationController
   def service
     @service ||= Service.where(id: params[:id], template: true).first
   end
+
+  def whitelist_query_limiting
+    Gitlab::QueryLimiting.whitelist('https://gitlab.com/gitlab-org/gitlab-ce/issues/42430')
+  end
 end

app/controllers/boards/issues_controller.rb

@@ -2,6 +2,7 @@ module Boards
   class IssuesController < Boards::ApplicationController
     include BoardsResponses
 
+    before_action :whitelist_query_limiting, only: [:index, :update]
     before_action :authorize_read_issue, only: [:index]
     before_action :authorize_create_issue, only: [:create]
     before_action :authorize_update_issue, only: [:update]
@@ -92,5 +93,10 @@ module Boards
         }
       )
     end
+
+    def whitelist_query_limiting
+      # Also see https://gitlab.com/gitlab-org/gitlab-ce/issues/42439
+      Gitlab::QueryLimiting.whitelist('https://gitlab.com/gitlab-org/gitlab-ce/issues/42428')
+    end
   end
 end

app/controllers/import/gitlab_projects_controller.rb

 class Import::GitlabProjectsController < Import::BaseController
+  before_action :whitelist_query_limiting, only: [:create]
   before_action :verify_gitlab_project_import_enabled
 
   def new
@@ -40,4 +41,8 @@ class Import::GitlabProjectsController < Import::BaseController
       :path, :namespace_id, :file
     )
   end
+
+  def whitelist_query_limiting
+    Gitlab::QueryLimiting.whitelist('https://gitlab.com/gitlab-org/gitlab-ce/issues/42437')
+  end
 end

app/controllers/projects/commits_controller.rb

@@ -4,6 +4,7 @@ class Projects::CommitsController < Projects::ApplicationController
   include ExtractsPath
   include RendersCommits
 
+  before_action :whitelist_query_limiting
   before_action :require_non_empty_project
   before_action :assign_ref_vars
   before_action :authorize_download_code!
@@ -65,4 +66,8 @@ class Projects::CommitsController < Projects::ApplicationController
     @commits = @commits.with_pipeline_status
     @commits = prepare_commits_for_rendering(@commits)
   end
+
+  def whitelist_query_limiting
+    Gitlab::QueryLimiting.whitelist('https://gitlab.com/gitlab-org/gitlab-ce/issues/42330')
+  end
 end

app/controllers/projects/cycle_analytics_controller.rb

@@ -3,6 +3,7 @@ class Projects::CycleAnalyticsController < Projects::ApplicationController
   include ActionView::Helpers::TextHelper
   include CycleAnalyticsParams
 
+  before_action :whitelist_query_limiting, only: [:show]
   before_action :authorize_read_cycle_analytics!
 
   def show
@@ -31,4 +32,8 @@ class Projects::CycleAnalyticsController < Projects::ApplicationController
       permissions: @cycle_analytics.permissions(user: current_user)
     }
   end
+
+  def whitelist_query_limiting
+    Gitlab::QueryLimiting.whitelist('https://gitlab.com/gitlab-org/gitlab-ce/issues/42671')
+  end
 end

app/controllers/projects/forks_controller.rb

@@ -2,6 +2,7 @@ class Projects::ForksController < Projects::ApplicationController
   include ContinueParams
 
   # Authorize
+  before_action :whitelist_query_limiting, only: [:create]
   before_action :require_non_empty_project
   before_action :authorize_download_code!
   before_action :authenticate_user!, only: [:new, :create]
@@ -54,4 +55,8 @@ class Projects::ForksController < Projects::ApplicationController
       render :error
     end
   end
+
+  def whitelist_query_limiting
+    Gitlab::QueryLimiting.whitelist('https://gitlab.com/gitlab-org/gitlab-ce/issues/42335')
+  end
 end

app/controllers/projects/issues_controller.rb

@@ -8,6 +8,7 @@ class Projects::IssuesController < Projects::ApplicationController
 
   prepend_before_action :authenticate_user!, only: [:new]
 
+  before_action :whitelist_query_limiting, only: [:create, :create_merge_request, :move, :bulk_update]
   before_action :check_issues_available!
   before_action :issue, except: [:index, :new, :create, :bulk_update]
   before_action :set_issuables_index, only: [:index]
@@ -247,4 +248,13 @@ class Projects::IssuesController < Projects::ApplicationController
     @finder_type = IssuesFinder
     super
   end
+
+  def whitelist_query_limiting
+    # Also see the following issues:
+    #
+    # 1. https://gitlab.com/gitlab-org/gitlab-ce/issues/42423
+    # 2. https://gitlab.com/gitlab-org/gitlab-ce/issues/42424
+    # 3. https://gitlab.com/gitlab-org/gitlab-ce/issues/42426
+    Gitlab::QueryLimiting.whitelist('https://gitlab.com/gitlab-org/gitlab-ce/issues/42422')
+  end
 end

app/controllers/projects/merge_requests/creations_controller.rb

@@ -4,6 +4,7 @@ class Projects::MergeRequests::CreationsController < Projects::MergeRequests::Ap
   include RendersCommits
 
   skip_before_action :merge_request
+  before_action :whitelist_query_limiting, only: [:create]
   before_action :authorize_create_merge_request!
   before_action :apply_diff_view_cookie!, only: [:diffs, :diff_for_path]
   before_action :build_merge_request, except: [:create]
@@ -125,4 +126,8 @@ class Projects::MergeRequests::CreationsController < Projects::MergeRequests::Ap
       @project.forked_from_project
     end
   end
+
+  def whitelist_query_limiting
+    Gitlab::QueryLimiting.whitelist('https://gitlab.com/gitlab-org/gitlab-ce/issues/42384')
+  end
 end

app/controllers/projects/merge_requests_controller.rb

@@ -7,6 +7,7 @@ class Projects::MergeRequestsController < Projects::MergeRequests::ApplicationCo
   include IssuableCollections
 
   skip_before_action :merge_request, only: [:index, :bulk_update]
+  before_action :whitelist_query_limiting, only: [:assign_related_issues, :update]
   before_action :authorize_update_issuable!, only: [:close, :edit, :update, :remove_wip, :sort]
   before_action :set_issuables_index, only: [:index]
   before_action :authenticate_user!, only: [:assign_related_issues]
@@ -339,4 +340,9 @@ class Projects::MergeRequestsController < Projects::MergeRequests::ApplicationCo
 
     access_denied! unless access_check
   end
+
+  def whitelist_query_limiting
+    # Also see https://gitlab.com/gitlab-org/gitlab-ce/issues/42441
+    Gitlab::QueryLimiting.whitelist('https://gitlab.com/gitlab-org/gitlab-ce/issues/42438')
+  end
 end

app/controllers/projects/network_controller.rb

@@ -2,6 +2,7 @@ class Projects::NetworkController < Projects::ApplicationController
   include ExtractsPath
   include ApplicationHelper
 
+  before_action :whitelist_query_limiting
   before_action :require_non_empty_project
   before_action :assign_ref_vars
   before_action :authorize_download_code!
@@ -35,4 +36,8 @@ class Projects::NetworkController < Projects::ApplicationController
     @options[:extended_sha1] = params[:extended_sha1]
     @commit = @repo.commit(@options[:extended_sha1])
   end
+
+  def whitelist_query_limiting
+    Gitlab::QueryLimiting.whitelist('https://gitlab.com/gitlab-org/gitlab-ce/issues/42333')
+  end
 end

app/controllers/projects/notes_controller.rb

@@ -2,6 +2,7 @@ class Projects::NotesController < Projects::ApplicationController
   include NotesActions
   include ToggleAwardEmoji
 
+  before_action :whitelist_query_limiting, only: [:create]
   before_action :authorize_read_note!
   before_action :authorize_create_note!, only: [:create]
   before_action :authorize_resolve_note!, only: [:resolve, :unresolve]
@@ -79,4 +80,8 @@ class Projects::NotesController < Projects::ApplicationController
 
     access_denied! unless can?(current_user, :create_note, noteable)
   end
+
+  def whitelist_query_limiting
+    Gitlab::QueryLimiting.whitelist('https://gitlab.com/gitlab-org/gitlab-ce/issues/42383')
+  end
 end

app/controllers/projects/pipelines_controller.rb

 class Projects::PipelinesController < Projects::ApplicationController
+  before_action :whitelist_query_limiting, only: [:create, :retry]
   before_action :pipeline, except: [:index, :new, :create, :charts]
   before_action :commit, only: [:show, :builds, :failures]
   before_action :authorize_read_pipeline!
@@ -166,4 +167,9 @@ class Projects::PipelinesController < Projects::ApplicationController
   def commit
     @commit ||= @pipeline.commit
   end
+
+  def whitelist_query_limiting
+    # Also see https://gitlab.com/gitlab-org/gitlab-ce/issues/42343
+    Gitlab::QueryLimiting.whitelist('https://gitlab.com/gitlab-org/gitlab-ce/issues/42339')
+  end
 end

app/controllers/projects_controller.rb

@@ -3,6 +3,7 @@ class ProjectsController < Projects::ApplicationController
   include ExtractsPath
   include PreviewMarkdown
 
+  before_action :whitelist_query_limiting, only: [:create]
   before_action :authenticate_user!, except: [:index, :show, :activity, :refs]
   before_action :redirect_git_extension, only: [:show]
   before_action :project, except: [:index, :new, :create]
@@ -405,4 +406,8 @@ class ProjectsController < Projects::ApplicationController
     #
     redirect_to request.original_url.sub(%r{\.git/?\Z}, '') if params[:format] == 'git'
   end
+
+  def whitelist_query_limiting
+    Gitlab::QueryLimiting.whitelist('https://gitlab.com/gitlab-org/gitlab-ce/issues/42440')
+  end
 end

app/controllers/registrations_controller.rb

 class RegistrationsController < Devise::RegistrationsController
   include Recaptcha::Verify
 
+  before_action :whitelist_query_limiting, only: [:destroy]
+
   def new
     redirect_to(new_user_session_path)
   end
@@ -83,4 +85,8 @@ class RegistrationsController < Devise::RegistrationsController
   def devise_mapping
     @devise_mapping ||= Devise.mappings[:user]
   end
+
+  def whitelist_query_limiting
+    Gitlab::QueryLimiting.whitelist('https://gitlab.com/gitlab-org/gitlab-ce/issues/42380')
+  end
 end

changelogs/unreleased/query-counts.yml

+---
+title: Track and act upon the number of executed queries
+merge_request:
+author:
+type: added

config/initializers/query_limiting.rb

+if Gitlab::QueryLimiting.enable?
+  require_dependency 'gitlab/query_limiting/active_support_subscriber'
+  require_dependency 'gitlab/query_limiting/transaction'
+  require_dependency 'gitlab/query_limiting/middleware'
+
+  Gitlab::Application.configure do |config|
+    config.middleware.use(Gitlab::QueryLimiting::Middleware)
+  end
+end

doc/development/README.md

@@ -75,6 +75,7 @@ comments: false
 - [Ordering table columns](ordering_table_columns.md)
 - [Verifying database capabilities](verifying_database_capabilities.md)
 - [Database Debugging and Troubleshooting](database_debugging.md)
+- [Query Count Limits](query_count_limits.md)
 
 ## Testing guides
 

doc/development/query_count_limits.md

+# Query Count Limits
+
+Each controller or API endpoint is allowed to execute up to 100 SQL queries. In
+a production environment we'll only log an error in case this threshold is
+exceeded, but in a test environment we'll raise an error instead.
+
+## Solving Failing Tests
+
+When a test fails because it executes more than 100 SQL queries there are two
+solutions to this problem:
+
+1. Reduce the number of SQL queries that are executed.
+2. Whitelist the controller or API endpoint.
+
+You should only resort to whitelisting when an existing controller or endpoint
+is to blame as in this case reducing the number of SQL queries can take a lot of
+effort. Newly added controllers and endpoints are not allowed to execute more
+than 100 SQL queries and no exceptions will be made for this rule. _If_ a large
+number of SQL queries is necessary to perform certain work it's best to have
+this work performed by Sidekiq instead of doing this directly in a web request.
+
+## Whitelisting
+
+In the event that you _have_ to whitelist a controller you'll first need to
+create an issue. This issue should (preferably in the title) mention the
+controller or endpoint and include the appropriate labels (`database`,
+`performance`, and at least a team specific label such as `Discussion`).
+
+Once the issue has been created you can whitelist the code in question. For
+Rails controllers it's best to create a `before_action` hook that runs as early
+as possible. The called method in turn should call
+`Gitlab::QueryLimiting.whitelist('issue URL here')`. For example:
+
+```ruby
+class MyController < ApplicationController
+  before_action :whitelist_query_limiting, only: [:show]
+
+  def index
+    # ...
+  end
+
+  def show
+    # ...
+  end
+
+  def whitelist_query_limiting
+    Gitlab::QueryLimiting.whitelist('https://gitlab.com/gitlab-org/...')
+  end
+end
+```
+
+By using a `before_action` you don't have to modify the controller method in
+question, reducing the likelihood of merge conflicts.
+
+For Grape API endpoints there unfortunately is not a reliable way of running a
+hook before a specific endpoint. This means that you have to add the whitelist
+call directly into the endpoint like so:
+
+```ruby
+get '/projects/:id/foo' do
+  Gitlab::QueryLimiting.whitelist('...')
+
+  # ...
+end
+```

lib/api/branches.rb

@@ -29,6 +29,8 @@ module API
         use :pagination
       end
       get ':id/repository/branches' do
+        Gitlab::QueryLimiting.whitelist('https://gitlab.com/gitlab-org/gitlab-ce/issues/42329')
+
         repository = user_project.repository
         branches = ::Kaminari.paginate_array(repository.branches.sort_by(&:name))
         merged_branch_names = repository.merged_branch_names(branches.map(&:name))

lib/api/issues.rb

@@ -161,6 +161,8 @@ module API
         use :issue_params
       end
       post ':id/issues' do
+        Gitlab::QueryLimiting.whitelist('https://gitlab.com/gitlab-org/gitlab-ce/issues/42320')
+
         authorize! :create_issue, user_project
 
         # Setting created_at time only allowed for admins and project owners
@@ -201,6 +203,8 @@ module API
                         :labels, :created_at, :due_date, :confidential, :state_event
       end
       put ':id/issues/:issue_iid' do
+        Gitlab::QueryLimiting.whitelist('https://gitlab.com/gitlab-org/gitlab-ce/issues/42322')
+
         issue = user_project.issues.find_by!(iid: params.delete(:issue_iid))
         authorize! :update_issue, issue
 
@@ -234,6 +238,8 @@ module API
         requires :to_project_id, type: Integer, desc: 'The ID of the new project'
       end
       post ':id/issues/:issue_iid/move' do
+        Gitlab::QueryLimiting.whitelist('https://gitlab.com/gitlab-org/gitlab-ce/issues/42323')
+
         issue = user_project.issues.find_by(iid: params[:issue_iid])
         not_found!('Issue') unless issue
 

lib/api/merge_requests.rb

@@ -152,6 +152,8 @@ module API
         use :optional_params
       end
       post ":id/merge_requests" do
+        Gitlab::QueryLimiting.whitelist('https://gitlab.com/gitlab-org/gitlab-ce/issues/42316')
+
         authorize! :create_merge_request, user_project
 
         mr_params = declared_params(include_missing: false)
@@ -256,6 +258,8 @@ module API
         at_least_one_of(*at_least_one_of_ce)
       end
       put ':id/merge_requests/:merge_request_iid' do
+        Gitlab::QueryLimiting.whitelist('https://gitlab.com/gitlab-org/gitlab-ce/issues/42318')
+
         merge_request = find_merge_request_with_access(params.delete(:merge_request_iid), :update_merge_request)
 
         mr_params = declared_params(include_missing: false)
@@ -283,6 +287,8 @@ module API
         optional :sha, type: String, desc: 'When present, must have the HEAD SHA of the source branch'
       end
       put ':id/merge_requests/:merge_request_iid/merge' do
+        Gitlab::QueryLimiting.whitelist('https://gitlab.com/gitlab-org/gitlab-ce/issues/42317')
+
         merge_request = find_project_merge_request(params[:merge_request_iid])
         merge_when_pipeline_succeeds = to_boolean(params[:merge_when_pipeline_succeeds])
 

lib/api/pipelines.rb

@@ -42,6 +42,8 @@ module API
         requires :ref, type: String,  desc: 'Reference'
       end
       post ':id/pipeline' do
+        Gitlab::QueryLimiting.whitelist('https://gitlab.com/gitlab-org/gitlab-ce/issues/42124')
+
         authorize! :create_pipeline, user_project
 
         new_pipeline = Ci::CreatePipelineService.new(user_project,

lib/api/projects.rb

@@ -210,6 +210,8 @@ module API
         optional :namespace, type: String, desc: 'The ID or name of the namespace that the project will be forked into'
       end
       post ':id/fork' do
+        Gitlab::QueryLimiting.whitelist('https://gitlab.com/gitlab-org/gitlab-ce/issues/42284')
+
         fork_params = declared_params(include_missing: false)
         namespace_id = fork_params[:namespace]
 

lib/api/triggers.rb

@@ -15,6 +15,8 @@ module API
         optional :variables, type: Hash, desc: 'The list of variables to be injected into build'
       end
       post ":id/(ref/:ref/)trigger/pipeline", requirements: { ref: /.+/ } do
+        Gitlab::QueryLimiting.whitelist('https://gitlab.com/gitlab-org/gitlab-ce/issues/42283')
+
         # validate variables
         params[:variables] = params[:variables].to_h
         unless params[:variables].all? { |key, value| key.is_a?(String) && value.is_a?(String) }

lib/api/users.rb

@@ -383,6 +383,8 @@ module API
         optional :hard_delete, type: Boolean, desc: "Whether to remove a user's contributions"
       end
       delete ":id" do
+        Gitlab::QueryLimiting.whitelist('https://gitlab.com/gitlab-org/gitlab-ce/issues/42279')
+
         authenticated_as_admin!
 
         user = User.find_by(id: params[:id])

lib/api/v3/branches.rb

@@ -14,6 +14,8 @@ module API
           success ::API::Entities::Branch
         end
         get ":id/repository/branches" do
+          Gitlab::QueryLimiting.whitelist('https://gitlab.com/gitlab-org/gitlab-ce/issues/42276')
+
           repository = user_project.repository
           branches = repository.branches.sort_by(&:name)
           merged_branch_names = repository.merged_branch_names(branches.map(&:name))

lib/api/v3/issues.rb

@@ -134,6 +134,8 @@ module API
           use :issue_params
         end
         post ':id/issues' do
+          Gitlab::QueryLimiting.whitelist('https://gitlab.com/gitlab-org/gitlab-ce/issues/42131')
+
           # Setting created_at time only allowed for admins and project owners
           unless current_user.admin? || user_project.owner == current_user
             params.delete(:created_at)
@@ -169,6 +171,8 @@ module API
                           :labels, :created_at, :due_date, :confidential, :state_event
         end
         put ':id/issues/:issue_id' do
+          Gitlab::QueryLimiting.whitelist('https://gitlab.com/gitlab-org/gitlab-ce/issues/42132')
+
           issue = user_project.issues.find(params.delete(:issue_id))
           authorize! :update_issue, issue
 
@@ -201,6 +205,8 @@ module API
           requires :to_project_id, type: Integer, desc: 'The ID of the new project'
         end
         post ':id/issues/:issue_id/move' do
+          Gitlab::QueryLimiting.whitelist('https://gitlab.com/gitlab-org/gitlab-ce/issues/42133')
+
           issue = user_project.issues.find_by(id: params[:issue_id])
           not_found!('Issue') unless issue
 

lib/api/v3/merge_requests.rb

@@ -91,6 +91,8 @@ module API
           use :optional_params
         end
         post ":id/merge_requests" do
+          Gitlab::QueryLimiting.whitelist('https://gitlab.com/gitlab-org/gitlab-ce/issues/42126')
+
           authorize! :create_merge_request, user_project
 
           mr_params = declared_params(include_missing: false)
@@ -167,6 +169,8 @@ module API
                             :remove_source_branch
           end
           put path do
+            Gitlab::QueryLimiting.whitelist('https://gitlab.com/gitlab-org/gitlab-ce/issues/42127')
+
             merge_request = find_merge_request_with_access(params.delete(:merge_request_id), :update_merge_request)
 
             mr_params = declared_params(include_missing: false)

lib/api/v3/pipelines.rb

@@ -19,6 +19,8 @@ module API
                               desc: 'Either running, branches, or tags'
         end
         get ':id/pipelines' do
+          Gitlab::QueryLimiting.whitelist('https://gitlab.com/gitlab-org/gitlab-ce/issues/42123')
+
           authorize! :read_pipeline, user_project
 
           pipelines = PipelinesFinder.new(user_project, scope: params[:scope]).execute

lib/api/v3/triggers.rb

@@ -16,6 +16,8 @@ module API
           optional :variables, type: Hash, desc: 'The list of variables to be injected into build'
         end
         post ":id/(ref/:ref/)trigger/builds", requirements: { ref: /.+/ } do
+          Gitlab::QueryLimiting.whitelist('https://gitlab.com/gitlab-org/gitlab-ce/issues/42121')
+
           # validate variables
           params[:variables] = params[:variables].to_h
           unless params[:variables].all? { |key, value| key.is_a?(String) && value.is_a?(String) }

lib/gitlab/query_limiting.rb

+module Gitlab
+  module QueryLimiting
+    # Returns true if we should enable tracking of query counts.
+    #
+    # This is only enabled in production/staging if we're running on GitLab.com.
+    # This ensures we don't produce any errors that users can't do anything
+    # about themselves.
+    def self.enable?
+      Gitlab.com? || Rails.env.development? || Rails.env.test?
+    end
+
+    # Allows the current request to execute any number of SQL queries.
+    #
+    # This method should _only_ be used when there's a corresponding issue to
+    # reduce the number of queries.
+    #
+    # The issue URL is only meant to push developers into creating an issue
+    # instead of blindly whitelisting offending blocks of code.
+    def self.whitelist(issue_url)
+      return unless enable_whitelist?
+
+      unless issue_url.start_with?('https://')
+        raise(
+          ArgumentError,
+          'You must provide a valid issue URL in order to whitelist a block of code'
+        )
+      end
+
+      Transaction&.current&.whitelisted = true
+    end
+
+    def self.enable_whitelist?
+      Rails.env.development? || Rails.env.test?
+    end
+  end
+end

lib/gitlab/query_limiting/active_support_subscriber.rb

+module Gitlab
+  module QueryLimiting
+    class ActiveSupportSubscriber < ActiveSupport::Subscriber
+      attach_to :active_record
+
+      def sql(*)
+        Transaction.current&.increment
+      end
+    end
+  end
+end

lib/gitlab/query_limiting/middleware.rb

+# 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

lib/gitlab/query_limiting/transaction.rb

+module Gitlab
+  module QueryLimiting
+    class Transaction
+      THREAD_KEY = :__gitlab_query_counts_transaction
+
+      attr_accessor :count, :whitelisted
+
+      # 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
+
+      # 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
+        @whitelisted = false
+      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)
+
+        if raise_error?
+          raise(error)
+        else
+          # Raven automatically logs to the Rails log if disabled, thus we don't
+          # need to manually log anything in case Sentry support is not enabled.
+          Raven.capture_exception(error)
+        end
+      end
+
+      def increment
+        @count += 1 unless whitelisted
+      end
+
+      def raise_error?
+        Rails.env.test?
+      end
+
+      def threshold_exceeded?
+        count > THRESHOLD
+      end
+
+      def error_message
+        header = 'Too many SQL queries were executed'
+        header += " in #{action}" if action
+
+        "#{header}: a maximum of #{THRESHOLD} is allowed but #{count} SQL queries were executed"
+      end
+    end
+  end
+end

spec/lib/gitlab/query_limiting/active_support_subscriber_spec.rb

+require 'spec_helper'
+
+describe Gitlab::QueryLimiting::ActiveSupportSubscriber do
+  describe '#sql' do
+    it 'increments the number of executed SQL queries' do
+      transaction = double(:transaction)
+
+      allow(Gitlab::QueryLimiting::Transaction)
+        .to receive(:current)
+        .and_return(transaction)
+
+      expect(transaction)
+        .to receive(:increment)
+        .at_least(:once)
+
+      User.count
+    end
+  end
+end

spec/lib/gitlab/query_limiting/middleware_spec.rb

+require 'spec_helper'
+
+describe Gitlab::QueryLimiting::Middleware do
+  describe '#call' do
+    it 'runs the application with query limiting in place' do
+      middleware = described_class.new(-> (env) { env })
+
+      expect_any_instance_of(Gitlab::QueryLimiting::Transaction)
+        .to receive(:act_upon_results)
+
+      expect(middleware.call({ number: 10 }))
+        .to eq({ number: 10 })
+    end
+  end
+
+  describe '#action_name' do
+    let(:middleware) { described_class.new(-> (env) { env }) }
+
+    context 'using a Rails request' do
+      it 'returns the name of the controller and action' do
+        env = {
+          described_class::CONTROLLER_KEY => double(
+            :controller,
+            action_name: 'show',
+            class: double(:class, name: 'UsersController'),
+            content_type: 'text/html'
+          )
+        }
+
+        expect(middleware.action_name(env)).to eq('UsersController#show')
+      end
+
+      it 'includes the content type if this is not text/html' do
+        env = {
+          described_class::CONTROLLER_KEY => double(
+            :controller,
+            action_name: 'show',
+            class: double(:class, name: 'UsersController'),
+            content_type: 'application/json'
+          )
+        }
+
+        expect(middleware.action_name(env))
+          .to eq('UsersController#show (application/json)')
+      end
+    end
+
+    context 'using a Grape API request' do
+      it 'returns the name of the request method and endpoint path' do
+        env = {
+          described_class::ENDPOINT_KEY => double(
+            :endpoint,
+            route: double(:route, request_method: 'GET', path: '/foo')
+          )
+        }
+
+        expect(middleware.action_name(env)).to eq('GET /foo')
+      end
+
+      it 'returns nil if the route can not be retrieved' do
+        endpoint = double(:endpoint)
+        env = { described_class::ENDPOINT_KEY => endpoint }
+
+        allow(endpoint)
+          .to receive(:route)
+          .and_raise(RuntimeError)
+
+        expect(middleware.action_name(env)).to be_nil
+      end
+    end
+  end
+end

spec/lib/gitlab/query_limiting/transaction_spec.rb

+require 'spec_helper'
+
+describe Gitlab::QueryLimiting::Transaction do
+  after do
+    Thread.current[described_class::THREAD_KEY] = nil
+  end
+
+  describe '.current' do
+    it 'returns nil when there is no transaction' do
+      expect(described_class.current).to be_nil
+    end
+
+    it 'returns the transaction when present' do
+      Thread.current[described_class::THREAD_KEY] = described_class.new
+
+      expect(described_class.current).to be_an_instance_of(described_class)
+    end
+  end
+
+  describe '.run' do
+    it 'runs a transaction and returns it and its return value' do
+      trans, ret = described_class.run do
+        10
+      end
+
+      expect(trans).to be_an_instance_of(described_class)
+      expect(ret).to eq(10)
+    end
+
+    it 'removes the transaction from the current thread upon completion' do
+      described_class.run do
+        10
+      end
+
+      expect(Thread.current[described_class::THREAD_KEY]).to be_nil
+    end
+  end
+
+  describe '#act_upon_results' do
+    context 'when the query threshold is not exceeded' do
+      it 'does nothing' do
+        trans = described_class.new
+
+        expect(trans).not_to receive(:raise)
+
+        trans.act_upon_results
+      end
+    end
+
+    context 'when the query threshold is exceeded' do
+      let(:transaction) do
+        trans = described_class.new
+        trans.count = described_class::THRESHOLD + 1
+
+        trans
+      end
+
+      it 'raises an error when this is enabled' do
+        expect { transaction.act_upon_results }
+          .to raise_error(described_class::ThresholdExceededError)
+      end
+
+      it 'reports the error in Sentry if raising an error is disabled' do
+        expect(transaction)
+          .to receive(:raise_error?)
+          .and_return(false)
+
+        expect(Raven)
+          .to receive(:capture_exception)
+          .with(an_instance_of(described_class::ThresholdExceededError))
+
+        transaction.act_upon_results
+      end
+    end
+  end
+
+  describe '#increment' do
+    it 'increments the number of executed queries' do
+      transaction = described_class.new
+
+      expect(transaction.count).to be_zero
+
+      transaction.increment
+
+      expect(transaction.count).to eq(1)
+    end
+  end
+
+  describe '#raise_error?' do
+    it 'returns true in a test environment' do
+      transaction = described_class.new
+
+      expect(transaction.raise_error?).to eq(true)
+    end
+
+    it 'returns false in a production environment' do
+      transaction = described_class.new
+
+      expect(Rails.env)
+        .to receive(:test?)
+        .and_return(false)
+
+      expect(transaction.raise_error?).to eq(false)
+    end
+  end
+
+  describe '#threshold_exceeded?' do
+    it 'returns false when the threshold is not exceeded' do
+      transaction = described_class.new
+
+      expect(transaction.threshold_exceeded?).to eq(false)
+    end
+
+    it 'returns true when the threshold is exceeded' do
+      transaction = described_class.new
+      transaction.count = described_class::THRESHOLD + 1
+
+      expect(transaction.threshold_exceeded?).to eq(true)
+    end
+  end
+
+  describe '#error_message' do
+    it 'returns the error message to display when the threshold is exceeded' do
+      transaction = described_class.new
+      transaction.count = max = described_class::THRESHOLD
+
+      expect(transaction.error_message).to eq(
+        "Too many SQL queries were executed: a maximum of #{max} " \
+        "is allowed but #{max} SQL queries were executed"
+      )
+    end
+
+    it 'includes the action name in the error message when present' do
+      transaction = described_class.new
+      transaction.count = max = described_class::THRESHOLD
+      transaction.action = 'UsersController#show'
+
+      expect(transaction.error_message).to eq(
+        "Too many SQL queries were executed in UsersController#show: " \
+        "a maximum of #{max} is allowed but #{max} SQL queries were executed"
+      )
+    end
+  end
+end

spec/lib/gitlab/query_limiting_spec.rb

+require 'spec_helper'
+
+describe Gitlab::QueryLimiting do
+  describe '.enable?' do
+    it 'returns true in a test environment' do
+      expect(described_class.enable?).to eq(true)
+    end
+
+    it 'returns true in a development environment' do
+      allow(Rails.env).to receive(:development?).and_return(true)
+
+      expect(described_class.enable?).to eq(true)
+    end
+
+    it 'returns true on GitLab.com' do
+      allow(Gitlab).to receive(:com?).and_return(true)
+
+      expect(described_class.enable?).to eq(true)
+    end
+
+    it 'returns true in a non GitLab.com' do
+      expect(Gitlab).to receive(:com?).and_return(false)
+      expect(Rails.env).to receive(:development?).and_return(false)
+      expect(Rails.env).to receive(:test?).and_return(false)
+
+      expect(described_class.enable?).to eq(false)
+    end
+  end
+
+  describe '.whitelist' do
+    it 'raises ArgumentError when an invalid issue URL is given' do
+      expect { described_class.whitelist('foo') }
+        .to raise_error(ArgumentError)
+    end
+
+    context 'without a transaction' do
+      it 'does nothing' do
+        expect { described_class.whitelist('https://example.com') }
+          .not_to raise_error
+      end
+    end
+
+    context 'with a transaction' do
+      let(:transaction) { Gitlab::QueryLimiting::Transaction.new }
+
+      before do
+        allow(Gitlab::QueryLimiting::Transaction)
+          .to receive(:current)
+          .and_return(transaction)
+      end
+
+      it 'does not increment the number of SQL queries executed in the block' do
+        before = transaction.count
+
+        described_class.whitelist('https://example.com')
+
+        2.times do
+          User.count
+        end
+
+        expect(transaction.count).to eq(before)
+      end
+    end
+  end
+end