yam-db-charmer 1.7.4.0 → 1.7.4.8

data/lib/db_charmer.rb

@@ -130,8 +130,9 @@ ActiveRecord::Base.send(:include, DbCharmer::ActiveRecord::MultiDbProxy::Instanc
 # Enable connection proxy for relations
 if DbCharmer.rails3?
   require 'db_charmer/rails3/active_record/relation_method'
+  require 'db_charmer/rails3/active_record/relation/arel_engine'
   require 'db_charmer/rails3/active_record/relation/connection_routing'
-  ActiveRecord::Base.extend(DbCharmer::ActiveRecord::RelationMethod)
+  ActiveRecord::Base.extend(DbCharmer::ActiveRecord::Relation::ArelEngine)
   ActiveRecord::Relation.send(:include, DbCharmer::ActiveRecord::Relation::ConnectionRouting)
 end
 

data/lib/db_charmer/active_record/class_attributes.rb

@@ -96,6 +96,15 @@ module DbCharmer
         raise "Mappings must be nil or respond to []" if mappings && (! mappings.respond_to?(:[]))
         @@db_charmer_database_remappings = mappings || { }
       end
+
+      # For sharded models, return the table name without the schema prefix
+      def table_name_without_schema
+        if self.respond_to?(:orig_table_name)
+          self.orig_table_name
+        else
+          self.table_name
+        end
+      end
     end
   end
 end

data/lib/db_charmer/active_record/connection_switching.rb

@@ -28,6 +28,10 @@ module DbCharmer
 
       #-----------------------------------------------------------------------------------------------------------------
       def hijack_connection!
+
+        # do not hijack anything if the model is not handled by db_charmer
+        return unless self.respond_to?(:db_charmer_extended)
+
         return if self.respond_to?(:connection_with_magic)
         class << self
           # Make sure we check our accessors before going to the default connection retrieval method

data/lib/db_charmer/active_record/db_magic.rb

@@ -3,9 +3,18 @@ module DbCharmer
     module DbMagic
 
       def db_magic(opt = {})
+
+        # define a method that can be used to check if
+        # a model is extended by db_charmer
+        create_true_method(:db_charmer_extended)
+    
         # Make sure we could use our connections management here
         hijack_connection!
 
+        if DbCharmer.rails3?
+          self.extend(DbCharmer::ActiveRecord::RelationMethod) unless self.respond_to?(:relation_with_db_charmer)
+        end
+
         # Should requested connections exist in the config?
         should_exist = opt.has_key?(:should_exist) ? opt[:should_exist] : DbCharmer.connections_should_exist?
 
@@ -33,6 +42,16 @@ module DbCharmer
         end
       end
 
+      def create_true_method(method_name)
+        unless self.respond_to?(method_name.to_sym)
+          (class << self; self; end).instance_eval do
+            define_method method_name.to_sym do
+              true
+            end
+          end
+        end
+      end
+
     private
 
       def setup_children_magic(opt)

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

@@ -31,6 +31,10 @@ module DbCharmer
 
           unless DbCharmer.rails31?
             def migrate_with_db_wrapper(direction)
+              # define a method that can be used to check if
+              # a model is extended by db_charmer
+              ::ActiveRecord::Base.create_true_method(:db_charmer_extended)
+
               if names = multi_db_names
                 names.each do |multi_db_name|
                   on_db(multi_db_name) do
@@ -42,6 +46,22 @@ module DbCharmer
               end
             end
 
+            # Run some migration method on all shards
+            def on_each_shard(sharded_connection)
+              db_magic :sharded_connection => sharded_connection
+              if names = multi_db_names
+                names.each do |multi_db_name|
+                  on_db(multi_db_name) do
+                    yield
+                  end
+                end
+              else
+                raise ArgumentError, 'no shard names defined'
+              end
+            ensure
+              self.multi_db_names = nil
+            end
+
             def on_db(db_name)
               s = "Switching connection to "
               if db_name.is_a?(Hash)

data/lib/db_charmer/active_record/sharding.rb

@@ -12,6 +12,16 @@ module DbCharmer
         on_db(conn, proxy_target, &block)
       end
 
+      # The block passed to this function should return an object responding to key_field
+      def create_with_allocated_shard(key_field = :id, &block)
+        group = sharded_connection.sharder.least_loaded_group
+        shard_info = shard_info_by_group_id(group_id)
+        conn = shard_connection_config(shard_info, group.id)
+        key_obj = on_db(conn, proxy_target, &block)
+        allocate_new_block_for_key_on_group(key_obj.send(key_field), group)
+        key_obj
+      end
+
       # Run on default shard (if supported by the sharding method)
       def on_default_shard(proxy_target = nil, &block)
         raise ArgumentError, "No sharded connection configured!" unless sharded_connection

data/lib/db_charmer/connection_factory.rb

@@ -12,10 +12,10 @@ module DbCharmer
       @@connection_classes = {}
     end
 
-    # Establishes connection or return an existing one from cache
-    def self.connect(connection_name, should_exist = true)
-      connection_name = connection_name.to_s
-      @@connection_classes[connection_name] ||= establish_connection(connection_name, should_exist)
+    # Uses the passed in connection class or establishes a new one connection if one exists it will return an existing one from cache
+    def self.connect(connection_or_name, should_exist = true)
+      connection_or_name = connection_or_name.to_s unless connection_or_name.respond_to?(:connection)
+      @@connection_classes[connection_or_name] ||= establish_connection(connection_or_name, should_exist)
     end
 
     # Establishes connection or return an existing one from cache (not using AR database configs)
@@ -33,9 +33,14 @@ module DbCharmer
     end
 
     # Establish connection with a specified name
-    def self.establish_connection(connection_name, should_exist = true)
-      abstract_class = generate_abstract_class(connection_name, should_exist)
-      DbCharmer::ConnectionProxy.new(abstract_class, connection_name)
+    def self.establish_connection(connection_or_name, should_exist = true)
+      abstract_class =
+          if connection_or_name.respond_to?(:connection)
+            connection_or_name
+          else
+            generate_abstract_class(connection_or_name, should_exist)
+          end
+      DbCharmer::ConnectionProxy.new(abstract_class, connection_or_name.to_s)
     end
 
     # Establish connection with a specified name (not using AR database configs)

data/lib/db_charmer/rails3/active_record/relation/arel_engine.rb

@@ -0,0 +1,22 @@
+module DbCharmer
+  module ActiveRecord
+    module Relation
+      module ArelEngine
+
+        def self.extended(base)
+          class << base
+            alias_method_chain :arel_engine, :db_charmer
+          end
+        end
+
+        # Use the model itself an engine for Arel, do not fall back to AR::Base
+        # That guarantees the correct connection will be used in Arel::Table
+        # for determining list of columns and stuff.
+        def arel_engine_with_db_charmer(*)
+          self
+        end
+
+      end
+    end
+  end
+end

data/lib/db_charmer/rails3/active_record/relation_method.rb

@@ -5,14 +5,13 @@ module DbCharmer
       def self.extended(base)
         class << base
           alias_method_chain :relation, :db_charmer
-          alias_method_chain :arel_engine, :db_charmer
         end
       end
 
       # Create a relation object and initialize its default connection
       def relation_with_db_charmer(*args, &block)
         relation_without_db_charmer(*args, &block).tap do |rel|
-          rel.db_charmer_connection = self.retrieve_connection
+          rel.db_charmer_connection = self.connection
           rel.db_charmer_enable_slaves = self.db_charmer_slaves.any?
           rel.db_charmer_connection_is_forced = !db_charmer_top_level_connection?
 
@@ -23,11 +22,6 @@ module DbCharmer
         end
       end
 
-      # Use the model itself an engine for Arel, do not fall back to AR::Base
-      def arel_engine_with_db_charmer(*)
-        self
-      end
-
     end
   end
 end

data/lib/db_charmer/sharding.rb

@@ -11,7 +11,7 @@ module DbCharmer
       @@sharded_connections[name] = DbCharmer::Sharding::Connection.new(config)
 
       # Enable multi-db migrations
-      if config[:method] == :db_block_schema_map
+      if @@sharded_connections[name].sharder.is_a?(DbCharmer::Sharding::Method::DbBlockSchemaMap)
         ::ActiveRecord::Base.extend(DbCharmer::Sharding::Method::SchemaTableNamePrefix)
       end
     end

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

@@ -14,6 +14,7 @@ module DbCharmer
         
         # Shard connection info model
         class Shard < ::ActiveRecord::Base
+          db_magic
           validates_presence_of :db_host
           validates_presence_of :db_port
           validates_presence_of :db_user
@@ -25,6 +26,7 @@ module DbCharmer
 
         # Table group info model
         class Group < ::ActiveRecord::Base
+          db_magic
           validates_presence_of :shard_id
           belongs_to :shard, :class_name => 'DbCharmer::Sharding::Method::DbBlockGroupMap::Shard'
         end

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

@@ -32,13 +32,11 @@ module DbCharmer
         def initialize(config)
           @name = config[:name] or raise(ArgumentError, "Missing required :name parameter!")
 
-          if config[:method] == :db_block_schema_map
-            if ::ActiveRecord::Base.configurations[DbCharmer.env]['adapter'] != 'postgresql'
-              raise(ArgumentError, 'DbBlockSchemaMap method can only be used with the postgresql adapter')
-            end
-          end
-          
           @connection = DbCharmer::ConnectionFactory.connect(config[:connection], true)
+          if self.is_a?(DbCharmer::Sharding::Method::DbBlockSchemaMap) && ::ActiveRecord::Base.configurations[DbCharmer.env]['adapter'] != 'postgresql'
+            raise(ArgumentError, 'DbBlockSchemaMap method can only be used with the postgresql adapter')
+          end
+
           @block_size = (config[:block_size] || 10000).to_i
 
           @map_table = config[:map_table] or raise(ArgumentError, "Missing required :map_table parameter!")

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

@@ -15,6 +15,7 @@ module DbCharmer
 
         # Shard connection info model
         class Shard < ::ActiveRecord::Base
+          db_magic
           validates_presence_of :db_host
           validates_presence_of :db_port
           validates_presence_of :db_user
@@ -26,6 +27,7 @@ module DbCharmer
 
         # Table group info model
         class Group < ::ActiveRecord::Base
+          db_magic
           validates_presence_of :shard_id
           belongs_to :shard, :class_name => 'DbCharmer::Sharding::Method::DbBlockSchemaMap::Shard'
         end
@@ -48,12 +50,13 @@ module DbCharmer
           # 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
+          slave_host = shard.respond_to?(:db_slave_host) ? shard.db_slave_host : ''
           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,
+            :slave_host => slave_host,
             :port => shard.db_port,
             :username => shard.db_user,
             :password => shard.db_pass,
@@ -64,6 +67,26 @@ module DbCharmer
           )
         end
 
+        #---------------------------------------------------------------------------------------------------------------
+        def allocate_new_block_for_key_on_group(key, group)
+          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}
+                           (start_id, end_id, group_id, block_size, created_at, updated_at) VALUES
+                           (#{start_id}, #{end_id}, #{group.id}, #{block_size}, NOW(), NOW())
+          SQL
+          connection.execute(sql, "Allocate new block")
+
+          # Increment the blocks counter on the shard
+          group_class.update_counters(group.id, :blocks_count => +1)
+
+          # Retry block search after creation
+          block_for_key(key)
+        end
+
+
         #---------------------------------------------------------------------------------------------------------------
         def create_shard(params)
           params = params.symbolize_keys
@@ -141,6 +164,9 @@ module DbCharmer
           #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] + '.'
+            if self.to_s!='ActiveRecord::Base' && @dbcharmer_table_name_prefix.blank?
+              @orig_table_name = table_name if defined?(:original_table_name)
+            end
           else
             new_prefix = '' if self.to_s=='ActiveRecord::Base' # this is for migrations
           end
@@ -160,15 +186,19 @@ module DbCharmer
             # 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
+              reset_cached_table_name(@orig_table_name)
             end
           end
         end
 
+        def orig_table_name
+          @orig_table_name ? @orig_table_name : table_name
+        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
+        def reset_cached_table_name(orig_table_name)
+          self.table_name = "#{full_table_name_prefix}#{orig_table_name}#{table_name_suffix}"
           @arel_table = nil
           @relation = nil
         end

data/lib/db_charmer/version.rb

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

data/lib/tasks/databases.rake

@@ -89,6 +89,10 @@ end
 # find all schema sharded databases and drop them
 def drop_schema_shard_databases(config)
   exit unless Rails.env=='test'
+
+  DbCharmer.env = 'test'
+  DbCharmer::ConnectionFactory.reset!
+  
   config.each do |name, sub_config|
     next unless sub_config.is_a?(Hash)
     next unless sub_config['database']
@@ -101,7 +105,14 @@ def drop_schema_shard_databases(config)
     # 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
+
+    # Rescue this as there may not be any tables in the admin db yet
+    begin
+      dbgm.drop_all_shard_databases
+    rescue ActiveRecord::StatementInvalid => e
+    end
+    
+    dbgm.connection.disconnect!  # or subsequent rake tasks may fail
   end
 end
 

data/lib/tasks/test.rake

@@ -14,6 +14,7 @@ namespace :db_charmer do
     end
 
     task :purge do
+      Rake::Task['db_charmer:drop:schema_shards'].invoke
       Rake::Task['db_charmer:drop'].invoke
       Rake::Task['db_charmer:create'].invoke
     end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification 
 name: yam-db-charmer
 version: !ruby/object:Gem::Version 
-  hash: 119
+  hash: 103
   prerelease: 
   segments: 
   - 1
   - 7
   - 4
-  - 0
-  version: 1.7.4.0
+  - 8
+  version: 1.7.4.8
 platform: ruby
 authors: 
 - Oleksiy Kovyrin
@@ -17,7 +17,7 @@ autorequire:
 bindir: bin
 cert_chain: []
 
-date: 2012-09-18 00:00:00 -07:00
+date: 2012-10-17 00:00:00 -07:00
 default_executable: 
 dependencies: 
 - !ruby/object:Gem::Dependency 
@@ -129,6 +129,7 @@ files:
 - lib/db_charmer/rails3/abstract_adapter/connection_name.rb
 - lib/db_charmer/rails3/active_record/log_subscriber.rb
 - lib/db_charmer/rails3/active_record/master_slave_routing.rb
+- lib/db_charmer/rails3/active_record/relation/arel_engine.rb
 - lib/db_charmer/rails3/active_record/relation/connection_routing.rb
 - lib/db_charmer/rails3/active_record/relation_method.rb
 - lib/db_charmer/rails31/active_record/migration/command_recorder.rb