From bc48c0aa4e3a820e63290ee8981c151f295eddc6 Mon Sep 17 00:00:00 2001
From: Ryan Senior <ryan@metabase.com>
Date: Mon, 15 Oct 2018 14:35:54 -0500
Subject: [PATCH] Changed query throttling to opt-in

This commit flips the logic for the query throttle. With this commit
query throttling is off by default. Users wanting to enable it will
use `MB_ENABLE_QUERY_THROTTLE=true`.
---
 .../middleware/add_query_throttle.clj         | 10 ++--
 .../middleware/add_query_throttle_test.clj    | 46 ++++++++++++++++---
 2 files changed, 44 insertions(+), 12 deletions(-)

diff --git a/src/metabase/query_processor/middleware/add_query_throttle.clj b/src/metabase/query_processor/middleware/add_query_throttle.clj
index 042904a3a6b76..09fd88cdf7d65 100644
--- a/src/metabase/query_processor/middleware/add_query_throttle.clj
+++ b/src/metabase/query_processor/middleware/add_query_throttle.clj
@@ -1,5 +1,5 @@
 (ns metabase.query-processor.middleware.add-query-throttle
-  "Middleware that constraints the number of concurrent queries, rejects queries by throwing an exception and
+  "Middleware that constrains the number of concurrent queries, rejects queries by throwing an exception and
   returning a 503 when we exceed our capacity"
   (:require [metabase.config :as config]
             [puppetlabs.i18n.core :refer [tru]])
@@ -44,8 +44,8 @@
       (throw-503-unavailable))))
 
 (defn maybe-add-query-throttle
-  "Adds the query throttle middleware as not as `MB_DISABLE_QUERY_THROTTLE` hasn't been set"
+  "Adds the query throttle middleware if `MB_ENABLE_QUERY_THROTTLE` has been set"
   [qp]
-  (if (config/config-bool :mb-disable-query-throttle)
-    qp
-    (throttle-queries query-semaphore qp)))
+  (if (config/config-bool :mb-enable-query-throttle)
+    (throttle-queries query-semaphore qp)
+    qp))
diff --git a/test/metabase/query_processor/middleware/add_query_throttle_test.clj b/test/metabase/query_processor/middleware/add_query_throttle_test.clj
index a132f753965b1..9abe9798be259 100644
--- a/test/metabase/query_processor/middleware/add_query_throttle_test.clj
+++ b/test/metabase/query_processor/middleware/add_query_throttle_test.clj
@@ -1,10 +1,13 @@
 (ns metabase.query-processor.middleware.add-query-throttle-test
-  (:require  [expectations :refer :all]
-             [metabase.query-processor.middleware
-              [add-query-throttle :as throttle :refer :all]
-              [catch-exceptions :as catch-exceptions]]
-             [metabase.test.data :as data]
-             [metabase.util :as u])
+  (:require [environ.core :as environ]
+            [expectations :refer :all]
+            [metabase.query-processor.middleware
+             [add-query-throttle :as throttle :refer :all]
+             [catch-exceptions :as catch-exceptions]]
+            [metabase.test
+             [data :as data]
+             [util :as tu]]
+            [metabase.util :as u])
   (:import java.util.concurrent.Semaphore))
 
 (defmacro ^:private exception-and-message [& body]
@@ -103,7 +106,7 @@
                                        (throw (Exception. "failure")))
           query-future              (future
                                       (u/ignore-exceptions
-                                       ((#'throttle/throttle-queries semaphore coordinate-then-fail) {:query "map"})))]
+                                        ((#'throttle/throttle-queries semaphore coordinate-then-fail) {:query "map"})))]
       {:beinning-permits       begin-num-permits
        :before-failure-permits (do
                                  @start-middleware-promise
@@ -112,3 +115,32 @@
                                  (deliver finish-middleware-promise true)
                                  @query-future
                                  (.availablePermits semaphore))})))
+
+;; Test the function that adds the middleware only when MB_ENABLE_QUERY_THROTTLE is set to true
+
+(defmacro ^:private with-query-throttle-value [enable-query-thottle-str & body]
+  `(with-redefs [environ/env {:mb-enable-query-throttle ~enable-query-thottle-str}]
+     ~@body))
+
+;; By default the query throttle should not be applied
+(expect
+  #'identity
+  (tu/throw-if-called throttle/throttle-queries
+    (with-query-throttle-value nil
+      (throttle/maybe-add-query-throttle #'identity))))
+
+;; The query throttle should not be applied if MB_ENABLE_QUERY_THROTTLE is false
+(expect
+  #'identity
+  (tu/throw-if-called throttle/throttle-queries
+    (with-query-throttle-value "false"
+      (throttle/maybe-add-query-throttle #'identity))))
+
+;; The query throttle should be applied if MB_ENABLE_QUERY_THROTTLE is true
+(expect
+  (let [called? (atom false)]
+    (with-redefs [throttle/throttle-queries (fn [& args]
+                                              (reset! called? true))]
+      (with-query-throttle-value "true"
+        (throttle/maybe-add-query-throttle #'identity)
+        @called?))))