From f60209d6585f6cdf1c6172e9b1009ed9e99b6300 Mon Sep 17 00:00:00 2001
From: "Marcos G. Zimmermann" <mgzmaster@gmail.com>
Date: Mon, 5 Aug 2024 16:00:49 -0300
Subject: [PATCH] feat: add sharding support to the collection

---
 lib/esse/active_record/collection.rb | 43 ++++++++++++++++++++++------
 spec/spec_helper.rb                  |  1 +
 2 files changed, 36 insertions(+), 8 deletions(-)

diff --git a/lib/esse/active_record/collection.rb b/lib/esse/active_record/collection.rb
index 9c027b7..81e2cbb 100644
--- a/lib/esse/active_record/collection.rb
+++ b/lib/esse/active_record/collection.rb
@@ -21,6 +21,12 @@ class Collection
       class_attribute :batch_contexts
       self.batch_contexts = {}
 
+      # Connects to a database or role (ex writing, reading, or another custom role) for the collection query
+      # @param [Symbol] role The role to connect to
+      # @param [Symbol] shard The shard to connect to
+      class_attribute :connect_with
+      self.connect_with = nil
+
       class << self
         def inspect
           return super unless self < Esse::ActiveRecord::Collection
@@ -40,6 +46,7 @@ def inherited(subclass)
 
           subclass.scopes = scopes.dup
           subclass.batch_contexts = batch_contexts.dup
+          subclass.connect_with = connect_with&.dup
         end
 
         def scope(name, proc = nil, override: false, &block)
@@ -57,6 +64,10 @@ def batch_context(name, proc = nil, override: false, &block)
 
           batch_contexts[name.to_sym] = proc
         end
+
+        def connected_to(**kwargs)
+          self.connect_with = kwargs
+        end
       end
 
       attr_reader :start, :finish, :batch_size, :params
@@ -74,23 +85,29 @@ def initialize(start: nil, finish: nil, batch_size: nil, **params)
       end
 
       def each
-        dataset.find_in_batches(**batch_options) do |rows|
-          kwargs = params.dup
-          self.class.batch_contexts.each do |name, proc|
-            kwargs[name] = proc.call(rows, **params)
+        with_connection do
+          dataset.find_in_batches(**batch_options) do |rows|
+            kwargs = params.dup
+            self.class.batch_contexts.each do |name, proc|
+              kwargs[name] = proc.call(rows, **params)
+            end
+            yield(rows, **kwargs)
           end
-          yield(rows, **kwargs)
         end
       end
 
       def each_batch_ids
-        dataset.select(:id).except(:includes, :preload, :eager_load).find_in_batches(**batch_options) do |rows|
-          yield(rows.map(&:id))
+        with_connection do
+          dataset.select(:id).except(:includes, :preload, :eager_load).find_in_batches(**batch_options) do |rows|
+            yield(rows.map(&:id))
+          end
         end
       end
 
       def count
-        dataset.except(:includes, :preload, :eager_load, :group, :order, :limit, :offset).count
+        with_connection do
+          dataset.except(:includes, :preload, :eager_load, :group, :order, :limit, :offset).count
+        end
       end
       alias_method :size, :count
 
@@ -127,6 +144,16 @@ def inspect
 
       protected
 
+      def with_connection
+        if self.class.connect_with&.any?
+          ::ActiveRecord::Base.connected_to(**self.class.connect_with) do
+            yield
+          end
+        else
+          yield
+        end
+      end
+
       def batch_options
         {
           batch_size: batch_size
diff --git a/spec/spec_helper.rb b/spec/spec_helper.rb
index 523cd17..1de0994 100644
--- a/spec/spec_helper.rb
+++ b/spec/spec_helper.rb
@@ -6,6 +6,7 @@
 require 'esse/rspec'
 
 require 'support/config_helpers'
+require 'support/sharding_hook'
 require 'support/webmock'
 require 'support/models'
 require 'pry'