activerecord-turntable 1.1.2 → 2.0.0.rc1

data/lib/active_record/turntable/connection_proxy.rb

@@ -11,49 +11,23 @@ module ActiveRecord::Turntable
     def initialize(cluster, options = {})
       @cluster      =  cluster
       @model_class  =  cluster.klass
-      @current_shard =  (cluster.master || cluster.shards.first[1])
-      @fixed_shard  = false
+      @default_current_shard =  (cluster.master || cluster.shards.first[1])
       @mixer = ActiveRecord::Turntable::Mixer.new(self)
     end
 
-    delegate :logger, :to => ActiveRecord::Base
+    delegate :logger, to: ActiveRecord::Base
+
+    delegate :shards_transaction, to: :cluster
 
     delegate :create_table, :rename_table, :drop_table, :add_column, :remove_colomn,
       :change_column, :change_column_default, :rename_column, :add_index,
       :remove_index, :initialize_schema_information,
-      :dump_schema_information, :execute_ignore_duplicate, :to => :master_connection
+      :dump_schema_information, :execute_ignore_duplicate, to: :master_connection
 
     def transaction(options = {}, &block)
       connection.transaction(options, &block)
     end
 
-    def shards_transaction(shards, options = {}, in_recursion = false, &block)
-      shards = in_recursion ? shards : Array.wrap(shards).dup
-      shard_or_object = shards.shift
-      shard = to_shard(shard_or_object)
-      if shards.present?
-        shard.connection.transaction(options) do
-          shards_transaction(shards, options, true, &block)
-        end
-      else
-        shard.connection.transaction(options) do
-          block.call
-        end
-      end
-    end
-
-    def to_shard(shard_or_object)
-      case shard_or_object
-      when ActiveRecord::Turntable::Shard
-        shard_or_object
-      when ActiveRecord::Base
-        shard_or_object.turntable_shard
-      else
-        raise ActiveRecord::Turntable::Error,
-                "transaction cannot call to object: #{shard_or_object}"
-      end
-    end
-
     def cache
       enable_query_cache!
       yield
@@ -92,13 +66,8 @@ module ActiveRecord::Turntable
       end
     end
 
-    # for 3.2.2
     def to_sql(arel, binds = [])
-      if master.connection.method(:to_sql).arity < 0
-        master.connection.to_sql(arel, binds)
-      else
-        master.connection.to_sql(arel)
-      end
+      master.connection.to_sql(arel, binds)
     end
 
     def cluster
@@ -114,7 +83,11 @@ module ActiveRecord::Turntable
     end
 
     def fixed_shard
-      @fixed_shard
+      fixed_shard_entry[object_id]
+    end
+
+    def fixed_shard=(shard)
+      fixed_shard_entry[object_id] = shard
     end
 
     def master
@@ -130,29 +103,31 @@ module ActiveRecord::Turntable
     end
 
     def current_shard
-      @current_shard
+      current_shard_entry[object_id] ||= @default_current_shard
     end
 
     def current_shard=(shard)
-      logger.debug { "Chainging #{@model_class}'s shard to #{shard.name}"}
-      @current_shard = shard
+      logger.debug { "Changing #{@model_class}'s shard to #{shard.name}"}
+      current_shard_entry[object_id] = shard
     end
 
     def connection
-      @current_shard.connection
+      current_shard.connection
     end
 
     def connection_pool
-      @current_shard.connection_pool
+      current_shard.connection_pool
     end
 
     def with_shard(shard)
+      shard = cluster.to_shard(shard)
+
       old_shard, old_fixed = current_shard, fixed_shard
       self.current_shard = shard
-      @fixed_shard = shard
+      self.fixed_shard = shard
       yield
     ensure
-      @fixed_shard = old_fixed
+      self.fixed_shard = old_fixed
       self.current_shard = old_shard
     end
 
@@ -182,45 +157,49 @@ module ActiveRecord::Turntable
       end
     end
 
+    def with_master_and_all(continue_on_error = false)
+      ([@cluster.master] + @cluster.shards.values).map do |shard|
+        begin
+          with_shard(shard) {
+            yield
+          }
+        rescue Exception => err
+          unless continue_on_error
+            raise err
+          end
+          err
+        end
+      end
+    end
+
     def with_master(&block)
       with_shard(@cluster.master) do
         yield
       end
     end
 
-    def connected?
-      connection_pool.connected?
-    end
+    delegate :connected?, :automatic_reconnect, :automatic_reconnect=, :checkout_timeout, :dead_connection_timeout,
+               :spec, :connections, :size, :reaper, :table_exists?, to: :connection_pool
 
-    if ActiveRecord::VERSION::STRING > '3.1'
-      %w(columns columns_hash column_defaults primary_keys).each do |name|
-        define_method(name.to_sym) do
-          master.connection_pool.send(name.to_sym)
-        end
+    %w(columns columns_hash column_defaults primary_keys).each do |name|
+      define_method(name.to_sym) do
+        master.connection_pool.send(name.to_sym)
       end
+    end
 
-      if ActiveRecord::VERSION::STRING < '3.2'
-        %w(table_exists?).each do |name|
-          define_method(name.to_sym) do |*args|
-            master.connection_pool.send(name.to_sym, *args)
-          end
-        end
-      else
-        %w(table_exists?).each do |name|
-          define_method(name.to_sym) do |*args|
-            master.connection_pool.with_connection do |c|
-              c.schema_cache.send(name.to_sym, *args)
-            end
-          end
+    %w(table_exists?).each do |name|
+      define_method(name.to_sym) do |*args|
+        master.connection_pool.with_connection do |c|
+          c.schema_cache.send(name.to_sym, *args)
         end
       end
+    end
 
-      def columns(*args)
-        if args.size > 0
-          master.connection_pool.columns[*args]
-        else
-          master.connection_pool.columns
-        end
+    def columns(*args)
+      if args.size > 0
+        master.connection_pool.columns[*args]
+      else
+        master.connection_pool.columns
       end
     end
 
@@ -239,5 +218,15 @@ module ActiveRecord::Turntable
     def spec
       @spec ||= master.connection_pool.spec
     end
+
+    private
+
+    def fixed_shard_entry
+      Thread.current[:turntable_fixed_shard] ||= ThreadSafe::Cache.new
+    end
+
+    def current_shard_entry
+      Thread.current[:turntable_current_shard] ||= ThreadSafe::Cache.new
+    end
   end
 end

data/lib/active_record/turntable/error.rb

@@ -1,8 +1,7 @@
 module ActiveRecord::Turntable
-  class Error < StandardError; end
-  class NotImplementedError < Error; end
-  class SequenceNotFoundError < Error; end
-  class CannotSpecifyShardError < Error; end
-  class MasterShardNotConnected < Error; end
-  class UnknownOperatorError < Error; end
+  class TurntableError < StandardError; end
+  class SequenceNotFoundError < TurntableError; end
+  class CannotSpecifyShardError < TurntableError; end
+  class MasterShardNotConnected < TurntableError; end
+  class UnknownOperatorError < TurntableError; end
 end

data/lib/active_record/turntable/helpers.rb

@@ -1,5 +1,9 @@
 module ActiveRecord::Turntable
   module Helpers
-    autoload :TestHelper, 'active_record/turntable/helpers/test_helper'
+    extend ActiveSupport::Autoload
+
+    eager_autoload do
+      autoload :TestHelper
+    end
   end
 end

data/lib/active_record/turntable/migration.rb

@@ -1,38 +1,15 @@
 module ActiveRecord::Turntable::Migration
   extend ActiveSupport::Concern
 
-  # AR < 3.1
-  def self.extended(base)
-    class << base
-      def announce_with_turntable(message)
-        announce_without_turntable("#{message} - #{get_current_shard}")
-      end
-
-      alias_method_chain :migrate, :turntable
-      alias_method_chain :announce, :turntable
-      include ShardDefinition
-    end
-    base.class_eval do
-      class_inheritable_accessor :target_shards
-    end
-    ::ActiveRecord::ConnectionAdapters::AbstractAdapter.send(:include, SchemaStatementsExt)
-  end
-
-  # AR >= 3.1
   included do
     extend ShardDefinition
-    class_attribute :target_shards
-    def announce_with_turntable(message)
-      announce_without_turntable("#{message} - #{get_current_shard}")
-    end
-
-    alias_method_chain :migrate, :turntable
+    class_attribute :target_shards, :current_shard
     alias_method_chain :announce, :turntable
+    alias_method_chain :exec_migration, :turntable
     ::ActiveRecord::ConnectionAdapters::AbstractAdapter.send(:include, SchemaStatementsExt)
     ::ActiveRecord::Migration::CommandRecorder.send(:include, CommandRecorder)
   end
 
-  # for all
   module ShardDefinition
     def clusters(*cluster_names)
       config = ActiveRecord::Base.turntable_config
@@ -55,32 +32,16 @@ module ActiveRecord::Turntable::Migration
     end
   end
 
-  def get_current_shard
-    "Shard: #{@@current_shard}" if @@current_shard
+  def target_shard?(shard_name)
+    target_shards.blank? or target_shards.include?(shard_name)
   end
 
-  def migrate_with_turntable(direction)
-    config = ActiveRecord::Base.configurations
-    @@current_shard = nil
-    shards = (self.class.target_shards||=[]).flatten.uniq.compact
-    if self.class.target_shards.blank?
-      return migrate_without_turntable(direction)
-    end
+  def announce_with_turntable(message)
+    announce_without_turntable("#{message} - Shard: #{current_shard}")
+  end
 
-    shards_conf = shards.map do |shard|
-      config[Rails.env||"development"]["shards"][shard]
-    end
-    seqs = config[Rails.env||"development"]["seq"]
-    shards_conf += seqs.values
-    shards_conf << config[Rails.env||"development"]
-    shards_conf.each_with_index do |conf, idx|
-      @@current_shard = (shards[idx] || seqs.keys[idx - shards.size] || "master")
-      ActiveRecord::Base.establish_connection(conf)
-      if !ActiveRecord::Base.connection.table_exists?(ActiveRecord::Migrator.schema_migrations_table_name())
-        ActiveRecord::Base.connection.initialize_schema_migrations_table
-      end
-      migrate_without_turntable(direction)
-    end
+  def exec_migration_with_turntable(*args)
+    exec_migration_without_turntable(*args) if target_shard?(current_shard)
   end
 
   module SchemaStatementsExt
@@ -132,4 +93,3 @@ module ActiveRecord::Turntable::Migration
   end
 
 end
-

data/lib/active_record/turntable/mixer/fader/calculate_shards_sum_result.rb

@@ -3,11 +3,22 @@ module ActiveRecord::Turntable
     class Fader
       class CalculateShardsSumResult < Fader
         def execute
-          @shards_query_hash.map do |shard, query|
+          results = @shards_query_hash.map do |shard, query|
             args = @args.dup
             args[1] = args[1].dup if args[1].present?
             shard.connection.send(@called_method, query, *@args, &@block)
-          end.inject(&:+)
+          end
+          merge_results(results)
+        end
+
+        private
+
+        def merge_results(results)
+          ActiveRecord::Result.new(
+            results.first.columns,
+            results[0].rows.zip(*results[1..-1].map{|r| r.rows}).map {|r| [r.map {|v| v.first}.inject(&:+)]},
+            results.first.column_types
+          )
         end
       end
     end

data/lib/active_record/turntable/mixer/fader/select_shards_merge_result.rb

@@ -3,17 +3,28 @@ module ActiveRecord::Turntable
     class Fader
       class SelectShardsMergeResult < Fader
         def execute
-          res = @shards_query_hash.map do |shard, query|
+          results = @shards_query_hash.map do |shard, query|
             args = @args.dup
             args[1] = args[1].dup if args[1].present?
             shard.connection.send(@called_method, query, *args, &@block)
-          end.flatten(1).compact
+          end
+          merge_results(results)
+        end
+
+        private
+
+        def merge_results(results)
+          if results.any? {|r| r.is_a?(ActiveRecord::Result) }
+            first_result = results.find {|r| r.present? }
+            return results.first unless first_result
 
-          case @called_method
-          when "select_value", "select_one"
-            res.first if res
+            ActiveRecord::Result.new(
+              first_result.columns,
+              results.map {|r| r.rows}.flatten(1),
+              first_result.column_types
+            )
           else
-            res
+            results.compact.inject(&:+)
           end
         end
       end

data/lib/active_record/turntable/mixer/fader/specified_shard.rb

@@ -4,7 +4,9 @@ module ActiveRecord::Turntable
       class SpecifiedShard < Fader
         def execute
           shard, query = @shards_query_hash.first
-          shard.connection.send(@called_method, query, *@args, &@block)
+          @proxy.with_shard(shard) do
+            shard.connection.send(@called_method, query, *@args, &@block)
+          end
         end
       end
     end

data/lib/active_record/turntable/mixer/fader.rb

@@ -2,16 +2,18 @@
 module ActiveRecord::Turntable
   class Mixer
     class Fader
-      # 単数shard
-      autoload :SpecifiedShard, "active_record/turntable/mixer/fader/specified_shard"
+      extend ActiveSupport::Autoload
 
-      # 複数shard
-      autoload :SelectShardsMergeResult, "active_record/turntable/mixer/fader/select_shards_merge_result"
-      autoload :InsertShardsMergeResult, "active_record/turntable/mixer/fader/insert_shards_merge_result"
-      autoload :UpdateShardsMergeResult, "active_record/turntable/mixer/fader/update_shards_merge_result"
-
-      # count
-      autoload :CalculateShardsSumResult, "active_record/turntable/mixer/fader/calculate_shards_sum_result"
+      eager_autoload do
+        # single shard
+        autoload :SpecifiedShard
+        # multiple shard merging
+        autoload :SelectShardsMergeResult
+        autoload :InsertShardsMergeResult
+        autoload :UpdateShardsMergeResult
+        # calcurations
+        autoload :CalculateShardsSumResult
+      end
 
       attr_reader :shards_query_hash
       attr_reader :called_method
@@ -27,7 +29,7 @@ module ActiveRecord::Turntable
       end
 
       def execute
-        raise ActiveRecord::Turntable::NotImplementedError, "Called abstract method"
+        raise NotImplementedError, "Called abstract method"
       end
     end
   end

data/lib/active_record/turntable/mixer.rb

@@ -3,7 +3,12 @@ require 'active_record/turntable/sql_tree_patch'
 
 module ActiveRecord::Turntable
   class Mixer
-    autoload :Fader, "active_record/turntable/mixer/fader"
+    extend ActiveSupport::Autoload
+
+    eager_autoload do
+      autoload :Fader
+    end
+
     delegate :logger, :to => ActiveRecord::Base
 
     NOT_USED_FOR_SHARDING_OPERATORS_REGEXP = /\A(NOT IN|IS|IS NOT|BETWEEN|LIKE|\!\=|<<|>>|<>|>\=|<=|[\*\+\-\/\%\|\&><])\z/
@@ -25,7 +30,7 @@ module ActiveRecord::Turntable
       begin
         tree = SQLTree[binded_query]
       rescue Exception => err
-        logger.warn { "[ActiveRecord::Turntable][BUG] Error on Parsing SQL: #{binded_query}, on_method: #{method_name}" }
+        logger.warn { "[ActiveRecord::Turntable] Error on Parsing SQL: #{binded_query}, on_method: #{method_name}" }
         raise err
       end
 
@@ -43,7 +48,7 @@ module ActiveRecord::Turntable
                            method, query, *args, &block)
       end
     rescue Exception => err
-      logger.warn { "[ActiveRecord::Turntable][BUG] Error on Building Fader: #{binded_query}, on_method: #{method_name}" }
+      logger.warn { "[ActiveRecord::Turntable] Error on Building Fader: #{binded_query}, on_method: #{method_name}" }
       raise err
     end
 
@@ -88,7 +93,7 @@ module ActiveRecord::Turntable
         []
       else
         raise ActiveRecord::Turntable::UnknownOperatorError,
-          "[ActiveRecord::Turntable][BUG] Found Unknown SQL Operator:'#{tree.operator if tree.respond_to?(:operaor)}', Please report this bug."
+          "[ActiveRecord::Turntable] Found Unknown SQL Operator:'#{tree.operator if tree.respond_to?(:operaor)}', Please report this bug."
       end
     end
 
@@ -107,21 +112,15 @@ module ActiveRecord::Turntable
       Hash[shards.map {|s| [s, query] }]
     end
 
-    if ActiveRecord::VERSION::STRING < '3.1'
-      def bind_sql(sql, binds)
-        sql
-      end
-    else
-      def bind_sql(sql, binds)
-        # TODO: substitution value should be determined by adapter
-        query = sql.is_a?(String) ? sql : @proxy.to_sql(sql, binds ? binds.dup : [])
-        query = if query.include?("\0") and binds.is_a?(Array) and binds[0].is_a?(Array) and binds[0][0].is_a?(ActiveRecord::ConnectionAdapters::Column)
-                  binds = binds.dup
-                  query.gsub("\0") { @proxy.master.connection.quote(*binds.shift.reverse) }
-                else
-                  query
-                end
-      end
+    def bind_sql(sql, binds)
+      # TODO: substitution value should be determined by adapter
+      query = sql.is_a?(String) ? sql : @proxy.to_sql(sql, binds ? binds.dup : [])
+      query = if query.include?("\0") and binds.is_a?(Array) and binds[0].is_a?(Array) and binds[0][0].is_a?(ActiveRecord::ConnectionAdapters::Column)
+                binds = binds.dup
+                query.gsub("\0") { @proxy.master.connection.quote(*binds.shift.reverse) }
+              else
+                query
+              end
     end
 
     def build_select_fader(tree, method, query, *args, &block)
@@ -137,10 +136,16 @@ module ActiveRecord::Turntable
 
       if shard_keys.size == 1 # shard
         return Fader::SpecifiedShard.new(@proxy,
-                                         { @proxy.cluster.select_shard(shard_keys.first) => query },
+                                         { @proxy.cluster.shard_for(shard_keys.first) => query },
                                          method, query, *args, &block)
-      elsif tree.group_by or tree.order_by or tree.limit.try(:value).to_i > 1
-        raise CannotSpecifyShardError, "cannot specify shard for query: #{binded_query}"
+      elsif SQLTree::Node::SelectDeclaration === tree.select.first and
+              tree.select.first.to_sql == '1 AS "one"'  # for `SELECT 1 AS one` (AR::Base.exists?)
+        return Fader::SelectShardsMergeResult.new(@proxy,
+                                                  build_shards_with_same_query(@proxy.shards.values, query),
+                                                  method, query, *args, &block
+                                                  )
+      elsif tree.group_by or tree.order_by or tree.limit.try(:value).to_i > 0
+        raise CannotSpecifyShardError, "cannot specify shard for query: #{tree.to_sql}"
       elsif shard_keys.present?
         if SQLTree::Node::SelectDeclaration === tree.select.first and
             SQLTree::Node::CountAggregrate === tree.select.first.expression
@@ -149,13 +154,17 @@ module ActiveRecord::Turntable
                                                      method, query, *args, &block)
         else
           return Fader::SelectShardsMergeResult.new(@proxy,
-                                                    Hash[shard_keys.map {|k| [@proxy.cluster.select_shard(k), query] }],
+                                                    Hash[shard_keys.map {|k| [@proxy.cluster.shard_for(k), query] }],
                                                     method, query, *args, &block
                                                     )
         end
       else # scan all shards
         if SQLTree::Node::SelectDeclaration === tree.select.first and
             SQLTree::Node::CountAggregrate === tree.select.first.expression
+
+          if raise_on_not_specified_shard_query?
+            raise CannotSpecifyShardError, "[Performance Notice] PLEASE FIX: #{tree.to_sql}"
+          end
           return Fader::CalculateShardsSumResult.new(@proxy,
                                                      build_shards_with_same_query(@proxy.shards.values, query),
                                                      method, query, *args, &block)
@@ -163,12 +172,15 @@ module ActiveRecord::Turntable
             SQLTree::Node::Expression::Value === tree.select.first.expression or
             SQLTree::Node::Expression::Variable === tree.select.first.expression
 
+          if raise_on_not_specified_shard_query?
+            raise CannotSpecifyShardError, "[Performance Notice] PLEASE FIX: #{tree.to_sql}"
+          end
           return Fader::SelectShardsMergeResult.new(@proxy,
                                                     build_shards_with_same_query(@proxy.shards.values, query),
                                                     method, query, *args, &block
                                                     )
         else
-          raise CannotSpecifyShardError, "cannot specify shard for query: #{binded_query}"
+          raise CannotSpecifyShardError, "cannot specify shard for query: #{tree.to_sql}"
         end
       end
     end
@@ -176,13 +188,23 @@ module ActiveRecord::Turntable
     def build_update_fader(tree, method, query, *args, &block)
       shard_keys = find_shard_keys(tree.where, @proxy.cluster.klass.table_name, @proxy.cluster.klass.turntable_shard_key.to_s)
       shards_with_query = if shard_keys.present?
-                            build_shards_with_same_query(shard_keys.map {|k| @proxy.cluster.select_shard(k) }, query)
+                            build_shards_with_same_query(shard_keys.map {|k| @proxy.cluster.shard_for(k) }, query)
                           else
                             build_shards_with_same_query(@proxy.shards.values, query)
                           end
-      Fader::UpdateShardsMergeResult.new(@proxy,
-                                         shards_with_query,
-                                         method, query, *args, &block)
+
+      if shards_with_query.size == 1
+        Fader::SpecifiedShard.new(@proxy,
+                                  shards_with_query,
+                                  method, query, *args, &block)
+      else
+        if raise_on_not_specified_shard_update?
+          raise CannotSpecifyShardError, "[Performance Notice] PLEASE FIX: #{tree.to_sql}"
+        end
+        Fader::UpdateShardsMergeResult.new(@proxy,
+                                           shards_with_query,
+                                           method, query, *args, &block)
+      end
     end
 
     def build_insert_fader(tree, method, query, *args, &block)
@@ -195,9 +217,17 @@ module ActiveRecord::Turntable
           "(#{val.map { |v| @proxy.connection.quote(v.value)}.join(', ')})"
         end.join(', ')
         sql.gsub!('("\0")') { value_sql }
-        shards_with_query[@proxy.cluster.select_shard(k)] = sql
+        shards_with_query[@proxy.cluster.shard_for(k)] = sql
       end
       Fader::InsertShardsMergeResult.new(@proxy, shards_with_query, method, query, *args, &block)
     end
+
+    def raise_on_not_specified_shard_query?
+      ActiveRecord::Base.turntable_config[:raise_on_not_specified_shard_query]
+    end
+
+    def raise_on_not_specified_shard_update?
+      ActiveRecord::Base.turntable_config[:raise_on_not_specified_shard_update]
+    end
   end
 end

data/lib/active_record/turntable/plugin.rb

@@ -0,0 +1,6 @@
+require 'active_record/turntable'
+
+module ActiveRecord::Turntable
+  module Plugin
+  end
+end

data/lib/active_record/turntable/pool_proxy.rb

@@ -14,34 +14,27 @@ module ActiveRecord::Turntable
       @proxy = proxy
     end
 
-    def connection
+    def proxy
       @proxy
     end
-
-    def spec
-      @proxy.spec
-    end
-
+    alias_method :connection, :proxy
 
     def with_connection
-      yield @proxy
+      yield proxy
     end
 
-    def connected?
-      @proxy.connected?
-    end
+    delegate :connected?, :automatic_reconnect, :automatic_reconnect=, :checkout_timeout, :dead_connection_timeout,
+               :spec, :connections, :size, :reaper, :table_exists?, to: :proxy
 
-    if ActiveRecord::VERSION::STRING > '3.1'
-      %w(columns_hash column_defaults primary_keys).each do |name|
-        define_method(name.to_sym) do
-          @proxy.send(name.to_sym)
-        end
+    %w(columns_hash column_defaults primary_keys).each do |name|
+      define_method(name.to_sym) do
+        @proxy.send(name.to_sym)
       end
+    end
 
-      %w(table_exists? columns).each do |name|
-        define_method(name.to_sym) do |*args|
-          @proxy.send(name.to_sym, *args)
-        end
+    %w(table_exists? columns).each do |name|
+      define_method(name.to_sym) do |*args|
+        @proxy.send(name.to_sym, *args)
       end
     end