Seems to be working somewhat, huzzah!

Rakefile

@@ -8,3 +8,27 @@ Rake::TestTask.new do |t|
 end
 
 task :default => :test
+
+task :create_db do
+  require 'rubygems'
+  require 'active_record'
+  ENV['RAILS_ENV'] = 'test'
+
+  ActiveRecord::Base.configurations = { 'test' => { :adapter => 'mysql', :host => 'localhost', :database => 'manage_test' } }
+  ActiveRecord::Base.establish_connection 'test'
+
+  def using_connection(&block)
+    ActiveRecord::Base.connection.instance_eval(&block)
+  end
+
+  databases = %w( vr_austin_master vr_austin_slave vr_dallas_master vr_dallas_slave )
+  databases.each do |db|
+    using_connection {
+      execute "drop database #{db}"
+      execute "create database #{db}"
+      execute "use #{db}"
+      execute "create table the_whole_burritos (id integer not null auto_increment, name varchar(30) not null, primary key(id))"
+      execute "insert into the_whole_burritos (id, name) values (1, '#{db}')"
+    }  
+  end
+end

lib/data_fabric.rb

@@ -77,7 +77,8 @@ def self.ensure_setup
   # Class methods injected into ActiveRecord::Base
   module ClassMethods
     def connection_topology(options)
-      ActiveRecord::Base.active_connections[name] = DataFabric::ConnectionProxy.new(self, options)
+      proxy = DataFabric::ConnectionProxy.new(self, options)
+      ActiveRecord::Base.active_connections[name] = proxy
     end
   end
 
@@ -92,17 +93,54 @@ def to_s
 
   class ConnectionProxy
     def initialize(model_class, options)
-      @model_class = model_class
-      @model_class.send :include, ActiveRecordConnectionMethods
-
+      @model_class = model_class      
       @replicated = options[:replicated]
       @shard_group = options[:shard_by]
       @prefix = options[:prefix]
       @current_role = 'slave' if @replicated
       @current_connection_name_builder = connection_name_builder
       @cached_connection = nil
-      @last_connection_name = nil
+      @current_connection_name = nil
+
+      @model_class.send :include, ActiveRecordConnectionMethods if @replicated
+    end
+
+    delegate :insert, :update, :delete, :create_table, :rename_table, :drop_table, :add_column, :remove_column, 
+      :change_column, :change_column_default, :rename_column, :add_index, :remove_index, :initialize_schema_information,
+      :dump_schema_information, :to => :master
+
+    def transaction(start_db_transaction = true, &block)
+      with_master { raw_connection.transaction(start_db_transaction, &block) }
+    end
+
+    def method_missing(method, *args, &block)
+      unless @cached_connection
+        raw_connection
+      end
+      @cached_connection.send(method, *args, &block) 
+    end
+
+    def connection_name
+      @current_connection_name_builder.join('_')
+    end
+
+    def disconnect!
+      @cached_connection.disconnect! if @cached_connection
+      @cached_connection = nil
+    end
+
+    def verify!(arg)
+      @cached_connection.verify!(0) if @cached_connection
+    end
+
+    def with_master
+      set_role('master')
+      yield
+    ensure
+      set_role('slave')
     end
+
+    private
 
     def connection_name_builder
       clauses = []
@@ -114,10 +152,6 @@ def connection_name_builder
       clauses
     end
 
-    def connection_name
-      @current_connection_name_builder.join('_')
-    end
-
     def raw_connection
       conn_name = connection_name
       unless already_connected_to? conn_name 
@@ -126,19 +160,21 @@ def raw_connection
           raise ArgumentError, "Unknown database config: #{conn_name}" unless config
           @model_class.establish_connection config
           conn = @model_class.connection
-          conn.reconnect! unless conn.active?
+#          conn.verify! 0
           conn
         end
+        @model_class.active_connections[@model_class.name] = self
       end
       @cached_connection
     end
 
     def already_connected_to?(conn_name)
-      conn_name == @last_connection_name and @cached_connection
+      conn_name == @current_connection_name and @cached_connection
     end
 
     def set_role(role)
-      if @replicated
+      if @replicated and @current_role != role
+#        puts "Role: #{role}"
         @current_role = role
         @cached_connection = nil
       end
@@ -150,25 +186,6 @@ def master
     ensure
       set_role('slave')
     end
-
-    def with_master
-      set_role('master')
-      yield
-    ensure
-      set_role('slave')
-    end
-
-    delegate :insert, :update, :delete, :create_table, :rename_table, :drop_table, :add_column, :remove_column, 
-      :change_column, :change_column_default, :rename_column, :add_index, :remove_index, :initialize_schema_information,
-      :dump_schema_information, :to => :master
-
-    def transaction(start_db_transaction = true, &block)
-      with_master { raw_connection.transaction(start_db_transaction, &block) }
-    end
-
-    def method_missing(method, *args, &block)
-      raw_connection.send(method, *args, &block)
-    end
   end
 
   module ActiveRecordConnectionMethods

test/connection_test.rb

@@ -15,11 +15,28 @@ class TheWholeEnchilada < ActiveRecord::Base
 
 class AdapterMock < ActiveRecord::ConnectionAdapters::AbstractAdapter
   # Minimum required to perform a find with no results
-  def columns(table_name, name=nil)
-    []
-  end
-  def select(sql, name = nil)
-    []
+   def columns(table_name, name=nil)
+     []
+   end
+   def select(sql, name=nil)
+     []
+   end
+   def execute(sql, name=nil)
+     0
+   end
+
+   def name
+     'fake-db'
+   end
+
+   def method_missing(name, *args)
+     raise ArgumentError, "#{self.class.name} missing '#{name}': #{args.inspect}"
+   end
+end
+
+class RawConnection
+  def method_missing(name, *args)
+      puts "#{self.class.name} missing '#{name}': #{args.inspect}"
   end
 end
 
@@ -51,18 +68,36 @@ def test_shard_connection_name
 
   def test_enchilada
     setup_configuration_for TheWholeEnchilada, 'fiveruns_city_dallas_test_slave'
+    setup_configuration_for TheWholeEnchilada, 'fiveruns_city_dallas_test_master'
     DataFabric.activate_shard :city, :dallas do
       assert_equal 'fiveruns_city_dallas_test_slave', TheWholeEnchilada.connection.connection_name
+
+      # Should use the slave
       assert_raises ActiveRecord::RecordNotFound do
         TheWholeEnchilada.find(1)
       end
+
+      # Should use the master
+      mmmm = TheWholeEnchilada.new
+      mmmm.instance_variable_set(:@attributes, { 'id' => 1 })
+      assert_raises ActiveRecord::RecordNotFound do
+        mmmm.reload
+      end
+      # ...but immediately set it back to default to the slave
+      assert_equal 'fiveruns_city_dallas_test_slave', TheWholeEnchilada.connection.connection_name
+
+      # Should use the master
+      TheWholeEnchilada.transaction do
+        mmmm.save!
+      end
     end
   end
 
   private
 
   def setup_configuration_for(clazz, name)
-    flexmock(clazz).should_receive(:mysql_connection).and_return(AdapterMock.new(nil))
-    ActiveRecord::Base.configurations = { name => { :adapter => 'mysql', :database => name, :host => 'localhost'} }
+    flexmock(clazz).should_receive(:mysql_connection).and_return(AdapterMock.new(RawConnection.new))
+    ActiveRecord::Base.configurations ||= HashWithIndifferentAccess.new
+    ActiveRecord::Base.configurations[name] = HashWithIndifferentAccess.new({ :adapter => 'mysql', :database => name, :host => 'localhost'})
   end
 end

test/database.yml

@@ -0,0 +1,27 @@
+fiveruns_city_austin_test_master:
+  adapter: mysql
+  host: localhost
+  database: vr_austin_master
+  username: root
+  password:
+
+fiveruns_city_austin_test_slave:
+  adapter: mysql
+  host: localhost
+  database: vr_austin_slave
+  username: root
+  password:
+
+fiveruns_city_dallas_test_master:
+  adapter: mysql
+  host: localhost
+  database: vr_dallas_master
+  username: root
+  password:
+
+fiveruns_city_dallas_test_slave:
+  adapter: mysql
+  host: localhost
+  database: vr_dallas_slave
+  username: root
+  password:

test/database_test.rb

@@ -0,0 +1,47 @@
+require File.join(File.dirname(__FILE__), 'test_helper')
+require 'flexmock/test_unit'
+require 'erb'
+
+class TheWholeBurrito < ActiveRecord::Base
+  connection_topology :prefix => 'fiveruns', :replicated => true, :shard_by => :city
+end
+
+class DatabaseTest < Test::Unit::TestCase
+
+  def setup
+    filename = File.join(File.dirname(__FILE__), "database.yml")
+    ActiveRecord::Base.configurations = YAML::load(ERB.new(IO.read(filename)).result)
+  end
+
+  def test_live_burrito
+    DataFabric.activate_shard :city, :dallas do
+      assert_equal 'fiveruns_city_dallas_test_slave', TheWholeBurrito.connection.connection_name
+
+      # Should use the slave
+      burrito = TheWholeBurrito.find(1)
+      assert_equal 'vr_dallas_slave', burrito.name
+
+      # Should use the master
+      burrito.reload
+      assert_equal 'vr_dallas_master', burrito.name
+
+      # ...but immediately set it back to default to the slave
+      assert_equal 'fiveruns_city_dallas_test_slave', TheWholeBurrito.connection.connection_name
+
+      # Should use the master
+      TheWholeBurrito.transaction do
+        burrito = TheWholeBurrito.find(1)
+        assert_equal 'vr_dallas_master', burrito.name
+        burrito.save!
+      end
+    end
+  end
+
+  private
+
+  def setup_configuration_for(clazz, name)
+    flexmock(clazz).should_receive(:mysql_connection).and_return(AdapterMock.new(RawConnection.new))
+    ActiveRecord::Base.configurations ||= HashWithIndifferentAccess.new
+    ActiveRecord::Base.configurations[name] = HashWithIndifferentAccess.new({ :adapter => 'mysql', :database => name, :host => 'localhost'})
+  end
+end

test/test_helper.rb

@@ -2,7 +2,12 @@
 RAILS_ENV='test'
 
 require 'rubygems'
+require 'ruby-debug'
 require 'active_support'
 Dependencies.load_paths << File.join(File.dirname(__FILE__), '../lib')
 
+require 'active_record'
+ActiveRecord::Base.logger = Logger.new(STDOUT)
+ActiveRecord::Base.logger.level = Logger::INFO
+
 require 'init'