slave_pools 0.1.0

data/lib/slave_pools/observer_extensions.rb

@@ -0,0 +1,19 @@
+module SlavePoolsModule
+  module ObserverExtensions
+    def self.included(base)
+      base.alias_method_chain :update, :masterdb
+    end
+    
+    # Send observed_method(object) if the method exists.
+    # currently replicating the update method instead of using the aliased method call to update_without_master
+    def update_with_masterdb(observed_method, object, &block) #:nodoc:
+      if object.class.connection.respond_to?(:with_master)
+        object.class.connection.with_master do
+          send(observed_method, object, &block) if respond_to?(observed_method) && !disabled_for?(object)
+        end
+      else
+        send(observed_method, object, &block) if respond_to?(observed_method) && !disabled_for?(object)
+      end
+    end
+  end
+end

data/lib/slave_pools/query_cache_compat.rb

@@ -0,0 +1,45 @@
+module SlavePoolsModule
+  # Implements the methods expected by the QueryCache module
+  module QueryCacheCompat
+
+    def select_all(*a, &b)
+      arel, name, binds = a
+      if query_cache_enabled && !locked?(arel)
+        # FIXME this still hits the +select_all+ method in AR connection's
+        # query_cache.rb. It'd be nice if we could avoid it somehow so
+        # +select_all+ and then +to_sql+ aren't called redundantly.
+        sql = to_sql(arel, binds)
+        @master.connection.send(:cache_sql, sql, binds) {send_to_current(:select_all, *[sql, name, binds], &b)}
+      else
+        send_to_current(:select_all, *a, &b)
+      end
+    end
+
+    def insert(*a, &b)
+      @master.connection.clear_query_cache if query_cache_enabled
+      send_to_master(:insert, *a, &b)
+    end
+
+    def update(*a, &b)
+      @master.connection.clear_query_cache if query_cache_enabled
+      send_to_master(:update, *a, &b)
+    end
+
+    def delete(*a, &b)
+      @master.connection.clear_query_cache if query_cache_enabled
+      send_to_master(:delete, *a, &b)
+    end
+
+    # Rails 3.2 changed query cacheing a little and affected slave_pools like this:
+    #
+    #   * ActiveRecord::Base.cache sets @query_cache_enabled for current connection
+    #   * ActiveRecord::QueryCache middleware (in call()) that Rails uses sets
+    #     @query_cache_enabled directly on ActiveRecord::Base.connection
+    #     (which could be master at that point)
+    #
+    # :`( So, let's just use the master connection for all query cacheing.
+    def query_cache_enabled
+      @master.connection.query_cache_enabled
+    end
+  end
+end

data/lib/slave_pools/slave_pool.rb

@@ -0,0 +1,29 @@
+module SlavePoolsModule
+  class SlavePool
+    
+    attr_accessor :name, :slaves,:pool_size, :current_index
+    
+    def initialize(name, slaves)
+      @name = name
+      @slaves = slaves
+      @pool_size = @slaves.length
+      @current_index = 0
+    end
+    
+    def current
+      @slaves[@current_index]
+    end
+    
+    def next
+      next_index! if @pool_size != 1
+      current
+    end
+    
+    protected
+    
+    def next_index!
+      @current_index = (@current_index + 1) % @pool_size
+    end    
+
+  end
+end

data/slave_pools.gemspec

@@ -0,0 +1,24 @@
+# -*- encoding: utf-8 -*-
+
+Gem::Specification.new do |s|
+  s.name = %q{slave_pools}
+  s.version = "0.1.0"
+
+  s.required_rubygems_version = Gem::Requirement.new(">= 1.2") if s.respond_to? :required_rubygems_version=
+  s.authors = ["Dan Drabik"]
+  s.description = "Connection proxy for ActiveRecord for single master / multiple slave database groups"
+  s.email = "dan@kickstarter.com"
+  s.extra_rdoc_files = ["LICENSE", "README.rdoc"]
+  s.files = ["lib/slave_pools.rb", "lib/slave_pools/active_record_extensions.rb", "lib/slave_pools/connection_proxy.rb", "lib/slave_pools/observer_extensions.rb", "lib/slave_pools/query_cache_compat.rb", "lib/slave_pools/slave_pool.rb", "LICENSE", "README.rdoc", "spec/config/database.yml", "spec/connection_proxy_spec.rb", "spec/slave_pool_spec.rb","spec/slave_pools_spec.rb", "spec/spec_helper.rb", "slave_pools.gemspec"]
+  s.has_rdoc = true
+  s.homepage = "https://github.com/kickstarter/slave_pools"
+  s.rdoc_options = ["--line-numbers", "--inline-source", "--title", "slave_pools", "--main", "README.rdoc"]
+  s.require_paths = ["lib"]
+  s.rubygems_version = %q{1.3.1}
+  s.summary = "Connection proxy for ActiveRecord for single master / multiple slave database groups"
+
+  s.add_dependency('activerecord', ["~> 3.2.12"])
+  s.add_development_dependency('mysql2', ["~> 0.3.11"])
+  s.add_development_dependency('rspec')
+  s.add_development_dependency('rake')
+end

data/spec/config/database.yml

@@ -0,0 +1,43 @@
+login: &login
+  adapter: mysql2
+  username: root
+  password:
+  host: localhost
+  encoding: utf8
+  read_timeout: 1
+
+readonly_login: &readonly_login
+  adapter: mysql2
+  username: read_only
+  password: readme
+  host: localhost
+  encoding: utf8
+  read_timeout: 1
+
+test:
+  database: test_db
+  <<: *login
+
+test_pool_secondary_name_db1:
+  database: test_db
+  <<: *readonly_login
+
+test_pool_secondary_name_db2:
+  database: test_db
+  <<: *readonly_login
+
+test_pool_secondary_name_db3:
+  database: test_db
+  <<: *readonly_login
+
+test_pool_default_name_db1:
+  database: test_db
+  <<: *readonly_login
+
+test_pool_default_name_db2:
+  database: test_db
+  <<: *readonly_login
+
+test_pool_default_name_fake_db:
+  database: fake_db_that_doesnt_exist #do not create this db, used to test the connection
+  <<: *readonly_login

data/spec/connection_proxy_spec.rb

@@ -0,0 +1,330 @@
+require File.expand_path(File.dirname(__FILE__) + '/spec_helper')
+
+describe SlavePools do
+
+  before(:each) do
+    ActiveRecord::Base.configurations = SLAVE_POOLS_SPEC_CONFIG
+    ActiveRecord::Base.establish_connection :test
+    ActiveRecord::Migration.verbose = false
+    ActiveRecord::Migration.create_table(:master_models, :force => true) {}
+    class MasterModel < ActiveRecord::Base; end
+    ActiveRecord::Migration.create_table(:foo_models, :force => true) {|t| t.string :bar}
+    class FooModel < ActiveRecord::Base; end
+    @sql = 'SELECT NOW()'
+  end
+
+  describe "standard setup" do
+    before(:each) do
+      SlavePoolsModule::ConnectionProxy.master_models = ['MasterModel']
+      SlavePoolsModule::ConnectionProxy.setup!
+      @proxy = ActiveRecord::Base.connection_proxy
+      @slave_pool_hash = @proxy.slave_pools
+      @slave_pool_array = @slave_pool_hash.values
+      @master = @proxy.master.retrieve_connection
+      # creates instance variables (@default_slave1, etc.) for each slave based on the order they appear in the slave_pool
+      # ruby 1.8.7 doesn't support ordered hashes, so we assign numbers to the slaves this way, and not the order in the yml file
+      # to prevent @default_slave1 from being different on different systems
+      ['default', 'secondary'].each do |pool_name|
+        @slave_pool_hash[pool_name.to_sym].slaves.each_with_index do |slave, i|
+          instance_variable_set("@#{pool_name}_slave#{i + 1}", slave.retrieve_connection)
+        end
+      end
+    end
+
+    it 'AR::B should respond to #connection_proxy' do
+      ActiveRecord::Base.connection_proxy.should be_kind_of(SlavePoolsModule::ConnectionProxy)
+    end
+
+    it 'FooModel#connection should return an instance of SlavePools::ConnectionProxy' do
+      FooModel.connection.should be_kind_of(SlavePoolsModule::ConnectionProxy)
+    end
+
+    it 'MasterModel#connection should not return an instance of SlavePools::ConnectionProxy' do
+      MasterModel.connection.should_not be_kind_of(SlavePoolsModule::ConnectionProxy)
+    end
+
+    it "should generate classes for each entry in the database.yml" do
+      defined?(SlavePoolsModule::DefaultDb1).should_not be_nil
+      defined?(SlavePoolsModule::DefaultDb2).should_not be_nil
+      defined?(SlavePoolsModule::SecondaryDb1).should_not be_nil
+      defined?(SlavePoolsModule::SecondaryDb2).should_not be_nil
+      defined?(SlavePoolsModule::SecondaryDb3).should_not be_nil
+    end
+
+    it "should not generate classes for an invalid DB in the database.yml" do
+      defined?(SlavePoolsModule::DefaultFakeDb).should be_nil
+    end
+
+    it 'should handle nested with_master-blocks correctly' do
+      @proxy.current.should_not == @proxy.master
+      @proxy.with_master do
+        @proxy.current.should == @proxy.master
+        @proxy.with_master do
+          @proxy.current.should == @proxy.master
+          @proxy.with_master do
+            @proxy.current.should == @proxy.master
+          end
+          @proxy.current.should == @proxy.master
+        end
+        @proxy.current.should == @proxy.master
+      end
+      @proxy.current.should_not == @proxy.master
+    end
+
+    it 'should perform transactions on the master' do
+      @master.should_receive(:select_all).exactly(5)
+      @default_slave1.should_receive(:select_all).exactly(0)
+      ActiveRecord::Base.transaction({}) do
+        5.times {@proxy.select_all(@sql)}
+      end
+    end
+
+    it 'should perform transactions on the master, and selects outside of transaction on the slave' do
+      @default_slave1.should_receive(:select_all).exactly(2) # before and after the transaction go to slaves
+      @master.should_receive(:select_all).exactly(5)
+      @proxy.select_all(@sql)
+      ActiveRecord::Base.transaction do
+        5.times {@proxy.select_all(@sql)}
+      end
+      @proxy.select_all(@sql)
+    end
+
+    it 'should not switch to the next reader on selects' do
+      @default_slave1.should_receive(:select_one).exactly(6)
+      @default_slave2.should_receive(:select_one).exactly(0)
+      6.times { @proxy.select_one(@sql) }
+    end
+
+    it '#next_slave! should switch to the next slave' do
+      @default_slave1.should_receive(:select_one).exactly(3)
+      @default_slave2.should_receive(:select_one).exactly(7)
+      3.times { @proxy.select_one(@sql) }
+      @proxy.next_slave!
+      7.times { @proxy.select_one(@sql) }
+    end
+
+    it 'should switch if next reader is explicitly called' do
+      @default_slave1.should_receive(:select_one).exactly(3)
+      @default_slave2.should_receive(:select_one).exactly(3)
+      6.times do
+        @proxy.select_one(@sql)
+        @proxy.next_slave!
+      end
+    end
+
+    it 'should not switch to the next reader when whithin a with_master-block' do
+      @master.should_receive(:select_one).twice
+      @default_slave1.should_not_receive(:select_one)
+      @default_slave2.should_not_receive(:select_one)
+      @proxy.with_master do
+        2.times { @proxy.select_one(@sql) }
+      end
+    end
+
+    it 'should send dangerous methods to the master' do
+      meths = [:insert, :update, :delete, :execute]
+      meths.each do |meth|
+        @default_slave1.stub!(meth).and_raise(RuntimeError)
+        @master.should_receive(meth).and_return(true)
+        @proxy.send(meth, @sql)
+      end
+    end
+
+    it "should send writes to the master, even if current gets called for a write" do
+      @proxy.instance_variable_set("@master_depth", 0)
+      @master.should_receive(:update).and_return(true)
+      @proxy.send(:send_to_current, :update, @sql, {})
+    end
+
+    it "should not allow master depth to get below 0" do
+      @proxy.instance_variable_set("@master_depth", -500)
+      @proxy.instance_variable_get("@master_depth").should == -500
+      @proxy.with_master {@sql}
+      @proxy.instance_variable_get("@master_depth").should == 0
+    end
+
+    it 'should dynamically generate safe methods' do
+      @proxy.should_not respond_to(:select_value)
+      @proxy.select_value(@sql)
+      @proxy.should respond_to(:select_value)
+    end
+
+    it 'should cache queries using select_all' do
+      ActiveRecord::Base.cache do
+        # next_slave will be called and switch to the SlaveDatabase2
+        @default_slave1.should_receive(:select_all).exactly(1).and_return([])
+        @default_slave2.should_not_receive(:select_all)
+        @master.should_not_receive(:select_all)
+        3.times { @proxy.select_all(@sql) }
+        @master.query_cache.keys.size.should == 1
+      end
+    end
+
+    it 'should invalidate the cache on insert, delete and update' do
+      ActiveRecord::Base.cache do
+        meths = [:insert, :update, :delete, :insert, :update]
+        meths.each do |meth|
+          @master.should_receive(meth).and_return(true)
+        end
+        @default_slave1.should_receive(:select_all).exactly(5).and_return([])
+        @default_slave2.should_receive(:select_all).exactly(0)
+        5.times do |i|
+          @proxy.select_all(@sql)
+          @proxy.select_all(@sql)
+          @master.query_cache.keys.size.should == 1
+          @proxy.send(meths[i])
+          @master.query_cache.keys.size.should == 0
+        end
+      end
+    end
+
+    describe "using querycache middleware" do
+      it 'should cache queries using select_all' do
+        mw = ActiveRecord::QueryCache.new lambda { |env|
+          @default_slave1.should_receive(:select_all).exactly(1).and_return([])
+          @default_slave2.should_not_receive(:select_all)
+          @master.should_not_receive(:select_all)
+          3.times { @proxy.select_all(@sql) }
+          @proxy.next_slave!
+          3.times { @proxy.select_all(@sql) }
+          @proxy.next_slave!
+          3.times { @proxy.select_all(@sql)}
+          @master.query_cache.keys.size.should == 1
+        }
+        mw.call({})
+      end
+
+      it 'should invalidate the cache on insert, delete and update' do
+        mw = ActiveRecord::QueryCache.new lambda { |env|
+          meths = [:insert, :update, :delete, :insert, :update]
+          meths.each do |meth|
+            @master.should_receive(meth).and_return(true)
+          end
+          @default_slave1.should_receive(:select_all).exactly(5).and_return([])
+          @default_slave2.should_receive(:select_all).exactly(0)
+          5.times do |i|
+            @proxy.select_all(@sql)
+            @proxy.select_all(@sql)
+            @master.query_cache.keys.size.should == 1
+            @proxy.send(meths[i])
+            @master.query_cache.keys.size.should == 0
+          end
+        }
+        mw.call({})
+      end
+    end
+
+    it 'should NOT rescue a non Mysql2::Error' do
+      @default_slave1.should_receive(:select_all).once.and_raise(RuntimeError.new('some error'))
+      @default_slave2.should_not_receive(:select_all)
+      @master.should_not_receive(:select_all)
+      lambda { @proxy.select_all(@sql) }.should raise_error
+    end
+
+    it 'should rescue a Mysql::Error fall back to the master' do
+      @default_slave1.should_receive(:select_all).once.and_raise(Mysql2::Error.new('connection error'))
+      @default_slave2.should_not_receive(:select_all)
+      @master.should_receive(:select_all).and_return(true)
+      lambda { @proxy.select_all(@sql) }.should_not raise_error
+    end
+
+    it 'should re-raise a Mysql::Error from a query timeout and not fall back to master' do
+      @default_slave1.should_receive(:select_all).once.and_raise(Mysql2::Error.new('Timeout waiting for a response from the last query. (waited 5 seconds)'))
+      @default_slave2.should_not_receive(:select_all)
+      @master.should_not_receive(:select_all)
+      lambda { @proxy.select_all(@sql) }.should raise_error
+    end
+
+    it 'should try to reconnect the master connection after the master has failed' do
+      @master.should_receive(:update).and_raise(RuntimeError)
+      lambda { @proxy.update(@sql) }.should raise_error
+      @master.should_receive(:reconnect!).and_return(true)
+      @master.should_receive(:insert).and_return(1)
+      @proxy.insert(@sql)
+    end
+
+    it 'should reload models from the master' do
+      foo = FooModel.create!(:bar => 'baz')
+      foo.bar = "not_saved"
+      @default_slave1.should_not_receive(:select_all)
+      @default_slave2.should_not_receive(:select_all)
+      foo.reload
+      # we didn't stub @master#select_all here, check that we actually hit the db
+      foo.bar.should == 'baz'
+    end
+
+    context "Using with_pool call" do
+
+      it "should switch to default pool if an invalid pool is specified" do
+        @default_slave1.should_receive(:select_one).exactly(3)
+        @secondary_slave1.should_not_receive(:select_one)
+        @secondary_slave2.should_not_receive(:select_one)
+        @secondary_slave3.should_not_receive(:select_one)
+        @proxy.with_pool('sfdsfsdf') do
+          3.times {@proxy.select_one(@sql)}
+        end
+      end
+
+      it "should switch to default pool if an no pool is specified" do
+        @default_slave1.should_receive(:select_one).exactly(1)
+        @proxy.with_pool do
+          @proxy.select_one(@sql)
+        end
+      end
+
+      it "should use a different pool if specified" do
+        @default_slave1.should_not_receive(:select_one)
+        @secondary_slave1.should_receive(:select_one).exactly(3)
+        @secondary_slave2.should_not_receive(:select_one)
+        @secondary_slave3.should_not_receive(:select_one)
+        @proxy.with_pool('secondary') do
+          3.times {@proxy.select_one(@sql)}
+        end
+      end
+
+      it "should different pool should use next_slave! to advance to the next DB" do
+        @default_slave1.should_not_receive(:select_one)
+        @secondary_slave1.should_receive(:select_one).exactly(2)
+        @secondary_slave2.should_receive(:select_one).exactly(1)
+        @secondary_slave3.should_receive(:select_one).exactly(1)
+        @proxy.with_pool('secondary') do
+          4.times do
+            @proxy.select_one(@sql)
+            @proxy.next_slave!
+          end
+        end
+      end
+
+      it "should switch to master if with_master is specified in an inner block" do
+        @master.should_receive(:select_one).exactly(5)
+        @default_slave1.should_receive(:select_one).exactly(0)
+        @secondary_slave1.should_receive(:select_one).exactly(0)
+        @secondary_slave2.should_receive(:select_one).exactly(0)
+        @secondary_slave3.should_receive(:select_one).exactly(0)
+        @proxy.with_pool('secondary') do
+          @proxy.with_master do
+            5.times do
+              @proxy.select_one(@sql)
+              @proxy.next_slave!
+            end
+          end
+        end
+      end
+
+      it "should switch to master if with_master is specified in an outer block (with master needs to trump with_pool)" do
+        @secondary_slave1.should_receive(:select_one).exactly(0)
+        @secondary_slave2.should_receive(:select_one).exactly(0)
+        @secondary_slave3.should_receive(:select_one).exactly(0)
+        @proxy.with_master do
+          @proxy.with_pool('secondary') do
+            5.times do
+              @proxy.select_one(@sql)
+              @proxy.next_slave!
+            end
+          end
+        end
+      end
+    end
+  end
+end
+