Sharding support (#12)

* feat: add sharding support to the collection

* fix: fixes to the db connection of specs

* fix: fix specs for rails older than 6.x

* chore: fix rubocop offenses

CHANGELOG.md

@@ -0,0 +1,12 @@
+# Changelog
+
+All notable changes to this project will be documented in this file.
+
+The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
+
+## 0.3.7 - 2024-08-05
+* Add `connected_to` to the collection for custom connection handling
+
+## 0.0.1
+The first release of the esse-active_record plugin
+* Added: Initial implementation of the plugin

Gemfile.lock

@@ -1,7 +1,7 @@
 PATH
   remote: .
   specs:
-    esse-active_record (0.3.6)
+    esse-active_record (0.3.7)
       activerecord (>= 4.2, < 8)
       esse (>= 0.3.0)
 

ci/Gemfile.rails-5.2.lock

@@ -1,7 +1,7 @@
 PATH
   remote: ..
   specs:
-    esse-active_record (0.3.6)
+    esse-active_record (0.3.7)
       activerecord (>= 4.2, < 8)
       esse (>= 0.3.0)
 

ci/Gemfile.rails-6.0.lock

@@ -1,7 +1,7 @@
 PATH
   remote: ..
   specs:
-    esse-active_record (0.3.6)
+    esse-active_record (0.3.7)
       activerecord (>= 4.2, < 8)
       esse (>= 0.3.0)
 

ci/Gemfile.rails-6.1.lock

@@ -1,7 +1,7 @@
 PATH
   remote: ..
   specs:
-    esse-active_record (0.3.6)
+    esse-active_record (0.3.7)
       activerecord (>= 4.2, < 8)
       esse (>= 0.3.0)
 

ci/Gemfile.rails-7.0.lock

@@ -1,7 +1,7 @@
 PATH
   remote: ..
   specs:
-    esse-active_record (0.3.6)
+    esse-active_record (0.3.7)
       activerecord (>= 4.2, < 8)
       esse (>= 0.3.0)
 

ci/Gemfile.rails-7.1.lock

@@ -1,7 +1,7 @@
 PATH
   remote: ..
   specs:
-    esse-active_record (0.3.6)
+    esse-active_record (0.3.7)
       activerecord (>= 4.2, < 8)
       esse (>= 0.3.0)
 

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

lib/esse/active_record/version.rb

@@ -2,6 +2,6 @@
 
 module Esse
   module ActiveRecord
-    VERSION = '0.3.6'
+    VERSION = '0.3.7'
   end
 end

spec/esse/active_record/collection_connected_to_spec.rb

@@ -0,0 +1,66 @@
+require 'spec_helper'
+
+RSpec.describe Esse::ActiveRecord::Collection, '.connected_to' do
+  let(:collection_class) { Class.new(described_class) }
+
+  describe '.connected_to' do
+    it 'sets the connect_with' do
+      collection_class.connected_to(role: :reading, shard: :default)
+      expect(collection_class.connect_with).to eq(role: :reading, shard: :default)
+    end
+  end
+
+  describe '#each using custom connection', :sharding do
+    let(:collection_class) do
+      klass = Class.new(described_class)
+      klass.base_scope = -> { State }
+      klass.connected_to(role: :reading)
+      klass
+    end
+
+    it 'uses the custom connection' do
+      expect(ActiveRecord::Base).to receive(:connected_to).with(role: :reading).and_call_original
+      il_state = State.create!(name: 'Illinois', abbr_name: 'IL')
+
+      instance = collection_class.new
+      expect { |b| instance.each(&b) }.to yield_successive_args([il_state])
+      il_state.destroy
+    end
+  end
+
+  describe '#each_batch_ids using custom connection', :sharding do
+    let(:collection_class) do
+      klass = Class.new(described_class)
+      klass.base_scope = -> { State }
+      klass.connected_to(role: :reading)
+      klass
+    end
+
+    it 'uses the custom connection' do
+      expect(ActiveRecord::Base).to receive(:connected_to).with(role: :reading).and_call_original
+      il_state = State.create!(name: 'Illinois', abbr_name: 'IL')
+
+      instance = collection_class.new
+      expect { |b| instance.each_batch_ids(&b) }.to yield_successive_args([il_state.id])
+      il_state.destroy
+    end
+  end
+
+  describe '#count using custom connection', :sharding do
+    let(:collection_class) do
+      klass = Class.new(described_class)
+      klass.base_scope = -> { State }
+      klass.connected_to(role: :reading)
+      klass
+    end
+
+    it 'uses the custom connection' do
+      expect(ActiveRecord::Base).to receive(:connected_to).with(role: :reading).and_call_original
+      il_state = State.create!(name: 'Illinois', abbr_name: 'IL')
+
+      instance = collection_class.new
+      expect(instance.count).to eq(1)
+      il_state.destroy
+    end
+  end
+end

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'

spec/support/models.rb

@@ -1,7 +1,32 @@
 if ENV['VERBOSE']
   ActiveRecord::Base.logger = Logger.new($stdout)
 end
-ActiveRecord::Base.establish_connection(adapter: 'sqlite3', database: ':memory:')
+
+db_env = ActiveRecord::ConnectionHandling::DEFAULT_ENV.call.to_s
+db_path = '/tmp/esse-active_record.db'
+db_config = {
+  adapter: 'sqlite3',
+  database: db_path,
+}
+
+if File.exist?(db_path)
+  File.delete(db_path)
+end
+
+if ::ActiveRecord.gem_version >= Gem::Version.new('6.0.0')
+  ActiveRecord::Base.configurations = { db_env => { primary: db_config, secondary: db_config } }
+  ActiveRecord::Base.establish_connection(:primary)
+  class ApplicationRecord < ActiveRecord::Base
+    self.abstract_class = true
+    connects_to database: { writing: :primary, reading: :secondary }
+  end
+else
+  ActiveRecord::Base.configurations = { db_env => db_config }
+  ActiveRecord::Base.establish_connection(db_config)
+  class ApplicationRecord < ActiveRecord::Base
+    self.abstract_class = true
+  end
+end
 
 ActiveRecord::Schema.define do
   self.verbose = false
@@ -27,7 +52,7 @@
   end
 end
 
-class Animal < ActiveRecord::Base
+class Animal < ApplicationRecord
 end
 
 class Dog < Animal
@@ -36,11 +61,11 @@ class Dog < Animal
 class Cat < Animal
 end
 
-class State < ActiveRecord::Base
+class State < ApplicationRecord
   has_many :counties
 end
 
-class County < ActiveRecord::Base
+class County < ApplicationRecord
   belongs_to :state
 end
 

spec/support/sharding_hook.rb

@@ -0,0 +1,21 @@
+# frozen_string_literal: true
+
+module ShardingHook
+  extend ActiveSupport::Concern
+
+  included do
+    around do |example|
+      if example.metadata[:sharding] &&
+          !ActiveRecord::Base.respond_to?(:connected_to)
+        skip 'ActiveRecord::Base.connected_to is not available in this version of Rails'
+        return
+      end
+
+      example.run
+    end
+  end
+end
+
+RSpec.configure do |config|
+  config.include ShardingHook
+end