activerecord-turntable 1.0.0

data/lib/active_record/turntable/cluster.rb

@@ -0,0 +1,70 @@
+module ActiveRecord::Turntable
+  class Cluster
+
+    DEFAULT_CONFIG = {
+      "shards" => [],
+      "algorithm" => "range",
+    }.with_indifferent_access
+
+    def initialize(klass, cluster_spec, options = {})
+      @klass = klass
+      @config = DEFAULT_CONFIG.merge(cluster_spec)
+      @options = options.with_indifferent_access
+      @shards = {}.with_indifferent_access
+
+      # setup master shard
+      @master_shard = MasterShard.new(klass)
+
+      # setup sequencer
+      if (seq = (@options[:seq] || @config[:seq]))
+        @seq_shard = SeqShard.new(seq)
+      end
+
+      # setup shards
+      @config[:shards].each do |spec|
+        @shards[spec["connection"]] ||= Shard.new(spec)
+      end
+
+      # setup algorithm
+      alg_name = "ActiveRecord::Turntable::Algorithm::#{@config["algorithm"].camelize}Algorithm"
+      @algorithm = alg_name.constantize.new(@config)
+
+      # setup proxy
+      @connection_proxy = ConnectionProxy.new(self, cluster_spec)
+    end
+
+    def klass
+      @klass
+    end
+
+    def master
+      @master_shard
+    end
+
+    def seq
+      @seq_shard || @master_shard
+    end
+
+    def shards
+      @shards
+    end
+
+    def connection_proxy
+      @connection_proxy
+    end
+
+    def select_shard(key)
+      @shards[@algorithm.calculate(key)]
+    rescue
+      raise ActiveRecord::Turntable::CannotSpecifyShardError,
+      "[#{klass}] cannot select_shard for key:#{key}"
+    end
+
+    def weighted_shards(key = nil)
+      key ||= @klass.current_sequence
+      Hash[@algorithm.calculate_used_shards_with_weight(key).map do |k,v|
+        [@shards[k], v]
+      end]
+    end
+  end
+end

data/lib/active_record/turntable/compatible.rb

@@ -0,0 +1,19 @@
+#
+#=ActiveRecord::Turntable::Compatible
+#
+# for ActiveRecord versions compatibility
+#
+module ActiveRecord::Turntable
+  module Compatible
+    extend ActiveSupport::Concern
+
+    included do
+      # class_attributes
+      unless respond_to?(:class_attribute)
+        class << self
+          alias_method :class_attribute, :class_inheritable_accessor
+        end
+      end
+    end
+  end
+end

data/lib/active_record/turntable/config.rb

@@ -0,0 +1,24 @@
+require 'active_support/core_ext/hash/indifferent_access'
+
+module ActiveRecord::Turntable
+  class Config
+    include Singleton
+
+    def self.[](key)
+      instance[key]
+    end
+
+    def [](key)
+      self.class.load!(ActiveRecord::Base.turntable_config_file) unless @config
+      @config[key]
+    end
+
+    def self.load!(config_file, env = (defined?(Rails) ? Rails.env : 'development'))
+      instance.load!(config_file, env)
+    end
+
+    def load!(config_file, env)
+      @config = YAML.load_file(config_file).with_indifferent_access[env]
+    end
+  end
+end

data/lib/active_record/turntable/connection_proxy.rb

@@ -0,0 +1,218 @@
+require 'active_record/turntable/connection_proxy/mixable'
+module ActiveRecord::Turntable
+  class ConnectionProxy
+    include Mixable
+
+    attr_writer :spec
+
+    def initialize(cluster, options = {})
+      @cluster      =  cluster
+      @model_class  =  cluster.klass
+      @current_shard =  (cluster.master || cluster.shards.first[1])
+      @fixed_shard  = false
+      @mixer = ActiveRecord::Turntable::Mixer.new(self)
+    end
+
+    delegate :logger, :to => ActiveRecord::Base
+
+    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
+
+    # delegate :insert_many, :to => :master # ar-extensions bulk insert support
+
+    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 method_missing(method, *args, &block)
+      if shard_fixed?
+        connection.send(method, *args, &block)
+      elsif mixable?(method, *args)
+        fader = @mixer.build_fader(method, *args, &block)
+        logger.debug { "[ActiveRecord::Turntable] Sending method: #{method}, " +
+          "sql: #{args.first}, " +
+          "shards: #{fader.shards_query_hash.keys.map(&:name)}" }
+        fader.execute
+      else
+        connection.send(method, *args, &block)
+      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
+    end
+
+    def cluster
+      @cluster
+    end
+
+    def shards
+      @cluster.shards
+    end
+
+    def shard_fixed?
+      !!fixed_shard
+    end
+
+    def fixed_shard
+      @fixed_shard
+    end
+
+    def master
+      @cluster.master
+    end
+
+    def master_connection
+      master.connection
+    end
+
+    def seq
+      @cluster.seq
+    end
+
+    def current_shard
+      @current_shard
+    end
+
+    def current_shard=(shard)
+      logger.debug { "Chainging #{@model_class}'s shard to #{shard.name}"}
+      @current_shard = shard
+    end
+
+    def connection
+      @current_shard.connection
+    end
+
+    def connection_pool
+      @current_shard.connection_pool
+    end
+
+    def with_shard(shard)
+      old_shard, old_fixed = current_shard, fixed_shard
+      self.current_shard = shard
+      @fixed_shard = shard
+      yield
+    ensure
+      @fixed_shard = old_fixed
+      self.current_shard = old_shard
+    end
+
+    def with_recursive_shards(connection_name, *klasses, &block)
+      with_shard(shards[connection_name]) do
+        if klasses.blank?
+          yield
+        else
+          current_klass = klasses.shift
+          current_klass.connection.with_recursive_shards(connection_name, *klasses, &block)
+        end
+      end
+    end
+
+    def with_all(continue_on_error = false)
+      @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
+
+    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
+      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
+        end
+      end
+
+      def columns(*args)
+        if args.size > 0
+          master.connection_pool.columns[*args]
+        else
+          master.connection_pool.columns
+        end
+      end
+    end
+
+    def pk_and_sequence_for(*args)
+      master.connection.send(:pk_and_sequence_for, *args)
+    end
+
+    def primary_key(*args)
+      master.connection.send(:primary_key, *args)
+    end
+
+    def supports_views?(*args)
+      master.connection.send(:supports_views?, *args)
+    end
+
+    def spec
+      @spec ||= master.connection_pool.spec
+    end
+  end
+end

data/lib/active_record/turntable/connection_proxy/mixable.rb

@@ -0,0 +1,39 @@
+module ActiveRecord::Turntable
+  class ConnectionProxy
+    module Mixable
+      extend ActiveSupport::Concern
+
+      included do
+        if ActiveRecord::VERSION::STRING < '3.1'
+          include Rails30
+        else
+          include Rails3x
+        end
+      end
+
+      module Rails3x
+        METHODS_REGEXP = /\A(insert|select|update|delete|exec_)/
+        EXCLUDE_QUERY_REGEXP = /\A\s*SHOW/i
+        QUERY_REGEXP = /\A\s*(INSERT|DELETE|UPDATE|SELECT)/i
+
+        def mixable?(method, *args)
+          (method.to_s =~ METHODS_REGEXP &&
+           args.first !~ EXCLUDE_QUERY_REGEXP) ||
+            (method.to_s == 'execute' && args.first =~ QUERY_REGEXP)
+        end
+      end
+
+      module Rails30
+        METHODS_REGEXP = /\A(insert|select|update|delete)/
+        EXCLUDE_QUERY_REGEXP = /\A\s*SHOW/i
+        QUERY_REGEXP = /\A\s*(INSERT|DELETE|UPDATE|SELECT)/i
+
+        def mixable?(method, *args)
+          (method.to_s =~ METHODS_REGEXP &&
+           args.first !~ EXCLUDE_QUERY_REGEXP) ||
+            (method.to_s == 'execute' && args.first =~ QUERY_REGEXP)
+        end
+      end
+    end
+  end
+end

data/lib/active_record/turntable/error.rb

@@ -0,0 +1,8 @@
+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
+end

data/lib/active_record/turntable/helpers.rb

@@ -0,0 +1,5 @@
+module ActiveRecord::Turntable
+  module Helpers
+    autoload :TestHelper, 'active_record/turntable/helpers/test_helper'
+  end
+end

data/lib/active_record/turntable/helpers/test_helper.rb

@@ -0,0 +1,25 @@
+module ActiveRecord::Turntable
+  module Helpers
+    module TestHelper
+      # all shards
+      def FabricateAll(name, overrides={}, &block)
+        obj = Fabrication::Fabricator.generate(name, {
+                                                 :save => true
+                                               }, overrides, &block)
+
+        default_pool = obj.class.connection_pool
+        connection_pools = obj.class.connection_handler.instance_variable_get(:@connection_pools)
+
+        ActiveRecord::Base.turntable_connections.each do |conn_name, conn|
+          new_obj = obj.dup
+          connection_pools[new_obj.class.name] = conn
+          new_obj.id = obj.id
+          new_obj.send(:create)
+        end
+        obj
+      ensure
+        connection_pools[obj.class.name] = default_pool
+      end
+    end
+  end
+end

data/lib/active_record/turntable/master_shard.rb

@@ -0,0 +1,28 @@
+module ActiveRecord::Turntable
+  class MasterShard < Shard
+    def initialize(klass)
+      (klass and klass.connection_pool) or
+        raise MasterShardNotConnected, "connection_pool is nil"
+      @klass = klass
+      @name  = 'master'
+    end
+
+    def connection_pool
+      if ActiveRecord::Base == @klass
+        ActiveRecord::Base.connection_pool
+      else
+        # use parentclass connection which is turntable disabled
+        klass = @klass.superclass
+        candidate_connection_pool = nil
+        while !candidate_connection_pool
+          if klass == ActiveRecord::Base or !klass.turntable_enabled?
+            candidate_connection_pool = klass.connection_pool
+          else
+            klass = klass.superclass
+          end
+        end
+        candidate_connection_pool
+      end
+    end
+  end
+end