activerecord-turntable 1.0.0

data/lib/active_record/turntable/migration.rb

@@ -0,0 +1,132 @@
+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
+    alias_method_chain :announce, :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
+      (self.target_shards ||= []) <<
+        if cluster_names.first == :all
+          config['clusters'].map do |name, cluster_conf|
+            cluster_conf["shards"].map {|shard| shard["connection"]}
+          end
+        else
+          cluster_names.map do |cluster_name|
+            config['clusters'][cluster_name]["shards"].map do |shard|
+              shard["connection"]
+            end
+          end
+        end
+    end
+
+    def shards(*connection_names)
+      (self.target_shards ||= []) << connection_names
+    end
+  end
+
+  def get_current_shard
+    "Shard: #{@@current_shard}" if @@current_shard
+  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
+
+    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
+  end
+
+  module SchemaStatementsExt
+    def create_sequence_for(table_name, options = { })
+      # TODO: pkname should be pulled from table definitions
+      pkname = "id"
+      sequence_table_name = ActiveRecord::Turntable::Sequencer.sequence_name(table_name, "id")
+      create_table(sequence_table_name, options)
+      execute "ALTER TABLE #{quote_table_name(sequence_table_name)} MODIFY id bigint(20) DEFAULT NULL auto_increment NOT NULL;"
+      execute "INSERT INTO #{quote_table_name(sequence_table_name)} (`id`) VALUES (0)"
+    end
+
+    def drop_sequence_for(table_name, options = { })
+      # TODO: pkname should be pulled from table definitions
+      pkname = "id"
+      sequence_table_name = ActiveRecord::Turntable::Sequencer.sequence_name(table_name, "id")
+      drop_table(sequence_table_name)
+    end
+
+    def rename_sequence_for(table_name, new_name)
+      # TODO: pkname should pulled from table definitions
+      seq_table_name = ActiveRecord::Turntable::Sequencer.sequence_name(table_name, "id")
+      new_seq_name = ActiveRecord::Turntable::Sequencer.sequence_name(new_name, "id")
+      rename_table(seq_table_name, new_seq_name)
+    end
+  end
+
+  module CommandRecorder
+    def create_sequence_for(*args)
+      record(:create_sequence_for, args)
+    end
+
+    def rename_sequence_for(*args)
+      record(:rename_sequence_for, args)
+    end
+
+    private
+
+    def invert_create_sequence_for(args)
+      [:drop_sequence_for, args]
+    end
+
+    def invert_rename_sequence_for(args)
+      [:rename_sequence_for, args.reverse]
+    end
+  end
+
+end
+

data/lib/active_record/turntable/mixer.rb

@@ -0,0 +1,203 @@
+require 'active_support/core_ext/object/try'
+require 'active_record/turntable/sql_tree_patch'
+
+module ActiveRecord::Turntable
+  class Mixer
+    autoload :Fader, "active_record/turntable/mixer/fader"
+    delegate :logger, :to => ActiveRecord::Base
+
+    NOT_USED_FOR_SHARDING_OPERATORS_REGEXP = /\A(NOT IN|IS|IS NOT|BETWEEN|LIKE|\!\=|<<|>>|<>|>\=|<=|[\*\+\-\/\%\|\&><])\z/
+
+    def initialize(proxy)
+      @proxy = proxy
+    end
+
+    def build_fader(method_name, query, *args, &block)
+      method = method_name.to_s
+      if @proxy.shard_fixed?
+        return SpecifiedShard.new(@proxy,
+                                  { @proxy.fixed_shard => query },
+                                  method, query, *args, &block)
+      end
+      binds = (method == 'insert') ? args[4] : args[1]
+      binded_query = bind_sql(query, binds)
+
+      begin
+        tree = SQLTree[binded_query]
+      rescue Exception => err
+        logger.warn { "[ActiveRecord::Turntable][BUG] Error on Parsing SQL: #{binded_query}, on_method: #{method_name}" }
+        raise err
+      end
+
+      case tree
+      when SQLTree::Node::SelectQuery
+        build_select_fader(tree, method, query, *args, &block)
+      when SQLTree::Node::UpdateQuery, SQLTree::Node::DeleteQuery
+        build_update_fader(tree, method, query, *args, &block)
+      when SQLTree::Node::InsertQuery
+        build_insert_fader(tree, method, query, *args, &block)
+      else
+        # send to master shard
+        Fader::SpecifiedShard.new(@proxy,
+                           { @proxy.master => query },
+                           method, query, *args, &block)
+      end
+    rescue Exception => err
+      logger.warn { "[ActiveRecord::Turntable][BUG] Error on Building Fader: #{binded_query}, on_method: #{method_name}" }
+      raise err
+    end
+
+    def find_shard_keys(tree, table_name, shard_key)
+      return [] unless tree.respond_to?(:operator)
+
+      case tree.operator
+      when "OR"
+        lkeys = find_shard_keys(tree.lhs, table_name, shard_key)
+        rkeys = find_shard_keys(tree.rhs, table_name, shard_key)
+        if lkeys.present? and rkeys.present?
+          lkeys + rkeys
+        else
+          []
+        end
+      when "AND"
+        lkeys = find_shard_keys(tree.lhs, table_name, shard_key)
+        rkeys = find_shard_keys(tree.rhs, table_name, shard_key)
+        if lkeys.present? or rkeys.present?
+          lkeys + rkeys
+        else
+          []
+        end
+      when "IN", "=", "=="
+        field = tree.lhs.respond_to?(:table) ? tree.lhs : nil
+        if tree.rhs.is_a?(SQLTree::Node::SubQuery)
+          if field.try(:table) == table_name and field.name == shard_key
+            find_shard_keys(tree.rhs.where, table_name, shard_key)
+          else
+            []
+          end
+        else
+          values = Array(tree.rhs)
+          if field.try(:table) == table_name and field.name == shard_key and
+              !tree.rhs.is_a?(SQLTree::Node::SubQuery)
+            values.map(&:value).compact
+          else
+            []
+          end
+        end
+      when NOT_USED_FOR_SHARDING_OPERATORS_REGEXP
+        []
+      else
+        raise ActiveRecord::Turntable::UnknownOperatorError,
+          "[ActiveRecord::Turntable][BUG] Found Unknown SQL Operator:'#{tree.operator if tree.respond_to?(:operaor)}', Please report this bug."
+      end
+    end
+
+    private
+
+    def divide_insert_values(tree, shard_key_name)
+      idx = tree.fields.find_index {|f| f.name == shard_key_name.to_s }
+      result = {}
+      tree.values.each do |val|
+        (result[val[idx].value] ||= []) << val
+      end
+      return result
+    end
+
+    def build_shards_with_same_query(shards, query)
+      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
+    end
+
+    def build_select_fader(tree, method, query, *args, &block)
+      shard_keys = if !tree.where and tree.from.size == 1 and SQLTree::Node::SubQuery === tree.from.first.table_reference.table
+                     find_shard_keys(tree.from.first.table_reference.table.where,
+                                     @proxy.cluster.klass.table_name,
+                                     @proxy.cluster.klass.turntable_shard_key.to_s)
+                   else
+                     find_shard_keys(tree.where,
+                                     @proxy.cluster.klass.table_name,
+                                     @proxy.cluster.klass.turntable_shard_key.to_s)
+                   end
+
+      if shard_keys.size == 1 # shard
+        return Fader::SpecifiedShard.new(@proxy,
+                                         { @proxy.cluster.select_shard(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 shard_keys.present?
+        if SQLTree::Node::SelectDeclaration === tree.select.first and
+            SQLTree::Node::CountAggregrate === tree.select.first.expression
+          return Fader::CalculateShardsSumResult.new(@proxy,
+                                                     build_shards_with_same_query(@proxy.shards.values, query),
+                                                     method, query, *args, &block)
+        else
+          return Fader::SelectShardsMergeResult.new(@proxy,
+                                                    Hash[shard_keys.map {|k| [@proxy.cluster.select_shard(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
+          return Fader::CalculateShardsSumResult.new(@proxy,
+                                                     build_shards_with_same_query(@proxy.shards.values, query),
+                                                     method, query, *args, &block)
+        elsif SQLTree::Node::AllFieldsDeclaration === tree.select.first or
+            SQLTree::Node::Expression::Value === tree.select.first.expression or
+            SQLTree::Node::Expression::Variable === tree.select.first.expression
+
+          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}"
+        end
+      end
+    end
+
+    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)
+                          else
+                            build_shards_with_same_query(@proxy.shards.values, query)
+                          end
+      Fader::UpdateShardsMergeResult.new(@proxy,
+                                         shards_with_query,
+                                         method, query, *args, &block)
+    end
+
+    def build_insert_fader(tree, method, query, *args, &block)
+      values_hash = divide_insert_values(tree, @proxy.cluster.klass.turntable_shard_key)
+      shards_with_query = {}
+      values_hash.each do |k,vs|
+        tree.values = [[SQLTree::Node::Expression::Variable.new("\\0")]]
+        sql = tree.to_sql
+        value_sql = vs.map do |val|
+          "(#{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
+      end
+      Fader::InsertShardsMergeResult.new(@proxy, shards_with_query, method, query, *args, &block)
+    end
+  end
+end

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

@@ -0,0 +1,34 @@
+# -*- coding: utf-8 -*-
+module ActiveRecord::Turntable
+  class Mixer
+    class Fader
+      # 単数shard
+      autoload :SpecifiedShard, "active_record/turntable/mixer/fader/specified_shard"
+
+      # 複数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"
+
+      attr_reader :shards_query_hash
+      attr_reader :called_method
+      attr_reader :query
+
+      def initialize(proxy, shards_query_hash, called_method, query, *args, &block)
+        @proxy = proxy
+        @shards_query_hash = shards_query_hash
+        @called_method = called_method
+        @query = query
+        @args = args
+        @block = block
+      end
+
+      def execute
+        raise ActiveRecord::Turntable::NotImplementedError, "Called abstract method"
+      end
+    end
+  end
+end

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

@@ -0,0 +1,15 @@
+module ActiveRecord::Turntable
+  class Mixer
+    class Fader
+      class CalculateShardsSumResult < Fader
+        def execute
+          @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
+      end
+    end
+  end
+end

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

@@ -0,0 +1,24 @@
+module ActiveRecord::Turntable
+  class Mixer
+    class Fader
+      class InsertShardsMergeResult < Fader
+        def execute
+          if @shards_query_hash.size == 1
+            @proxy.shards_transaction(@shards_query_hash.keys) do
+              shard, query = @shards_query_hash.first
+              shard.connection.send(@called_method, query, *@args, &@block)
+            end
+          else
+            @proxy.shards_transaction(@shards_query_hash.keys) do
+              @shards_query_hash.each do |shard, query|
+                args    = @args.dup
+                args[4] = args[4].dup if args[4].present?
+                shard.connection.send(@called_method, query, *args, &@block)
+              end
+            end
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,22 @@
+module ActiveRecord::Turntable
+  class Mixer
+    class Fader
+      class SelectShardsMergeResult < Fader
+        def execute
+          res = @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
+
+          case @called_method
+          when "select_value", "select_one"
+            res.first if res
+          else
+            res
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,12 @@
+module ActiveRecord::Turntable
+  class Mixer
+    class Fader
+      class SpecifiedShard < Fader
+        def execute
+          shard, query = @shards_query_hash.first
+          shard.connection.send(@called_method, query, *@args, &@block)
+        end
+      end
+    end
+  end
+end