DEV: Add configurable rate limit for Data Explorer API query runs (#238)

Data Explorer can run arbitrary SQL queries which can be costly for us if over-used. Because of that we want to add the ability to rate limit the query run endpoint, in particular when requested programmatically using API.

This commit introduces a rate limit to the `QueryController#run` endpoint. It heavily leans on the existing `RateLimiter` implementation, and the ability of `ApplicationController` to turn rate limit exceptions into nicely formatted JSON responses.

The rate limit (per 10 seconds) can be configured through the global setting `max_data_explorer_api_reqs_per_10_seconds`, and defaults to 2.

Handling can be configured through `max_data_explorer_api_req_mode`, and can be set to warn, block, or both warn and block. We will default to warn for now and monitor the logs for a while.

app/controllers/discourse_data_explorer/query_controller.rb

@@ -135,6 +135,8 @@ module ::DiscourseDataExplorer
     # explain - string. (Optional - pass explain=true in the request) Postgres query plan, UNIX newlines.
     # rows - array of array of strings. Results of the query. In the same order as 'columns'.
     def run
+      rate_limit_query_runs!
+
       check_xhr unless params[:download]
 
       query = Query.find(params[:id].to_i)
@@ -227,6 +229,22 @@ module ::DiscourseDataExplorer
 
     private
 
+    def rate_limit_query_runs!
+      return if !is_api? && !is_user_api?
+
+      RateLimiter.new(
+        nil,
+        "api-query-run-10-sec",
+        GlobalSetting.max_data_explorer_api_reqs_per_10_seconds,
+        10.seconds,
+      ).performed!
+    rescue RateLimiter::LimitExceeded => e
+      if GlobalSetting.max_data_explorer_api_req_mode.include?("warn")
+        Discourse.warn("Query run 10 second rate limit exceeded", query_id: params[:id])
+      end
+      raise e if GlobalSetting.max_data_explorer_api_req_mode.include?("block")
+    end
+
     def set_group
       @group = Group.find_by(name: params["group_name"])
     end

plugin.rb

@@ -38,6 +38,14 @@ after_initialize do
   require_relative "lib/discourse_data_explorer/queries"
   require_relative "lib/discourse_data_explorer/query_group_bookmarkable"
 
+  GlobalSetting.add_default(:max_data_explorer_api_reqs_per_10_seconds, 2)
+
+  # Available options:
+  #   - warn
+  #   - warn+block
+  #   - block
+  GlobalSetting.add_default(:max_data_explorer_api_req_mode, "warn")
+
   add_to_class(:guardian, :user_is_a_member_of_group?) do |group|
     return false if !current_user
     return true if current_user.admin?

spec/requests/query_controller_spec.rb

@@ -170,6 +170,97 @@ describe DiscourseDataExplorer::QueryController do
         expect(response_json["errors"].first).to match(/ValidationError/)
       end
 
+      context "when rate limited" do
+        def unlimited_request(query_id, headers = {})
+          post "/admin/plugins/explorer/queries/#{query_id}/run.json",
+               params: {
+                 params: {}.to_json,
+               },
+               headers: headers
+
+          expect(response.status).to eq(200)
+        end
+
+        def limited_request(query_id, headers = {})
+          post "/admin/plugins/explorer/queries/#{query_id}/run.json",
+               params: {
+                 params: {}.to_json,
+               },
+               headers: headers
+
+          expect(response.status).to eq(429)
+          expect(response.parsed_body["extras"]).to eq(
+            { "wait_seconds" => 9, "time_left" => "9 seconds" },
+          )
+        end
+
+        it "limits query runs from API when using block mode" do
+          global_setting :max_data_explorer_api_reqs_per_10_seconds, 1
+          global_setting :max_data_explorer_api_req_mode, "block"
+
+          RateLimiter.enable
+          RateLimiter.clear_all!
+
+          admin = Fabricate(:admin)
+          api_key = Fabricate(:api_key, user: admin)
+
+          query = make_query("SELECT 23 as my_value")
+
+          headers = { HTTP_API_KEY: api_key.key, HTTP_API_USERNAME: admin.username }
+
+          now = Time.now
+          freeze_time(now)
+
+          unlimited_request(query.id, headers)
+
+          freeze_time(now + 1.second)
+
+          limited_request(query.id, headers)
+
+          freeze_time(now + 10.seconds)
+
+          unlimited_request(query.id, headers)
+        end
+
+        it "does not limit query runs from API when using warn mode" do
+          global_setting :max_data_explorer_api_reqs_per_10_seconds, 1
+          global_setting :max_data_explorer_api_req_mode, "warn"
+
+          RateLimiter.enable
+          RateLimiter.clear_all!
+
+          admin = Fabricate(:admin)
+          api_key = Fabricate(:api_key, user: admin)
+
+          query = make_query("SELECT 23 as my_value")
+
+          headers = { HTTP_API_KEY: api_key.key, HTTP_API_USERNAME: admin.username }
+
+          freeze_time
+
+          unlimited_request(query.id, headers)
+
+          Discourse.expects(:warn).once
+
+          unlimited_request(query.id, headers)
+        end
+
+        it "does not limit query runs from UI" do
+          global_setting :max_data_explorer_api_reqs_per_10_seconds, 1
+          global_setting :max_data_explorer_api_req_mode, "block"
+
+          RateLimiter.enable
+          RateLimiter.clear_all!
+
+          query = make_query("SELECT 23 as my_value")
+
+          freeze_time
+
+          unlimited_request(query.id)
+          unlimited_request(query.id)
+        end
+      end
+
       it "doesn't allow you to modify the database #1" do
         p = create_post