Merge branch 'query-counts' into 'master'

Track and act upon the number of executed SQL queries

See merge request gitlab-org/gitlab-ce!16466

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