FIX: Ensure that we revert back to default connection after running jobs.

app/jobs/base.rb

@@ -126,23 +126,23 @@ module Jobs
         begin
           exception = {}
 
-          begin
-            RailsMultisite::ConnectionManagement.establish_connection(db: db)
-            I18n.locale = SiteSetting.default_locale || "en"
-            I18n.ensure_all_loaded!
+          RailsMultisite::ConnectionManagement.with_connection(db: db) do
             begin
-              execute(opts)
+              I18n.locale = SiteSetting.default_locale || "en"
+              I18n.ensure_all_loaded!
+              begin
+                execute(opts)
+              rescue => e
+                exception[:ex] = e
+                exception[:other] = { problem_db: db }
+              end
             rescue => e
               exception[:ex] = e
+              exception[:message] = "While establishing database connection to #{db}"
               exception[:other] = { problem_db: db }
+            ensure
+              total_db_time += Instrumenter.stats.duration_ms
             end
-          rescue => e
-            exception[:ex] = e
-            exception[:message] = "While establishing database connection to #{db}"
-            exception[:other] = { problem_db: db }
-          ensure
-            ActiveRecord::Base.connection_handler.clear_active_connections!
-            total_db_time += Instrumenter.stats.duration_ms
           end
 
           exceptions << exception unless exception.empty?

lib/scheduler/defer.rb

@@ -63,18 +63,21 @@ module Scheduler
     # using non_block to match Ruby #deq
     def do_work(non_block = false)
       db, job, desc = @queue.deq(non_block)
+
       begin
-        RailsMultisite::ConnectionManagement.establish_connection(db: db) if db
-        job.call
+        if db
+          RailsMultisite::ConnectionManagement.with_connection(db: db) do
+            job.call
+          end
+        else
+          job.call
+        end
       rescue => ex
         Discourse.handle_job_exception(ex, message: "Running deferred code '#{desc}'")
       end
     rescue => ex
       Discourse.handle_job_exception(ex, message: "Processing deferred code queue")
-    ensure
-      ActiveRecord::Base.connection_handler.clear_active_connections!
     end
-
   end
 
   class Defer

lib/scheduler/manager.rb

@@ -77,9 +77,7 @@ module Scheduler
           @mutex.synchronize { info.write! }
 
           if @manager.enable_stats
-            begin
-              RailsMultisite::ConnectionManagement.establish_connection(db: "default")
-
+            RailsMultisite::ConnectionManagement.with_connection(db: "default") do
               stat = SchedulerStat.create!(
                 name: klass.to_s,
                 hostname: hostname,
@@ -87,8 +85,6 @@ module Scheduler
                 started_at: Time.zone.now,
                 live_slots_start: GC.stat[:heap_live_slots]
               )
-            ensure
-              ActiveRecord::Base.connection_handler.clear_active_connections!
             end
           end
 
@@ -106,16 +102,13 @@ module Scheduler
         info.prev_result = failed ? "FAILED" : "OK"
         info.current_owner = nil
         if stat
-          begin
-            RailsMultisite::ConnectionManagement.establish_connection(db: "default")
+          RailsMultisite::ConnectionManagement.with_connection(db: "default") do
             stat.update!(
               duration_ms: duration,
               live_slots_finish: GC.stat[:heap_live_slots],
               success: !failed,
               error: error
             )
-          ensure
-            ActiveRecord::Base.connection_handler.clear_active_connections!
           end
         end
         attempts(3) do

spec/multisite/jobs.rb

@@ -0,0 +1,13 @@
+require 'rails_helper'
+
+RSpec.describe "Running Sidekiq Jobs in Multisite" do
+  it 'should revert back to the default connection' do
+    expect(RailsMultisite::ConnectionManagement.current_db)
+      .to eq('default')
+
+    Jobs::DestroyOldDeletionStubs.new.perform({})
+
+    expect(RailsMultisite::ConnectionManagement.current_db)
+      .to eq('default')
+  end
+end