diff --git a/lib/multi_db/connection_proxy.rb b/lib/multi_db/connection_proxy.rb index 4b4e0a2..a86d38b 100644 --- a/lib/multi_db/connection_proxy.rb +++ b/lib/multi_db/connection_proxy.rb @@ -59,6 +59,14 @@ def setup!(scheduler = Scheduler) master.logger.info("** multi_db with master and #{slaves.length} slave#{"s" if slaves.length > 1} loaded.") end + def with_master + if defined?(ActiveRecord::Base.connection_proxy) + ActiveRecord::Base.connection_proxy.with_master{ yield } + else + yield + end + end + protected # Slave entries in the database.yml must be named like this diff --git a/spec/connection_proxy_spec.rb b/spec/connection_proxy_spec.rb index 6c40bd8..5d8fbf2 100644 --- a/spec/connection_proxy_spec.rb +++ b/spec/connection_proxy_spec.rb @@ -70,6 +70,14 @@ class FooModel < ActiveRecord::Base; end @proxy.select_all(@sql) end end + + it 'should force connection on the master' do + @master.should_receive(:select_all).exactly(1) # makes sure the first one goes to a slave + @proxy.select_all(@sql) + MultiDb::ConnectionProxy.with_master do + @proxy.select_all(@sql) + end + end it 'should switch to the next reader on selects' do @slave1.should_receive(:select_one).exactly(2)