yam-db-charmer 1.7.01

data/Rakefile

@@ -0,0 +1,4 @@
+require 'rake'
+require 'bundler'
+
+Bundler::GemHelper.install_tasks

data/db-charmer.gemspec

@@ -0,0 +1,29 @@
+# -*- encoding: utf-8 -*-
+$:.push File.expand_path('../lib', __FILE__)
+require 'db_charmer/version'
+
+Gem::Specification.new do |s|
+  s.name          = 'yam-db-charmer'
+  s.version       = DbCharmer::Version::STRING
+  s.platform      = Gem::Platform::RUBY
+
+  s.authors       = [ 'Oleksiy Kovyrin', 'Mike Ihbe']
+  s.email         = 'mihbe@yammer-inc.com'
+  s.homepage      = 'http://kovyrin.github.com/db-charmer'
+  s.summary       = 'ActiveRecord Connections Magic (slaves, multiple connections, etc)'
+  s.description   = 'DbCharmer is a Rails plugin (and gem) that could be used to manage AR model connections, implement master/slave query schemes, sharding and other magic features many high-scale applications need.'
+
+  s.rdoc_options = [ '--charset=UTF-8' ]
+
+  s.files            = `git ls-files`.split("\n")
+  s.require_paths    = [ 'lib' ]
+  s.extra_rdoc_files = [ 'LICENSE', 'README.rdoc' ]
+
+  # Dependencies
+  s.add_dependency 'activesupport', '< 3.1'
+  s.add_dependency 'activerecord', '< 3.1'
+
+  s.add_development_dependency 'rspec'
+  s.add_development_dependency 'yard'
+  s.add_development_dependency 'actionpack'
+end

data/init.rb

@@ -0,0 +1 @@
+require 'db_charmer'

data/lib/db_charmer/action_controller/force_slave_reads.rb

@@ -0,0 +1,69 @@
+module DbCharmer
+  module ActionController
+    module ForceSlaveReads
+
+      module ClassMethods
+        @@db_charmer_force_slave_reads_actions = {}
+        def force_slave_reads(params)
+          @@db_charmer_force_slave_reads_actions[self.name] = {
+            :except => params[:except] ? [*params[:except]].map(&:to_s) : [],
+            :only => params[:only] ? [*params[:only]].map(&:to_s) : []
+          }
+        end
+
+        def force_slave_reads_options
+          @@db_charmer_force_slave_reads_actions[self.name]
+        end
+
+        def force_slave_reads_action?(name = nil)
+          name = name.to_s
+
+          options = force_slave_reads_options
+          # If no options were defined for this controller, all actions are not forced to use slaves
+          return false unless options
+
+          # Actions where force_slave_reads mode was turned off
+          return false if options[:except].include?(name)
+
+          # Only for these actions force_slave_reads was turned on
+          return options[:only].include?(name) if options[:only].any?
+
+          # If :except is not empty, we're done with the checks and rest of the actions are should force slave reads
+          # Otherwise, all the actions are not in force_slave_reads mode
+          options[:except].any?
+        end
+      end
+
+      module InstanceMethods
+        DISPATCH_METHOD = (DbCharmer.rails3?) ? :process_action : :perform_action
+
+        def self.included(base)
+          base.alias_method_chain DISPATCH_METHOD, :forced_slave_reads
+        end
+
+        def force_slave_reads!
+          @db_charmer_force_slave_reads = true
+        end
+
+        def dont_force_slave_reads!
+          @db_charmer_force_slave_reads = false
+        end
+
+        def force_slave_reads?
+          @db_charmer_force_slave_reads || self.class.force_slave_reads_action?(params[:action])
+        end
+
+      protected
+
+        class_eval <<-EOF
+          def #{DISPATCH_METHOD}_with_forced_slave_reads(*args, &block)
+            DbCharmer.with_controller(self) do
+              #{DISPATCH_METHOD}_without_forced_slave_reads(*args, &block)
+            end
+          end
+        EOF
+      end
+
+    end
+  end
+end

data/lib/db_charmer/active_record/association_preload.rb

@@ -0,0 +1,23 @@
+module DbCharmer
+  module ActiveRecord
+    module AssociationPreload
+      ASSOCIATION_TYPES = [ :has_one, :has_many, :belongs_to, :has_and_belongs_to_many ]
+
+      def self.extended(base)
+        ASSOCIATION_TYPES.each do |association_type|
+          base.class_eval <<-EOF, __FILE__, __LINE__ + 1
+            def self.preload_#{association_type}_association(records, reflection, preload_options = {})
+              if self.db_charmer_top_level_connection? || reflection.options[:polymorphic] ||
+                    self.db_charmer_default_connection != reflection.klass.db_charmer_default_connection
+                return super(records, reflection, preload_options)
+              end
+              reflection.klass.on_db(self) do
+                super(records, reflection, preload_options)
+              end
+            end
+          EOF
+        end
+      end
+    end
+  end
+end

data/lib/db_charmer/active_record/class_attributes.rb

@@ -0,0 +1,101 @@
+module DbCharmer
+  module ActiveRecord
+    module ClassAttributes
+      @@db_charmer_opts = {}
+      def db_charmer_opts=(opts)
+        @@db_charmer_opts[self.name] = opts
+      end
+
+      def db_charmer_opts
+        @@db_charmer_opts[self.name] || {}
+      end
+
+      #-----------------------------------------------------------------------------
+      @@db_charmer_connection_proxies = {}
+      def db_charmer_connection_proxy=(proxy)
+        @@db_charmer_connection_proxies[self.name] = proxy
+      end
+
+      def db_charmer_connection_proxy
+        @@db_charmer_connection_proxies[self.name]
+      end
+
+      #-----------------------------------------------------------------------------
+      @@db_charmer_default_connections = {}
+      def db_charmer_default_connection=(conn)
+        @@db_charmer_default_connections[self.name] = conn
+      end
+
+      def db_charmer_default_connection
+        @@db_charmer_default_connections[self.name]
+      end
+
+      #-----------------------------------------------------------------------------
+      @@db_charmer_slaves = {}
+      def db_charmer_slaves=(slaves)
+        @@db_charmer_slaves[self.name] = slaves
+      end
+
+      def db_charmer_slaves
+        @@db_charmer_slaves[self.name] || []
+      end
+
+      def db_charmer_random_slave
+        return nil unless db_charmer_slaves.any?
+        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)
+        @@db_charmer_connection_levels[self.name] = level
+      end
+
+      def db_charmer_connection_level
+        @@db_charmer_connection_levels[self.name] || 0
+      end
+
+      def db_charmer_top_level_connection?
+        db_charmer_connection_level.zero?
+      end
+
+      #-----------------------------------------------------------------------------
+      @@db_charmer_database_remappings = Hash.new
+      def db_charmer_remapped_connection
+        return nil if (db_charmer_connection_level || 0) > 0
+        name = :master
+        proxy = db_charmer_connection_proxy
+        name = proxy.db_charmer_connection_name.to_sym if proxy
+
+        remapped = @@db_charmer_database_remappings[name]
+        remapped ? DbCharmer::ConnectionFactory.connect(remapped, true) : nil
+      end
+
+      def db_charmer_database_remappings
+        @@db_charmer_database_remappings
+      end
+
+      def db_charmer_database_remappings=(mappings)
+        raise "Mappings must be nil or respond to []" if mappings && (! mappings.respond_to?(:[]))
+        @@db_charmer_database_remappings = mappings || { }
+      end
+    end
+  end
+end

data/lib/db_charmer/active_record/connection_switching.rb

@@ -0,0 +1,81 @@
+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
+          # Make sure we check our accessors before going to the default connection retrieval method
+          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, should_exist = true)
+        new_conn = coerce_to_connection_proxy(conn, should_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!
+
+        # self.reset_column_information
+      end
+
+    end
+  end
+end

data/lib/db_charmer/active_record/db_magic.rb

@@ -0,0 +1,85 @@
+module DbCharmer
+  module ActiveRecord
+    module DbMagic
+
+      def db_magic(opt = {})
+        # Make sure we could use our connections management here
+        hijack_connection!
+
+        # Should requested connections exist in the config?
+        should_exist = opt.has_key?(:should_exist) ? opt[:should_exist] : DbCharmer.connections_should_exist?
+
+        # Main connection management
+        setup_connection_magic(opt[:connection], should_exist)
+
+        # Set up slaves pool
+        opt[:slaves] ||= []
+        opt[:slaves] = [ opt[:slaves] ].flatten
+        opt[:slaves] << opt[:slave] if opt[:slave]
+
+        # 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)
+
+        # Setup inheritance magic
+        setup_children_magic(opt)
+
+        # Setup sharding if needed
+        if opt[:sharded]
+          raise ArgumentError, "Can't use sharding on a model with slaves!" if opt[:slaves].any?
+          setup_sharding_magic(opt[:sharded])
+        end
+      end
+
+    private
+
+      def setup_children_magic(opt)
+        self.db_charmer_opts = opt.clone
+
+        def self.inherited(child)
+          child.db_magic(self.db_charmer_opts)
+          super
+        end
+      end
+
+      def setup_sharding_magic(config)
+        # Add sharding-specific methods
+        self.extend(DbCharmer::ActiveRecord::Sharding)
+
+        # Get configuration
+        name = config[:sharded_connection] or raise ArgumentError, "No :sharded_connection!"
+        # Assign sharded connection
+        self.sharded_connection = DbCharmer::Sharding.sharded_connection(name)
+
+        # Setup model default connection
+        setup_connection_magic(sharded_connection.default_connection)
+      end
+
+      def setup_connection_magic(conn, should_exist = true)
+        switch_connection_to(conn, should_exist)
+        self.db_charmer_default_connection = conn
+      end
+
+      def setup_slaves_magic(slaves, force_slave_reads, should_exist = true)
+        self.db_charmer_force_slave_reads = force_slave_reads
+
+        # Initialize the slave connections list
+        self.db_charmer_slaves = slaves.collect do |slave|
+          coerce_to_connection_proxy(slave, should_exist)
+        end
+        return if db_charmer_slaves.empty?
+
+        # Enable on_slave/on_master methods
+        self.extend(DbCharmer::ActiveRecord::MultiDbProxy::MasterSlaveClassMethods)
+
+        # Enable automatic master/slave queries routing (we have specialized versions on those modules for rails2/3)
+        self.extend(DbCharmer::ActiveRecord::MasterSlaveRouting::ClassMethods)
+        self.send(:include, DbCharmer::ActiveRecord::MasterSlaveRouting::InstanceMethods)
+      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