switchman 0.0.1 → 1.0.0

data/lib/switchman/active_record/finder_methods.rb

@@ -20,6 +20,31 @@ module Switchman
           super
         end
       end
+
+      def find_or_instantiator_by_attributes(match, attributes, *args)
+        primary_shard.activate { super }
+      end
+
+      def exists?(id = false)
+        id = id.id if ActiveRecord::Base === id
+        return false if id.nil?
+
+        join_dependency = construct_join_dependency_for_association_find
+        relation = construct_relation_for_association_find(join_dependency)
+        relation = relation.except(:select, :order).select("1 AS one").limit(1)
+
+        case id
+          when Array, Hash
+            relation = relation.where(id)
+          else
+            relation = relation.where(table[primary_key].eq(id)) if id
+        end
+
+        activate { return true if connection.select_value(relation, "#{name} Exists") }
+        false
+      rescue ThrowResult
+        false
+      end
     end
   end
 end

data/lib/switchman/active_record/log_subscriber.rb

@@ -17,7 +17,8 @@ module Switchman
         name  = '%s (%.1fms)' % [payload[:name], event.duration]
         sql   = payload[:sql].squeeze(' ')
         binds = nil
-        connection = ObjectSpace._id2ref(payload[:connection_id])
+        shard = payload[:shard]
+        shard = "  [shard #{shard[:id]} #{shard[:env]}]" if shard
 
         unless (payload[:binds] || []).empty?
           binds = "  " + payload[:binds].map { |col,v|
@@ -36,7 +37,7 @@ module Switchman
           name = color(name, self.class::MAGENTA, true)
         end
 
-        debug "  #{name}  #{sql}#{binds}  [shard #{connection.shard.id} #{::Shackles.environment}]"
+        debug "  #{name}  #{sql}#{binds}#{shard}"
       end
     end
   end

data/lib/switchman/active_record/query_cache.rb

@@ -3,10 +3,81 @@ module Switchman
     # needs to be included in the same class as ::ActiveRecord::ConnectionAdapters::QueryCache
     # *after* that module is included
     module QueryCache
+      # thread local accessors to replace @query_cache_enabled
+      def query_cache_enabled
+        Thread.current[:query_cache_enabled]
+      end
+
+      def query_cache_enabled=(value)
+        Thread.current[:query_cache_enabled] = value
+      end
+
+      # basically wholesale repeat of the methods from the original (see
+      # https://github.com/rails/rails/blob/master/activerecord/lib/active_record/connection_adapters/abstract/query_cache.rb),
+      # but with self.query_cache_enabled and self.query_cache_enabled= instead
+      # of @query_cache_enabled.
+
+      def enable_query_cache!
+        self.query_cache_enabled = true
+      end
+
+      def disable_query_cache!
+        self.query_cache_enabled = false
+      end
+
+      def cache
+        old, self.query_cache_enabled = query_cache_enabled, true
+        yield
+      ensure
+        self.query_cache_enabled = old
+        clear_query_cache unless self.query_cache_enabled
+      end
+
+      def uncached
+        old, self.query_cache_enabled = query_cache_enabled, false
+        yield
+      ensure
+        self.query_cache_enabled = old
+      end
+
+      def select_all(arel, name = nil, binds = [])
+        if self.query_cache_enabled && !locked?(arel)
+          sql = to_sql(arel, binds)
+          cache_sql(sql, binds) { super(sql, name, binds) }
+        else
+          super
+        end
+      end
+
+      # no reason to define these on the including class directly. the super
+      # works just as well from a method on the included module
+      [:insert, :update, :delete].each do |method_name|
+        class_eval <<-end_code, __FILE__, __LINE__ + 1
+          def #{method_name}(*args)
+            clear_query_cache if self.query_cache_enabled
+            super
+          end
+        end_code
+      end
+
+      private
       def cache_sql(sql, *args, &block)
         # have to include the shard id in the cache key because of switching dbs on the same connection
         super("#{self.shard.id}::#{sql}", *args, &block)
       end
+
+      def self.included(base)
+        base.class_eval do
+          # when we call insert, update, and delete, we want it to find the
+          # definitions from this module (which will then find the definitions
+          # from ActiveRecord::ConnectionAdapters::DatabaseStatements as
+          # 'super'), not the ones defined on base by
+          # ActiveRecord::ConnectionAdapters::QueryCache.
+          remove_method :insert
+          remove_method :update
+          remove_method :delete
+        end
+      end
     end
   end
 end

data/lib/switchman/active_record/query_methods.rb

@@ -6,10 +6,12 @@ module Switchman
       #   An array or relation of shards
       #   An AR object (query runs against that object's associated_shards)
       # shard_source_value is one of:
-      #   :implicit - inferred from current shard when relation was created, or primary key where clause
-      #   :explicit - explicit set on the relation
-      #   :to_a - a special value that Relation#to_a uses when querying multiple shards to
-      #           remove primary keys from conditions that aren't applicable to the current shard
+      #   :implicit    - inferred from current shard when relation was created, or primary key where clause
+      #   :explicit    - explicit set on the relation
+      #   :association - a special value that scopes from associations use to use slightly different logic
+      #                  for foreign key transposition
+      #   :to_a        - a special value that Relation#to_a uses when querying multiple shards to
+      #                  remove primary keys from conditions that aren't applicable to the current shard
       attr_accessor :shard_value, :shard_source_value
 
       def shard(value, source = :explicit)
@@ -43,12 +45,11 @@ module Switchman
       end
 
       def build_where(opts, other = [])
-        source_shard = Shard.current(klass.shard_category)
         case opts
         when Hash, Arel::Nodes::Node
           predicates = super
           infer_shards_from_primary_key(predicates) if shard_source_value == :implicit && shard_value.is_a?(Shard)
-          predicates = transpose_predicates(predicates, source_shard, primary_shard) if shard_source_value != :explicit
+          predicates = transpose_predicates(predicates, nil, primary_shard) if shard_source_value != :explicit
           predicates
         else
           super
@@ -73,11 +74,24 @@ module Switchman
         end
       end
 
+      # the shard value as an array or a relation
+      def all_shards
+        case shard_value
+        when Shard, DefaultShard
+          [shard_value]
+        when ::ActiveRecord::Base
+          shard_value.respond_to?(:associated_shards) ? shard_value.associated_shards : [shard_value.shard]
+        else
+          shard_value
+        end
+      end
+
       private
       def infer_shards_from_primary_key(predicates)
         primary_key = predicates.detect do |predicate|
           predicate.is_a?(Arel::Nodes::Binary) && predicate.left.is_a?(Arel::Attributes::Attribute) &&
-          predicate.left.relation.engine == klass && klass.primary_key == predicate.left.name
+            predicate.left.relation.is_a?(Arel::Table) && predicate.left.relation.engine == klass &&
+            klass.primary_key == predicate.left.name
         end
         if primary_key
           case primary_key.right
@@ -115,45 +129,66 @@ module Switchman
         end
       end
 
-      def transposable_attribute_type(attribute)
-        return false unless attribute.is_a?(Arel::Attributes::Attribute)
-        if sharded_primary_key?(attribute)
-          return :primary
-        elsif sharded_foreign_key?(attribute)
-          return :foreign
+      def transposable_attribute_type(relation, column)
+        if sharded_primary_key?(relation, column)
+          :primary
+        elsif sharded_foreign_key?(relation, column)
+          :foreign
         end
       end
 
-      def sharded_foreign_key?(attribute)
-        @@foreign_keys ||= {}
-        @@foreign_keys[attribute.relation.table_name] ||= {}
-        if @@foreign_keys[attribute.relation.table_name].has_key?(attribute.name)
-          @@foreign_keys[attribute.relation.table_name][attribute.name]
-        else
-          models = attribute.relation.engine.descendants.select{|d| d.table_name == attribute.relation.table_name}
-          models << attribute.relation.engine unless attribute.relation.engine == ::ActiveRecord::Base
+      def models_for_table(table_name)
+        @@models_for_table ||= {}
+        @@models_for_table[table_name] ||= ::ActiveRecord::Base.descendants.select { |d| d.table_name == table_name }
+      end
 
-          @@foreign_keys[attribute.relation.table_name][attribute.name] = models.any?{|m| m.sharded_column?(attribute.name)}
-        end
+      def sharded_foreign_key?(relation, column)
+        models_for_table(relation.table_name).any? { |m| m.sharded_column?(column) }
       end
 
-      def sharded_primary_key?(attribute)
-        attribute.relation.engine.primary_key == attribute.name
+      def sharded_primary_key?(relation, column)
+        relation.engine.primary_key == column
       end
 
+      def source_shard_for_foreign_key(relation, column)
+        reflection = nil
+        models_for_table(relation.table_name).each do |model|
+          reflection = model.send(:reflection_for_integer_attribute, column)
+          break if reflection
+        end
+        return Shard.current(klass.shard_category) if reflection.options[:polymorphic]
+        Shard.current(reflection.klass.shard_category)
+      end
+
+      def relation_and_column(attribute)
+        column = attribute.name
+        attribute = attribute.relation if attribute.relation.is_a?(Arel::Nodes::TableAlias)
+        [attribute.relation, column]
+      end
       # semi-private
       public
       def transpose_predicates(predicates, source_shard, target_shard, remove_nonlocal_primary_keys = false)
         predicates.map do |predicate|
-          next predicate unless predicate.is_a?(Arel::Nodes::Binary) && type = transposable_attribute_type(predicate.left)
+          next predicate unless predicate.is_a?(Arel::Nodes::Binary)
+          next predicate unless predicate.left.is_a?(Arel::Attributes::Attribute)
+          relation, column = relation_and_column(predicate.left)
+          next predicate unless (type = transposable_attribute_type(relation, column))
 
           remove = true if type == :primary && remove_nonlocal_primary_keys && predicate.left.relation.engine == klass
+          current_source_shard =
+              if source_shard
+                source_shard
+              elsif type == :primary
+                Shard.current(klass.shard_category)
+              elsif type == :foreign
+                source_shard_for_foreign_key(relation, column)
+              end
 
           new_right_value = case predicate.right
           when Array
             local_ids = []
             predicate.right.each do |value|
-              local_id = Shard.relative_id_for(value, source_shard, target_shard)
+              local_id = Shard.relative_id_for(value, current_source_shard, target_shard)
               local_ids << local_id unless remove && local_id > Shard::IDS_PER_SHARD
             end
             local_ids
@@ -161,13 +196,13 @@ module Switchman
             # look for a bind param with a matching column name
             if @bind_params && idx = @bind_params.find_index{|b| b.is_a?(Array) && b.first.try(:name) == predicate.left}
               column, value = @bind_params[idx]
-              local_id = Shard.relative_id_for(value, source_shard, target_shard)
+              local_id = Shard.relative_id_for(value, current_source_shard, target_shard)
               local_id = [] if remove && local_id > Shard::IDS_PER_SHARD
               @bind_params[idx] = [column, local_id]
             end
             predicate.right
           else
-            local_id = Shard.relative_id_for(predicate.right, source_shard, target_shard)
+            local_id = Shard.relative_id_for(predicate.right, current_source_shard, target_shard)
             local_id = [] if remove && local_id > Shard::IDS_PER_SHARD
             local_id
           end

data/lib/switchman/active_record/relation.rb

@@ -3,7 +3,12 @@ module Switchman
     module Relation
       def self.included(klass)
         klass::SINGLE_VALUE_METHODS.concat [ :shard, :shard_source ]
-        %w{initialize exec_queries update_all delete_all}.each do |method|
+
+        %w{exec_queries update_all delete_all}.each do |method|
+          klass.alias_method_chain(method, :deshackles)
+        end
+
+        %w{initialize exec_queries update_all delete_all new create create!}.each do |method|
           klass.alias_method_chain(method, :sharding)
         end
       end
@@ -23,6 +28,28 @@ module Switchman
         relation
       end
 
+      def new_with_sharding(*args, &block)
+        primary_shard.activate(klass.shard_category) { new_without_sharding(*args, &block) }
+      end
+
+      def create_with_sharding(*args, &block)
+        primary_shard.activate(klass.shard_category) { create_without_sharding(*args, &block) }
+      end
+
+      def create_with_sharding!(*args, &block)
+        primary_shard.activate(klass.shard_category) { create_without_sharding!(*args, &block) }
+      end
+
+      def exec_queries_with_deshackles(*args)
+        if self.lock_value
+          db = Shard.current(shard_category).database_server
+          if ::Shackles.environment != db.shackles_environment
+            return db.unshackle { exec_queries_without_deshackles(*args) }
+          end
+        end
+        exec_queries_without_deshackles(*args)
+      end
+
       def exec_queries_with_sharding
         return @records if loaded?
         results = self.activate{|relation| relation.send(:exec_queries_without_sharding) }
@@ -36,6 +63,15 @@ module Switchman
 
       %w{update_all delete_all}.each do |method|
         class_eval <<-RUBY
+          def #{method}_with_deshackles(*args)
+            db = Shard.current(shard_category).database_server
+            if ::Shackles.environment != db.shackles_environment
+              db.unshackle { #{method}_without_deshackles(*args) }
+            else
+              #{method}_without_deshackles(*args)
+            end
+          end
+
           def #{method}_with_sharding(*args)
             self.activate{|relation| relation.#{method}_without_sharding(*args)}
           end
@@ -43,22 +79,15 @@ module Switchman
       end
 
       def activate(&block)
-        case shard_value
-        when DefaultShard, Shard.current(klass.shard_category)
-          yield(self, shard_value)
-        when Shard
-          shard_value.activate(klass.shard_category) { yield(self, shard_value) }
-        when Array, ::ActiveRecord::Relation, ::ActiveRecord::Base
-          # TODO: implement local limit to avoid querying extra shards
-          if shard_value.is_a?(::ActiveRecord::Base)
-            if shard_value.respond_to?(:associated_shards)
-              shards = shard_value.associated_shards
-            else
-              shards = [shard_value.shard]
-            end
+        shards = all_shards
+        if (Array === shards && shards.length == 1)
+          if shards.first == DefaultShard || shards.first == Shard.current(klass.shard_category)
+            yield(self, shards.first)
           else
-            shards = shard_value
+            shards.first.activate(klass.shard_category) { yield(self, shard_value) }
           end
+        else
+          # TODO: implement local limit to avoid querying extra shards
           Shard.with_each_shard(shards, [klass.shard_category]) do
             shard(Shard.current(klass.shard_category), :to_a).activate(&block)
           end

data/lib/switchman/active_record/spawn_methods.rb

@@ -0,0 +1,74 @@
+module Switchman
+  module ActiveRecord
+    module SpawnMethods
+      def merge(r)
+        return self unless r
+        return to_a & r if r.is_a?(Array)
+
+        # have to figure out shard stuff *before* conditions are merged
+        if shard_value != r.shard_value
+          if (r.shard_source_value == :implicit)
+            final_shard_value = shard_value
+            final_primary_shard = primary_shard
+            final_shard_source_value = shard_source_value
+          elsif (shard_source_value == :implicit)
+            final_shard_value = r.shard_value
+            final_primary_shard = r.primary_shard
+            final_shard_source_value = r.shard_source_value
+          else
+            final_shard_source_value = [:explicit, :association].detect do |source_value|
+              shard_source_value == source_value || r.shard_source_value == source_value
+            end
+            raise "unknown shard_source_value" unless final_shard_source_value
+
+            # have to merge shard_value
+            lhs_shard_value = all_shards
+            rhs_shard_value = r.all_shards
+            if (::ActiveRecord::Relation === lhs_shard_value &&
+                ::ActiveRecord::Relation === rhs_shard_value)
+              final_shard_value = lhs_shard_value.merge(rhs_shard_value)
+              final_primary_shard = Shard.default
+            else
+              final_shard_value = lhs_shard_value.to_a & rhs_shard_value.to_a
+              return none if final_shard_value.length == 0
+              final_primary_shard = final_shard_value.first
+              final_shard_value = final_shard_value.first if final_shard_value.length == 1
+            end
+          end
+        elsif shard_source_value != r.shard_source_value
+          final_shard_source_value = [:explicit, :association, :implicit].detect do |source_value|
+            shard_source_value == source_value || r.shard_source_value == source_value
+          end
+          raise "unknown shard_source_value" unless final_shard_source_value
+
+          result = super
+          result.shard_source_value = final_shard_source_value
+          return result
+        else
+          # nothing fancy
+          return super
+        end
+
+        # change the primary shard if necessary before merging
+        result = if primary_shard != final_primary_shard && r.primary_shard != final_primary_shard
+          lhs = shard(final_primary_shard)
+          r = r.shard(final_primary_shard)
+          lhs.merge(r)
+        elsif primary_shard != final_primary_shard
+          lhs = shard(final_primary_shard)
+          lhs.merge(r)
+        elsif r.primary_shard != final_primary_shard
+          r = r.shard(final_primary_shard)
+          super(r)
+        else
+          super
+        end
+
+        result.shard_value = final_shard_value
+        result.shard_source_value = final_shard_source_value
+
+        result
+      end
+    end
+  end
+end

data/lib/switchman/active_support/cache.rb

@@ -0,0 +1,17 @@
+module Switchman
+  module ActiveSupport
+    module Cache
+      module Store
+        def initialize_with_sharding(options = nil)
+          options ||= {}
+          options[:namespace] ||= lambda { Shard.current.default? ? nil : "shard_#{Shard.current.id}" }
+          initialize_without_sharding(options)
+        end
+
+        def self.included(klass)
+          klass.alias_method_chain(:initialize, :sharding)
+        end
+      end
+    end
+  end
+end

data/lib/switchman/connection_pool_proxy.rb

@@ -1,4 +1,11 @@
 module Switchman
+  module ConnectionError
+    def self.===(other)
+      return true if defined?(PG::Error) && PG::Error === other
+      false
+    end
+  end
+
   class ConnectionPoolProxy
     delegate :spec, :connected?, :default_schema, :with_connection,
              :to => :current_pool
@@ -19,8 +26,12 @@ module Switchman
       Shard.current(@category)
     end
 
+    def active_shackles_environment
+      ::Rails.env.test? ? :master : active_shard.database_server.shackles_environment
+    end
+
     def current_pool
-      pool = self.default_pool if active_shard.default?
+      pool = self.default_pool if active_shard.database_server == Shard.default.database_server && active_shackles_environment == :master && active_shard.database_server.shareable?
       pool = @connection_pools[pool_key] ||= create_pool unless pool
       pool.shard = active_shard
       pool
@@ -28,7 +39,24 @@ module Switchman
 
     def connection
       pool = current_pool
-      pool.connection
+      begin
+        pool.connection
+      rescue ConnectionError
+        raise if active_shard.database_server == Shard.default.database_server && active_shackles_environment == :master
+        configs = active_shard.database_server.config(active_shackles_environment)
+        raise unless configs.is_a?(Array)
+        configs.each_with_index do |config, idx|
+          pool = create_pool(config.dup)
+          begin
+            connection = pool.connection
+          rescue ConnectionError
+            raise if idx == configs.length - 1
+            next
+          end
+          @connection_pools[pool_key] = pool
+          break connection
+        end
+      end
     end
 
     %w{release_connection disconnect! clear_reloadable_connections! verify_active_connections! clear_stale_cached_connections!}.each do |method|
@@ -42,12 +70,26 @@ module Switchman
     protected
 
     def pool_key
-      active_shard.database_server.shareable? ? active_shard.database_server.id : active_shard
+      [active_shackles_environment,
+        active_shard.database_server.shareable? ? active_shard.database_server.pool_key : active_shard]
     end
 
-    def create_pool
+    def create_pool(config = nil)
       shard = active_shard
-      config = shard.database_server.config.dup
+      unless config
+        if shard != Shard.default
+          config = shard.database_server.config(active_shackles_environment)
+          config = config.first if config.is_a?(Array)
+          config = config.dup
+        else
+          config = default_pool.spec.config
+          if config[active_shackles_environment].is_a?(Hash)
+            config = config.merge(config[active_shackles_environment])
+          else
+            config = config.dup
+          end
+        end
+      end
       spec = ::ActiveRecord::Base::ConnectionSpecification.new(config, "#{config[:adapter]}_connection")
       # unfortunately the AR code that does this require logic can't really be
       # called in isolation