lib/replica.rb in replica-1.0.1 vs lib/replica.rb in replica-1.0.2

@@ -9,31 +9,34 @@
       #   Account.with_slave do
       #     Account.first
       #   end
       # the first account will be found on the slave DB
       #
+      # For one-liners you can simply do
+      #   Account.with_slave.first
+      #
       # this is the same as:
       #   Account.with_replica(:slave) do
       #     Account.first
       #   end
       def with_slave(&block)
         with_replica(:slave, &block)
       end
-      
+
       # See with_slave
       def with_master(&block)
         with_replica(nil, &block)
       end
 
       def with_slave_if(condition, &block)
         condition ? with_slave(&block) : with_master(&block)
       end
-      
+
       def with_slave_unless(condition, &block)
         with_slave_if(!condition, &block)
       end
-      
+
       # Name of the connection pool. Used by ConnectionHandler to retrieve the current connection pool.
       def connection_pool_name # :nodoc:
         replica = current_replica_name
         if replica
           "#{name}_#{replica}"
@@ -42,83 +45,101 @@
         else
           superclass.connection_pool_name
         end
       end
 
-      # Specify which database to use. 
+      # Specify which database to use.
       #
       # Example:
       # database.yml
       #   test_slave:
       #     adapter: mysql
       #     ...
       #
       # Account.with_replica(:slave) { Account.count }
+      # Account.with_replica(:slave).count
       #
       def with_replica(replica_name, &block)
+        if block_given?
+          with_replica_block(replica_name, &block)
+        else
+          Proxy.new(self, replica_name)
+        end
+      end
+
+      def with_replica_block(replica_name, &block)
         old_replica_name = current_replica_name
         begin
           self.current_replica_name = replica_name
         rescue ActiveRecord::AdapterNotSpecified => e
           self.current_replica_name = old_replica_name
           logger.warn("Failed to establish replica connection: #{e.message} - defaulting to master")
-        end        
+        end
         yield
       ensure
         self.current_replica_name = old_replica_name
       end
 
       private
-      
+
         def current_replica_name
           Thread.current[replica_key]
         end
 
         def current_replica_name=(new_replica_name)
           Thread.current[replica_key] = new_replica_name
 
           establish_replica_connection(new_replica_name) unless connected_to_replica?
         end
-        
+
         def establish_replica_connection(replica_name)
-           name = replica_name ? "#{RAILS_ENV}_#{replica_name}" : RAILS_ENV
-           spec = configurations[name]
-           raise AdapterNotSpecified.new("No database defined by #{name} in database.yml") if spec.nil?
-           
-           connection_handler.establish_connection(connection_pool_name, ConnectionSpecification.new(spec, "#{spec['adapter']}_connection"))
+          name = replica_name ? "#{RAILS_ENV}_#{replica_name}" : RAILS_ENV
+          spec = configurations[name]
+          raise AdapterNotSpecified.new("No database defined by #{name} in database.yml") if spec.nil?
+
+          connection_handler.establish_connection(connection_pool_name, ConnectionSpecification.new(spec, "#{spec['adapter']}_connection"))
         end
-        
+
         def connected_to_replica?
           connection_handler.connection_pools.has_key?(connection_pool_name)
         end
-        
+
         def replica_key
           @replica_key ||= "#{name}_replica"
         end
-        
+
+        class Proxy
+          def initialize(target, replica)
+            @target = target
+            @replica = replica
+          end
+
+          def method_missing(method, *args)
+            @target.with_replica_block(@replica) { @target.send(method, *args) }
+          end
+        end
     end
   end
   Base.extend(Base::Replica)
-  
+
   # The only difference here is that we use klass.connection_pool_name
   # instead of klass.name as the pool key
   module ConnectionAdapters # :nodoc:
     class ConnectionHandler # :nodoc:
-      
+
       def retrieve_connection_pool(klass)
         pool = @connection_pools[klass.connection_pool_name]
         return pool if pool
         return nil if ActiveRecord::Base == klass
         retrieve_connection_pool klass.superclass
       end
-      
+
       def remove_connection(klass)
-         pool = @connection_pools[klass.connection_pool_name]
-         @connection_pools.delete_if { |key, value| value == pool }
-         pool.disconnect! if pool
-         pool.spec.config if pool
-       end
-      
+        pool = @connection_pools[klass.connection_pool_name]
+        @connection_pools.delete_if { |key, value| value == pool }
+        pool.disconnect! if pool
+        pool.spec.config if pool
+      end
+
     end
   end
 end
-