switchman 0.0.1

data/db/migrate/20130328212039_create_switchman_shards.rb

@@ -0,0 +1,9 @@
+class CreateSwitchmanShards < ActiveRecord::Migration
+  def change
+    create_table :switchman_shards do |t|
+      t.string :name
+      t.string :database_server_id
+      t.boolean :default, :default => false, :null => false
+    end
+  end
+end

data/db/migrate/20130328224244_create_default_shard.rb

@@ -0,0 +1,9 @@
+class CreateDefaultShard < ActiveRecord::Migration
+  def up
+    unless Switchman::Shard.default.is_a?(Switchman::Shard)
+      Switchman::Shard.reset_column_information
+      Switchman::Shard.create!(:default => true)
+      Switchman::Shard.default(true)
+    end
+  end
+end

data/lib/switchman.rb

@@ -0,0 +1,9 @@
+require "shackles"
+require "switchman/engine"
+
+module Switchman
+  def self.config
+    # TODO: load from yaml
+    {}
+  end
+end

data/lib/switchman/active_record/abstract_adapter.rb

@@ -0,0 +1,11 @@
+module Switchman
+  module ActiveRecord
+    module AbstractAdapter
+      attr_writer :shard
+
+      def shard
+        @shard || Shard.default
+      end
+    end
+  end
+end

data/lib/switchman/active_record/association.rb

@@ -0,0 +1,108 @@
+module Switchman
+  module ActiveRecord
+    module Association
+      def self.included(klass)
+        %w{build_record 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
+          @owner.shard
+        else
+          Shard.default
+        end
+      end
+
+      def build_record_with_sharding(*args)
+        self.shard.activate { build_record_without_sharding(*args) }
+      end
+
+      def load_target_with_sharding
+        self.shard.activate { load_target_without_sharding }
+      end
+
+      def scoped_with_sharding
+        shard_value = @reflection.options[:multishard] ? @owner : self.shard
+        self.shard.activate { scoped_without_sharding.shard(shard_value, :association) }
+      end
+    end
+
+    module Builder
+      module Association
+        def self.included(klass)
+          klass.descendants.each{|d| d.valid_options += [:multishard]}
+        end
+      end
+    end
+
+    module Preloader
+      module Association
+        def self.included(klass)
+          klass.send(:remove_method, :associated_records_by_owner)
+          klass.send(:remove_method, :owners_by_key)
+          klass.send(:remove_method, :scoped)
+        end
+
+        def associated_records_by_owner
+          owners_map = owners_by_key
+
+          if klass.nil? || owners_map.empty?
+            records = []
+          else
+            # Some databases impose a limit on the number of ids in a list (in Oracle it's 1000)
+            # Make several smaller queries if necessary or make one query if the adapter supports it
+            records = Shard.partition_by_shard(owners) do |partitioned_owners|
+              sliced_owners = partitioned_owners.each_slice(model.connection.in_clause_length || partitioned_owners.size)
+              sliced_owners.map do |slice|
+                relative_owner_keys = slice.map do |owner|
+                  key = owner[owner_key_name]
+                  if key && owner.class.sharded_column?(owner_key_name)
+                    key = Shard.relative_id_for(key, owner.shard, Shard.current(owner.class.shard_category))
+                  end
+                  key && key.to_s
+                end
+                relative_owner_keys.compact!
+                records_for(relative_owner_keys)
+              end
+            end
+            records.flatten!
+          end
+
+          # Each record may have multiple owners, and vice-versa
+          records_by_owner = Hash[owners.map { |owner| [owner, []] }]
+          records.each do |record|
+            owner_key = record[association_key_name]
+            owner_key = Shard.global_id_for(owner_key, record.shard) if owner_key && record.class.sharded_column?(association_key_name)
+
+            owners_map[owner_key.to_s].each do |owner|
+              records_by_owner[owner] << record
+            end
+          end
+          records_by_owner
+        end
+
+        def owners_by_key
+          @owners_by_key ||= owners.group_by do |owner|
+            key = owner[owner_key_name]
+            key = Shard.global_id_for(key, owner.shard) if key && owner.class.sharded_column?(owner_key_name)
+            key && key.to_s
+          end
+        end
+
+        def scoped
+          build_scope
+        end
+      end
+    end
+
+    module CollectionProxy
+      def shard(*args)
+        scoped.shard(*args)
+      end
+    end
+  end
+end
+

data/lib/switchman/active_record/attribute_methods.rb

@@ -0,0 +1,104 @@
+module Switchman
+  module ActiveRecord
+    module AttributeMethods
+      module ClassMethods
+
+        def sharded_primary_key?
+          self != Shard && shard_category != :unsharded && integral_id?
+        end
+
+        def sharded_foreign_key?(column_name)
+          reflection = reflection_for_integer_attribute(column_name.to_s)
+          return false unless reflection
+          reflection.options[:polymorphic] || reflection.klass.sharded_primary_key?
+        end
+
+        def sharded_column?(column_name)
+          column_name = column_name.to_s
+          @sharded_column_values ||= {}
+          unless @sharded_column_values.has_key?(column_name)
+            @sharded_column_values[column_name] = (column_name == primary_key && sharded_primary_key?) || sharded_foreign_key?(column_name)
+          end
+          @sharded_column_values[column_name]
+        end
+
+        protected
+
+        def reflection_for_integer_attribute(attr_name)
+          columns_hash[attr_name] && columns_hash[attr_name].type == :integer &&
+              reflections.find { |_, r| r.belongs_to? && r.foreign_key == attr_name }.try(:last)
+        rescue ::ActiveRecord::StatementInvalid
+          # this is for when models are referenced in initializers before migrations have been run
+          nil
+        end
+
+        def define_method_global_attribute(attr_name)
+          if sharded_column?(attr_name)
+            generated_attribute_methods.module_eval <<-RUBY, __FILE__, __LINE__ + 1
+              def __temp__
+                Shard.global_id_for(original_#{attr_name}, shard)
+              end
+              alias_method 'global_#{attr_name}', :__temp__
+              undef_method :__temp__
+            RUBY
+          end
+        end
+
+        def define_method_local_attribute(attr_name)
+          if sharded_column?(attr_name)
+            generated_attribute_methods.module_eval <<-RUBY, __FILE__, __LINE__ + 1
+              def __temp__
+                Shard.local_id_for(original_#{attr_name}).first
+              end
+              alias_method 'local_#{attr_name}', :__temp__
+              undef_method :__temp__
+            RUBY
+          end
+        end
+
+        def shard_category_code_for_reflection(reflection)
+          if reflection
+            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"
+            else
+              # otherwise we can just return a symbol for the statically known type of the association
+              reflection.klass.shard_category.inspect
+            end
+          else
+            shard_category.inspect
+          end
+        end
+
+        def define_method_original_attribute(attr_name)
+          if sharded_column?(attr_name)
+            reflection = reflection_for_integer_attribute(attr_name)
+            generated_attribute_methods.module_eval <<-RUBY, __FILE__, __LINE__ + 1
+              # rename the original method to original_
+              alias_method 'original_#{attr_name}', '#{attr_name}'
+              # and replace with one that transposes the id
+              def __temp__
+                Shard.relative_id_for(original_#{attr_name}, shard, Shard.current(#{shard_category_code_for_reflection(reflection)}))
+              end
+              alias_method '#{attr_name}', :__temp__
+              undef_method :__temp__
+
+              alias_method 'original_#{attr_name}=', '#{attr_name}='
+              def __temp__(new_value)
+                self.original_#{attr_name} = Shard.relative_id_for(new_value, Shard.current(#{shard_category_code_for_reflection(reflection)}), shard)
+              end
+              alias_method '#{attr_name}=', :__temp__
+              undef_method :__temp__
+            RUBY
+          end
+        end
+      end
+
+      def self.included(klass)
+        klass.extend(ClassMethods)
+        klass.attribute_method_prefix "global_", "local_", "original_"
+      end
+    end
+  end
+end

data/lib/switchman/active_record/base.rb

@@ -0,0 +1,95 @@
+module Switchman
+  module ActiveRecord
+    module Base
+      module ClassMethods
+        attr_writer :shard_category
+        delegate :shard, :to => :scoped
+
+        def shard_category
+          @shard_category || :default
+        end
+
+        def shard_category=(category)
+          categories = Shard.const_get(:CATEGORIES)
+          if categories[shard_category]
+            categories[shard_category].delete(self)
+            categories.delete(shard_category) if categories[shard_category].empty?
+          end
+          # TODO: de-initialize the proxy
+          categories[category] ||= []
+          categories[category] << self
+          @shard_category = category
+          connection_handler.initialize_categories(superclass)
+        end
+
+        def integral_id?
+          if @integral_id == nil
+            @integral_id = columns_hash[primary_key].type == :integer
+          end
+          @integral_id
+        end
+
+        def transaction(*args, &block)
+          return super if !block || !current_scope
+          super(*args) do
+            current_scope.activate(&block)
+          end
+        end
+      end
+
+      def self.included(klass)
+        klass.extend(ClassMethods)
+        klass.set_callback(:initialize, :before) { @shard = Shard.current(self.class.shard_category) }
+      end
+
+      def shard
+        @shard || Shard.current(self.class.shard_category) || Shard.default
+      end
+
+      def shard=(new_shard)
+        raise ::ActiveRecord::ReadOnlyRecord if !self.new_record? || @shard_set_in_stone
+        if shard != new_shard
+          # TODO: adjust foreign keys
+          @shard = new_shard
+        end
+      end
+
+      def save(*args)
+        @shard_set_in_stone = true
+        shard.activate(self.class.shard_category) { super }
+      end
+
+      def save!(*args)
+        @shard_set_in_stone = true
+        shard.activate(self.class.shard_category) { super }
+      end
+
+      def destroy
+        shard.activate(self.class.shard_category) { super }
+      end
+
+      def clone
+        result = super
+        # TODO: adjust foreign keys
+        # don't use the setter, cause the foreign keys are already
+        # relative to this shard
+        result.instance_variable_set(:@shard, self.shard)
+        result
+      end
+
+      def transaction(&block)
+        shard.activate do
+          super
+        end
+      end
+
+      def hash
+        global_id.hash
+      end
+
+      def to_param
+        Shard.short_id_for(self.id) if persisted?
+      end
+    end
+  end
+end

data/lib/switchman/active_record/calculations.rb

@@ -0,0 +1,63 @@
+module Switchman
+  module ActiveRecord
+    module Calculations
+      def self.included(klass)
+        %w{execute_simple_calculation pluck}.each do |method|
+          klass.alias_method_chain(method, :sharding)
+        end
+      end
+
+      def pluck_with_sharding(column_name)
+        target_shard = Shard.current(klass.shard_category)
+        self.activate do |relation, shard|
+          results = relation.pluck_without_sharding(column_name)
+          if klass.sharded_column?(column_name.to_s)
+            results = results.map{|result| Shard.relative_id_for(result, shard, target_shard)}
+          end
+          results
+        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)
+        else
+          result = self.activate{ |relation| relation.send(:execute_simple_calculation_without_sharding, operation, column_name, distinct) }
+          if result.is_a?(Array)
+            case operation
+            when "count", "sum"
+              result = result.sum
+            when "minimum"
+              result = result.min
+            when "maximum"
+              result = result.max
+            end
+          end
+        end
+        result
+      end
+
+      def calculate_average(column_name, distinct)
+        # See activerecord#execute_simple_calculation
+        relation = reorder(nil)
+        column = aggregate_column(column_name)
+        relation.select_values = [operation_over_aggregate_column(column, "average", distinct).as("average"),
+                                  operation_over_aggregate_column(column, "count", distinct).as("count")]
+
+        initial_results = relation.activate{ |rel| @klass.connection.select_all(rel) }
+        if initial_results.is_a?(Array)
+          initial_results.each do |r|
+            r["average"] = type_cast_calculated_value(r["average"], nil, "average")
+            r["count"] = type_cast_calculated_value(r["count"], nil, "count")
+          end
+          result = initial_results.map{|r| r["average"] * r["count"]}.sum / initial_results.map{|r| r["count"]}.sum
+        else
+          result = type_cast_calculated_value(initial_results["average"], nil, "average")
+        end
+        result
+      end
+    end
+  end
+end

data/lib/switchman/active_record/connection_handler.rb

@@ -0,0 +1,147 @@
+require_dependency 'switchman/connection_pool_proxy'
+require_dependency 'switchman/shard'
+
+module Switchman
+  module ActiveRecord
+    module ConnectionHandler
+      def self.make_sharing_automagic(config, shard)
+        key = config[:adapter] == 'postgresql' ? :schema_search_path : :database
+
+        # we may not be able to connect to this shard yet, cause it might be an empty database server
+        shard_name = shard.name rescue nil
+        return unless shard_name
+
+        config[:shard_name] ||= shard_name
+        if !config[key] || config[key] == shard_name
+          # this may truncate the schema_search_path if it was not specified in database.yml
+          # but that's what our old behavior was anyway, so I guess it's okay
+          config[key] = '%{shard_name}'
+        end
+      end
+
+      def self.included(klass)
+        klass.alias_method_chain(:establish_connection, :sharding)
+        klass.alias_method_chain(:remove_connection, :sharding)
+      end
+
+      def establish_connection_with_sharding(name, spec)
+        establish_connection_without_sharding(name, spec)
+
+        # this is the first place that the adapter would have been required; but now we
+        # need this addition ASAP since it will be called when loading the default shard below
+        if defined?(::ActiveRecord::ConnectionAdapters::PostgreSQLAdapter)
+          require "switchman/active_record/postgresql_adapter"
+          ::ActiveRecord::ConnectionAdapters::PostgreSQLAdapter.send(:include, ActiveRecord::PostgreSQLAdapter)
+        end
+
+        model = name.constantize
+        pool = connection_pools[spec]
+
+        first_time = !Shard.instance_variable_get(:@default)
+        if first_time
+          # Have to cache the default shard before we insert sharding, otherwise the first access
+          # to sharding will recurse onto itself trying to access column information
+          Shard.default
+
+          # automatically change config to allow for sharing connections with simple config
+          ConnectionHandler.make_sharing_automagic(spec.config, Shard.default)
+          ConnectionHandler.make_sharing_automagic(Shard.default.database_server.config, Shard.default)
+        end
+        @shard_connection_pools ||= { Shard.default.database_server.id => pool }
+
+        proxy = ConnectionPoolProxy.new(model.shard_category,
+                                        pool,
+                                        @shard_connection_pools)
+        connection_pools[spec] = proxy
+
+        initialize_categories(model)
+        @class_to_pool[name] = proxy
+
+        # reload the default shard if we just got a new connection
+        # to where the Shards table is
+        # DON'T do it if we're not the current connection handler - that means
+        # we're in the middle of switching environments, and we don't want to
+        # establish a connection with incorrect settings
+        if (model == ::ActiveRecord::Base || model == Shard) && self == ::ActiveRecord::Base.connection_handler && !first_time
+          Shard.default(true) unless first_time
+          proxy.disconnect!
+        end
+
+        if first_time
+          # do the change for other database servers, now that we can switch shards
+          if Shard.default.is_a?(Shard)
+            DatabaseServer.all.each do |server|
+              next if server == Shard.default.database_server
+              shard = server.shards.where(:name => nil).first
+              shard ||= Shard.new(:database_server => server)
+              ConnectionHandler.make_sharing_automagic(server.config, shard)
+            end
+          end
+          # we may have established some connections above trying to infer the shard's name.
+          # close them, so that someone that doesn't expect them doesn't try to fork
+          # without closing them
+          self.clear_all_connections!
+        end
+      end
+
+      def remove_connection_with_sharding(model)
+        uninitialize_ar(model) if @class_to_pool[model.name].is_a?(ConnectionPoolProxy)
+        remove_connection_without_sharding(model)
+      end
+
+      private
+
+      def uninitialize_ar(model = ::ActiveRecord::Base)
+        # take the proxies out
+        @class_to_pool.each_key do |model_name|
+          pool_model = model_name.constantize
+          # only de-proxify models that inherit from what we're uninitializing
+          next unless pool_model == model || pool_model < model
+          proxy = @class_to_pool[model_name]
+          next unless proxy.is_a?(ConnectionPoolProxy)
+
+          # make sure we're switched back to the default shard for the
+          # connection that will remain
+          if proxy.connected?
+            Shard.default.activate { proxy.connection }
+          end
+          connection_pools[proxy.spec] = proxy.default_pool
+          @class_to_pool[model_name] = proxy.default_pool
+        end
+
+        # prune dups that were created for implementing shard categories
+        @class_to_pool.each_key do |model_name|
+          next if model_name == ::ActiveRecord::Base.name
+          pool_model = model_name.constantize
+          @class_to_pool.delete(model_name) if retrieve_connection_pool(pool_model.superclass) == @class_to_pool[model_name]
+        end
+      end
+
+      # semi-private
+      public
+      def initialize_categories(model = ::ActiveRecord::Base)
+        # now set up pools for models that inherit from this model, but with a different
+        # sharding category
+        Shard.const_get(:CATEGORIES).each do |category, models|
+          next if category == :default
+          next if category == model.shard_category
+
+          this_proxy = nil
+          Array(models).each do |category_model|
+            category_model = category_model.constantize if category_model.is_a? String
+            next unless category_model < model
+
+            # don't replace existing connections
+            next if @class_to_pool[category_model.name]
+
+            default_pool = retrieve_connection_pool(model)
+            default_pool = default_pool.default_pool if default_pool.is_a?(ConnectionPoolProxy)
+            # look for an existing compatible proxy for this category
+            this_proxy ||= ConnectionPoolProxy.new(category_model.shard_category, default_pool, @shard_connection_pools)
+            @class_to_pool[category_model.name] = this_proxy
+          end
+        end
+      end
+    end
+  end
+end