db-charmer 1.6.19 → 1.7.0.pre1

data/lib/db_charmer/{association_preload.rb → active_record/association_preload.rb}

@@ -1,6 +1,6 @@
 module DbCharmer
-  module AssociationPreload
-    module ClassMethods
+  module ActiveRecord
+    module AssociationPreload
       ASSOCIATION_TYPES = [ :has_one, :has_many, :belongs_to, :has_and_belongs_to_many ]
 
       def self.extended(base)

data/lib/db_charmer/{active_record_extensions.rb → active_record/class_attributes.rb}

@@ -1,32 +1,6 @@
 module DbCharmer
-  module ActiveRecordExtensions
-    module ClassMethods
-
-      def establish_real_connection_if_exists(name, should_exist = false)
-        name = name.to_s
-
-        # Check environment name
-        config = configurations[DbCharmer.env]
-        unless config
-          error = "Invalid environment name (does not exist in database.yml): #{DbCharmer.env}. Please set correct Rails.env or DbCharmer.env."
-          raise ArgumentError, error
-        end
-
-        # Check connection name
-        config = config[name]
-        unless config
-          if should_exist
-            raise ArgumentError, "Invalid connection name (does not exist in database.yml): #{DbCharmer.env}/#{name}"
-          end
-          return # No need to establish connection - they do not want us to
-        end
-
-        # Pass connection name with config
-        config[:connection_name] = name
-        establish_connection(config)
-      end
-
-      #-----------------------------------------------------------------------------
+  module ActiveRecord
+    module ClassAttributes
       @@db_charmer_opts = {}
       def db_charmer_opts=(opts)
         @@db_charmer_opts[self.name] = opts
@@ -71,6 +45,23 @@ module DbCharmer
         db_charmer_slaves[rand(db_charmer_slaves.size)]
       end
 
+      #-----------------------------------------------------------------------------
+      @@db_charmer_force_slave_reads = {}
+      def db_charmer_force_slave_reads=(force)
+        @@db_charmer_force_slave_reads[self.name] = force
+      end
+
+      def db_charmer_force_slave_reads
+        @@db_charmer_force_slave_reads[self.name]
+      end
+
+      # Slave reads are used in two cases:
+      #  - per-model slave reads are enabled (see db_magic method for more details)
+      #  - global slave reads enforcing is enabled (in a controller action)
+      def db_charmer_force_slave_reads?
+        db_charmer_force_slave_reads || DbCharmer.force_slave_reads?
+      end
+
       #-----------------------------------------------------------------------------
       @@db_charmer_connection_levels = Hash.new(0)
       def db_charmer_connection_level=(level)
@@ -105,18 +96,6 @@ module DbCharmer
         raise "Mappings must be nil or respond to []" if mappings && (! mappings.respond_to?(:[]))
         @@db_charmer_database_remappings = mappings || { }
       end
-
-      #-----------------------------------------------------------------------------
-      def hijack_connection!
-        return if self.respond_to?(:connection_with_magic)
-        class << self
-          def connection_with_magic
-            db_charmer_remapped_connection || db_charmer_connection_proxy || connection_without_magic
-          end
-          alias_method_chain :connection, :magic
-        end
-      end
-
     end
   end
 end

data/lib/db_charmer/active_record/connection_switching.rb

@@ -0,0 +1,77 @@
+module DbCharmer
+  module ActiveRecord
+    module ConnectionSwitching
+      def establish_real_connection_if_exists(name, should_exist = false)
+        name = name.to_s
+
+        # Check environment name
+        config = configurations[DbCharmer.env]
+        unless config
+          error = "Invalid environment name (does not exist in database.yml): #{DbCharmer.env}. Please set correct Rails.env or DbCharmer.env."
+          raise ArgumentError, error
+        end
+
+        # Check connection name
+        config = config[name]
+        unless config
+          if should_exist
+            raise ArgumentError, "Invalid connection name (does not exist in database.yml): #{DbCharmer.env}/#{name}"
+          end
+          return # No need to establish connection - they do not want us to
+        end
+
+        # Pass connection name with config
+        config[:connection_name] = name
+        establish_connection(config)
+      end
+
+      #-----------------------------------------------------------------------------------------------------------------
+      def hijack_connection!
+        return if self.respond_to?(:connection_with_magic)
+        class << self
+          def connection_with_magic
+            db_charmer_remapped_connection || db_charmer_connection_proxy || connection_without_magic
+          end
+          alias_method_chain :connection, :magic
+        end
+      end
+
+      #-----------------------------------------------------------------------------------------------------------------
+      def coerce_to_connection_proxy(conn, should_exist = true)
+        return nil if conn.nil?
+
+        if conn.kind_of?(Symbol) || conn.kind_of?(String)
+          return DbCharmer::ConnectionFactory.connect(conn, should_exist)
+        end
+
+        if conn.kind_of?(Hash)
+          conn = conn.symbolize_keys
+          raise ArgumentError, "Missing required :connection_name parameter" unless conn[:connection_name]
+          return DbCharmer::ConnectionFactory.connect_to_db(conn[:connection_name], conn)
+        end
+
+        if conn.respond_to?(:db_charmer_connection_proxy)
+          return conn.db_charmer_connection_proxy
+        end
+
+        if conn.kind_of?(::ActiveRecord::ConnectionAdapters::AbstractAdapter) || conn.kind_of?(DbCharmer::Sharding::StubConnection)
+          return conn
+        end
+
+        raise "Unsupported connection type: #{conn.class}"
+      end
+
+      #-----------------------------------------------------------------------------------------------------------------
+      def switch_connection_to(conn, require_config_to_exist = true)
+        new_conn = coerce_to_connection_proxy(conn, require_config_to_exist)
+
+        if db_charmer_connection_proxy.is_a?(DbCharmer::Sharding::StubConnection)
+          db_charmer_connection_proxy.set_real_connection(new_conn)
+        end
+
+        self.db_charmer_connection_proxy = new_conn
+        self.hijack_connection!
+      end
+    end
+  end
+end

data/lib/db_charmer/{db_magic.rb → active_record/db_magic.rb}

@@ -1,6 +1,7 @@
 module DbCharmer
-  module DbMagic
-    module ClassMethods
+  module ActiveRecord
+    module DbMagic
+
       def db_magic(opt = {})
         # Make sure we could use our connections management here
         hijack_connection!
@@ -13,8 +14,14 @@ module DbCharmer
 
         # Set up slaves pool
         opt[:slaves] ||= []
+        opt[:slaves] = [ opt[:slaves] ].flatten
         opt[:slaves] << opt[:slave] if opt[:slave]
-        setup_slaves_magic(opt[:slaves], should_exist) if opt[:slaves].any?
+
+        # Forced reads are enabled for all models by default, could be disabled by the user
+        forced_slave_reads = opt.has_key?(:force_slave_reads) ? opt[:force_slave_reads] : true
+
+        # Setup all the slaves related magic if needed
+        setup_slaves_magic(opt[:slaves], forced_slave_reads, should_exist) if opt[:slaves].any?
 
         # Setup inheritance magic
         setup_children_magic(opt)
@@ -39,7 +46,7 @@ module DbCharmer
 
       def setup_sharding_magic(config)
         # Add sharding-specific methods
-        self.extend(DbCharmer::Sharding::ClassMethods)
+        self.extend(DbCharmer::ActiveRecord::Sharding)
 
         # Get configuration
         name = config[:sharded_connection] or raise ArgumentError, "No :sharded_connection!"
@@ -55,15 +62,18 @@ module DbCharmer
         self.db_charmer_default_connection = conn
       end
 
-      def setup_slaves_magic(slaves, should_exist = true)
+      def setup_slaves_magic(slaves, force_slave_reads, should_exist = true)
         self.db_charmer_slaves = slaves.collect do |slave|
           coerce_to_connection_proxy(slave, should_exist)
         end
 
-        self.extend(DbCharmer::FinderOverrides::ClassMethods)
-        self.send(:include, DbCharmer::FinderOverrides::InstanceMethods)
-        self.extend(DbCharmer::MultiDbProxy::MasterSlaveClassMethods)
+        self.db_charmer_force_slave_reads = force_slave_reads
+
+        self.extend(DbCharmer::ActiveRecord::FinderOverrides::ClassMethods)
+        self.send(:include, DbCharmer::ActiveRecord::FinderOverrides::InstanceMethods)
+        self.extend(DbCharmer::ActiveRecord::MultiDbProxy::MasterSlaveClassMethods)
       end
+
     end
   end
 end

data/lib/db_charmer/active_record/finder_overrides.rb

@@ -0,0 +1,61 @@
+module DbCharmer
+  module ActiveRecord
+    module FinderOverrides
+
+      module ClassMethods
+        SLAVE_METHODS = [ :find_by_sql, :count_by_sql, :calculate ]
+        MASTER_METHODS = [ :update, :create, :delete, :destroy, :delete_all, :destroy_all, :update_all, :update_counters ]
+
+        SLAVE_METHODS.each do |slave_method|
+          class_eval <<-EOF, __FILE__, __LINE__ + 1
+            def #{slave_method}(*args, &block)
+              first_level_on_slave do
+                super(*args, &block)
+              end
+            end
+          EOF
+        end
+
+        MASTER_METHODS.each do |master_method|
+          class_eval <<-EOF, __FILE__, __LINE__ + 1
+            def #{master_method}(*args, &block)
+              on_master do
+                super(*args, &block)
+              end
+            end
+          EOF
+        end
+
+        def find(*args, &block)
+          options = args.last
+          if options.is_a?(Hash) && options[:lock]
+            on_master { super(*args, &block) }
+          else
+            super(*args, &block)
+          end
+        end
+
+      private
+
+        def first_level_on_slave
+          first_level = db_charmer_top_level_connection? && on_master.connection.open_transactions.zero?
+          if first_level && db_charmer_force_slave_reads?
+            on_slave { yield }
+          else
+            yield
+          end
+        end
+
+      end
+
+      module InstanceMethods
+        def reload(*args, &block)
+          self.class.on_master do
+            super(*args, &block)
+          end
+        end
+      end
+
+    end
+  end
+end

data/lib/db_charmer/active_record/migration/multi_db_migrations.rb

@@ -0,0 +1,71 @@
+module DbCharmer
+  module ActiveRecord
+    module Migration
+      module MultiDbMigrations
+
+        def self.extended(base)
+          class << base
+            alias_method_chain :migrate, :db_wrapper
+          end
+        end
+
+        @@multi_db_names = {}
+        def multi_db_names
+          @@multi_db_names[self.name] || @@multi_db_names['ActiveRecord::Migration']
+        end
+
+        def multi_db_names=(names)
+          @@multi_db_names[self.name] = names
+        end
+
+        def migrate_with_db_wrapper(direction)
+          if names = multi_db_names
+            names.each do |multi_db_name|
+              on_db(multi_db_name) do
+                migrate_without_db_wrapper(direction)
+              end
+            end
+          else
+            migrate_without_db_wrapper(direction)
+          end
+        end
+
+        def on_db(db_name)
+          name = db_name.is_a?(Hash) ? db_name[:connection_name] : db_name.inspect
+          announce "Switching connection to #{name}"
+          # Switch connection
+          old_proxy = ::ActiveRecord::Base.db_charmer_connection_proxy
+          db_name = nil if db_name == :default
+          ::ActiveRecord::Base.switch_connection_to(db_name, DbCharmer.connections_should_exist?)
+          # Yield the block
+          yield
+        ensure
+          # Switch it back
+          ::ActiveRecord::Base.verify_active_connections!
+          announce "Switching connection back"
+          ::ActiveRecord::Base.switch_connection_to(old_proxy)
+        end
+
+        def db_magic(opts = {})
+          # Collect connections from all possible options
+          conns = [ opts[:connection], opts[:connections] ]
+          conns << shard_connections(opts[:sharded_connection]) if opts[:sharded_connection]
+
+          # Get a unique set of connections
+          conns = conns.flatten.compact.uniq
+          raise ArgumentError, "No connection name - no magic!" unless conns.any?
+
+          # Save connections
+          self.multi_db_names = conns
+        end
+
+        # Return a list of connections to shards in a sharded connection
+        def shard_connections(conn_name)
+          conn = DbCharmer::Sharding.sharded_connection(conn_name)
+          conn.shard_connections
+        end
+
+      end
+    end
+  end
+end

data/lib/db_charmer/active_record/multi_db_proxy.rb

@@ -0,0 +1,65 @@
+module DbCharmer
+  module ActiveRecord
+    module MultiDbProxy
+      # Simple proxy class that switches connections and then proxies all the calls
+      # This class is used to implement chained on_db calls
+      class OnDbProxy < BlankSlate
+        def initialize(proxy_target, slave)
+          @proxy_target = proxy_target
+          @slave = slave
+        end
+
+      private
+
+        def method_missing(meth, *args, &block)
+          # Switch connection and proxy the method call
+          @proxy_target.on_db(@slave) do |proxy_target|
+            res = proxy_target.__send__(meth, *args, &block)
+
+            # If result is a scope/association, return a new proxy for it, otherwise return the result itself
+            (res.proxy?) ? OnDbProxy.new(res, @slave) : res
+          end
+        end
+      end
+
+      module ClassMethods
+        def on_db(con, proxy_target = nil)
+          proxy_target ||= self
+
+          # Chain call
+          return OnDbProxy.new(proxy_target, con) unless block_given?
+
+          # Block call
+          begin
+            self.db_charmer_connection_level += 1
+            old_proxy = db_charmer_connection_proxy
+            switch_connection_to(con, DbCharmer.connections_should_exist?)
+            yield(proxy_target)
+          ensure
+            switch_connection_to(old_proxy)
+            self.db_charmer_connection_level -= 1
+          end
+        end
+      end
+
+      module InstanceMethods
+        def on_db(con, proxy_target = nil, &block)
+          proxy_target ||= self
+          self.class.on_db(con, proxy_target, &block)
+        end
+      end
+
+      module MasterSlaveClassMethods
+        def on_slave(con = nil, proxy_target = nil, &block)
+          con ||= db_charmer_random_slave
+          raise ArgumentError, "No slaves found in the class and no slave connection given" unless con
+          on_db(con, proxy_target, &block)
+        end
+
+        def on_master(proxy_target = nil, &block)
+          on_db(db_charmer_default_connection, proxy_target, &block)
+        end
+      end
+    end
+  end
+end

data/lib/db_charmer/active_record/named_scope/scope_proxy.rb

@@ -0,0 +1,26 @@
+module DbCharmer
+  module ActiveRecord
+    module NamedScope
+      module ScopeProxy
+
+        def proxy?
+          true
+        end
+
+        def on_db(con, proxy_target = nil, &block)
+          proxy_target ||= self
+          proxy_scope.on_db(con, proxy_target, &block)
+        end
+
+        def on_slave(con = nil, &block)
+          proxy_scope.on_slave(con, self, &block)
+        end
+
+        def on_master(&block)
+          proxy_scope.on_master(self, &block)
+        end
+
+      end
+    end
+  end
+end