From 896841bb32d907c9d7879a960ded5caf35d39420 Mon Sep 17 00:00:00 2001 From: "Marcos G. Zimmermann" 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 +++++++++++++++---- .../collection_connected_to_spec.rb | 29 +++++++++++++ spec/spec_helper.rb | 1 + spec/support/sharding_hook.rb | 21 +++++++++ 4 files changed, 86 insertions(+), 8 deletions(-) create mode 100644 spec/esse/active_record/collection_connected_to_spec.rb create mode 100644 spec/support/sharding_hook.rb 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/esse/active_record/collection_connected_to_spec.rb b/spec/esse/active_record/collection_connected_to_spec.rb new file mode 100644 index 0000000..0eabec5 --- /dev/null +++ b/spec/esse/active_record/collection_connected_to_spec.rb @@ -0,0 +1,29 @@ +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: true do + let(:collection_class) do + klass = Class.new(described_class) + klass.base_scope = -> { State } + klass.connected_to(role: :reading, shard: :default) + klass + end + + it 'uses the custom connection' do + expect(ActiveRecord::Base).to receive(:connected_to).with(role: :reading, shard: :default).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]]) + end + end +end 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' diff --git a/spec/support/sharding_hook.rb b/spec/support/sharding_hook.rb new file mode 100644 index 0000000..0dbf333 --- /dev/null +++ b/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