1 files changed, 15 insertions, 7 deletions

diff --git a/lib/gitlab/query_limiting.rb b/lib/gitlab/query_limiting.rb
index 5e46e26e14e..03386dca141 100644
--- a/lib/gitlab/query_limiting.rb
+++ b/lib/gitlab/query_limiting.rb
@@ -6,28 +6,36 @@ module Gitlab
     #
     # This is only enabled in development and test to ensure we don't produce
     # any errors that users of other environments can't do anything about themselves.
-    def self.enable?
+    def self.enabled_for_env?
       Rails.env.development? || Rails.env.test?
     end
 
+    def self.enabled?
+      enabled_for_env? &&
+        !Gitlab::SafeRequestStore[:query_limiting_disabled]
+    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?
-
+    # instead of blindly disabling for offending blocks of code.
+    def self.disable!(issue_url)
       unless issue_url.start_with?('https://')
         raise(
           ArgumentError,
-          'You must provide a valid issue URL in order to whitelist a block of code'
+          'You must provide a valid issue URL in order to allow a block of code'
         )
       end
 
-      Transaction&.current&.whitelisted = true
+      Gitlab::SafeRequestStore[:query_limiting_disabled] = true
+    end
+
+    # Enables query limiting for the request.
+    def self.enable!
+      Gitlab::SafeRequestStore[:query_limiting_disabled] = nil
     end
   end
 end