Enable shard selection for processes

Author: alxckn

lib/active_job/queue_adapters/solid_queue_adapter.rb

@@ -33,8 +33,10 @@ def enqueue_all(active_jobs) # :nodoc:
       private
 
       def select_shard
-        if @db_shard
-          ActiveRecord::Base.connected_to(shard: @db_shard) { yield }
+        shard = @db_shard || SolidQueue.primary_shard
+
+        if shard
+          ActiveRecord::Base.connected_to(shard: shard) { yield }
         else
           yield
         end

lib/solid_queue.rb

@@ -41,6 +41,8 @@ module SolidQueue
   mattr_accessor :clear_finished_jobs_after, default: 1.day
   mattr_accessor :default_concurrency_control_period, default: 3.minutes
 
+  mattr_accessor :primary_shard, :active_shard
+
   delegate :on_start, :on_stop, to: Supervisor
 
   def on_worker_start(...)

lib/solid_queue/processes/base.rb

@@ -6,6 +6,23 @@ class Base
       include Callbacks # Defines callbacks needed by other concerns
       include AppExecutor, Registrable, Interruptible, 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/integration/jobs_lifecycle_test.rb

@@ -30,6 +30,22 @@ class JobsLifecycleTest < ActiveSupport::TestCase
     assert_equal 2, SolidQueue::Job.finished.count
   end
 
+  test "enqueue and run jobs from different shards" do
+    AddToBufferJob.perform_later "hey"
+    ShardTwoJob.perform_later "ho"
+
+    change_active_shard_to(:queue_shard_two) do
+      @dispatcher.start
+      @worker.start
+
+      wait_for_jobs_to_finish_for(2.seconds)
+    end
+
+    assert_equal [ "ho" ], JobBuffer.values.sort
+    assert_equal 1, SolidQueue::ReadyExecution.count
+    assert_equal 1, ActiveRecord::Base.connected_to(shard: :queue_shard_two) { SolidQueue::Job.finished.count }
+  end
+
   test "enqueue and run jobs that fail without retries" do
     RaisingJob.perform_later(ExpectedTestError, "A")
     RaisingJob.perform_later(ExpectedTestError, "B")

test/test_helpers/jobs_test_helper.rb

@@ -34,4 +34,13 @@ 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