switchman 0.0.1 → 1.0.0

checksums.yaml

@@ -1,7 +1,15 @@
 ---
-SHA1:
-  metadata.gz: cc14f7783fcf89a84b17781a64cc331aa29f0447
-  data.tar.gz: 7921d4093ae4c1194b03724fa10013d9c71dbcae
+!binary "U0hBMQ==":
+  metadata.gz: !binary |-
+    ZGY5N2I3ZDdiMzViMzAzYmRjOTQxMGFjMmNmNjIwYzhjNTNiMDM2NQ==
+  data.tar.gz: !binary |-
+    NTc3ZGM5YTY5NTI3MTc4ODNiMjdlNDRmM2QwNjU5N2JmNzA2ZDg4Nw==
 SHA512:
-  metadata.gz: ec8355813399e1a4d0f01c48c299498735ee1bd913fe2083624e161d26e32533c9d5a0695f0ddfa726648f63292fd9a5e1f58ecbdacfc6d0f70274d148cae52b
-  data.tar.gz: b7f54166cf2d7eab085cfe199bd47be4d1566860de6805a67bc3614e86bc04eebef4d80b892b2c5a26a71d5604d5b80a62efb260deee6dbee6f7a8f5190acf34
+  metadata.gz: !binary |-
+    M2RjZDFiMzVhZDgyNzFhMTg2YjBmY2M2OTY0Nzg2OTc0NjkxZjkzOTdjYzU4
+    ZmE0OGY2YmI0OGUwM2FhOWFhYjFjMDk3MGI5OThmZDZlZWNiMWZjMTEzMDZl
+    ZDM4MmIyMTJhNDZlZTk1NDg3ODhiZjczMDYxNTQ1ZDE1NDY4ODU=
+  data.tar.gz: !binary |-
+    OTc4N2ZiNjNlY2YzMjk5NTMxN2ZhZWRiNzJmZDY0NWU0YmFmNGU0NzRlYmRi
+    YmRlMDcyMzY5MTU4ZWI3YWEwMDM1YTc1ZDZhYTAwMzU5OTY4MDMxZTNjYjE4
+    YWFlNGU5Y2ZlMzBhZGYwZDRjNjY3OTI5YThkZjc5NWFlNGQyMzA=

data/app/models/switchman/shard.rb

@@ -14,6 +14,7 @@ module Switchman
           :unsharded => [Shard]
       }
     private_constant :CATEGORIES
+    @shard_category = :unsharded
 
     attr_accessible :name, :database_server, :default
 
@@ -22,7 +23,6 @@ module Switchman
 
     after_save :clear_cache
 
-
     class << self
       def categories
         CATEGORIES.keys
@@ -35,10 +35,12 @@ module Switchman
           # is up and running everywhere else).  This includes for looking up the
           # default shard itself. This also needs to be a local so that this method
           # can be re-entrant
-          default = DefaultShard.new
+          default = DefaultShard.instance
 
           # the first time we need a dummy dummy for re-entrancy to avoid looping on ourselves
           @default ||= default
+          # forget our current shard activations - it might have "another" default shard serialized there
+          active_shards.clear
 
           # Now find the actual record, if it exists; rescue the fake default if the table doesn't exist
           @default = Shard.find_by_default(true) || default rescue default
@@ -77,7 +79,7 @@ module Switchman
         cached_shards[id] ||= Shard.default.activate do
           # can't simply cache the AR object since Shard has a custom serializer
           # that calls this method
-          attributes = Rails.cache.fetch(['shard', id].join('/')) do
+          attributes = ::Rails.cache.fetch(['shard', id].join('/')) do
             shard = find_by_id(id)
             shard.try(:attributes) || :nil
           end
@@ -99,16 +101,35 @@ module Switchman
         @cached_shards = {}
       end
 
-      # options
-      #  :parallel - true/false to execute in parallel, or a integer of how many
-      #              sub-processes per database server. Note that parallel
-      #              invocation currently uses forking, so should be used sparingly
-      #              because errors are not raised, and you cannot get results back
-      def with_each_shard(scope = nil, categories = nil, options = {})
+      # ==== Parameters
+      #
+      # * +shards+ - an array or relation of Shards to iterate over
+      # * +categories+ - an array of categories to activate
+      # * +options+ -
+      #    :parallel - true/false to execute in parallel, or a integer of how many
+      #                sub-processes per database server. Note that parallel
+      #                invocation currently uses forking, so should be used sparingly
+      #                because errors are not raised, and you cannot get results back
+      #    :exception - :ignore, :raise, :defer (wait until the end and raise the first
+      #                error), or a proc
+      def with_each_shard(*args)
+        raise ArgumentError("wrong number of arguments (#{args.length} for 0...3)") if args.length > 3
+
         unless default.is_a?(Shard)
           return Array(yield)
         end
 
+        options = args.extract_options!
+        if args.length == 1
+          if Array === args.first && args.first.first.is_a?(Symbol)
+            categories = args.first
+          else
+            scope = args.first
+          end
+        else
+          scope, categories = args
+        end
+
         parallel = case options[:parallel]
                      when true
                        1
@@ -117,6 +138,7 @@ module Switchman
                      else
                        options[:parallel]
                    end
+        options.delete(:parallel)
         scope ||= Shard.order("database_server_id IS NOT NULL, database_server_id, id")
 
         if parallel > 0
@@ -125,7 +147,7 @@ module Switchman
             database_servers = scope.reorder('database_server_id').select(:database_server_id).uniq.
                 map(&:database_server).compact.uniq
             scopes = Hash[database_servers.map do |server|
-              server_scope = server.shards(scope)
+              server_scope = server.shards.merge(scope)
               if parallel == 1
                 subscopes = [server_scope]
               else
@@ -166,13 +188,20 @@ module Switchman
             if scopes.length == 1 && subscopes.length == 1
               exception_pipe.first.close
               exception_pipe.last.close
-              return with_each_shard(subscopes.first, categories) { yield }
+              return with_each_shard(subscopes.first, categories, options) { yield }
             end
             subscopes.each_with_index do |subscope, idx|
+              if subscopes.length > 1
+                name = "#{server.id} #{idx + 1}"
+              else
+                name = server.id
+              end
+
               details = Open4.pfork4(lambda do
                 begin
                   ::ActiveRecord::Base.clear_all_connections!
-                  with_each_shard(subscope, categories) { yield }
+                  $0 = [$0, ARGV, name].flatten.join(' ')
+                  with_each_shard(subscope, categories, options) { yield }
                 rescue Exception => e
                   exception_pipe.last.write(Marshal.dump(e))
                   exception_pipe.last.flush
@@ -183,11 +212,6 @@ module Switchman
               details[1].close
               fds.concat details[2..3]
               pids << details[0]
-              if subscopes.length > 1
-                name = "#{server.id} #{idx + 1}"
-              else
-                name = server.id
-              end
               fd_to_name_map[details[2]] = name
               fd_to_name_map[details[3]] = name
             end
@@ -239,16 +263,32 @@ module Switchman
         end
 
         result = []
+        exception = nil
         scope.each do |shard|
           # shard references a database server that isn't configured in this environment
           next unless shard.database_server
           close_connections_if_needed.call(shard)
           shard.activate(*categories) do
-            result.concat Array(yield)
+            begin
+              result.concat Array(yield)
+            rescue
+              case options[:exception]
+              when :ignore
+              when :defer
+                exception ||= $!
+              when Proc
+                options[:exception].call
+              when :raise
+                raise
+              else
+                raise
+              end
+            end
           end
           previous_shard = shard
         end
         close_connections_if_needed.call(Shard.current)
+        raise exception if exception
         result
       end
 
@@ -475,11 +515,16 @@ module Switchman
       local_id + self.id * IDS_PER_SHARD
     end
 
+    def ==(rhs)
+      return true if rhs.is_a?(DefaultShard) && default?
+      super
+    end
+
     private
 
     def clear_cache
       Shard.default.activate do
-        Rails.cache.delete(['shard', id].join('/'))
+        ::Rails.cache.delete(['shard', id].join('/'))
       end
     end
 

data/lib/switchman/action_controller/caching.rb

@@ -0,0 +1,26 @@
+module Switchman
+  module ActionController
+    module Caching
+      module ConfigMethods
+        # always go through Rails.cache, which will give you the cache store
+        # appropriate to the current shard.
+        def cache_store
+          ::Rails.cache
+        end
+
+        # disallow assigning to ActionController::Base.cache_store or
+        # ActionController::Base#cache_store for the same reasons we disallow
+        # assigning to Rails.cache
+        def cache_store=(cache)
+          raise NoMethodError
+        end
+      end
+
+      include ConfigMethods
+
+      def self.included(base)
+        base.extend(ConfigMethods)
+      end
+    end
+  end
+end

data/lib/switchman/active_record/abstract_adapter.rb

@@ -1,3 +1,5 @@
+require 'switchman/sharded_instrumenter'
+
 module Switchman
   module ActiveRecord
     module AbstractAdapter
@@ -6,6 +8,15 @@ module Switchman
       def shard
         @shard || Shard.default
       end
+
+      def initialize_with_shard(*args)
+        initialize_without_shard(*args)
+        @instrumenter = Switchman::ShardedInstrumenter.new(@instrumenter, self)
+      end
+
+      def self.included(klass)
+        klass.alias_method_chain :initialize, :shard
+      end
     end
   end
 end

data/lib/switchman/active_record/association.rb

@@ -2,14 +2,14 @@ module Switchman
   module ActiveRecord
     module Association
       def self.included(klass)
-        %w{build_record load_target scoped}.each do |method|
+        %w{build_record creation_attributes load_target scoped}.each do |method|
           klass.alias_method_chain(method, :sharding)
         end
       end
 
       def shard
-        # polymorphic associations assume the same shard as the owning item
         if @reflection.options[:polymorphic] || @reflection.klass.shard_category == @owner.class.shard_category
+          # polymorphic associations assume the same shard as the owning item
           @owner.shard
         else
           Shard.default
@@ -26,7 +26,41 @@ module Switchman
 
       def scoped_with_sharding
         shard_value = @reflection.options[:multishard] ? @owner : self.shard
-        self.shard.activate { scoped_without_sharding.shard(shard_value, :association) }
+        @owner.shard.activate { scoped_without_sharding.shard(shard_value, :association) }
+      end
+
+      def creation_attributes_with_sharding
+        attributes = creation_attributes_without_sharding
+
+        # translate keys
+        if reflection.macro.in?([:has_one, :has_many]) && !options[:through]
+          attributes[reflection.foreign_key] = Shard.relative_id_for(owner[reflection.active_record_primary_key], owner.shard, self.shard)
+        end
+        attributes
+      end
+    end
+
+    module BelongsToAssociation
+      def self.included(klass)
+        klass.send(:alias_method_chain, :replace_keys, :sharding)
+      end
+
+      def replace_keys_with_sharding(record)
+        if record && record.class.sharded_column?(reflection.association_primary_key(record.class))
+          foreign_id = record[reflection.association_primary_key(record.class)]
+          owner[reflection.foreign_key] = Shard.relative_id_for(foreign_id, record.shard, owner.shard)
+        else
+          replace_keys_without_sharding(record)
+        end
+      end
+
+      def shard
+        if @owner.class.sharded_column?(@reflection.foreign_key) &&
+            foreign_id = @owner[@reflection.foreign_key]
+          Shard.shard_for(foreign_id, @owner.shard)
+        else
+          super
+        end
       end
     end
 

data/lib/switchman/active_record/attribute_methods.rb

@@ -25,8 +25,9 @@ module Switchman
         protected
 
         def reflection_for_integer_attribute(attr_name)
+          attr_name = attr_name.to_s
           columns_hash[attr_name] && columns_hash[attr_name].type == :integer &&
-              reflections.find { |_, r| r.belongs_to? && r.foreign_key == attr_name }.try(:last)
+              reflections.find { |_, r| r.belongs_to? && r.foreign_key.to_s == attr_name }.try(:last)
         rescue ::ActiveRecord::StatementInvalid
           # this is for when models are referenced in initializers before migrations have been run
           nil
@@ -61,7 +62,7 @@ module Switchman
             if reflection.options[:polymorphic]
               # a polymorphic association has to be discovered at runtime. This code ends up being something like
               # context_type.try(:constantize).try(:shard_category) || :default
-              "#{reflection.foreign_type}.try(:constantize).try(:shard_category) || :default"
+              "read_attribute(:#{reflection.foreign_type}).try(:constantize).try(:shard_category) || :default"
             else
               # otherwise we can just return a symbol for the statically known type of the association
               reflection.klass.shard_category.inspect

data/lib/switchman/active_record/base.rb

@@ -2,7 +2,6 @@ module Switchman
   module ActiveRecord
     module Base
       module ClassMethods
-        attr_writer :shard_category
         delegate :shard, :to => :scoped
 
         def shard_category
@@ -24,22 +23,43 @@ module Switchman
 
         def integral_id?
           if @integral_id == nil
-            @integral_id = columns_hash[primary_key].type == :integer
+            @integral_id = columns_hash[primary_key].try(:type) == :integer
           end
           @integral_id
         end
 
-        def transaction(*args, &block)
-          return super if !block || !current_scope
-          super(*args) do
-            current_scope.activate(&block)
+        def transaction(*args)
+          if current_scope
+            current_scope.activate do
+              db = Shard.current(shard_category).database_server
+              if ::Shackles.environment != db.shackles_environment
+                db.unshackle { super }
+              else
+                super
+              end
+            end
+          else
+            db = Shard.current(shard_category).database_server
+            if ::Shackles.environment != db.shackles_environment
+              db.unshackle { super }
+            else
+              super
+            end
           end
         end
       end
 
       def self.included(klass)
         klass.extend(ClassMethods)
-        klass.set_callback(:initialize, :before) { @shard = Shard.current(self.class.shard_category) }
+        klass.set_callback(:initialize, :before) do
+          unless @shard
+            if self.class.sharded_primary_key?
+              @shard = Shard.shard_for(self[self.class.primary_key], Shard.current(self.class.shard_category))
+            else
+              @shard = Shard.current(self.class.shard_category)
+            end
+          end
+        end
       end
 
       def shard
@@ -49,23 +69,25 @@ module Switchman
       def shard=(new_shard)
         raise ::ActiveRecord::ReadOnlyRecord if !self.new_record? || @shard_set_in_stone
         if shard != new_shard
-          # TODO: adjust foreign keys
+          attributes.each do |attr, value|
+            self[attr] = Shard.relative_id_for(value, shard, new_shard) if self.class.sharded_column?(attr)
+          end
           @shard = new_shard
         end
       end
 
       def save(*args)
         @shard_set_in_stone = true
-        shard.activate(self.class.shard_category) { super }
+        self.class.with_scope(self.class.shard(shard), :overwrite) { super }
       end
 
       def save!(*args)
         @shard_set_in_stone = true
-        shard.activate(self.class.shard_category) { super }
+        self.class.with_scope(self.class.shard(shard), :overwrite) { super }
       end
 
       def destroy
-        shard.activate(self.class.shard_category) { super }
+        self.class.with_scope(self.class.shard(shard), :overwrite) { super }
       end
 
       def clone
@@ -77,9 +99,9 @@ module Switchman
         result
       end
 
-      def transaction(&block)
-        shard.activate do
-          super
+      def transaction(options={}, &block)
+        shard.activate(self.class.shard_category) do
+          self.class.transaction(options, &block)
         end
       end
 
@@ -88,7 +110,14 @@ module Switchman
       end
 
       def to_param
-        Shard.short_id_for(self.id) if persisted?
+        short_id = Shard.short_id_for(id)
+        short_id && short_id.to_s
+      end
+
+      def initialize_dup(*args)
+        copy = super
+        @shard_set_in_stone = false
+        copy
       end
     end
   end

data/lib/switchman/active_record/calculations.rb

@@ -18,11 +18,10 @@ module Switchman
         end
       end
 
-      # TODO: grouped calculations
       def execute_simple_calculation_with_sharding(operation, column_name, distinct)
         operation = operation.to_s.downcase
         if operation == "average"
-          result = calculate_average(column_name, distinct)
+          result = calculate_simple_average(column_name, distinct)
         else
           result = self.activate{ |relation| relation.send(:execute_simple_calculation_without_sharding, operation, column_name, distinct) }
           if result.is_a?(Array)
@@ -39,7 +38,7 @@ module Switchman
         result
       end
 
-      def calculate_average(column_name, distinct)
+      def calculate_simple_average(column_name, distinct)
         # See activerecord#execute_simple_calculation
         relation = reorder(nil)
         column = aggregate_column(column_name)
@@ -58,6 +57,149 @@ module Switchman
         end
         result
       end
+
+      # See activerecord#execute_grouped_calculation
+      def execute_grouped_calculation(operation, column_name, distinct)
+        opts = grouped_calculation_options(operation.to_s.downcase, column_name, distinct)
+
+        relation = build_grouped_calculation_relation(opts)
+        target_shard = Shard.current(:default)
+
+        rows = relation.activate do |rel, shard|
+          calculated_data = @klass.connection.select_all(rel)
+
+          if opts[:association]
+            key_ids     = calculated_data.collect { |row| row[opts[:group_aliases].first] }
+            key_records = opts[:association].klass.base_class.where(:id => key_ids)
+            key_records = Hash[key_records.map { |r| [Shard.relative_id_for(r, shard, target_shard), r] }]
+          end
+
+          calculated_data.map do |row|
+            row[opts[:aggregate_alias]] = type_cast_calculated_value(
+                row[opts[:aggregate_alias]], column_for(opts[:column_name]), opts[:operation])
+            row['count'] = row['count'].to_i if opts[:operation] == 'average'
+
+            opts[:group_columns].each do |aliaz, column|
+              if opts[:associated] && (aliaz == opts[:group_aliases].first)
+                row[aliaz] = key_records[Shard.relative_id_for(row[aliaz], shard, target_shard)]
+              elsif column && @klass.sharded_column?(column.name)
+                row[aliaz] = Shard.relative_id_for(type_cast_calculated_value(row[aliaz], column), shard, target_shard)
+              end
+            end
+            row
+          end
+        end
+
+        compact_grouped_calculation_rows(rows, opts)
+      end
+
+      private
+      def grouped_calculation_options(operation, column_name, distinct)
+        opts = {:operation => operation, :column_name => column_name, :distinct => distinct}
+
+        opts[:aggregate_alias] = aggregate_alias_for(operation, column_name)
+        group_attrs = @group_values
+        if group_attrs.first.respond_to?(:to_sym)
+          association  = @klass.reflect_on_association(group_attrs.first.to_sym)
+          associated   = group_attrs.size == 1 && association && association.macro == :belongs_to # only count belongs_to associations
+          group_fields = Array(associated ? association.foreign_key : group_attrs)
+        else
+          group_fields = group_attrs
+        end
+
+        group_aliases = group_fields.map { |field| column_alias_for(field) }
+        group_columns = group_aliases.zip(group_fields).map { |aliaz,field|
+          [aliaz, column_for(field)]
+        }
+
+        opts.merge!(:association => association, :associated => associated,
+         :group_aliases => group_aliases, :group_columns => group_columns,
+         :group_fields => group_fields)
+
+        opts
+      end
+
+      def aggregate_alias_for(operation, column_name)
+        if operation == 'count' && column_name == :all
+          'count_all'
+        elsif operation == 'average'
+          'average'
+        else
+          column_alias_for(operation, column_name)
+        end
+      end
+
+      def build_grouped_calculation_relation(opts)
+        group = @klass.connection.adapter_name == 'FrontBase' ? opts[:group_aliases] : opts[:group_fields]
+
+        select_values = [
+            operation_over_aggregate_column(
+                aggregate_column(opts[:column_name]),
+                opts[:operation],
+                opts[:distinct]).as(opts[:aggregate_alias])
+        ]
+        if opts[:operation ]== 'average'
+          # include count in average so we can recalculate the average
+          # across all shards if needed
+          select_values << operation_over_aggregate_column(
+              aggregate_column(opts[:column_name]),
+              'count', opts[:distinct]).as('count')
+        end
+
+        select_values += @select_values unless @having_values.empty?
+        select_values.concat opts[:group_fields].zip(opts[:group_aliases]).map { |field,aliaz|
+          if field.respond_to?(:as)
+            field.as(aliaz)
+          else
+            "#{field} AS #{aliaz}"
+          end
+        }
+
+        relation = except(:group).group(group)
+        relation.select_values = select_values
+        relation
+      end
+
+      def compact_grouped_calculation_rows(rows, opts)
+        result = ::ActiveSupport::OrderedHash.new
+        rows.each do |row|
+          key = opts[:group_columns].map { |aliaz, column| row[aliaz] }
+          key = key.first if key.size == 1
+          value = row[opts[:aggregate_alias]]
+
+          if opts[:operation] == 'average'
+            if result.has_key?(key)
+              old_value, old_count = result[key]
+              new_count = old_count + row['count']
+              new_value = ((old_value * old_count) + (value * row['count'])) / new_count
+              result[key] = [new_value, new_count]
+            else
+              result[key] = [value, row['count']]
+            end
+          else
+            if result.has_key?(key)
+              case opts[:operation]
+              when "count", "sum"
+                result[key] += value
+              when "minimum"
+                result[key] = value if value < result[key]
+              when "maximum"
+                result[key] = value if value > result[key]
+              end
+            else
+              result[key] = value
+            end
+          end
+        end
+
+        if opts[:operation] == 'average'
+          result = Hash[result.map{|k, v| [k, v.first]}]
+        end
+
+        result
+      end
+
+
     end
   end
 end

data/lib/switchman/active_record/connection_handler.rb

@@ -54,6 +54,17 @@ module Switchman
                                         @shard_connection_pools)
         connection_pools[spec] = proxy
 
+        if first_time
+          if Shard.default.database_server.config[:prefer_slave]
+            Shard.default.database_server.shackle!
+          end
+
+          if Shard.default.is_a?(DefaultShard) && Shard.default.database_server.config[:slave]
+            Shard.default.database_server.shackle!
+            Shard.default(true)
+          end
+        end
+
         initialize_categories(model)
         @class_to_pool[name] = proxy
 

data/lib/switchman/active_record/connection_pool.rb

@@ -64,12 +64,14 @@ module Switchman
         }.map { |thread| thread.object_id }
         keys.each do |key|
           conns = @reserved_connections[key]
-          ActiveSupport::Deprecation.warn(<<-eowarn) if conn.in_use?
+          conns.each do |conn|
+            ::ActiveSupport::Deprecation.warn(<<-eowarn) if conn.in_use?
 Database connections will not be closed automatically, please close your
 database connection at the end of the thread by calling `close` on your
 connection.  For example: ActiveRecord::Base.connection.close
-          eowarn
-          conns.each { |conn| checkin conn }
+            eowarn
+            checkin conn
+          end
           @reserved_connections.delete(key)
         end
       end