active_record_shards 2.5.8 → 2.5.9

data/lib/active_record_shards/association_collection_connection_selection.rb

@@ -39,12 +39,7 @@ module ActiveRecordShards
           reflection = @association_collection.proxy_reflection
         end
 
-        case @which
-          when :slave
-            reflection.klass.on_slave_block { @association_collection.send(method, *args, &block) }
-          when :master
-            reflection.klass.on_master_block { @association_collection.send(method, *args, &block) }
-        end
+        reflection.klass.on_cx_switch_block(@which) { @association_collection.send(method, *args, &block) }
       end
     end
   end

data/lib/active_record_shards/connection_switcher.rb

@@ -55,12 +55,7 @@ module ActiveRecordShards
 
     def on_master_or_slave(which, &block)
       if block_given?
-        case which
-          when :master
-            on_master_block(&block)
-          when :slave
-            on_slave_block(&block)
-        end
+        on_cx_switch_block(which, &block)
       else
         MasterSlaveProxy.new(self, which)
       end
@@ -88,9 +83,13 @@ module ActiveRecordShards
     alias_method :with_slave_if, :on_slave_if
     alias_method :with_slave_unless, :on_slave_unless
 
-    def on_slave_block(&block)
+    def on_cx_switch_block(which, &block)
       old_options = current_shard_selection.options
-      switch_connection(:slave => (@disallow_slave.nil? || @disallow_slave == 0))
+      switch_to_slave = (which == :slave && (@disallow_slave.nil? || @disallow_slave == 0))
+      switch_connection(:slave => switch_to_slave)
+
+      @disallow_slave = (@disallow_slave || 0) + 1 if which == :master
+
       # setting read-only scope on ActiveRecord::Base never made any sense, anyway
       if self == ActiveRecord::Base
         yield
@@ -98,18 +97,10 @@ module ActiveRecordShards
         with_scope({:find => {:readonly => on_slave?}}, :merge, &block)
       end
     ensure
+      @disallow_slave -= 1 if which == :master
       switch_connection(old_options)
     end
 
-    # to make this re-entrant we need to keep a counter.
-    def on_master_block(&block)
-      @disallow_slave ||= 0
-      @disallow_slave += 1
-      yield
-    ensure
-      @disallow_slave -= 1
-    end
-
     # Name of the connection pool. Used by ConnectionHandler to retrieve the current connection pool.
     def connection_pool_name # :nodoc:
       name = @connection_pool_name_override || current_shard_selection.shard_name(self)

data/lib/active_record_shards/default_slave_patches.rb

@@ -24,6 +24,11 @@ module ActiveRecordShards
           self.class.quote_value(*args, &block)
         end
 
+        def reload_with_slave_off
+          self.class.on_master { reload_without_slave_off }
+        end
+        alias_method_chain :reload, :slave_off
+
         class << self
           def transaction_with_slave_off(*args, &block)
             if on_slave_by_default?

data/test/connection_switching_test.rb

@@ -375,6 +375,11 @@ class ConnectionSwitchenTest < ActiveSupport::TestCase
           assert_equal 2, count
         end
 
+        should "reload() on the master" do
+          account = Account.find(1000)
+          assert_equal 'master_name', account.reload.name
+        end
+
         should "Allow override using on_master" do
           model = Account.on_master.find(1000)
           assert_equal "master_name", model.name