Clean up

Author: alxckn

app/models/solid_queue/record.rb

@@ -4,7 +4,16 @@ module SolidQueue
   class Record < ActiveRecord::Base
     self.abstract_class = true
 
-    connects_to(**SolidQueue.connects_to) if SolidQueue.connects_to
+    def self.connects_to_and_set_active_shard
+      connects_to(**SolidQueue.connects_to)
+
+      if SolidQueue.connects_to.key?(:shards) &&
+           SolidQueue.connects_to[:shards].key?(SolidQueue.active_shard)
+        self.default_shard = SolidQueue.active_shard
+      end
+    end
+
+    connects_to_and_set_active_shard if SolidQueue.connects_to
 
     def self.non_blocking_lock
       if SolidQueue.use_skip_locked

lib/active_job/queue_adapters/solid_queue_adapter.rb

@@ -32,13 +32,13 @@ def enqueue_all(active_jobs) # :nodoc:
 
       private
 
-      def select_shard
+      def select_shard(&block)
         shard = @db_shard || SolidQueue.primary_shard
 
         if shard
-          ActiveRecord::Base.connected_to(shard: shard) { yield }
+          ActiveRecord::Base.connected_to(shard: shard) { block.call }
         else
-          yield
+          block.call
         end
       end
     end

lib/solid_queue/engine.rb

@@ -45,5 +45,15 @@ class Engine < ::Rails::Engine
         SolidQueue::Processes::Base.include SolidQueue::Processes::OgInterruptible
       end
     end
+
+    initializer "solid_queue.shard_configuration" do
+      ActiveSupport.on_load(:solid_queue) do
+        # Record the name of the primary shard, which should be used for
+        # adapter less jobs
+        if SolidQueue.connects_to.key?(:shards) && SolidQueue.primary_shard.nil?
+          SolidQueue.primary_shard = SolidQueue.connects_to[:shards].keys.first
+        end
+      end
+    end
   end
 end

lib/solid_queue/processes/base.rb

@@ -6,23 +6,6 @@ class Base
       include Callbacks # Defines callbacks needed by other concerns
       include AppExecutor, Registrable, Procline
 
-      after_boot -> do
-        if SolidQueue.connects_to.key?(:shards)
-          # Record the name of the primary shard, which should be used for
-          # adapter less jobs
-          if SolidQueue.primary_shard.nil?
-            SolidQueue.primary_shard = SolidQueue.connects_to[:shards].keys.first
-          end
-
-          # Move active_shard to first position in connects_to[:shards] Hash to
-          # make it the default
-          if SolidQueue.active_shard.present? &&
-               SolidQueue.connects_to[:shards].key?(SolidQueue.active_shard)
-            SolidQueue::Record.default_shard = SolidQueue.active_shard
-          end
-        end
-      end
-
       attr_reader :name
 
       def initialize(*)

test/test_helper.rb

@@ -28,7 +28,7 @@ class ExpectedTestError < RuntimeError; end
 
 
 class ActiveSupport::TestCase
-  include ConfigurationTestHelper, ProcessesTestHelper, JobsTestHelper
+  include ConfigurationTestHelper, ProcessesTestHelper, JobsTestHelper, MultishardingTestHelper
 
   setup do
     @_on_thread_error = SolidQueue.on_thread_error

test/test_helpers/jobs_test_helper.rb

@@ -34,13 +34,4 @@ def assert_claimed_jobs(count = 1)
       assert_equal count, SolidQueue::ClaimedExecution.count
     end
   end
-
-  def change_active_shard_to(new_shard_name, &block)
-    old_shard_name = SolidQueue.active_shard
-    SolidQueue.active_shard = new_shard_name
-    block.call
-  ensure
-    SolidQueue.active_shard = old_shard_name
-    SolidQueue::Record.connects_to(**SolidQueue.connects_to) if SolidQueue.connects_to
-  end
 end

test/test_helpers/multisharding_test_helper.rb

@@ -0,0 +1,17 @@
+module MultishardingTestHelper
+  private
+
+  def connected_to_shard_two(&block)
+    ActiveRecord::Base.connected_to(shard: :queue_shard_two) { block.call }
+  end
+
+  def change_active_shard_to(new_shard_name, &block)
+    old_shard_name = SolidQueue.active_shard
+    SolidQueue.active_shard = new_shard_name
+    SolidQueue::Record.connects_to_and_set_active_shard
+    block.call
+  ensure
+    SolidQueue.active_shard = old_shard_name
+    SolidQueue::Record.connects_to_and_set_active_shard
+  end
+end

test/unit/multisharding_test.rb

@@ -3,26 +3,22 @@
 class MultishardingTest < ActiveSupport::TestCase
   test "jobs are enqueued in the right shard" do
     assert_difference -> { SolidQueue::Job.count }, 1 do
-      assert_difference -> do
-                          ActiveRecord::Base.connected_to(
-                            shard: :queue_shard_two
-                          ) { SolidQueue::Job.count }
-                        end,
-                        1 do
+      assert_difference -> { connected_to_shard_two { SolidQueue::Job.count } }, 1 do
         AddToBufferJob.perform_later "hey!"
         ShardTwoJob.perform_later "coucou!"
       end
     end
   end
 
+  test "jobs are enqueued in the right shard no matter the primary shard" do
+    assert_difference -> { SolidQueue::Job.count }, 1 do
+      change_active_shard_to(:queue_shard_two) { AddToBufferJob.perform_later "hey!" }
+    end
+  end
+
   test "jobs are enqueued for later in the right shard" do
     assert_difference -> { SolidQueue::ScheduledExecution.count }, 1 do
-      assert_difference -> do
-                          ActiveRecord::Base.connected_to(
-                            shard: :queue_shard_two
-                          ) { SolidQueue::ScheduledExecution.count }
-                        end,
-                        1 do
+      assert_difference -> { connected_to_shard_two { SolidQueue::ScheduledExecution.count } }, 1 do
         AddToBufferJob.set(wait: 1).perform_later "hey!"
         ShardTwoJob.set(wait: 1).perform_later "coucou!"
       end
@@ -38,12 +34,7 @@ class MultishardingTest < ActiveSupport::TestCase
     ]
 
     assert_difference -> { SolidQueue::Job.count }, 2 do
-      assert_difference -> do
-                          ActiveRecord::Base.connected_to(
-                            shard: :queue_shard_two
-                          ) { SolidQueue::Job.count }
-                        end,
-                        2 do
+      assert_difference -> { connected_to_shard_two { SolidQueue::Job.count } }, 2 do
         ActiveJob.perform_all_later(active_jobs)
       end
     end