slave_pools 0.1.2 → 1.0.0.rc1

data/lib/slave_pools/hijack.rb

@@ -0,0 +1,24 @@
+module SlavePools
+  # The hijack is added to ActiveRecord::Base but only applies to
+  # its descendants. The Base.connection is left in place.
+  module Hijack
+    def self.extended(base)
+      # hijack models that have already been loaded
+      base.send(:descendants).each do |child|
+        child.hijack_connection
+      end
+    end
+
+    # hijack models that get loaded later
+    def inherited(child)
+      super
+      child.hijack_connection
+    end
+
+    def hijack_connection
+      class << self
+        alias_method :connection, :connection_proxy
+      end
+    end
+  end
+end

data/lib/slave_pools/observer_extensions.rb

@@ -1,19 +1,19 @@
-module SlavePoolsModule
+module SlavePools
   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)
+          update_without_masterdb(observed_method, object)
         end
       else
-        send(observed_method, object, &block) if respond_to?(observed_method) && !disabled_for?(object)
+        update_without_masterdb(observed_method, object)
       end
     end
   end
-end
+end

data/lib/slave_pools/pool.rb

@@ -0,0 +1,21 @@
+module SlavePools
+  class Pool
+    attr_reader :name, :slaves, :current, :size
+
+    def initialize(name, connections)
+      @name    = name
+      @slaves  = connections
+      @size    = connections.size
+      self.reset
+    end
+
+    def reset
+      @cycle   = slaves.cycle
+      self.next
+    end
+
+    def next
+      @current = @cycle.next
+    end
+  end
+end

data/lib/slave_pools/pools.rb

@@ -0,0 +1,60 @@
+require 'delegate'
+
+module SlavePools
+  class Pools < ::SimpleDelegator
+    include Enumerable
+
+    def initialize
+      pools = {}
+      pool_configurations.group_by{|_, name, _| name }.each do |name, set|
+        pools[name.to_sym] = SlavePools::Pool.new(
+          name,
+          set.map{ |conn_name, _, replica_name|
+            connection_class(name, replica_name, conn_name)
+          }
+        )
+      end
+
+      super pools
+    end
+
+    private
+
+    # finds valid pool configs
+    def pool_configurations
+      ActiveRecord::Base.configurations.map do |name, config|
+        next unless name.to_s =~ /#{SlavePools.config.environment}_pool_(.*)_name_(.*)/
+        next unless connection_valid?(config)
+        [name, $1, $2]
+      end.compact
+    end
+
+    # generates a unique ActiveRecord::Base subclass for a single replica
+    def connection_class(pool_name, replica_name, connection_name)
+      class_name = "#{pool_name.camelize}#{replica_name.camelize}"
+
+      SlavePools.module_eval %Q{
+        class #{class_name} < ActiveRecord::Base
+          self.abstract_class = true
+          establish_connection :#{connection_name}
+          def self.connection_config
+            configurations[#{connection_name.to_s.inspect}]
+          end
+        end
+      }, __FILE__, __LINE__
+      SlavePools.const_get(class_name)
+    end
+
+    # tests a connection to be sure it's configured
+    def connection_valid?(db_config)
+      ActiveRecord::Base.establish_connection(db_config)
+      return ActiveRecord::Base.connection && ActiveRecord::Base.connected?
+    rescue => e
+      SlavePools.log :error, "Could not connect to #{db_config.inspect}"
+      SlavePools.log :error, e.to_s
+      return false
+    ensure
+      ActiveRecord::Base.establish_connection(SlavePools.config.environment)
+    end
+  end
+end

data/lib/slave_pools/query_cache.rb

@@ -0,0 +1,35 @@
+module SlavePools
+  # duck-types with ActiveRecord::ConnectionAdapters::QueryCache
+  # but relies on ActiveRecord::Base.query_cache for state so we
+  # don't fragment the cache across multiple connections
+  #
+  # we could use more of ActiveRecord's QueryCache if it only
+  # used accessors for its internal ivars.
+  module QueryCache
+    query_cache_methods = ActiveRecord::ConnectionAdapters::QueryCache.instance_methods(false)
+
+    # these methods can all use the master connection
+    (query_cache_methods - [:select_all]).each do |method_name|
+      module_eval <<-END, __FILE__, __LINE__ + 1
+        def #{method_name}(*a, &b)
+          ActiveRecord::Base.connection.#{method_name}(*a, &b)
+        end
+      END
+    end
+
+    # select_all is trickier. it needs to use the master
+    # connection for cache logic, but ultimately pass its query
+    # through to whatever connection is current.
+    def select_all(arel, name = nil, binds = [])
+      if query_cache_enabled && !locked?(arel)
+        sql = to_sql(arel, binds)
+        cache_sql(sql, binds) { route_to(current, :select_all, sql, name, binds) }
+      else
+        route_to(current, :select_all, arel, name, binds)
+      end
+    end
+
+    # these can use the unsafe delegation built into ConnectionProxy
+    # [:insert, :update, :delete]
+  end
+end

data/lib/slave_pools/version.rb

@@ -0,0 +1,3 @@
+module SlavePools
+  VERSION = "1.0.0.rc1"
+end

data/lib/slave_pools.rb

@@ -1,43 +1,71 @@
 require 'active_record'
-require 'slave_pools/slave_pool'
+require 'slave_pools/config'
+require 'slave_pools/pool'
+require 'slave_pools/pools'
 require 'slave_pools/active_record_extensions'
+require 'slave_pools/hijack'
 require 'slave_pools/observer_extensions'
-require 'slave_pools/query_cache_compat'
+require 'slave_pools/query_cache'
 require 'slave_pools/connection_proxy'
 
-#wrapper class to make the calls more succinct
+require 'slave_pools/engine' if defined? Rails
+ActiveRecord::Observer.send :include, SlavePools::ObserverExtensions
+ActiveRecord::Base.send :include, SlavePools::ActiveRecordExtensions
 
-class SlavePools
+module SlavePools
+  class << self
 
-  def self.setup!
-    SlavePoolsModule::ConnectionProxy.setup!
-  end
+    def config
+      @config ||= SlavePools::Config.new
+    end
 
-  def self.active?
-    ActiveRecord::Base.respond_to?('connection_proxy')
-  end
+    def setup!
+      if pools.empty?
+        log :info, "No pools found for #{config.environment}. Loading a default pool with master instead."
+        pools['default'] = SlavePools::Pool.new('default', [ActiveRecord::Base])
+      end
 
-  def self.next_slave!
-    ActiveRecord::Base.connection_proxy.next_slave! if active?
-  end
+      ConnectionProxy.generate_safe_delegations
+
+      ActiveRecord::Base.send(:extend, SlavePools::Hijack)
+      ActiveRecord::Base.connection_proxy = self.proxy
 
-  def self.with_pool(pool_name = 'default')
-    if active?
-      ActiveRecord::Base.connection_proxy.with_pool(pool_name) { yield }
-    else
-      yield
+      log :info, "Proxy loaded with: #{pools.keys.join(', ')}"
     end
-  end
 
-  def self.with_master
-    if active?
-      ActiveRecord::Base.connection_proxy.with_master { yield }
-    else
-      yield
+    def proxy
+      Thread.current[:slave_pools_proxy] ||= SlavePools::ConnectionProxy.new(
+        ActiveRecord::Base,
+        SlavePools.pools
+      )
+    end
+
+    def current
+      proxy.current
+    end
+
+    def next_slave!
+      proxy.next_slave!
+    end
+
+    def with_pool(*a)
+      proxy.with_pool(*a){ yield }
+    end
+
+    def with_master
+      proxy.with_master{ yield }
     end
-  end
 
-  def self.current
-    ActiveRecord::Base.connection_proxy.current if active?
+    def pools
+      Thread.current[:slave_pools] ||= SlavePools::Pools.new
+    end
+
+    def log(level, message)
+      logger.send(level, "[SlavePools] #{message}")
+    end
+
+    def logger
+      ActiveRecord::Base.logger
+    end
   end
 end

data/spec/config/test_model.rb

@@ -0,0 +1,19 @@
+class TestModel < ActiveRecord::Base
+  has_many :test_subs
+end
+
+class TestSub < ActiveRecord::Base
+  belongs_to :test_model
+end
+
+class TestModelObserver < ActiveRecord::Observer
+  
+  def after_create(test_model)
+    ActiveRecord::Base.logger.info "in observer"
+    TestSub.create(:test_model=>test_model)
+    TestModel.first
+  end
+end
+
+
+