activerecord-turntable 1.0.0

data/lib/active_record/turntable/active_record_ext/persistence.rb

@@ -0,0 +1,95 @@
+module ActiveRecord::Turntable::ActiveRecordExt
+  module Persistence
+    if ActiveRecord::VERSION::STRING < '3.1'
+      ::ActiveRecord::Persistence.class_eval do
+        def destroy
+          klass = self.class
+          if persisted?
+            condition_scope = klass.unscoped.where(klass.arel_table[klass.primary_key].eq(id))
+            if klass.turntable_enabled? and klass.primary_key != klass.turntable_shard_key.to_s
+              condition_scope = condition_scope.where(klass.turntable_shard_key => self.send(turntable_shard_key))
+            end
+            condition_scope.delete_all
+          end
+
+          @destroyed = true
+          freeze
+        end
+
+        private
+
+        # overrides ActiveRecord::Persistence's original method so that
+        def update(attribute_names = @attributes.keys)
+          klass = self.class
+          attributes_with_values = arel_attributes_values(false, false, attribute_names)
+          return 0 if attributes_with_values.empty?
+          condition_scope = klass.unscoped.where(klass.arel_table[klass.primary_key].eq(id))
+          if klass.turntable_enabled? and klass.primary_key != klass.turntable_shard_key.to_s
+            condition_scope = condition_scope.where(klass.turntable_shard_key => self.send(turntable_shard_key))
+          end
+          condition_scope.arel.update(attributes_with_values)
+        end
+      end
+
+    else
+      ::ActiveRecord::Persistence.class_eval do
+        def destroy
+          klass = self.class
+          destroy_associations
+
+          if persisted?
+            ActiveRecord::IdentityMap.remove(self) if ActiveRecord::IdentityMap.enabled?
+            pk         = klass.primary_key
+            column     = klass.columns_hash[pk]
+            substitute = connection.substitute_at(column, 0)
+
+            relation = klass.unscoped.where(klass.arel_table[pk].eq(substitute))
+            if klass.turntable_enabled? and klass.primary_key != klass.turntable_shard_key.to_s
+              relation = relation.where(klass.turntable_shard_key => self.send(turntable_shard_key))
+            end
+            relation.bind_values = [[column, id]]
+            relation.delete_all
+          end
+
+          @destroyed = true
+          freeze
+        end
+
+        def destroy_without_callbacks
+          klass = self.class
+          destroy_associations
+
+          if persisted?
+            ActiveRecord::IdentityMap.remove(self) if ActiveRecord::IdentityMap.enabled?
+            pk         = klass.primary_key
+            column     = klass.columns_hash[pk]
+            substitute = connection.substitute_at(column, 0)
+
+            relation = klass.unscoped.where(klass.arel_table[pk].eq(substitute))
+            if klass.turntable_enabled? and klass.primary_key != klass.turntable_shard_key.to_s
+              relation = relation.where(klass.turntable_shard_key => self.send(turntable_shard_key))
+            end
+            relation.bind_values = [[column, id]]
+            relation.delete_all
+          end
+
+          @destroyed = true
+          freeze
+        end
+
+        private
+        def update(attribute_names = @attributes.keys)
+          attributes_with_values = arel_attributes_values(false, false, attribute_names)
+          return 0 if attributes_with_values.empty?
+          klass = self.class
+          condition_scope = klass.unscoped.where(klass.arel_table[klass.primary_key].eq(id))
+          if klass.turntable_enabled? and klass.primary_key != klass.turntable_shard_key.to_s
+            condition_scope = condition_scope.where(klass.turntable_shard_key => self.send(turntable_shard_key))
+          end
+          stmt = condition_scope.arel.compile_update(attributes_with_values)
+          klass.connection.update stmt.to_sql
+        end
+      end
+    end
+  end
+end

data/lib/active_record/turntable/active_record_ext/schema_dumper.rb

@@ -0,0 +1,107 @@
+# -*- coding: utf-8 -*-
+module ActiveRecord::Turntable
+  module ActiveRecordExt
+    module SchemaDumper
+      extend ActiveSupport::Concern
+
+      included do
+        alias_method_chain :table, :turntable
+      end
+
+      private
+
+      def table_with_turntable(table, stream)
+        columns = @connection.columns(table)
+        begin
+          tbl = StringIO.new
+
+          # first dump primary key column
+          if @connection.respond_to?(:pk_and_sequence_for)
+            pk, _ = @connection.pk_and_sequence_for(table)
+          elsif @connection.respond_to?(:primary_key)
+            pk = @connection.primary_key(table)
+          end
+
+          # turntable sequencer dump
+          if table =~ /\A(.*)_id_seq\z/
+            tbl.print "  create_sequence_for #{$1.inspect}"
+          else
+            tbl.print "  create_table #{table.inspect}"
+          end
+
+          if columns.detect { |c| c.name == pk }
+            if pk != 'id'
+              tbl.print %Q(, :primary_key => "#{pk}")
+            end
+          else
+            tbl.print ", :id => false"
+          end
+          tbl.print ", :force => true"
+          tbl.puts " do |t|"
+
+          # then dump all non-primary key columns
+          column_specs = columns.map do |column|
+            raise StandardError, "Unknown type '#{column.sql_type}' for column '#{column.name}'" if @types[column.type].nil?
+            next if column.name == pk
+            spec = {}
+            spec[:name]      = column.name.inspect
+
+            # AR has an optimisation which handles zero-scale decimals as integers.  This
+            # code ensures that the dumper still dumps the column as a decimal.
+            spec[:type]      = if column.type == :integer && [/^numeric/, /^decimal/].any? { |e| e.match(column.sql_type) }
+                                 'decimal'
+                               else
+                                 column.type.to_s
+                               end
+            spec[:limit]     = column.limit.inspect if column.limit != @types[column.type][:limit] && spec[:type] != 'decimal'
+            spec[:precision] = column.precision.inspect if column.precision
+            spec[:scale]     = column.scale.inspect if column.scale
+            spec[:null]      = 'false' unless column.null
+            spec[:default]   = default_string(column.default) if column.has_default?
+            (spec.keys - [:name, :type]).each{ |k| spec[k].insert(0, "#{k.inspect} => ")}
+            spec
+          end.compact
+
+          # find all migration keys used in this table
+          keys = [:name, :limit, :precision, :scale, :default, :null] & column_specs.map{ |k| k.keys }.flatten
+
+          # figure out the lengths for each column based on above keys
+          lengths = keys.map{ |key| column_specs.map{ |spec| spec[key] ? spec[key].length + 2 : 0 }.max }
+
+          # the string we're going to sprintf our values against, with standardized column widths
+          format_string = lengths.map{ |len| "%-#{len}s" }
+
+          # find the max length for the 'type' column, which is special
+          type_length = column_specs.map{ |column| column[:type].length }.max
+
+          # add column type definition to our format string
+          format_string.unshift "    t.%-#{type_length}s "
+
+          format_string *= ''
+
+          column_specs.each do |colspec|
+            values = keys.zip(lengths).map{ |key, len| colspec.key?(key) ? colspec[key] + ", " : " " * len }
+            values.unshift colspec[:type]
+            tbl.print((format_string % values).gsub(/,\s*$/, ''))
+            tbl.puts
+          end
+
+          tbl.puts "  end"
+          tbl.puts
+
+          indexes(table, tbl)
+
+          tbl.rewind
+          stream.print tbl.read
+        rescue => e
+          stream.puts "# Could not dump table #{table.inspect} because of following #{e.class}"
+          stream.puts "#   #{e.message}"
+          stream.puts
+        end
+
+        stream
+      end
+
+    end
+  end
+end

data/lib/active_record/turntable/active_record_ext/sequencer.rb

@@ -0,0 +1,28 @@
+module ActiveRecord::Turntable::ActiveRecordExt
+  module Sequencer
+    extend ActiveSupport::Concern
+
+    included do
+      include DatabaseStatements
+      alias_method_chain :prefetch_primary_key?, :turntable
+    end
+
+    module DatabaseStatements
+      def default_sequence_name(table_name, pk = nil)
+        ActiveRecord::Turntable::Sequencer.sequence_name(table_name, pk)
+      end
+    end
+
+    def prefetch_primary_key_with_turntable?(table_name = nil)
+      ActiveRecord::Turntable::Sequencer.has_sequencer?(table_name)
+    end
+
+    def next_sequence_value(sequence_name)
+      ActiveRecord::Turntable::Sequencer.sequences[sequence_name].next_sequence_value(sequence_name)
+    end
+
+    def current_sequence_value(sequence_name)
+      ActiveRecord::Turntable::Sequencer.sequences[sequence_name].current_sequence_value(sequence_name)
+    end
+  end
+end

data/lib/active_record/turntable/active_record_ext/transactions.rb

@@ -0,0 +1,33 @@
+module ActiveRecord::Turntable
+  module ActiveRecordExt
+    module Transactions
+      def with_transaction_returning_status
+        if self.class.turntable_enabled?
+          status = nil
+          if self.new_record? and self.turntable_shard_key.to_s == self.class.primary_key and
+              self.id.nil? and connection.prefetch_primary_key?(self.class.table_name)
+            self.id = connection.next_sequence_value(self.class.sequence_name)
+          end
+          self.class.connection.shards_transaction([self.turntable_shard]) do
+            add_to_transaction
+            status = yield
+            raise ActiveRecord::Rollback unless status
+          end
+          status
+        else
+          super
+        end
+      end
+
+      def add_to_transaction
+        if self.class.turntable_enabled?
+          if self.turntable_shard.connection.add_transaction_record(self)
+            remember_transaction_record_state
+          end
+        else
+          super
+        end
+      end
+    end
+  end
+end

data/lib/active_record/turntable/algorithm.rb

@@ -0,0 +1,7 @@
+module ActiveRecord::Turntable
+  module Algorithm
+    autoload :Base, "active_record/turntable/algorithm/base"
+    autoload :RangeAlgorithm, "active_record/turntable/algorithm/range_algorithm"
+    autoload :RangeBsearchAlgorithm, "active_record/turntable/algorithm/range_bsearch_algorithm"
+  end
+end

data/lib/active_record/turntable/algorithm/base.rb

@@ -0,0 +1,11 @@
+module ActiveRecord::Turntable::Algorithm
+  class Base
+    def initialize(config)
+      @config = config
+    end
+
+    def calculate(key)
+      raise ActiveRecord::Turntable::NotImplementedError, "not implemented"
+    end
+  end
+end

data/lib/active_record/turntable/algorithm/range_algorithm.rb

@@ -0,0 +1,37 @@
+# -*- coding: utf-8 -*-
+module ActiveRecord::Turntable::Algorithm
+  class RangeAlgorithm < Base
+    def initialize(config)
+      @config = config
+    end
+
+    def calculate(key)
+      idx = calculate_idx(key)
+      @config["shards"][idx]["connection"]
+    rescue
+      raise ActiveRecord::Turntable::CannotSpecifyShardError, "cannot specify shard for key:#{key}"
+    end
+
+    def calculate_idx(key)
+      @config["shards"].find_index {|h| h["less_than"] > key }
+    end
+
+    # { connection_name => weight, ... }
+    def calculate_used_shards_with_weight(sequence_value)
+      idx = calculate_idx(sequence_value)
+      last_connection = calculate(sequence_value)
+      shards = @config["shards"][0..idx]
+      weighted_hash = Hash.new {|h,k| h[k]=0}
+      prev_max = 0
+      shards.each_with_index do |h,idx|
+        weighted_hash[h["connection"]] += if idx < shards.size - 1
+                                            h["less_than"] - prev_max - 1
+                                          else
+                                            sequence_value - prev_max
+                                          end
+        prev_max = h["less_than"] - 1
+      end
+      return weighted_hash
+    end
+  end
+end

data/lib/active_record/turntable/algorithm/range_bsearch_algorithm.rb

@@ -0,0 +1,41 @@
+# -*- coding: utf-8 -*-
+require 'bsearch'
+module ActiveRecord::Turntable::Algorithm
+  class RangeBsearchAlgorithm < Base
+    def initialize(config)
+      @config = config
+      @config["shards"].sort_by! {|a| a["less_than"]}
+    end
+
+    def calculate(key)
+      idx = calculate_idx(key)
+      @config["shards"][idx]["connection"]
+    rescue
+      raise ActiveRecord::Turntable::CannotSpecifyShardError, "cannot specify shard for key:#{key}"
+    end
+
+    def calculate_idx(key)
+      @config["shards"].bsearch_lower_boundary { |h|
+        h["less_than"] <=> key
+      }
+    end
+
+    # { connection_name => weight, ... }
+    def calculate_used_shards_with_weight(sequence_value)
+      idx = calculate_idx(sequence_value)
+      last_connection = calculate(sequence_value)
+      shards = @config["shards"][0..idx]
+      weighted_hash = Hash.new {|h,k| h[k]=0}
+      prev_max = 0
+      shards.each_with_index do |h,idx|
+        weighted_hash[h["connection"]] += if idx < shards.size - 1
+                                            h["less_than"] - prev_max - 1
+                                          else
+                                            sequence_value - prev_max
+                                          end
+        prev_max = h["less_than"] - 1
+      end
+      return weighted_hash
+    end
+  end
+end

data/lib/active_record/turntable/base.rb

@@ -0,0 +1,130 @@
+module ActiveRecord::Turntable
+  module Base
+    extend ActiveSupport::Concern
+
+    included do
+      include Compatible
+      class_attribute :turntable_connections,
+                        :turntable_enabled, :turntable_sequencer_enabled
+
+      self.turntable_connections = {}
+      self.turntable_enabled = false
+      self.turntable_sequencer_enabled = false
+      class << self
+        delegate :shards_transaction, :to => :connection
+      end
+    end
+
+    module ClassMethods
+      def turntable(cluster_name, shard_key_name, options = {})
+        class_attribute :turntable_shard_key,
+                          :turntable_cluster, :turntable_cluster_name
+
+        self.turntable_enabled = true
+        self.turntable_cluster_name = cluster_name
+        self.turntable_shard_key = shard_key_name
+        self.turntable_cluster = Cluster.new(
+                                   self,
+                                   turntable_config[:clusters][cluster_name],
+                                   options
+                                 )
+        turntable_replace_connection_pool
+      end
+
+      def force_transaction_all_shards!(options={}, &block)
+        force_connect_all_shards!
+        shards = turntable_connections.values
+        shards += [ActiveRecord::Base.connection_pool]
+        recursive_transaction(shards, options, &block)
+      end
+
+      def recursive_transaction(pools, options, &block)
+        pool = pools.shift
+        if pools.present?
+          pool.connection.transaction(options) do
+            recursive_transaction(pools, options, &block)
+          end
+        else
+          pool.connection.transaction(options, &block)
+        end
+      end
+
+      def force_connect_all_shards!
+        conf = configurations[Rails.env]
+        shards = conf["shards"]
+        shards = shards.merge(conf["seq"]) if conf["seq"]
+        shards.each do |name, config|
+          turntable_connections[name] ||=
+            ActiveRecord::ConnectionAdapters::ConnectionPool.new(spec_for(config))
+        end
+      end
+
+      def turntable_replace_connection_pool
+        ch = connection_handler
+        cp = turntable_cluster.connection_proxy
+        if ActiveRecord::VERSION::STRING >= '3.2.0'
+          ch.connection_pools[cp.spec] = PoolProxy.new(cp)
+          ch.instance_variable_get(:@class_to_pool)[name] = ch.connection_pools[cp.spec]
+        else
+          ch.connection_pools[name] = PoolProxy.new(cp)
+        end
+      end
+
+      def spec_for(config)
+        begin
+          require "active_record/connection_adapters/#{config['adapter']}_adapter"
+        rescue LoadError => e
+          raise "Please install the #{config['adapter']} adapter: `gem install activerecord-#{config['adapter']}-adapter` (#{e})"
+        end
+        adapter_method = "#{config['adapter']}_connection"
+        ActiveRecord::Base::ConnectionSpecification.new(config, adapter_method)
+      end
+
+      def clear_all_connections!
+        turntable_connections.values.each do |pool|
+          pool.disconnect!
+        end
+      end
+
+      def sequencer
+        class_attribute :turntable_sequencer
+        self.turntable_sequencer_enabled = true
+        self.turntable_sequencer = ActiveRecord::Turntable::Sequencer.build(self)
+      end
+
+      def turntable_enabled?
+        turntable_enabled
+      end
+
+      def sequencer_enabled?
+        turntable_sequencer_enabled
+      end
+
+      def current_sequence
+        connection.current_sequence_value(self.sequence_name) if sequencer_enabled?
+      end
+
+      def current_last_shard
+        turntable_cluster.select_shard(current_sequence) if sequencer_enabled?
+      end
+
+      def weighted_random_shard_with(*klasses, &block)
+        shards_weight = self.turntable_cluster.weighted_shards
+        sum = shards_weight.values.inject(&:+)
+        idx = rand(sum)
+        shard, weight = shards_weight.find {|k,v|
+          (idx -= v) < 0
+        }
+        self.connection.with_recursive_shards(shard.name, *klasses, &block)
+      end
+    end
+
+    def shards_transaction(options = {}, &block)
+      self.class.shards_transaction(options, &block)
+    end
+
+    def turntable_shard
+      turntable_cluster.select_shard(self.send(turntable_shard_key))
+    end
+  end
+end