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.

1 files changed, 65 insertions, 0 deletions

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