ar-octopus 0.8.5 → 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,6 +1,4 @@
 begin
-  require 'rails/railtie'
-
   module Octopus
     class Railtie < Rails::Railtie
       rake_tasks do

data/lib/octopus/relation_proxy.rb

@@ -16,8 +16,43 @@ module Octopus
       @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)
-      run_on_shard { @ar_relation.public_send(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)

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

@@ -11,20 +11,24 @@ module Octopus
 
     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
+
     def initialize(shard, klass)
       @current_shard = shard
       @klass = klass
     end
 
     def using(shard)
-      fail "Nonexistent Shard Name: #{shard}" if @klass.connection.instance_variable_get(:@shards)[shard].nil?
+      fail "Nonexistent Shard Name: #{shard}" if @klass.connection.shards[shard].nil?
       @current_shard = shard
       self
     end
 
     # Transaction Method send all queries to a specified shard.
     def transaction(options = {}, &block)
-      run_on_shard { @klass = klass.transaction(options, &block) }
+      run_on_shard { klass.transaction(options, &block) }
     end
 
     def connection
@@ -40,10 +44,13 @@ module Octopus
     end
 
     def method_missing(method, *args, &block)
-      result = run_on_shard { @klass.send(method, *args, &block) }
+      result = run_on_shard { @klass.__send__(method, *args, &block) }
       if result.respond_to?(:all)
-        @klass = result
-        return self
+        return ::Octopus::ScopeProxy.new(current_shard, result)
+      end
+
+      if result.respond_to?(:current_shard)
+        result.current_shard = current_shard
       end
 
       result

data/lib/octopus/shard_tracking.rb

@@ -20,7 +20,8 @@ module Octopus
         define_method with do |*args, &block|
           run_on_shard { send(without, *args, &block) }
         end
-        alias_method_chain name.to_sym, :octopus
+        alias_method without.to_sym, name.to_sym
+        alias_method name.to_sym, with.to_sym
       end
     end
 
@@ -31,8 +32,12 @@ module Octopus
         # 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
-        r = Octopus::RelationProxy.new(cs, r) if ActiveRecord::Relation === r and not Octopus::RelationProxy === r
-        r
+
+        if (ActiveRecord::Relation === r || ActiveRecord::QueryMethods::WhereChain === r) && !(Octopus::RelationProxy === r)
+          Octopus::RelationProxy.new(cs, r)
+        else
+          r
+        end
       else
         yield
       end

data/lib/octopus/slave_group.rb

@@ -3,11 +3,11 @@ module Octopus
     def initialize(slaves)
       slaves = HashWithIndifferentAccess.new(slaves)
       slaves_list = slaves.values
-      @load_balancer = Octopus::LoadBalancing::RoundRobin.new(slaves_list)
+      @load_balancer = Octopus.load_balancer.new(slaves_list)
     end
 
-    def next
-      @load_balancer.next
+    def next(options)
+      @load_balancer.next options
     end
   end
 end

data/lib/octopus/version.rb

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

data/lib/octopus.rb

@@ -11,12 +11,12 @@ module Octopus
   end
 
   def self.rails_env
-    @rails_env ||= self.rails? ? Rails.env.to_s : 'shards'
+    @rails_env ||= defined?(::Rails.env) ? Rails.env.to_s : 'shards'
   end
 
   def self.config
     @config ||= begin
-      file_name = Octopus.directory + '/config/shards.yml'
+      file_name = File.join(Octopus.directory, 'config/shards.yml').to_s
 
       if File.exist?(file_name) || File.symlink?(file_name)
         config ||= HashWithIndifferentAccess.new(YAML.load(ERB.new(File.read(file_name)).result))[Octopus.env]
@@ -28,13 +28,25 @@ module Octopus
     end
   end
 
+  def self.load_balancer=(balancer)
+    @load_balancer = balancer
+  end
+
+  def self.load_balancer
+    @load_balancer ||= Octopus::LoadBalancing::RoundRobin
+  end
+
+  def self.master_shard
+    ((config && config[:master_shard]) || :master).to_sym
+  end
+
   # Public: Whether or not Octopus is configured and should hook into the
   # current environment. Checks the environments config option for the Rails
   # environment by default.
   #
   # Returns a boolean
   def self.enabled?
-    if defined?(::Rails)
+    if defined?(::Rails.env)
       Octopus.environments.include?(Rails.env.to_s)
     else
       # TODO: This doens't feel right but !Octopus.config.blank? is breaking a
@@ -46,7 +58,7 @@ module Octopus
   # Returns the Rails.root_to_s when you are using rails
   # Running the current directory in a generic Ruby process
   def self.directory
-    @directory ||= defined?(Rails) ?  Rails.root.to_s : Dir.pwd
+    @directory ||= defined?(::Rails.root) ? Rails.root.to_s : Dir.pwd
   end
 
   # This is the default way to do Octopus Setup
@@ -78,26 +90,42 @@ module Octopus
     robust_environments.include? rails_env
   end
 
-  def self.rails3?
-    ActiveRecord::VERSION::MAJOR <= 3
+  def self.rails4?
+    ActiveRecord::VERSION::MAJOR == 4
   end
 
-  def self.rails4?
-    ActiveRecord::VERSION::MAJOR >= 4
+  def self.rails42?
+    rails4? && ActiveRecord::VERSION::MINOR == 2
   end
 
-  def self.rails41?
-    rails4? && ActiveRecord::VERSION::MINOR >= 1
+  def self.rails50?
+    ActiveRecord::VERSION::MAJOR == 5 && ActiveRecord::VERSION::MINOR == 0
+  end
+  
+  def self.atleast_rails50?
+    ActiveRecord::VERSION::MAJOR >= 5
+  end
+
+  def self.rails51?
+    ActiveRecord::VERSION::MAJOR == 5 && ActiveRecord::VERSION::MINOR == 1
+  end
+
+  def self.rails52?
+    ActiveRecord::VERSION::MAJOR == 5 && ActiveRecord::VERSION::MINOR == 2
+  end
+
+  def self.atleast_rails51?
+    ActiveRecord::VERSION::MAJOR > 5 || (ActiveRecord::VERSION::MAJOR == 5 && ActiveRecord::VERSION::MINOR >= 1)
   end
 
-  def self.rails?
-    defined?(Rails)
+  def self.atleast_rails52?
+    ActiveRecord::VERSION::MAJOR > 5 || (ActiveRecord::VERSION::MAJOR == 5 && ActiveRecord::VERSION::MINOR > 1)
   end
 
   attr_writer :logger
 
   def self.logger
-    if defined?(Rails)
+    if defined?(Rails.logger)
       @logger ||= Rails.logger
     else
       @logger ||= Logger.new($stderr)
@@ -119,32 +147,57 @@ module Octopus
     end
   end
 
+  def self.using_group(group, &block)
+    conn = ActiveRecord::Base.connection
+
+    if conn.is_a?(Octopus::Proxy)
+      conn.send_queries_to_group(group, &block)
+    else
+      yield
+    end
+  end
+
+  def self.using_all(&block)
+    conn = ActiveRecord::Base.connection
+
+    if conn.is_a?(Octopus::Proxy)
+      conn.send_queries_to_all_shards(&block)
+    else
+      yield
+    end
+  end
+
   def self.fully_replicated(&_block)
-    old_fully_replicated = Thread.current['octopus.fully_replicated']
-    Thread.current['octopus.fully_replicated'] = true
+    old_fully_replicated = Thread.current[Octopus::ProxyConfig::FULLY_REPLICATED_KEY]
+    Thread.current[Octopus::ProxyConfig::FULLY_REPLICATED_KEY] = true
     yield
   ensure
-    Thread.current['octopus.fully_replicated'] = old_fully_replicated
+    Thread.current[Octopus::ProxyConfig::FULLY_REPLICATED_KEY] = old_fully_replicated
   end
 end
 
+require 'octopus/exception'
+
 require 'octopus/shard_tracking'
 require 'octopus/shard_tracking/attribute'
 require 'octopus/shard_tracking/dynamic'
 
 require 'octopus/model'
+require 'octopus/result_patch'
 require 'octopus/migration'
 require 'octopus/association'
 require 'octopus/collection_association'
-require 'octopus/has_and_belongs_to_many_association' unless Octopus.rails41?
 require 'octopus/association_shard_tracking'
 require 'octopus/persistence'
 require 'octopus/log_subscriber'
 require 'octopus/abstract_adapter'
 require 'octopus/singular_association'
+require 'octopus/finder_methods'
+require 'octopus/query_cache_for_shards' unless Octopus.rails4?
 
-require 'octopus/railtie' if defined?(::Rails)
+require 'octopus/railtie' if defined?(::Rails::Railtie)
 
+require 'octopus/proxy_config'
 require 'octopus/proxy'
 require 'octopus/collection_proxy'
 require 'octopus/relation_proxy'

data/spec/config/shards.yml

@@ -86,6 +86,18 @@ production_replicated:
       <<: *mysql
 
 
+production_fully_replicated:
+  replicated: true
+  fully_replicated: true
+
+  shards:
+    slave1:
+      database: octopus_shard_2
+      <<: *mysql
+    slave2:
+      database: octopus_shard_3
+      <<: *mysql
+
 replicated_with_one_slave:
   replicated: true
   shards:

data/spec/migrations/10_create_users_using_replication.rb

@@ -1,4 +1,4 @@
-class CreateUsersUsingReplication < ActiveRecord::Migration
+class CreateUsersUsingReplication < BaseOctopusMigrationClass
   def self.up
     Cat.create!(:name => 'Replication')
   end

data/spec/migrations/11_add_field_in_all_slaves.rb

@@ -1,4 +1,4 @@
-class AddFieldInAllSlaves < ActiveRecord::Migration
+class AddFieldInAllSlaves < BaseOctopusMigrationClass
   using(:slave1, :slave2, :slave3, :slave4)
 
   def self.up

data/spec/migrations/12_create_users_using_block.rb

@@ -1,4 +1,4 @@
-class CreateUsersUsingBlock < ActiveRecord::Migration
+class CreateUsersUsingBlock < BaseOctopusMigrationClass
   def self.up
     Octopus.using(:brazil) do
       User.create!(:name => 'UsingBlock1')

data/spec/migrations/13_create_users_using_block_and_using.rb

@@ -1,4 +1,4 @@
-class CreateUsersUsingBlockAndUsing < ActiveRecord::Migration
+class CreateUsersUsingBlockAndUsing < BaseOctopusMigrationClass
   using(:brazil)
 
   def self.up

data/spec/migrations/14_create_users_on_shards_of_a_group_with_versions.rb

@@ -1,4 +1,4 @@
-class CreateUsersOnShardsOfAGroupWithVersions < ActiveRecord::Migration
+class CreateUsersOnShardsOfAGroupWithVersions < BaseOctopusMigrationClass
   using_group(:country_shards)
 
   def self.up

data/spec/migrations/15_create_user_on_shards_of_default_group_with_versions.rb

@@ -1,4 +1,4 @@
-class CreateUserOnShardsOfDefaultGroupWithVersions < ActiveRecord::Migration
+class CreateUserOnShardsOfDefaultGroupWithVersions < BaseOctopusMigrationClass
   def self.up
     User.create!(:name => 'Default Group')
   end

data/spec/migrations/1_create_users_on_master.rb

@@ -1,4 +1,4 @@
-class CreateUsersOnMaster < ActiveRecord::Migration
+class CreateUsersOnMaster < BaseOctopusMigrationClass
   def self.up
     User.create!(:name => 'Master')
   end

data/spec/migrations/2_create_users_on_canada.rb

@@ -1,4 +1,4 @@
-class CreateUsersOnCanada < ActiveRecord::Migration
+class CreateUsersOnCanada < BaseOctopusMigrationClass
   using(:canada)
 
   def self.up