From 39743eaf194b18af622e357ef6bc8edb3db50f10 Mon Sep 17 00:00:00 2001
From: Jan Provaznik <jprovaznik@gitlab.com>
Date: Mon, 3 Sep 2018 12:01:34 +0200
Subject: [PATCH] Disable query limit counter in sidekiq

Inside Sidekiq the request is whitelisted to
avoid increasing query counter.
---
 spec/support/sidekiq.rb | 20 ++++++++++++++++++++
 1 file changed, 20 insertions(+)

diff --git a/spec/support/sidekiq.rb b/spec/support/sidekiq.rb
index d143014692d..6c4e11910d3 100644
--- a/spec/support/sidekiq.rb
+++ b/spec/support/sidekiq.rb
@@ -1,7 +1,27 @@
 require 'sidekiq/testing/inline'
 
+# If Sidekiq::Testing.inline! is used, SQL transactions done inside
+# Sidekiq worker are included in the SQL query limit (in a real
+# deployment sidekiq worker is executed separately). To avoid
+# increasing SQL limit counter, the request is marked as whitelisted
+# during Sidekiq block
+class DisableQueryLimit
+  def call(worker_instance, msg, queue)
+    transaction = Gitlab::QueryLimiting::Transaction.current
+
+    if !transaction.respond_to?(:whitelisted) || transaction.whitelisted
+      yield
+    else
+      transaction.whitelisted = true
+      yield
+      transaction.whitelisted = false
+    end
+  end
+end
+
 Sidekiq::Testing.server_middleware do |chain|
   chain.add Gitlab::SidekiqStatus::ServerMiddleware
+  chain.add DisableQueryLimit
 end
 
 RSpec.configure do |config|