yam-db-charmer 1.7.01

data/lib/db_charmer/sharding/method/db_block_map.rb

@@ -0,0 +1,211 @@
+# This is a more sophisticated sharding method based on a database-backed
+# blocks map that holds block-shard associations. It automatically
+# creates new blocks for new keys and assigns them to shards.
+#
+module DbCharmer
+  module Sharding
+    module Method
+      class DbBlockMap
+        # Sharder name
+        attr_accessor :name
+
+        # Mapping db connection
+        attr_accessor :connection, :connection_name
+
+        # Mapping table name
+        attr_accessor :map_table
+
+        # Shards table name
+        attr_accessor :shards_table
+
+        # Sharding keys block size
+        attr_accessor :block_size
+
+        def initialize(config)
+          @name = config[:name] or raise(ArgumentError, "Missing required :name parameter!")
+          @connection = DbCharmer::ConnectionFactory.connect(config[:connection], true)
+          @block_size = (config[:block_size] || 10000).to_i
+
+          @map_table = config[:map_table] or raise(ArgumentError, "Missing required :map_table parameter!")
+          @shards_table = config[:shards_table] or raise(ArgumentError, "Missing required :shards_table parameter!")
+
+          # Local caches
+          @shard_info_cache = {}
+
+          @blocks_cache = Rails.cache
+          @blocks_cache_prefix = config[:blocks_cache_prefix] || "#{@name}_block:"
+        end
+
+        def shard_for_key(key)
+          block = block_for_key(key)
+
+          begin
+            # Auto-allocate new blocks
+            block ||= allocate_new_block_for_key(key)
+          rescue ::ActiveRecord::StatementInvalid => e
+            raise unless e.message.include?('Duplicate entry')
+            block = block_for_key(key)
+          end
+
+          raise ArgumentError, "Invalid key value, no shards found for this key and could not create a new block!" unless block
+
+          # Bail if no shard found
+          shard_id = block['shard_id'].to_i
+          shard_info = shard_info_by_id(shard_id)
+          raise ArgumentError, "Invalid shard_id: #{shard_id}" unless shard_info
+
+          # Get config
+          shard_connection_config(shard_info)
+        end
+
+        class ShardInfo < ::ActiveRecord::Base
+          validates_presence_of :db_host
+          validates_presence_of :db_port
+          validates_presence_of :db_user
+          validates_presence_of :db_pass
+          validates_presence_of :db_name
+        end
+
+        # Returns a block for a key
+        def block_for_key(key, cache = true)
+          # Cleanup the cache if asked to
+          key_range = [ block_start_for_key(key), block_end_for_key(key) ]
+          block_cache_key = "%d-%d" % key_range
+
+          if cache
+            cached_block = get_cached_block(block_cache_key)
+            return cached_block if cached_block
+          end
+
+          # Fetch cached value or load from db
+          block = begin
+            sql = "SELECT * FROM #{map_table} WHERE start_id = #{key_range.first} AND end_id = #{key_range.last} LIMIT 1"
+            connection.select_one(sql, 'Find a shard block')
+          end
+
+          set_cached_block(block_cache_key, block)
+
+          return block
+        end
+
+        def get_cached_block(block_cache_key)
+          @blocks_cache.read("#{@blocks_cache_prefix}#{block_cache_key}")
+        end
+
+        def set_cached_block(block_cache_key, block)
+          @blocks_cache.write("#{@blocks_cache_prefix}#{block_cache_key}", block)
+        end
+
+        # Load shard info
+        def shard_info_by_id(shard_id, cache = true)
+          # Cleanup the cache if asked to
+          @shard_info_cache[shard_id] = nil unless cache
+
+          # Either load from cache or from db
+          @shard_info_cache[shard_id] ||= begin
+            prepare_shard_model
+            ShardInfo.find_by_id(shard_id)
+          end
+        end
+
+        def allocate_new_block_for_key(key)
+          # Can't find any shards to use for blocks allocation!
+          return nil unless shard = least_loaded_shard
+
+          # Figure out block limits
+          start_id = block_start_for_key(key)
+          end_id = block_end_for_key(key)
+
+          # Try to insert a new mapping (ignore duplicate key errors)
+          sql = <<-SQL
+            INSERT INTO #{map_table}
+                   SET start_id = #{start_id},
+                       end_id = #{end_id},
+                       shard_id = #{shard.id},
+                       block_size = #{block_size},
+                       created_at = NOW(),
+                       updated_at = NOW()
+          SQL
+          connection.execute(sql, "Allocate new block")
+
+          # Increment the blocks counter on the shard
+          ShardInfo.update_counters(shard.id, :blocks_count => +1)
+
+          # Retry block search after creation
+          block_for_key(key)
+        end
+
+        def least_loaded_shard
+          prepare_shard_model
+
+          # Select shard
+          shard = ShardInfo.all(:conditions => { :enabled => true, :open => true }, :order => 'blocks_count ASC', :limit => 1).first
+          raise "Can't find any shards to use for blocks allocation!" unless shard
+          return shard
+        end
+
+        def block_start_for_key(key)
+          block_size.to_i * (key.to_i / block_size.to_i)
+        end
+
+        def block_end_for_key(key)
+          block_size.to_i + block_start_for_key(key)
+        end
+
+        # Create configuration (use mapping connection as a template)
+        def shard_connection_config(shard)
+          # Format connection name
+          shard_name = "db_charmer_db_block_map_#{name}_shard_%05d" % shard.id
+
+          # Here we get the mapping connection's configuration
+          # They do not expose configs so we hack in and get the instance var
+          # FIXME: Find a better way, maybe move config method to our ar extenstions
+          connection.instance_variable_get(:@config).clone.merge(
+            # Name for the connection factory
+            :connection_name => shard_name,
+            # Connection params
+            :host => shard.db_host,
+            :port => shard.db_port,
+            :username => shard.db_user,
+            :password => shard.db_pass,
+            :database => shard.db_name
+          )
+        end
+
+        def create_shard(params)
+          params = params.symbolize_keys
+          [ :db_host, :db_port, :db_user, :db_pass, :db_name ].each do |arg|
+            raise ArgumentError, "Missing required parameter: #{arg}" unless params[arg]
+          end
+
+          # Prepare model
+          prepare_shard_model
+
+          # Create the record
+          ShardInfo.create! do |shard|
+            shard.db_host = params[:db_host]
+            shard.db_port = params[:db_port]
+            shard.db_user = params[:db_user]
+            shard.db_pass = params[:db_pass]
+            shard.db_name = params[:db_name]
+          end
+        end
+
+        def shard_connections
+          # Find all shards
+          prepare_shard_model
+          shards = ShardInfo.all(:conditions => { :enabled => true })
+          # Map them to connections
+          shards.map { |shard| shard_connection_config(shard) }
+        end
+
+        # Prepare model for working with our shards table
+        def prepare_shard_model
+          ShardInfo.set_table_name(shards_table)
+          ShardInfo.switch_connection_to(connection)
+        end
+
+      end
+    end
+  end
+end

data/lib/db_charmer/sharding/method/hash_map.rb

@@ -0,0 +1,23 @@
+module DbCharmer
+  module Sharding
+    module Method
+      class HashMap
+        attr_accessor :map
+
+        def initialize(config)
+          @map = config[:map].clone or raise ArgumentError, "No :map defined!"
+        end
+
+        def shard_for_key(key)
+          res = map[key] || map[:default]
+          raise ArgumentError, "Invalid key value, no shards found for this key!" unless res
+          return res
+        end
+
+        def support_default_shard?
+          map.has_key?(:default)
+        end
+      end
+    end
+  end
+end

data/lib/db_charmer/sharding/method/range.rb

@@ -0,0 +1,33 @@
+module DbCharmer
+  module Sharding
+    module Method
+      class Range
+        attr_accessor :ranges
+
+        def initialize(config)
+          @ranges = config[:ranges] ? config[:ranges].clone : raise(ArgumentError, "No :ranges defined!")
+        end
+
+        def shard_for_key(key)
+          return ranges[:default] if key == :default
+
+          ranges.each do |range, shard|
+            next if range == :default
+            return shard if range.member?(key.to_i)
+          end
+
+          return ranges[:default] if ranges[:default]
+          raise ArgumentError, "Invalid key value, no shards found for this key!"
+        end
+
+        def support_default_shard?
+          ranges.has_key?(:default)
+        end
+
+        def shard_connections
+          ranges.values.uniq
+        end
+      end
+    end
+  end
+end

data/lib/db_charmer/sharding/method.rb

@@ -0,0 +1,10 @@
+module DbCharmer
+  module Sharding
+    module Method
+      autoload :Range, 'db_charmer/sharding/method/range'
+      autoload :HashMap, 'db_charmer/sharding/method/hash_map'
+      autoload :DbBlockMap, 'db_charmer/sharding/method/db_block_map'
+      autoload :DbBlockGroupMap, 'db_charmer/sharding/method/db_block_group_map'
+    end
+  end
+end

data/lib/db_charmer/sharding/stub_connection.rb

@@ -0,0 +1,60 @@
+# This is a simple proxy class used as a default connection on sharded models
+#
+# The idea is to proxy all utility method calls to a real connection (set by
+# the +set_real_connection+ method when we switch shards) and fail on real
+# database querying calls forcing users to switch shard connections.
+#
+module DbCharmer
+  module Sharding
+    class StubConnection
+      attr_accessor :sharded_connection
+
+      def initialize(sharded_connection)
+        @sharded_connection = sharded_connection
+        @real_conn = nil
+      end
+
+      def set_real_connection(real_conn)
+        @real_conn = real_conn
+      end
+
+      def db_charmer_connection_name
+        "StubConnection"
+      end
+
+      def real_connection
+        # Return memoized real connection
+        return @real_conn if @real_conn
+
+        # If sharded connection supports shards enumeration, get the first shard
+        conn = sharded_connection.shard_connections.try(:first)
+
+        # If we do not have real connection yet, try to use the default one (if it is supported by the sharder)
+        conn ||= sharded_connection.sharder.shard_for_key(:default) if sharded_connection.support_default_shard?
+
+        # Get connection proxy for our real connection
+        return nil unless conn
+        @real_conn = ::ActiveRecord::Base.coerce_to_connection_proxy(conn, DbCharmer.connections_should_exist?)
+      end
+
+      def method_missing(meth, *args, &block)
+        # Fail on database statements
+        if ::ActiveRecord::ConnectionAdapters::DatabaseStatements.instance_methods.member?(meth.to_s)
+          raise ::ActiveRecord::ConnectionNotEstablished, "You have to switch connection on your model before using it!"
+        end
+
+        # Fail if no connection has been established yet
+        unless real_connection
+          raise ::ActiveRecord::ConnectionNotEstablished, "No real connection to proxy this method to!"
+        end
+
+        if real_connection.kind_of?(DbCharmer::Sharding::StubConnection)
+          raise ::ActiveRecord::ConnectionNotEstablished, "You have to switch connection on your model before using it!"
+        end
+
+        # Proxy the call to our real connection target
+        real_connection.__send__(meth, *args, &block)
+      end
+    end
+  end
+end

data/lib/db_charmer/sharding.rb

@@ -0,0 +1,18 @@
+module DbCharmer
+  module Sharding
+    autoload :Connection, 'db_charmer/sharding/connection'
+    autoload :StubConnection, 'db_charmer/sharding/stub_connection'
+    autoload :Method, 'db_charmer/sharding/method'
+
+    @@sharded_connections = {}
+
+    def self.register_connection(config)
+      name = config[:name] or raise ArgumentError, "No :name in connection!"
+      @@sharded_connections[name] = DbCharmer::Sharding::Connection.new(config)
+    end
+
+    def self.sharded_connection(name)
+      @@sharded_connections[name] or raise ArgumentError, "Invalid sharded connection name!"
+    end
+  end
+end

data/lib/db_charmer/version.rb

@@ -0,0 +1,10 @@
+module DbCharmer
+  module Version
+    MAJOR = 1
+    MINOR = 7
+    PATCH = "01"
+    BUILD = nil
+
+    STRING = [MAJOR, MINOR, PATCH, BUILD].compact.join('.')
+  end
+end

data/lib/db_charmer.rb

@@ -0,0 +1,192 @@
+# In Rails 2.2 they did not add it to the autoload so it won't work w/o this require
+require 'active_record/version' unless defined?(::ActiveRecord::VERSION::MAJOR)
+
+module DbCharmer
+  # Configure autoload
+  autoload :Sharding, 'db_charmer/sharding'
+  autoload :Version,  'db_charmer/version'
+  module ActionController
+    autoload :ForceSlaveReads, 'db_charmer/action_controller/force_slave_reads'
+  end
+
+  # Used in all Rails3-specific places
+  def self.rails3?
+    ::ActiveRecord::VERSION::MAJOR > 2
+  end
+
+  # Used in all Rails2-specific places
+  def self.rails2?
+    ::ActiveRecord::VERSION::MAJOR == 2
+  end
+
+  # Accessors
+  @@connections_should_exist = true
+  mattr_accessor :connections_should_exist
+
+  # Try to detect current environment or use development by default
+  if defined?(Rails)
+    @@env = Rails.env
+  elsif ENV['RAILS_ENV']
+    @@env = ENV['RAILS_ENV']
+  elsif ENV['RACK_ENV']
+    @@env = ENV['RACK_ENV']
+  else
+    @@env = 'development'
+  end
+  mattr_accessor :env
+
+  def self.connections_should_exist?
+    !! connections_should_exist
+  end
+
+  # Extend ActionController to support forcing slave reads
+  def self.enable_controller_magic!
+    ::ActionController::Base.extend(DbCharmer::ActionController::ForceSlaveReads::ClassMethods)
+    ::ActionController::Base.send(:include, DbCharmer::ActionController::ForceSlaveReads::InstanceMethods)
+  end
+
+  def self.logger
+    return Rails.logger if defined?(Rails)
+    @logger ||= Logger.new(STDERR)
+  end
+
+  def self.with_remapped_databases(mappings, &proc)
+    old_mappings = ::ActiveRecord::Base.db_charmer_database_remappings
+    begin
+      ::ActiveRecord::Base.db_charmer_database_remappings = mappings
+      if mappings[:master] || mappings['master']
+        with_all_hijacked(&proc)
+      else
+        proc.call
+      end
+    ensure
+      ::ActiveRecord::Base.db_charmer_database_remappings = old_mappings
+    end
+  end
+
+  def self.hijack_new_classes?
+    @@hijack_new_classes
+  end
+
+private
+
+  @@hijack_new_classes = false
+  def self.with_all_hijacked
+    old_hijack_new_classes = @@hijack_new_classes
+    begin
+      @@hijack_new_classes = true
+      subclasses_method = DbCharmer.rails3? ? :descendants : :subclasses
+      ::ActiveRecord::Base.send(subclasses_method).each do |subclass|
+        subclass.hijack_connection!
+      end
+      yield
+    ensure
+      @@hijack_new_classes = old_hijack_new_classes
+    end
+  end
+end
+
+# Add useful methods to global object
+require 'db_charmer/core_extensions'
+
+require 'db_charmer/connection_factory'
+require 'db_charmer/connection_proxy'
+require 'db_charmer/force_slave_reads'
+
+# Add our custom class-level attributes to AR models
+require 'db_charmer/active_record/class_attributes'
+ActiveRecord::Base.extend(DbCharmer::ActiveRecord::ClassAttributes)
+
+# Enable connections switching in AR
+require 'db_charmer/active_record/connection_switching'
+ActiveRecord::Base.extend(DbCharmer::ActiveRecord::ConnectionSwitching)
+
+# Enable AR logging extensions
+if DbCharmer.rails3?
+  require 'db_charmer/rails3/abstract_adapter/connection_name'
+  require 'db_charmer/rails3/active_record/log_subscriber'
+  ActiveRecord::LogSubscriber.send(:include, DbCharmer::ActiveRecord::LogSubscriber)
+  ActiveRecord::ConnectionAdapters::AbstractAdapter.send(:include, DbCharmer::AbstractAdapter::ConnectionName)
+else
+  require 'db_charmer/rails2/abstract_adapter/log_formatting'
+  ActiveRecord::ConnectionAdapters::AbstractAdapter.send(:include, DbCharmer::AbstractAdapter::LogFormatting)
+end
+
+# Enable connection proxy in AR
+require 'db_charmer/active_record/multi_db_proxy'
+ActiveRecord::Base.extend(DbCharmer::ActiveRecord::MultiDbProxy::ClassMethods)
+ActiveRecord::Base.extend(DbCharmer::ActiveRecord::MultiDbProxy::MasterSlaveClassMethods)
+ActiveRecord::Base.send(:include, DbCharmer::ActiveRecord::MultiDbProxy::InstanceMethods)
+
+# Enable connection proxy for relations
+if DbCharmer.rails3?
+  require 'db_charmer/rails3/active_record/relation_method'
+  require 'db_charmer/rails3/active_record/relation/connection_routing'
+  ActiveRecord::Base.extend(DbCharmer::ActiveRecord::RelationMethod)
+  ActiveRecord::Relation.send(:include, DbCharmer::ActiveRecord::Relation::ConnectionRouting)
+end
+
+# Enable connection proxy for scopes (rails 2.x only)
+if DbCharmer.rails2?
+  require 'db_charmer/rails2/active_record/named_scope/scope_proxy'
+  ActiveRecord::NamedScope::Scope.send(:include, DbCharmer::ActiveRecord::NamedScope::ScopeProxy)
+end
+
+# Enable connection proxy for associations
+# WARNING: Inject methods to association class right here (they proxy include calls somewhere else, so include does not work)
+module ActiveRecord
+  module Associations
+    class AssociationProxy
+      def proxy?
+        true
+      end
+
+      def on_db(con, proxy_target = nil, &block)
+        proxy_target ||= self
+        @reflection.klass.on_db(con, proxy_target, &block)
+      end
+
+      def on_slave(con = nil, &block)
+        @reflection.klass.on_slave(con, self, &block)
+      end
+
+      def on_master(&block)
+        @reflection.klass.on_master(self, &block)
+      end
+    end
+  end
+end
+
+# Enable multi-db migrations
+require 'db_charmer/active_record/migration/multi_db_migrations'
+ActiveRecord::Migration.extend(DbCharmer::ActiveRecord::Migration::MultiDbMigrations)
+
+# Enable the magic
+if DbCharmer.rails3?
+  require 'db_charmer/rails3/active_record/master_slave_routing'
+else
+  require 'db_charmer/rails2/active_record/master_slave_routing'
+end
+
+require 'db_charmer/active_record/sharding'
+require 'db_charmer/active_record/db_magic'
+ActiveRecord::Base.extend(DbCharmer::ActiveRecord::DbMagic)
+
+# Setup association preload magic
+require 'db_charmer/active_record/association_preload'
+ActiveRecord::Base.extend(DbCharmer::ActiveRecord::AssociationPreload)
+
+# Open up really useful API method
+ActiveRecord::AssociationPreload::ClassMethods.send(:public, :preload_associations)
+
+class ::ActiveRecord::Base
+  class << self
+    def inherited_with_hijacking(subclass)
+      out = inherited_without_hijacking(subclass)
+      hijack_connection! if DbCharmer.hijack_new_classes?
+      out
+    end
+
+    alias_method_chain :inherited, :hijacking
+  end
+end