Track and act upon the number of executed queriesquery-counts

This ensures that we have more visibility in the number of SQL queries
that are executed in web requests. The current threshold is hardcoded to
100 as we will rarely (maybe once or twice) change it.

In production and development we use Sentry if enabled, in the test
environment we raise an error. This feature is also only enabled in
production/staging when running on GitLab.com as it's not very useful to
other users.

4 files changed, 300 insertions, 0 deletions

diff --git a/spec/lib/gitlab/query_limiting/active_support_subscriber_spec.rb b/spec/lib/gitlab/query_limiting/active_support_subscriber_spec.rb
new file mode 100644
index 00000000000..b49bc5c328c
--- /dev/null
+++ b/spec/lib/gitlab/query_limiting/active_support_subscriber_spec.rb
@@ -0,0 +1,19 @@
+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
diff --git a/spec/lib/gitlab/query_limiting/middleware_spec.rb b/spec/lib/gitlab/query_limiting/middleware_spec.rb
new file mode 100644
index 00000000000..a04bcdecb4b
--- /dev/null
+++ b/spec/lib/gitlab/query_limiting/middleware_spec.rb
@@ -0,0 +1,72 @@
+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
diff --git a/spec/lib/gitlab/query_limiting/transaction_spec.rb b/spec/lib/gitlab/query_limiting/transaction_spec.rb
new file mode 100644
index 00000000000..b4231fcd0fa
--- /dev/null
+++ b/spec/lib/gitlab/query_limiting/transaction_spec.rb
@@ -0,0 +1,144 @@
+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
diff --git a/spec/lib/gitlab/query_limiting_spec.rb b/spec/lib/gitlab/query_limiting_spec.rb
new file mode 100644
index 00000000000..2eddab0b8c3
--- /dev/null
+++ b/spec/lib/gitlab/query_limiting_spec.rb
@@ -0,0 +1,65 @@
+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