yam-db-charmer 1.7.01 → 1.7.4.0

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

@@ -43,7 +43,7 @@ module DbCharmer
             # Auto-allocate new blocks
             block ||= allocate_new_block_for_key(key)
           rescue ::ActiveRecord::StatementInvalid => e
-            raise unless e.message.include?('Duplicate entry')
+            raise unless e.message.include?('Duplicate entry') || e.message.include?('duplicate key')
             block = block_for_key(key)
           end
 
@@ -118,13 +118,8 @@ module DbCharmer
 
           # 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()
+            INSERT INTO #{map_table} (start_id, end_id, shard_id, block_size, created_at, updated_at)
+            VALUES (#{start_id}, #{end_id}, #{shard.id}, #{block_size}, NOW(), NOW())
           SQL
           connection.execute(sql, "Allocate new block")
 
@@ -201,7 +196,7 @@ module DbCharmer
 
         # Prepare model for working with our shards table
         def prepare_shard_model
-          ShardInfo.set_table_name(shards_table)
+          ShardInfo.table_name = shards_table
           ShardInfo.switch_connection_to(connection)
         end
 

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

@@ -0,0 +1,179 @@
+require File.join(File.dirname(__FILE__), '/db_block_group_map_base')
+
+# This is a more sophisticated sharding method based on a two layer database-backed
+# blocks map that holds block-shard associations. Record blocks are mapped to groups
+# (= database schemas) and groups are mapped to shards (= databases).
+#
+# It automatically creates new blocks for new keys and assigns them to existing groups.
+# Warning: make sure to create at least one shard and one group before inserting any records.
+#
+module DbCharmer
+  module Sharding
+    module Method
+      class DbBlockSchemaMap
+        include DbBlockGroupMapBase
+
+        # Shard connection info model
+        class Shard < ::ActiveRecord::Base
+          validates_presence_of :db_host
+          validates_presence_of :db_port
+          validates_presence_of :db_user
+          validates_presence_of :db_pass, :unless => Proc.new { |shard| shard.db_pass=='' }
+          validates_presence_of :schema_name_prefix
+
+          has_many :groups, :class_name => 'DbCharmer::Sharding::Method::DbBlockSchemaMap::Group'
+        end
+
+        # Table group info model
+        class Group < ::ActiveRecord::Base
+          validates_presence_of :shard_id
+          belongs_to :shard, :class_name => 'DbCharmer::Sharding::Method::DbBlockSchemaMap::Shard'
+        end
+
+        def connection_name(shard_id)
+          "db_charmer_db_block_schema_map_#{name}_s%d" % shard_id
+        end
+
+        def schema_name(schema_name_prefix, group_id)
+          "%s_%05d" % [ schema_name_prefix, group_id ]
+        end
+        
+        #---------------------------------------------------------------------------------------------------------------
+        # Create configuration (use mapping connection as a template)
+        def shard_connection_config(shard, group_id)
+          # Format connection name
+          connection_name = connection_name(shard.id)
+          schema_name = schema_name(shard.schema_name_prefix, group_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.config.clone.merge(
+            # Name for the connection factory
+            :connection_name => connection_name,
+            # Connection params
+            :host => shard.db_host,
+            :slave_host => shard.db_slave_host,
+            :port => shard.db_port,
+            :username => shard.db_user,
+            :password => shard.db_pass,
+            :database => shard.db_name,
+            :schema_name => schema_name,
+            :shard_id => shard.id,
+            :sharder_name => self.name
+          )
+        end
+
+        #---------------------------------------------------------------------------------------------------------------
+        def create_shard(params)
+          params = params.symbolize_keys
+          [ :db_host, :db_port, :db_user, :db_name, :schema_name_prefix ].each do |arg|
+            raise ArgumentError, "Missing required parameter: #{arg}" unless params[arg]
+          end
+
+          # Prepare model
+          prepare_shard_models
+
+          # Create the record
+          Shard.create! do |shard|
+            shard.db_host       = params[:db_host]
+            shard.db_slave_host = (params[:db_slave_host] || '') if shard.respond_to? :db_slave_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]
+            shard.schema_name_prefix = params[:schema_name_prefix]
+          end
+        end
+
+        def create_group(shard, open, enabled)
+          # Prepare model
+          prepare_shard_models
+
+          # Create the record
+          group = Group.create! do |group|
+            group.shard_id = shard.id
+            group.open     = open
+            group.enabled  = enabled
+          end
+
+          old_connection = Group.connection
+          conn_config = shard_connection_config(shard, group.id)
+          Group.switch_connection_to(conn_config)
+
+          schema_name = schema_name(shard.schema_name_prefix, group.id)
+
+          # create schema only if it doesn't already exist
+          sql = "SELECT schema_name FROM information_schema.schemata WHERE schema_name='#{schema_name}'"
+          existing_schema = Group.connection.execute(sql)
+          
+          unless existing_schema.first
+            sql = "CREATE SCHEMA #{schema_name}"
+            Group.connection.execute(sql)
+          end
+          
+          Group.switch_connection_to(old_connection)
+        end
+
+        def shard_connections
+          # Find all groups
+          prepare_shard_models
+          groups = Group.all(:conditions => { :enabled => true }, :include => :shard)
+          # Map them to shards
+          groups.map { |group| shard_connection_config(group.shard, group.id) }
+        end
+
+        # Prepare model for working with our shards table
+        def prepare_shard_models
+          Shard.switch_connection_to(connection)
+          Shard.table_name = shards_table
+
+          Group.switch_connection_to(connection)
+          Group.table_name = groups_table
+        end
+      end
+
+      # To be mixed in AR#base.
+      # Allows a schema name to be set at the same time a db connection is selected.
+      module SchemaTableNamePrefix
+        def set_schema_table_name_prefix(con)
+          
+          #Rails.logger.debug "set_schema_table_name_prefix: self=#{self} @dbcharmer_table_name_prefix=#{@dbcharmer_table_name_prefix}"
+          if con.is_a?(Hash) && con[:schema_name]
+            new_prefix = con[:schema_name] + '.'
+          else
+            new_prefix = '' if self.to_s=='ActiveRecord::Base' # this is for migrations
+          end
+          
+          # remove the old dbcharmer prefix first
+          if !@dbcharmer_table_name_prefix.blank? && @dbcharmer_table_name_prefix != new_prefix
+            self.table_name_prefix = self.table_name_prefix.gsub(/#{Regexp.escape(@dbcharmer_table_name_prefix)}/, '')
+          end
+
+          # Set the new table_name_prefix
+          if new_prefix && @dbcharmer_table_name_prefix != new_prefix
+            @dbcharmer_table_name_prefix = new_prefix
+            self.table_name_prefix = "#{new_prefix}#{self.table_name_prefix}"
+
+            # Reset all forms of table_name that were memoized in Rails.
+            # Don't do it in the context of migrations where the current class
+            # is AR::Base and there is no actual table name.
+            unless self.to_s=='ActiveRecord::Base'
+              #Rails.logger.debug "set_schema_table_name_prefix: resetting"
+              reset_cached_table_name
+            end
+          end
+        end
+
+        # Rails memoizes table_name and table_name_prefix at the time the models are loaded.
+        # This method forces refreshing those names
+        def reset_cached_table_name
+          self.reset_table_name
+          @arel_table = nil
+          @relation = nil
+        end
+      end
+      
+    end
+  end
+end

data/lib/db_charmer/sharding/method.rb

@@ -4,7 +4,9 @@ module DbCharmer
       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 :DbBlockGroupMapBase, 'db_charmer/sharding/method/db_block_group_map_base'
       autoload :DbBlockGroupMap, 'db_charmer/sharding/method/db_block_group_map'
+      autoload :DbBlockSchemaMap, 'db_charmer/sharding/method/db_block_schema_map'
     end
   end
 end

data/lib/db_charmer/sharding.rb

@@ -9,10 +9,38 @@ module DbCharmer
     def self.register_connection(config)
       name = config[:name] or raise ArgumentError, "No :name in connection!"
       @@sharded_connections[name] = DbCharmer::Sharding::Connection.new(config)
+
+      # Enable multi-db migrations
+      if config[:method] == :db_block_schema_map
+        ::ActiveRecord::Base.extend(DbCharmer::Sharding::Method::SchemaTableNamePrefix)
+      end
     end
 
+    # name is the is the sharded_connection name passed to db_magic in the model
     def self.sharded_connection(name)
       @@sharded_connections[name] or raise ArgumentError, "Invalid sharded connection name!"
     end
+
+    # Return the DbCharmer::Sharding::Connection that matches name
+    # name is the the connection name calculated in the sharder for a specific database
+    def self.sharded_connection_by_connection_name(name)
+      connection = @@sharded_connections.detect do |c|
+        c[1].config[:connection] == name.to_sym
+      end
+      connection[1] if connection
+    end
+
+    def self.sharder_for_connection_name(sharder_name)
+      connection = sharded_connection(sharder_name)
+      return connection.sharder if connection
+    end
+
+    # Return shard record for the given config hash.
+    # connection_config is the config hash taken from a PostgreSQLAdapter instance.
+    def self.shard_for_connection_name(connection_config)
+      sharder = sharder_for_connection_name(connection_config[:sharder_name])
+      return unless sharder
+      sharder.shard_class.where(:id => connection_config[:shard_id]).first
+    end
   end
 end

data/lib/db_charmer/version.rb

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

data/lib/db_charmer.rb

@@ -1,6 +1,10 @@
 # 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)
 
+# Load the rake tasks so that they become available in the application
+path = File.join(File.dirname(__FILE__),'tasks/*.rake')
+Dir[path].each { |ext|  load ext } if defined?(Rake)
+
 module DbCharmer
   # Configure autoload
   autoload :Sharding, 'db_charmer/sharding'
@@ -14,6 +18,11 @@ module DbCharmer
     ::ActiveRecord::VERSION::MAJOR > 2
   end
 
+  # Used in all Rails3.1-specific places
+  def self.rails31?
+    rails3? && ::ActiveRecord::VERSION::MINOR >= 1
+  end
+
   # Used in all Rails2-specific places
   def self.rails2?
     ::ActiveRecord::VERSION::MAJOR == 2
@@ -134,32 +143,43 @@ 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
+association_proxy_class = DbCharmer.rails31? ? ActiveRecord::Associations::CollectionProxy : ActiveRecord::Associations::AssociationProxy
+association_proxy_class.class_eval do
+  def proxy?
+    true
+  end
 
-      def on_db(con, proxy_target = nil, &block)
-        proxy_target ||= self
-        @reflection.klass.on_db(con, proxy_target, &block)
-      end
+  if DbCharmer.rails31?
+    def on_db(con, proxy_target = nil, &block)
+      proxy_target ||= self
+      @association.klass.on_db(con, proxy_target, &block)
+    end
 
-      def on_slave(con = nil, &block)
-        @reflection.klass.on_slave(con, self, &block)
-      end
+    def on_slave(con = nil, &block)
+      @association.klass.on_slave(con, self, &block)
+    end
 
-      def on_master(&block)
-        @reflection.klass.on_master(self, &block)
-      end
+    def on_master(&block)
+      @association.klass.on_master(self, &block)
+    end
+  else
+    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
 
-# Enable multi-db migrations
 require 'db_charmer/active_record/migration/multi_db_migrations'
-ActiveRecord::Migration.extend(DbCharmer::ActiveRecord::Migration::MultiDbMigrations)
+ActiveRecord::Migration.send(:include, DbCharmer::ActiveRecord::Migration::MultiDbMigrations)
 
 # Enable the magic
 if DbCharmer.rails3?
@@ -173,11 +193,19 @@ 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)
+if DbCharmer.rails31?
+  require 'db_charmer/rails31/active_record/preloader/association'
+  ActiveRecord::Associations::Preloader::Association.send(:include, DbCharmer::ActiveRecord::Preloader::Association)
+  require 'db_charmer/rails31/active_record/preloader/has_and_belongs_to_many'
+  ActiveRecord::Associations::Preloader::HasAndBelongsToMany.send(:include, DbCharmer::ActiveRecord::Preloader::HasAndBelongsToMany)
+else
+  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)
+end
 
-# Open up really useful API method
-ActiveRecord::AssociationPreload::ClassMethods.send(:public, :preload_associations)
 
 class ::ActiveRecord::Base
   class << self

data/lib/tasks/databases.rake

@@ -23,7 +23,7 @@ namespace :db_charmer do
 
   desc 'Create the databases defined in config/database.yml for the current RAILS_ENV'
   task :create => "db:load_config" do
-    create_core_and_sub_database(ActiveRecord::Base.configurations[RAILS_ENV])
+    create_core_and_sub_database(ActiveRecord::Base.configurations[Rails.env])
   end
 
   def create_core_and_sub_database(config)
@@ -45,11 +45,16 @@ namespace :db_charmer do
         local_database?(config) { drop_core_and_sub_database(config) }
       end
     end
+
+    task :schema_shards => :environment do
+      config = ::ActiveRecord::Base.configurations[Rails.env || 'development']
+      drop_schema_shard_databases(config)
+    end      
   end
 
   desc 'Drops the database for the current RAILS_ENV'
   task :drop => "db:load_config" do
-    config = ::ActiveRecord::Base.configurations[RAILS_ENV || 'development']
+    config = ::ActiveRecord::Base.configurations[Rails.env || 'development']
     begin
       drop_core_and_sub_database(config)
     rescue Exception => e
@@ -68,15 +73,36 @@ namespace :db_charmer do
 end
 
 def drop_core_and_sub_database(config)
+  exit unless Rails.env=='test'
   drop_database(config)
   config.each_value do | sub_config |
     next unless sub_config.is_a?(Hash)
     next unless sub_config['database']
     begin
       drop_database(sub_config)
-    rescue
-      $stderr.puts "#{config['database']} not exists"
+    rescue => e
+      $stderr.puts "#{e.to_s}, #{config['database']} not exists"
     end
   end
 end
 
+# find all schema sharded databases and drop them
+def drop_schema_shard_databases(config)
+  exit unless Rails.env=='test'
+  config.each do |name, sub_config|
+    next unless sub_config.is_a?(Hash)
+    next unless sub_config['database']
+
+    # find the database connection for the schema admin db
+    next unless sub_config['shard_db_name_prefix']
+    connection = DbCharmer::Sharding.sharded_connection_by_connection_name(name)
+    next unless connection
+
+    # itereate through entries in the shards_info table to find the
+    # databases that will be dropped
+    dbgm = DbCharmer::Sharding::Method::DbBlockSchemaMap.new(connection.config)
+    dbgm.drop_all_shard_databases
+  end
+end
+
+