ar-octopus 0.8.1 → 0.10.2

data/lib/octopus/proxy_config.rb

@@ -0,0 +1,251 @@
+module Octopus
+  class ProxyConfig
+    CURRENT_MODEL_KEY = 'octopus.current_model'.freeze
+    CURRENT_SHARD_KEY = 'octopus.current_shard'.freeze
+    CURRENT_GROUP_KEY = 'octopus.current_group'.freeze
+    CURRENT_SLAVE_GROUP_KEY = 'octopus.current_slave_group'.freeze
+    CURRENT_LOAD_BALANCE_OPTIONS_KEY = 'octopus.current_load_balance_options'.freeze
+    BLOCK_KEY = 'octopus.block'.freeze
+    FULLY_REPLICATED_KEY = 'octopus.fully_replicated'.freeze
+
+    attr_accessor :config, :sharded, :shards, :shards_slave_groups, :slave_groups,
+                  :adapters, :replicated, :slaves_load_balancer, :slaves_list, :shards_slave_groups,
+                  :slave_groups, :groups, :entire_sharded, :shards_config
+
+    def initialize(config)
+      initialize_shards(config)
+      initialize_replication(config) if !config.nil? && config['replicated']
+    end
+
+    def current_model
+      Thread.current[CURRENT_MODEL_KEY]
+    end
+
+    def current_model=(model)
+      Thread.current[CURRENT_MODEL_KEY] = model.is_a?(ActiveRecord::Base) ? model.class : model
+    end
+
+    def current_shard
+      Thread.current[CURRENT_SHARD_KEY] ||= Octopus.master_shard
+    end
+
+    def current_shard=(shard_symbol)
+      if shard_symbol.is_a?(Array)
+        self.current_slave_group = nil
+        shard_symbol.each { |symbol| fail "Nonexistent Shard Name: #{symbol}" if shards[symbol].nil? }
+      elsif shard_symbol.is_a?(Hash)
+        hash = shard_symbol
+        shard_symbol = hash[:shard]
+        slave_group_symbol = hash[:slave_group]
+        load_balance_options = hash[:load_balance_options]
+
+        if shard_symbol.nil? && slave_group_symbol.nil?
+          fail 'Neither shard or slave group must be specified'
+        end
+
+        if shard_symbol.present?
+          fail "Nonexistent Shard Name: #{shard_symbol}" if shards[shard_symbol].nil?
+        end
+
+        if slave_group_symbol.present?
+          if (shards_slave_groups.try(:[], shard_symbol).present? && shards_slave_groups[shard_symbol][slave_group_symbol].nil?) ||
+              (shards_slave_groups.try(:[], shard_symbol).nil? && @slave_groups[slave_group_symbol].nil?)
+            fail "Nonexistent Slave Group Name: #{slave_group_symbol} in shards config: #{shards_config.inspect}"
+          end
+        end
+        self.current_slave_group = slave_group_symbol
+        self.current_load_balance_options = load_balance_options
+      else
+        fail "Nonexistent Shard Name: #{shard_symbol}" if shards[shard_symbol].nil?
+      end
+
+      Thread.current[CURRENT_SHARD_KEY] = shard_symbol
+    end
+
+    def current_group
+      Thread.current[CURRENT_GROUP_KEY]
+    end
+
+    def current_group=(group_symbol)
+      # TODO: Error message should include all groups if given more than one bad name.
+      [group_symbol].flatten.compact.each do |group|
+        fail "Nonexistent Group Name: #{group}" unless has_group?(group)
+      end
+
+      Thread.current[CURRENT_GROUP_KEY] = group_symbol
+    end
+
+    def current_slave_group
+      Thread.current[CURRENT_SLAVE_GROUP_KEY]
+    end
+
+    def current_slave_group=(slave_group_symbol)
+      Thread.current[CURRENT_SLAVE_GROUP_KEY] = slave_group_symbol
+      Thread.current[CURRENT_LOAD_BALANCE_OPTIONS_KEY] = nil if slave_group_symbol.nil?
+    end
+
+    def current_load_balance_options
+      Thread.current[CURRENT_LOAD_BALANCE_OPTIONS_KEY]
+    end
+
+    def current_load_balance_options=(options)
+      Thread.current[CURRENT_LOAD_BALANCE_OPTIONS_KEY] = options
+    end
+
+    def block
+      Thread.current[BLOCK_KEY]
+    end
+
+    def block=(block)
+      Thread.current[BLOCK_KEY] = block
+    end
+
+    def fully_replicated?
+      @fully_replicated || Thread.current[FULLY_REPLICATED_KEY]
+    end
+
+    # Public: Whether or not a group exists with the given name converted to a
+    # string.
+    #
+    # Returns a boolean.
+    def has_group?(group)
+      @groups.key?(group.to_s)
+    end
+
+    # Public: Retrieves names of all loaded shards.
+    #
+    # Returns an array of shard names as symbols
+    def shard_names
+      shards.keys
+    end
+
+    def shard_name
+      current_shard.is_a?(Array) ? current_shard.first : current_shard
+    end
+
+    # Public: Retrieves the defined shards for a given group.
+    #
+    # Returns an array of shard names as symbols or nil if the group is not
+    # defined.
+    def shards_for_group(group)
+      @groups.fetch(group.to_s, nil)
+    end
+
+    def initialize_shards(config)
+      @original_config = config
+
+      self.shards = HashWithIndifferentAccess.new
+      self.shards_slave_groups = HashWithIndifferentAccess.new
+      self.slave_groups = HashWithIndifferentAccess.new
+      self.groups = {}
+      self.config = ActiveRecord::Base.connection_pool_without_octopus.spec.config
+
+      unless config.nil?
+        self.entire_sharded = config['entire_sharded']
+        self.shards_config = config[Octopus.rails_env]
+      end
+
+      self.shards_config ||= []
+
+      shards_config.each do |key, value|
+        if value.is_a?(String)
+          value = resolve_string_connection(value).merge(:octopus_shard => key)
+          initialize_adapter(value['adapter'])
+          shards[key.to_sym] = connection_pool_for(value, "#{value['adapter']}_connection")
+        elsif value.is_a?(Hash) && value.key?('adapter')
+          value.merge!(:octopus_shard => key)
+          initialize_adapter(value['adapter'])
+          shards[key.to_sym] = connection_pool_for(value, "#{value['adapter']}_connection")
+
+          slave_group_configs = value.select do |_k, v|
+            structurally_slave_group? v
+          end
+
+          if slave_group_configs.present?
+            slave_groups = HashWithIndifferentAccess.new
+            slave_group_configs.each do |slave_group_name, slave_configs|
+              slaves = HashWithIndifferentAccess.new
+              slave_configs.each do |slave_name, slave_config|
+                shards[slave_name.to_sym] = connection_pool_for(slave_config, "#{value['adapter']}_connection")
+                slaves[slave_name.to_sym] = slave_name.to_sym
+              end
+              slave_groups[slave_group_name.to_sym] = Octopus::SlaveGroup.new(slaves)
+            end
+            @shards_slave_groups[key.to_sym] = slave_groups
+            @sharded = true
+          end
+        elsif value.is_a?(Hash)
+          @groups[key.to_s] = []
+
+          value.each do |k, v|
+            fail 'You have duplicated shard names!' if shards.key?(k.to_sym)
+
+            initialize_adapter(v['adapter'])
+            config_with_octopus_shard = v.merge(:octopus_shard => k)
+
+            shards[k.to_sym] = connection_pool_for(config_with_octopus_shard, "#{v['adapter']}_connection")
+            @groups[key.to_s] << k.to_sym
+          end
+
+          if structurally_slave_group? value
+            slaves = Hash[@groups[key.to_s].map { |v| [v, v] }]
+            @slave_groups[key.to_sym] = Octopus::SlaveGroup.new(slaves)
+          end
+        end
+      end
+
+      shards[:master] ||= ActiveRecord::Base.connection_pool_without_octopus if Octopus.master_shard == :master
+    end
+
+    def initialize_replication(config)
+      @replicated = true
+      if config.key?('fully_replicated')
+        @fully_replicated = config['fully_replicated']
+      else
+        @fully_replicated = true
+      end
+
+      @slaves_list = shards.keys.map(&:to_s).sort
+      @slaves_list.delete('master')
+      @slaves_load_balancer = Octopus.load_balancer.new(@slaves_list)
+    end
+
+    def reinitialize_shards
+      initialize_shards(@original_config)
+    end
+
+    private
+
+    def connection_pool_for(config, adapter)
+      if Octopus.rails4?
+        spec = ActiveRecord::ConnectionAdapters::ConnectionSpecification.new(config.dup, adapter )
+      else
+        name = adapter["octopus_shard"]
+        spec = ActiveRecord::ConnectionAdapters::ConnectionSpecification.new(name, config.dup, adapter)
+      end
+
+      ActiveRecord::ConnectionAdapters::ConnectionPool.new(spec)
+    end
+
+    def resolve_string_connection(spec)
+      resolver = ActiveRecord::ConnectionAdapters::ConnectionSpecification::Resolver.new({})
+      HashWithIndifferentAccess.new(resolver.spec(spec).config)
+    end
+
+    def structurally_slave?(config)
+      config.is_a?(Hash) && config.key?('adapter')
+    end
+
+    def structurally_slave_group?(config)
+      config.is_a?(Hash) && config.values.any? { |v| structurally_slave? v }
+    end
+
+    def initialize_adapter(adapter)
+      begin
+        require "active_record/connection_adapters/#{adapter}_adapter"
+      rescue LoadError
+        raise "Please install the #{adapter} adapter: `gem install activerecord-#{adapter}-adapter` (#{$ERROR_INFO})"
+      end
+    end
+  end
+end

data/lib/octopus/query_cache_for_shards.rb

@@ -0,0 +1,24 @@
+# query cache methods are moved to ConnectionPool for Rails >= 5.0
+module Octopus
+  module ConnectionPool
+    module QueryCacheForShards
+      %i(enable_query_cache! disable_query_cache!).each do |method|
+        define_method(method) do
+          if(Octopus.enabled? && (shards = ActiveRecord::Base.connection.shards)['master'] == self)
+            shards.each do |shard_name, v|
+              if shard_name == 'master'
+                super()
+              else
+                v.public_send(method)
+              end
+            end
+          else
+            super()
+          end
+        end
+      end
+    end
+  end
+end
+
+ActiveRecord::ConnectionAdapters::ConnectionPool.send(:prepend, Octopus::ConnectionPool::QueryCacheForShards)

data/lib/octopus/railtie.rb

@@ -1,10 +1,8 @@
 begin
-  require "rails/railtie"
-
   module Octopus
     class Railtie < Rails::Railtie
       rake_tasks do
-        Dir[File.join(File.dirname(__FILE__), "../tasks/*.rake")].each { |ext| load ext }
+        Dir[File.join(File.dirname(__FILE__), '../tasks/*.rake')].each { |ext| load ext }
       end
     end
   end

data/lib/octopus/relation_proxy.rb

@@ -0,0 +1,70 @@
+module Octopus
+  class RelationProxy < BasicObject
+    include ::Octopus::ShardTracking::Attribute
+
+    module CaseFixer
+      def ===(other)
+        other = other.ar_relation while ::Octopus::RelationProxy === other
+        super
+      end
+    end
+
+    attr_accessor :ar_relation
+
+    def initialize(shard, ar_relation)
+      @current_shard = shard
+      @ar_relation = ar_relation
+    end
+
+    def respond_to?(*args)
+      method_missing(:respond_to?, *args)
+    end
+
+    # methods redefined in ActiveRecord that should run_on_shard
+    ENUM_METHODS = (::Enumerable.instance_methods - ::Object.instance_methods).reject do |m|
+      ::ActiveRecord::Relation.instance_method(m).source_location rescue nil
+    end + [:each, :map, :index_by]
+    # `find { ... }` etc. should run_on_shard, `find(id)` should be sent to relation
+    ENUM_WITH_BLOCK_METHODS = [:find, :select, :none?, :any?, :one?, :many?, :sum]
+    BATCH_METHODS = [:find_each, :find_in_batches, :in_batches]
+    WHERE_CHAIN_METHODS = [:not]
+
+    def method_missing(method, *args, &block)
+      if !block && BATCH_METHODS.include?(method)
+        ::Enumerator.new do |yielder|
+          run_on_shard do
+            @ar_relation.public_send(method, *args) do |batch_item|
+              yielder << batch_item
+            end
+          end
+        end
+      elsif ENUM_METHODS.include?(method) || block && ENUM_WITH_BLOCK_METHODS.include?(method)
+        run_on_shard { @ar_relation.to_a }.public_send(method, *args, &block)
+      elsif WHERE_CHAIN_METHODS.include?(method)
+        ::Octopus::ScopeProxy.new(@current_shard, run_on_shard { @ar_relation.public_send(method, *args) } )
+      elsif block
+        @ar_relation.public_send(method, *args, &block)
+      else
+        run_on_shard do
+          if method == :load_records
+            @ar_relation.send(method, *args)
+          else
+            @ar_relation.public_send(method, *args)
+          end
+        end
+      end
+    end
+
+    def ==(other)
+      case other
+      when ::Octopus::RelationProxy
+        method_missing(:==, other.ar_relation)
+      else
+        method_missing(:==, other)
+      end
+    end
+    alias_method :eql?, :==
+  end
+end
+
+ActiveRecord::Relation.extend(Octopus::RelationProxy::CaseFixer)

data/lib/octopus/result_patch.rb

@@ -0,0 +1,19 @@
+module Octopus::ResultPatch
+  attr_accessor :current_shard
+
+  private
+
+  def hash_rows
+    if current_shard.blank?
+      super
+    else
+      foo = super
+      foo.each { |f| f.merge!('current_shard' => current_shard) }
+      foo
+    end
+  end
+end
+
+class ActiveRecord::Result
+  prepend Octopus::ResultPatch
+end

data/lib/octopus/scope_proxy.rb

@@ -1,50 +1,68 @@
-class Octopus::ScopeProxy
-  attr_accessor :shard, :klass
+module Octopus
+  class ScopeProxy < BasicObject
+    include ::Octopus::ShardTracking::Attribute
 
-  def initialize(shard, klass)
-    @shard = shard
-    @klass = klass
-  end
+    module CaseFixer
+      def ===(other)
+        other = other.klass while ::Octopus::ScopeProxy === other
+        super
+      end
+    end
 
-  def using(shard)
-    raise "Nonexistent Shard Name: #{shard}" if @klass.connection.instance_variable_get(:@shards)[shard].nil?
-    @shard = shard
-    return self
-  end
+    attr_accessor :klass
+
+    # Dup and clone should be delegated to the class.
+    # We want to dup the query, not the scope proxy.
+    delegate :dup, :clone, to: :klass
 
-  # Transaction Method send all queries to a specified shard.
-  def transaction(options = {}, &block)
-    @klass.connection.run_queries_on_shard(@shard) do
-      @klass = @klass.connection().transaction(options, &block)
+    def initialize(shard, klass)
+      @current_shard = shard
+      @klass = klass
     end
-  end
 
-  def connection
-    @klass.connection().current_shard = @shard
-    @klass.connection()
-  end
+    def using(shard)
+      fail "Nonexistent Shard Name: #{shard}" if @klass.connection.shards[shard].nil?
+      @current_shard = shard
+      self
+    end
 
-  def method_missing(method, *args, &block)
-    result = @klass.connection.run_queries_on_shard(@shard) do
-      @klass.send(method, *args, &block)
+    # Transaction Method send all queries to a specified shard.
+    def transaction(options = {}, &block)
+      run_on_shard { klass.transaction(options, &block) }
     end
 
-    if result.respond_to?(:scoped)
-      @klass = result
-      return self
+    def connection
+      @klass.connection_proxy.current_shard = @current_shard
+
+      if @klass.custom_octopus_connection && @klass.allowed_shard?(@current_shard)
+        # Force use of proxy, given we called 'using' explicitly to get here
+        @klass.connection_proxy.current_model = @klass
+        @klass.connection_proxy
+      else
+        @klass.connection
+      end
     end
 
-    result
-  end
+    def method_missing(method, *args, &block)
+      result = run_on_shard { @klass.__send__(method, *args, &block) }
+      if result.respond_to?(:all)
+        return ::Octopus::ScopeProxy.new(current_shard, result)
+      end
 
-  def as_json(options = nil)
-    method_missing(:as_json, options)
-  end
+      if result.respond_to?(:current_shard)
+        result.current_shard = current_shard
+      end
 
-  # Delegates to method_missing (instead of @klass) so that User.using(:blah).where(:name => "Mike")
-  # gets run in the correct shard context when #== is evaluated.
-  def ==(*args)
-    method_missing(:==, *args)
+      result
+    end
+
+    # Delegates to method_missing (instead of @klass) so that User.using(:blah).where(:name => "Mike")
+    # gets run in the correct shard context when #== is evaluated.
+    def ==(other)
+      method_missing(:==, other)
+    end
+    alias_method :eql?, :==
   end
-  alias :eql? :==
 end
+
+ActiveRecord::Relation.extend(Octopus::ScopeProxy::CaseFixer)

data/lib/octopus/shard_tracking/attribute.rb

@@ -0,0 +1,22 @@
+# Adds current_shard as an attribute; provide a default
+# implementation of set_current_shard which considers
+# only the current ActiveRecord::Base.connection_proxy
+module Octopus
+  module ShardTracking
+    module Attribute
+      def self.included(base)
+        base.send(:include, Octopus::ShardTracking)
+      end
+
+      attr_accessor :current_shard
+
+      def set_current_shard
+        return unless Octopus.enabled?
+
+        if ActiveRecord::Base.connection_proxy.block
+          self.current_shard = ActiveRecord::Base.connection_proxy.current_shard
+        end
+      end
+    end
+  end
+end

data/lib/octopus/shard_tracking/dynamic.rb

@@ -0,0 +1,11 @@
+require 'octopus/shard_tracking'
+
+module Octopus
+  module ShardTracking
+    module Dynamic
+      def self.included(base)
+        base.send(:include, Octopus::ShardTracking)
+      end
+    end
+  end
+end

data/lib/octopus/shard_tracking.rb

@@ -0,0 +1,46 @@
+module Octopus
+  module ShardTracking
+    def self.included(base)
+      base.extend(ClassMethods)
+    end
+
+    module ClassMethods
+      # If the class which includes this module responds to the class
+      # method sharded_methods, then automagically alias_method_chain
+      # a sharding-friendly version of each of those methods into existence
+      def sharded_methods(*methods)
+        methods.each { |m| create_sharded_method(m) }
+      end
+
+      def create_sharded_method(name)
+        name.to_s =~ /([^!?]+)([!?])?/
+        method, punctuation = [Regexp.last_match[1], Regexp.last_match[2]]
+        with = :"#{method}_with_octopus#{punctuation}"
+        without = :"#{method}_without_octopus#{punctuation}"
+        define_method with do |*args, &block|
+          run_on_shard { send(without, *args, &block) }
+        end
+        alias_method without.to_sym, name.to_sym
+        alias_method name.to_sym, with.to_sym
+      end
+    end
+
+    # Adds run_on_shard method, but does not implement current_shard method
+    def run_on_shard(&block)
+      if (cs = current_shard)
+        r = ActiveRecord::Base.connection_proxy.run_queries_on_shard(cs, &block)
+        # Use a case statement to avoid any path through ActiveRecord::Delegation's
+        # respond_to? code. We want to avoid the respond_to? code because it can have
+        # the side effect of causing a call to load_target
+
+        if (ActiveRecord::Relation === r || ActiveRecord::QueryMethods::WhereChain === r) && !(Octopus::RelationProxy === r)
+          Octopus::RelationProxy.new(cs, r)
+        else
+          r
+        end
+      else
+        yield
+      end
+    end
+  end
+end

data/lib/octopus/singular_association.rb

@@ -0,0 +1,9 @@
+module Octopus
+  module SingularAssociation
+    def self.included(base)
+      base.sharded_methods :reader, :writer, :create, :create!, :build
+    end
+  end
+end
+
+ActiveRecord::Associations::SingularAssociation.send(:include, Octopus::SingularAssociation)

data/lib/octopus/slave_group.rb

@@ -0,0 +1,13 @@
+module Octopus
+  class SlaveGroup
+    def initialize(slaves)
+      slaves = HashWithIndifferentAccess.new(slaves)
+      slaves_list = slaves.values
+      @load_balancer = Octopus.load_balancer.new(slaves_list)
+    end
+
+    def next(options)
+      @load_balancer.next options
+    end
+  end
+end

data/lib/octopus/version.rb

@@ -1,3 +1,3 @@
 module Octopus
-  VERSION = '0.8.1'
+  VERSION = '0.10.2'
 end