switchman 0.0.1

data/lib/switchman/active_record/connection_pool.rb

@@ -0,0 +1,117 @@
+module Switchman
+  module ActiveRecord
+    module ConnectionPool
+      def self.included(klass)
+        klass.alias_method_chain(:checkout_new_connection, :sharding)
+        klass.send(:remove_method, :connection)
+        klass.send(:remove_method, :active_connection?)
+        klass.send(:remove_method, :release_connection)
+        klass.send(:remove_method, :clear_stale_cached_connections!)
+      end
+
+      attr_writer :shard
+
+      def shard
+        @shard || Shard.default
+      end
+
+      def default_schema
+        raise "Not postgres!" unless self.spec.config[:adapter] == 'postgresql'
+        connection unless @schemas
+        # default shard will not switch databases immediately, so it won't be set yet
+        @schemas ||= connection.schemas
+        @schemas.first
+      end
+
+      def checkout_new_connection_with_sharding
+        # TODO: this might be a threading issue
+        spec.config[:shard_name] = self.shard.name
+
+        conn = checkout_new_connection_without_sharding
+        conn.shard = self.shard
+        conn
+      end
+
+      def connection
+        conns = synchronize { @reserved_connections[current_connection_id] ||= [] }
+        conn = conns.find { |conn| !conn.open_transactions || conn.shard == self.shard || conn.adapter_name == 'PostgreSQL' }
+        unless conn
+          conn = checkout
+          yield conn if block_given?
+          conns << conn
+        end
+        switch_database(conn) if conn.shard != self.shard
+        conn
+      end
+
+      # Is there an open connection that is being used for the current thread?
+      def active_connection?
+        synchronize do
+          @reserved_connections.fetch(current_connection_id) {
+            return false
+          }.any? { |conn| conn.in_use? }
+        end
+      end
+
+      def release_connection
+        conns = synchronize { @reserved_connections.delete(current_connection_id) }
+        conns.each { |conn| checkin conn } if conns
+      end
+
+      def clear_stale_cached_connections!
+        keys = @reserved_connections.keys - Thread.list.find_all { |t|
+          t.alive?
+        }.map { |thread| thread.object_id }
+        keys.each do |key|
+          conns = @reserved_connections[key]
+          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 }
+          @reserved_connections.delete(key)
+        end
+      end
+
+      def release(conn)
+        synchronize do
+          thread_id = nil
+
+          if @reserved_connections[current_connection_id].include?(conn)
+            thread_id = current_connection_id
+            @reserved_connections[thread_id].delete(conn)
+          else
+            thread_id = @reserved_connections.keys.find { |k|
+              @reserved_connections[k].include?(conn)
+            }
+            if thread_id
+              @reserved_connections[thread_id].delete(conn)
+            end
+          end
+
+          @reserved_connections.delete thread_id if thread_id && @reserved_connections[thread_id].empty?
+        end
+      end
+
+      def switch_database(conn)
+        if !@schemas && conn.adapter_name == 'PostgreSQL' && !self.shard.database_server.config[:shard_name]
+          @schemas = conn.schemas
+        end
+
+        spec.config[:shard_name] = self.shard.name
+        case conn.adapter_name
+          when 'MySQL', 'Mysql2'
+            conn.execute("USE #{spec.config[:database]}")
+          when 'PostgreSQL'
+            conn.schema_search_path = spec.config[:schema_search_path]
+          when 'SQLite'
+            # This is an artifact of the adapter modifying the path to be an absolute path when it is instantiated; just let it slide
+          else
+            raise("Cannot switch databases on same DatabaseServer with adapter type: #{conn.adapter_name}. Limit one Shard per DatabaseServer.")
+        end
+        conn.shard = shard
+      end
+    end
+  end
+end

data/lib/switchman/active_record/finder_methods.rb

@@ -0,0 +1,25 @@
+module Switchman
+  module ActiveRecord
+    module FinderMethods
+      # find_one uses binds, so we can't depend on QueryMethods
+      # catching it
+      def find_one(id)
+        local_id, shard = Shard.local_id_for(id)
+
+        return super(local_id) if shard_source_value != :implicit
+
+        if shard
+          begin
+            old_shard_value = shard_value
+            self.shard_value = shard
+            super(local_id)
+          ensure
+            self.shard_value = old_shard_value
+          end
+        else
+          super
+        end
+      end
+    end
+  end
+end

data/lib/switchman/active_record/log_subscriber.rb

@@ -0,0 +1,43 @@
+module Switchman
+  module ActiveRecord
+    module LogSubscriber
+      def self.included(klass)
+        klass.send(:remove_method, :sql)
+      end
+
+      # sadly, have to completely replace this
+      def sql(event)
+        self.class.runtime += event.duration
+        return unless logger.debug?
+
+        payload = event.payload
+
+        return if 'SCHEMA' == payload[:name]
+
+        name  = '%s (%.1fms)' % [payload[:name], event.duration]
+        sql   = payload[:sql].squeeze(' ')
+        binds = nil
+        connection = ObjectSpace._id2ref(payload[:connection_id])
+
+        unless (payload[:binds] || []).empty?
+          binds = "  " + payload[:binds].map { |col,v|
+            if col
+              [col.name, v]
+            else
+              [nil, v]
+            end
+          }.inspect
+        end
+
+        if odd?
+          name = color(name, self.class::CYAN, true)
+          sql  = color(sql, nil, true)
+        else
+          name = color(name, self.class::MAGENTA, true)
+        end
+
+        debug "  #{name}  #{sql}#{binds}  [shard #{connection.shard.id} #{::Shackles.environment}]"
+      end
+    end
+  end
+end

data/lib/switchman/active_record/postgresql_adapter.rb

@@ -0,0 +1,13 @@
+module Switchman
+  module ActiveRecord
+    module PostgreSQLAdapter
+      def self.included(klass)
+        klass::NATIVE_DATABASE_TYPES[:primary_key] = "bigserial primary key".freeze
+      end
+
+      def schemas
+        select_values("SELECT * FROM unnest(current_schemas(false))")
+      end
+    end
+  end
+end

data/lib/switchman/active_record/query_cache.rb

@@ -0,0 +1,12 @@
+module Switchman
+  module ActiveRecord
+    # needs to be included in the same class as ::ActiveRecord::ConnectionAdapters::QueryCache
+    # *after* that module is included
+    module QueryCache
+      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
+    end
+  end
+end

data/lib/switchman/active_record/query_methods.rb

@@ -0,0 +1,184 @@
+module Switchman
+  module ActiveRecord
+    module QueryMethods
+      # shard_value is one of:
+      #   A shard
+      #   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
+      attr_accessor :shard_value, :shard_source_value
+
+      def shard(value, source = :explicit)
+        relation = clone
+        relation.shard_value = value
+        relation.shard_source_value = source
+        if (primary_shard != relation.primary_shard || source == :to_a)
+          relation.where_values = relation.transpose_predicates(relation.where_values, primary_shard, relation.primary_shard, source == :to_a) if !relation.where_values.empty?
+          relation.having_values = relation.transpose_predicates(relation.having_values, primary_shard, relation.primary_shard, source == :to_a) if !relation.having_values.empty?
+        end
+        relation
+      end
+
+      # replace these with versions that call build_where on the
+      # result relation, not the source relation (so build_where
+      # is able to implicitly change the shard_value)
+      def where(opts, *rest)
+        return self if opts.blank?
+
+        relation = clone
+        relation.where_values += relation.build_where(opts, rest)
+        relation
+      end
+
+      def having(opts, *rest)
+        return self if opts.blank?
+
+        relation = clone
+        relation.having_values += relation.build_where(opts, rest)
+        relation
+      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
+        else
+          super
+        end
+      end
+
+      # the shard that where_values are relative to. if it's multiple shards, they're stored
+      # relative to the first shard
+      def primary_shard
+        case shard_value
+        when Shard, DefaultShard
+          shard_value
+        # associated_shards
+        when ::ActiveRecord::Base
+          shard_value.shard
+        when Array
+          shard_value.first
+        when ::ActiveRecord::Relation
+          Shard.default
+        else
+          raise ArgumentError("invalid shard value #{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
+        end
+        if primary_key
+          case primary_key.right
+          when Array
+            id_shards = Set.new
+            primary_key.right.each do |value|
+              local_id, id_shard = Shard.local_id_for(value)
+              id_shard ||= Shard.current(klass.shard_category) if local_id
+              id_shards << id_shard if id_shard
+            end
+            if id_shards.empty?
+              return
+            elsif id_shards.length == 1
+              id_shard = id_shards.first
+            # prefer to not change the shard
+            elsif id_shards.include?(primary_shard)
+              id_shards.delete(primary_shard)
+              self.shard_value = [primary_shard] + id_shards.to_a
+              return
+            else
+              id_shards = id_shards.to_a
+              self.where_values = transpose_predicates(where_values, primary_shard, id_shards.first) if !where_values.empty?
+              self.having_values = transpose_predicates(having_values, primary_shard, id_shards.first) if !having_values.empty?
+              self.shard_value = id_shards
+              return
+            end
+          else
+            local_id, id_shard = Shard.local_id_for(primary_key.right)
+            id_shard ||= Shard.current(klass.shard_category) if local_id
+          end
+          return if !id_shard || id_shard == primary_shard
+          self.where_values = transpose_predicates(where_values, primary_shard, id_shard) if !where_values.empty?
+          self.having_values = transpose_predicates(having_values, primary_shard, id_shard) if !having_values.empty?
+          self.shard_value = id_shard
+        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
+        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
+
+          @@foreign_keys[attribute.relation.table_name][attribute.name] = models.any?{|m| m.sharded_column?(attribute.name)}
+        end
+      end
+
+      def sharded_primary_key?(attribute)
+        attribute.relation.engine.primary_key == attribute.name
+      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)
+
+          remove = true if type == :primary && remove_nonlocal_primary_keys && predicate.left.relation.engine == klass
+
+          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_ids << local_id unless remove && local_id > Shard::IDS_PER_SHARD
+            end
+            local_ids
+          when Arel::Nodes::BindParam
+            # 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 = [] 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 = [] if remove && local_id > Shard::IDS_PER_SHARD
+            local_id
+          end
+
+          if new_right_value == predicate.right
+            predicate
+          else
+            predicate.class.new(predicate.left, new_right_value)
+          end
+        end
+      end
+    end
+  end
+end

data/lib/switchman/active_record/relation.rb

@@ -0,0 +1,69 @@
+module Switchman
+  module ActiveRecord
+    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|
+          klass.alias_method_chain(method, :sharding)
+        end
+      end
+
+      def initialize_with_sharding(klass, table)
+        initialize_without_sharding(klass, table)
+        self.shard_value = Shard.current(klass.shard_category)
+        self.shard_source_value = :implicit
+      end
+
+      def merge(*args)
+        relation = super
+        if relation.shard_value != self.shard_value && relation.shard_source_value == :implicit
+          relation.shard_value = self.shard_value
+          relation.shard_source_value = self.shard_source_value
+        end
+        relation
+      end
+
+      def exec_queries_with_sharding
+        return @records if loaded?
+        results = self.activate{|relation| relation.send(:exec_queries_without_sharding) }
+        case shard_value
+        when Array, ::ActiveRecord::Relation, ::ActiveRecord::Base
+          @records = results
+          @loaded = true
+        end
+        results
+      end
+
+      %w{update_all delete_all}.each do |method|
+        class_eval <<-RUBY
+          def #{method}_with_sharding(*args)
+            self.activate{|relation| relation.#{method}_without_sharding(*args)}
+          end
+        RUBY
+      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
+          else
+            shards = shard_value
+          end
+          Shard.with_each_shard(shards, [klass.shard_category]) do
+            shard(Shard.current(klass.shard_category), :to_a).activate(&block)
+          end
+        end
+      end
+    end
+  end
+end