ar-multidb 0.4.2 → 0.6.0

data/spec/lib/multidb/balancer_spec.rb

@@ -0,0 +1,393 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Multidb::Balancer do
+  let(:config) { configuration_with_replicas }
+  let(:configuration) {
+    c = config.with_indifferent_access
+    Multidb::Configuration.new(c.except(:multidb), c[:multidb] || {})
+  }
+  let(:balancer) { global_config ? Multidb.balancer : described_class.new(configuration) }
+  let(:global_config) { false }
+
+  before do
+    ActiveRecord::Base.establish_connection(config) if global_config
+  end
+
+  describe '#initialize' do
+    subject { balancer }
+
+    context 'when configuration has no multidb config' do
+      let(:config) { configuration_with_replicas.except('multidb') }
+
+      it 'sets @candidates to have only default set of candidates' do
+        expect(subject.instance_variable_get(:@candidates).keys).to contain_exactly('default')
+      end
+
+      it 'sets @default_candidate to be the fist candidate for the default @candidates' do
+        candidates = subject.instance_variable_get(:@candidates)
+
+        expect(subject.instance_variable_get(:@default_candidate)).to eq(candidates['default'].first)
+      end
+
+      it 'sets @default_configuration to be the configuration' do
+        expect(subject.instance_variable_get(:@default_configuration)).to eq(configuration)
+      end
+
+      it 'sets fallback to false' do
+        expect(subject.fallback).to eq(false)
+      end
+
+      context 'when rails ENV is development' do
+        before do
+          stub_const('Rails', class_double('Rails', env: 'development'))
+        end
+
+        it 'sets fallback to true' do
+          expect(subject.fallback).to eq(true)
+        end
+      end
+
+      context 'when rails ENV is test' do
+        before do
+          stub_const('Rails', class_double('Rails', env: 'test'))
+        end
+
+        it 'sets fallback to true' do
+          expect(subject.fallback).to eq(true)
+        end
+      end
+    end
+
+    context 'when configuration has fallback: true' do
+      let(:config) { configuration_with_replicas.merge('multidb' => { 'fallback' => true }) }
+
+      it 'sets @candidates to have only default set of candidates' do
+        expect(subject.instance_variable_get(:@candidates).keys).to contain_exactly('default')
+      end
+
+      it 'sets @default_candidate to be the fist candidate for the default @candidates' do
+        candidates = subject.instance_variable_get(:@candidates)
+
+        expect(subject.instance_variable_get(:@default_candidate)).to eq(candidates['default'].first)
+      end
+
+      it 'sets @default_configuration to be the configuration' do
+        expect(subject.instance_variable_get(:@default_configuration)).to eq(configuration)
+      end
+
+      it 'sets fallback to true' do
+        expect(subject.fallback).to eq(true)
+      end
+    end
+
+    context 'when configuration has default multidb configuration' do
+      let(:config) {
+        extra = { multidb: { databases: {
+          default: {
+            adapter: 'sqlite3',
+            database: 'spec/test-default.sqlite'
+          }
+        } } }
+        configuration_with_replicas.merge(extra)
+      }
+
+      it 'set @candidates default to that configuration and not @default_candidate' do
+        candidates        = subject.instance_variable_get(:@candidates)
+        default_candidate = subject.instance_variable_get(:@default_candidate)
+
+        expect(candidates[:default].first).not_to eq default_candidate
+      end
+    end
+
+    context 'when configuration is nil' do
+      let(:configuration) { nil }
+
+      it 'set @candidates to an empty hash' do
+        expect(subject.instance_variable_get(:@candidates)).to eq({})
+      end
+    end
+  end
+
+  describe '#append' do
+    subject { balancer.append(appended_config) }
+
+    let(:config) { configuration_with_replicas.except('multidb') }
+
+    context 'with a basic configuration' do
+      let(:appended_config) { { replica4: { database: 'spec/test-replica4.sqlite' } } }
+
+      it 'registers the new candidate set in @candidates' do
+        expect { subject }.to change {
+          balancer.instance_variable_get(:@candidates)
+        }.to include('replica4')
+      end
+
+      it 'makes it available for use' do
+        subject
+
+        balancer.use(:replica4) do
+          expect(balancer.current_connection).to have_database 'test-replica4.sqlite'
+        end
+      end
+
+      it 'returns the connection name' do
+        subject
+
+        balancer.use(:replica4) do
+          expect(balancer.current_connection_name).to eq :replica4
+        end
+      end
+    end
+
+    context 'with an alias' do
+      let(:appended_config) {
+        {
+          replica2: {
+            database: 'spec/test-replica4.sqlite'
+          },
+          replica_alias: {
+            alias: 'replica2'
+          }
+        }
+      }
+
+      it 'aliases replica4 as replica2' do
+        subject
+
+        candidates = balancer.instance_variable_get(:@candidates)
+
+        expect(candidates['replica2']).to eq(candidates['replica_alias'])
+      end
+    end
+  end
+
+  describe '#disconnect!' do
+    subject { balancer.disconnect! }
+
+    it 'calls disconnect! on all the candidates' do
+      candidate1 = instance_double(Multidb::Candidate, disconnect!: nil)
+      candidate2 = instance_double(Multidb::Candidate, disconnect!: nil)
+
+      candidates = { 'replica1' => [candidate1], 'replica2' => [candidate2] }
+
+      balancer.instance_variable_set(:@candidates, candidates)
+
+      subject
+
+      expect(candidate1).to have_received(:disconnect!)
+      expect(candidate2).to have_received(:disconnect!)
+    end
+  end
+
+  describe '#get' do
+    subject { balancer.get(name) }
+
+    let(:name) { :replica1 }
+    let(:candidates) { balancer.instance_variable_get(:@candidates) }
+
+    context 'when there is only one candidate' do
+      it 'returns the candidate' do
+        is_expected.to eq candidates['replica1'].first
+      end
+    end
+
+    context 'when there is more than one candidate' do
+      it 'returns a random candidate' do
+        returned = Set.new
+        100.times do
+          returned << balancer.get(:replica3)
+        end
+
+        expect(returned).to match_array candidates['replica3']
+      end
+    end
+
+    context 'when the name has no configuration' do
+      let(:name) { :other }
+
+      context 'when fallback is false' do
+        it 'raises an error' do
+          expect { subject }.to raise_error(ArgumentError, /No such database connection/)
+        end
+      end
+
+      context 'when fallback is true' do
+        before do
+          balancer.fallback = true
+        end
+
+        it 'returns the default connection' do
+          is_expected.to eq candidates[:default].first
+        end
+      end
+
+      context 'when given a block' do
+        it 'yields the candidate' do
+          expect { |y|
+            balancer.get(:replica1, &y)
+          }.to yield_with_args(candidates[:replica1].first)
+        end
+      end
+    end
+  end
+
+  describe '#use' do
+    context 'with an undefined connection' do
+      it 'raises exception' do
+        expect {
+          balancer.use(:something) { nil }
+        }.to raise_error(ArgumentError)
+      end
+    end
+
+    context 'with a configured connection' do
+      let(:global_config) { true }
+
+      it 'returns default connection on :default' do
+        balancer.use(:default) do
+          expect(balancer.current_connection).to have_database 'test.sqlite'
+        end
+      end
+
+      it 'returns results instead of relation' do
+        foobar_class = Class.new(ActiveRecord::Base) do
+          self.table_name = 'foo_bars'
+        end
+
+        res = balancer.use(:replica1) do
+          ActiveRecord::Migration.verbose = false
+          ActiveRecord::Schema.define(version: 1) { create_table :foo_bars }
+          foobar_class.where(id: 42)
+        end
+
+        expect(res).to eq []
+      end
+    end
+
+    it 'returns replica connection' do
+      balancer.use(:replica1) do
+        expect(balancer.current_connection).to have_database 'test-replica1.sqlite'
+      end
+    end
+
+    it 'returns supports nested replica connection' do
+      balancer.use(:replica1) do
+        balancer.use(:replica2) do
+          expect(balancer.current_connection).to have_database 'test-replica2.sqlite'
+        end
+      end
+    end
+
+    it 'returns preserves state when nesting' do
+      balancer.use(:replica1) do
+        balancer.use(:replica2) do
+          expect(balancer.current_connection).to have_database 'test-replica2.sqlite'
+        end
+
+        expect(balancer.current_connection).to have_database 'test-replica1.sqlite'
+      end
+    end
+
+    it 'returns the parent connection for aliases' do
+      expect(balancer.use(:replica1)).not_to eq balancer.use(:replica_alias)
+      expect(balancer.use(:replica2)).to eq balancer.use(:replica_alias)
+    end
+
+    context 'when there are multiple candidates' do
+      it 'returns random candidate' do
+        names = []
+        100.times do
+          balancer.use(:replica3) do
+            list = balancer.current_connection.execute('pragma database_list')
+            names.push(File.basename(list.first&.[]('file')))
+          end
+        end
+        expect(names.uniq).to match_array %w[test-replica3-1.sqlite test-replica3-2.sqlite]
+      end
+    end
+  end
+
+  describe '#current_connection' do
+    subject { balancer.current_connection }
+
+    context 'when no alternate connection is active' do
+      let(:global_config) { true }
+
+      it 'returns main connection by default' do
+        is_expected.to have_database 'test.sqlite'
+
+        is_expected.to eq ActiveRecord::Base.retrieve_connection
+      end
+    end
+
+    context 'when an alternate connection is active' do
+      before do
+        Thread.current[:multidb] = { connection: 'a different connection' }
+      end
+
+      it 'returns the thread local connection' do
+        is_expected.to eq 'a different connection'
+      end
+    end
+  end
+
+  describe '#current_connection_name' do
+    subject { balancer.current_connection_name }
+
+    context 'when no alternate connection is active' do
+      it 'returns default connection name for default connection' do
+        is_expected.to eq :default
+      end
+    end
+
+    context 'when an alternate connection is active' do
+      before do
+        Thread.current[:multidb] = { connection_name: :replica1 }
+      end
+
+      it 'returns the thread local connection' do
+        is_expected.to eq :replica1
+      end
+    end
+  end
+
+  describe 'class delegates' do
+    let(:balancer) {
+      instance_double('Multidb::Balancer',
+                      use: nil,
+                      current_connection: nil,
+                      current_connection_name: nil,
+                      disconnect!: nil)
+    }
+
+    before do
+      Multidb.instance_variable_set(:@balancer, balancer)
+    end
+
+    it 'delegates use to the balancer' do
+      described_class.use(:name)
+
+      expect(balancer).to have_received(:use).with(:name)
+    end
+
+    it 'delegates current_connection to the balancer' do
+      described_class.current_connection
+
+      expect(balancer).to have_received(:current_connection)
+    end
+
+    it 'delegates current_connection_name to the balancer' do
+      described_class.current_connection_name
+
+      expect(balancer).to have_received(:current_connection_name)
+    end
+
+    it 'delegates disconnect! to the balancer' do
+      described_class.disconnect!
+
+      expect(balancer).to have_received(:disconnect!)
+    end
+  end
+end

data/spec/lib/multidb/candidate_spec.rb

@@ -0,0 +1,105 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Multidb::Candidate do
+  subject(:candidate) { described_class.new(name, target) }
+
+  let(:name) { :default }
+  let(:config) { configuration_with_replicas.with_indifferent_access.except(:multidb) }
+  let(:target) { config }
+
+  describe '#initialize' do
+    context 'when target is a config hash' do
+      let(:target) { config }
+
+      it 'sets the connection_handler to a new AR connection handler' do
+        handler = subject.instance_variable_get(:@connection_handler)
+        expect(handler).to an_instance_of(ActiveRecord::ConnectionAdapters::ConnectionHandler)
+      end
+
+      it 'merges the name: primary into the hash', rails: '< 6.1' do
+        handler = instance_double('ActiveRecord::ConnectionAdapters::ConnectionHandler')
+        allow(ActiveRecord::ConnectionAdapters::ConnectionHandler).to receive(:new).and_return(handler)
+        allow(handler).to receive(:establish_connection)
+
+        subject
+
+        expect(handler).to have_received(:establish_connection).with(hash_including(name: 'primary'))
+      end
+    end
+
+    context 'when target is a connection handler' do
+      let(:target) { ActiveRecord::ConnectionAdapters::ConnectionHandler.new }
+
+      it 'sets the connection_handler to the passed handler' do
+        handler = subject.instance_variable_get(:@connection_handler)
+        expect(handler).to eq(target)
+      end
+    end
+
+    context 'when target is anything else' do
+      let(:target) { 'something else' }
+
+      it 'raises an ArgumentError' do
+        expect { subject }.to raise_error(ArgumentError, /Connection handler not passed/)
+      end
+    end
+
+    it 'sets the name to the name' do
+      expect(subject.name).to eq name
+    end
+  end
+
+  describe '#connection' do
+    let(:pool) {
+      instance_double('ActiveRecord::ConnectionAdapters::ConnectionPool').tap do |o|
+        allow(o).to receive(:with_connection).and_yield('a connection')
+      end
+    }
+    let(:target) {
+      ActiveRecord::ConnectionAdapters::ConnectionHandler.new.tap do |o|
+        allow(o).to receive(:retrieve_connection)
+        allow(o).to receive(:retrieve_connection_pool).and_return(pool)
+      end
+    }
+
+    context 'when given a block' do
+      it 'calls retrieve_connection_pool' do
+        subject.connection { |_| nil }
+
+        expect(target).to have_received(:retrieve_connection_pool).with(Multidb::Candidate::SPEC_NAME)
+      end
+
+      it 'yields a connection object' do
+        expect { |y|
+          subject.connection(&y)
+        }.to yield_with_args('a connection')
+      end
+    end
+
+    context 'when not given a block' do
+      it 'calls retrieve_connection on the handler' do
+        subject.connection
+
+        expect(target).to have_received(:retrieve_connection).with(Multidb::Candidate::SPEC_NAME)
+      end
+    end
+  end
+
+  describe '#disconnect!' do
+    subject { candidate.disconnect! }
+
+    let(:target) {
+      ActiveRecord::ConnectionAdapters::ConnectionHandler.new.tap do |o|
+        allow(o).to receive(:clear_all_connections!)
+      end
+    }
+
+    it 'calls clear_all_connections! on the handler' do
+      subject
+
+      expect(target).to have_received(:clear_all_connections!)
+    end
+  end
+end

data/spec/lib/multidb/configuration_spec.rb

@@ -0,0 +1,23 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Multidb::Configuration do
+  subject { described_class.new(config.except(:multidb), config[:multidb]) }
+
+  let(:config) { configuration_with_replicas.with_indifferent_access }
+
+  describe '#initialize' do
+    it 'sets the default_handler to the AR connection handler' do
+      expect(subject.default_handler).to eq(ActiveRecord::Base.connection_handler)
+    end
+
+    it 'sets the default_adapter to the main configuration' do
+      expect(subject.default_adapter).to eq config.except(:multidb)
+    end
+
+    it 'sets the raw_configuration to the multidb configuration' do
+      expect(subject.raw_configuration).to eq config[:multidb]
+    end
+  end
+end

data/spec/lib/multidb/log_subscriber_extension_spec.rb

@@ -0,0 +1,89 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Multidb::LogSubscriberExtension do
+  before do
+    ActiveRecord::Base.establish_connection(configuration_with_replicas)
+  end
+
+  let(:klass) {
+    klass = Class.new do
+      def sql(event)
+        event
+      end
+
+      def debug(msg)
+        msg
+      end
+
+      def color(text, _color, _bold)
+        text
+      end
+    end
+
+    klass.tap do |o|
+      o.prepend described_class
+    end
+  }
+
+  let(:instance) { klass.new }
+
+  it 'prepends the extension into the ActiveRecord::LogSubscriber' do
+    expect(ActiveRecord::LogSubscriber.included_modules).to include(described_class)
+  end
+
+  describe '#sql' do
+    subject { instance.sql(event) }
+
+    let(:event) { instance_double('Event', payload: {}) }
+
+    it 'sets the :default db_name into the event payload' do
+      expect { subject }.to change { event.payload }.to include(db_name: :default)
+    end
+
+    context 'when a replica is active' do
+      it 'sets the db_name into the event payload to the replica' do
+        expect {
+          Multidb.use(:replica1) { subject }
+        }.to change { event.payload }.to include(db_name: :replica1)
+      end
+    end
+
+    context 'when there is no name returned from the balancer' do
+      before do
+        allow(Multidb.balancer).to receive(:current_connection_name)
+      end
+
+      it 'does not change the payload' do
+        expect { subject }.not_to change { event.payload }
+      end
+    end
+  end
+
+  describe '#debug' do
+    subject { instance.debug('message') }
+
+    it 'prepends the db name to the message' do
+      is_expected.to include('[DB: default]')
+    end
+
+    context 'when a replica is active' do
+      it 'prepends the replica dbname to the message' do
+        Multidb.use(:replica1) {
+          is_expected.to include('[DB: replica1')
+        }
+      end
+    end
+
+    context 'when there is no name returned from the balancer' do
+      before do
+        allow(Multidb.balancer).to receive(:current_connection_name)
+      end
+
+      it 'does not prepend to the message' do
+        is_expected.to eq('message')
+      end
+    end
+  end
+end

data/spec/lib/multidb/model_extensions_spec.rb

@@ -0,0 +1,61 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Multidb::ModelExtensions do
+  it 'includes the Multidb::Connection module into the class methods of ActiveRecord::Base' do
+    expect(ActiveRecord::Base.singleton_class.included_modules).to include Multidb::Connection
+  end
+
+  describe Multidb::Connection do
+    describe '.establish_connection' do
+      subject { ActiveRecord::Base.establish_connection(configuration_with_replicas) }
+
+      it 'initializes multidb' do
+        allow(Multidb).to receive(:init)
+
+        subject
+
+        expect(Multidb).to have_received(:init)
+      end
+    end
+
+    describe '.connection' do
+      subject { klass.connection }
+
+      let(:klass) {
+        Class.new do
+          def self.connection
+            'AR connection'
+          end
+
+          include Multidb::ModelExtensions
+        end
+      }
+
+      context 'when multidb is not initialized' do
+        it 'calls AR::Base.connection' do
+          is_expected.to eq('AR connection')
+        end
+      end
+
+      context 'when multidb is initialized' do
+        let(:balancer) { instance_double('Multidb::Balancer', current_connection: 'Multidb connection') }
+
+        before do
+          Multidb.instance_variable_set(:@balancer, balancer)
+        end
+
+        it 'calls current_connection on the balancer' do
+          subject
+
+          expect(balancer).to have_received(:current_connection)
+        end
+
+        it 'returns the balancer connection' do
+          is_expected.to eq('Multidb connection')
+        end
+      end
+    end
+  end
+end