Disable query limit counter in sidekiq

Inside Sidekiq the request is whitelisted to
avoid increasing query counter.

1 files changed, 20 insertions, 0 deletions

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|