active_record_shards 3.22.0 → 4.0.0.beta1

data/lib/active_record_shards/connection_switcher.rb

@@ -1,34 +1,16 @@
 # frozen_string_literal: true
-
 require 'active_record_shards/shard_support'
 
 module ActiveRecordShards
   module ConnectionSwitcher
-    SHARD_NAMES_CONFIG_KEY = 'shard_names'
+    SHARD_NAMES_CONFIG_KEY = 'shard_names'.freeze
 
     def self.extended(base)
-      if ActiveRecord::VERSION::MAJOR >= 5
-        base.singleton_class.send(:alias_method, :load_schema_without_default_shard!, :load_schema!)
-        base.singleton_class.send(:alias_method, :load_schema!, :load_schema_with_default_shard!)
-      else
-        base.singleton_class.send(:alias_method, :columns_without_default_shard, :columns)
-        base.singleton_class.send(:alias_method, :columns, :columns_with_default_shard)
-      end
+      base.singleton_class.send(:alias_method, :load_schema_without_default_shard!, :load_schema!)
+      base.singleton_class.send(:alias_method, :load_schema!, :load_schema_with_default_shard!)
 
       base.singleton_class.send(:alias_method, :table_exists_without_default_shard?, :table_exists?)
       base.singleton_class.send(:alias_method, :table_exists?, :table_exists_with_default_shard?)
-
-      base.singleton_class.send(:alias_method, :reset_primary_key_without_default_shard, :reset_primary_key)
-      base.singleton_class.send(:alias_method, :reset_primary_key, :reset_primary_key_with_default_shard)
-    end
-
-    def default_shard=(new_default_shard)
-      ActiveRecordShards::ShardSelection.default_shard = new_default_shard
-      switch_connection(shard: new_default_shard)
-    end
-
-    def on_primary_db(&block)
-      on_shard(nil, &block)
     end
 
     def on_shard(shard)
@@ -39,9 +21,9 @@ module ActiveRecordShards
       switch_connection(old_options)
     end
 
-    def on_first_shard(&block)
+    def on_first_shard
       shard_name = shard_names.first
-      on_shard(shard_name, &block)
+      on_shard(shard_name) { yield }
     end
 
     def shards
@@ -62,93 +44,88 @@ module ActiveRecordShards
       switch_connection(old_options)
     end
 
-    def on_replica_if(condition, &block)
-      condition ? on_replica(&block) : yield
+    def on_slave_if(condition, &block)
+      condition ? on_slave(&block) : yield
     end
-    alias_method :on_slave_if, :on_replica_if
 
-    def on_replica_unless(condition, &block)
-      on_replica_if(!condition, &block)
+    def on_slave_unless(condition, &block)
+      on_slave_if(!condition, &block)
     end
-    alias_method :on_slave_unless, :on_replica_unless
 
-    def on_primary_if(condition, &block)
-      condition ? on_primary(&block) : yield
+    def on_master_if(condition, &block)
+      condition ? on_master(&block) : yield
     end
-    alias_method :on_master_if, :on_primary_if
 
-    def on_primary_unless(condition, &block)
-      on_primary_if(!condition, &block)
+    def on_master_unless(condition, &block)
+      on_master_if(!condition, &block)
     end
-    alias_method :on_master_unless, :on_primary_unless
 
-    def on_primary_or_replica(which, &block)
+    def on_master_or_slave(which, &block)
       if block_given?
         on_cx_switch_block(which, &block)
       else
-        PrimaryReplicaProxy.new(self, which)
+        MasterSlaveProxy.new(self, which)
       end
     end
-    alias_method :on_master_or_slave, :on_primary_or_replica
 
-    # Executes queries using the replica database. Fails over to primary if no replica is found.
-    # if you want to execute a block of code on the replica you can go:
-    #   Account.on_replica do
+    # Executes queries using the slave database. Fails over to master if no slave is found.
+    # if you want to execute a block of code on the slave you can go:
+    #   Account.on_slave do
     #     Account.first
     #   end
-    # the first account will be found on the replica DB
+    # the first account will be found on the slave DB
     #
     # For one-liners you can simply do
-    #   Account.on_replica.first
-    def on_replica(&block)
-      on_primary_or_replica(:replica, &block)
+    #   Account.on_slave.first
+    def on_slave(&block)
+      on_master_or_slave(:slave, &block)
     end
-    alias_method :on_slave, :on_replica
 
-    def on_primary(&block)
-      on_primary_or_replica(:primary, &block)
+    def on_master(&block)
+      on_master_or_slave(:master, &block)
     end
-    alias_method :on_master, :on_primary
-
-    # just to ease the transition from replica to active_record_shards
-    alias_method :with_slave, :on_replica
-    alias_method :with_slave_if, :on_replica_if
-    alias_method :with_slave_unless, :on_replica_unless
 
     def on_cx_switch_block(which, force: false, construct_ro_scope: nil, &block)
-      @disallow_replica ||= 0
-      @disallow_replica += 1 if [:primary, :master].include?(which)
+      @disallow_slave ||= 0
+      @disallow_slave += 1 if which == :master
 
-      ActiveRecordShards::Deprecation.warn('the `:master` option should be replaced with `:primary`!') if which == :master
-
-      switch_to_replica = force || @disallow_replica.zero?
+      switch_to_slave = force || @disallow_slave.zero?
       old_options = current_shard_selection.options
 
-      switch_connection(replica: switch_to_replica)
+      switch_connection(slave: switch_to_slave)
 
       # we avoid_readonly_scope to prevent some stack overflow problems, like when
       # .columns calls .with_scope which calls .columns and onward, endlessly.
-      if self == ActiveRecord::Base || !switch_to_replica || construct_ro_scope == false || ActiveRecordShards.disable_replica_readonly_records == true
+      if self == ActiveRecord::Base || !switch_to_slave || construct_ro_scope == false
         yield
       else
         readonly.scoping(&block)
       end
     ensure
-      @disallow_replica -= 1 if [:primary, :master].include?(which)
+      @disallow_slave -= 1 if which == :master
       switch_connection(old_options) if old_options
     end
 
+    def connection_specification_name
+      name = current_shard_selection.resolve_connection_name(sharded: is_sharded?, configurations: configurations)
+
+      unless configurations[name] || name == "primary"
+        raise ActiveRecord::AdapterNotSpecified, "No database defined by #{name} in your database config. (configurations: #{configurations.inspect})"
+      end
+
+      name
+    end
+
     def supports_sharding?
       shard_names.any?
     end
 
-    def on_replica?
-      current_shard_selection.on_replica?
+    def on_slave?
+      current_shard_selection.on_slave?
     end
-    alias_method :on_slave?, :on_replica?
 
     def current_shard_selection
-      Thread.current[:shard_selection] ||= ShardSelection.new
+      Thread.current[:shard_selection] ||= ShardSelection.new(shard_names.first)
     end
 
     def current_shard_id
@@ -156,99 +133,78 @@ module ActiveRecordShards
     end
 
     def shard_names
-      unless config_for_env.fetch(SHARD_NAMES_CONFIG_KEY, []).all? { |shard_name| shard_name.is_a?(Integer) }
-        raise "All shard names must be integers: #{config_for_env[SHARD_NAMES_CONFIG_KEY].inspect}."
+      unless config = configurations[shard_env]
+        raise "Did not find #{shard_env} in configurations, did you forget to add it to your database config? (configurations: #{configurations.inspect})"
       end
-
-      config_for_env[SHARD_NAMES_CONFIG_KEY] || []
-    end
-
-    def reset_primary_key_with_default_shard
-      with_default_shard { reset_primary_key_without_default_shard }
+      unless config[SHARD_NAMES_CONFIG_KEY]
+        raise "No shards configured for #{shard_env}"
+      end
+      unless config[SHARD_NAMES_CONFIG_KEY].all? { |shard_name| shard_name.is_a?(Integer) }
+        raise "All shard names must be integers: #{config[SHARD_NAMES_CONFIG_KEY].inspect}."
+      end
+      config[SHARD_NAMES_CONFIG_KEY]
     end
 
     private
 
-    def config_for_env
-      @_ars_config_for_env ||= {}
-      @_ars_config_for_env[shard_env] ||= begin
-        unless config = configurations[shard_env]
-          raise "Did not find #{shard_env} in configurations, did you forget to add it to your database config? (configurations: #{configurations.to_h.keys.inspect})"
-        end
-
-        config
-      end
-    end
-    alias_method :check_config_for_env, :config_for_env
-
     def switch_connection(options)
       if options.any?
-        if options.key?(:replica)
-          current_shard_selection.on_replica = options[:replica]
-        end
-
         if options.key?(:shard)
-          check_config_for_env
-
+          unless config = configurations[shard_env]
+            raise "Did not find #{shard_env} in configurations, did you forget to add it to your database config? (configurations: #{configurations.inspect})"
+          end
+          unless config['shard_names'].include?(options[:shard])
+            raise "Did not find shard #{options[:shard]} in configurations"
+          end
           current_shard_selection.shard = options[:shard]
         end
 
+        if options.key?(:slave)
+          current_shard_selection.on_slave = options[:slave]
+        end
+
         ensure_shard_connection
       end
     end
 
     def shard_env
-      ActiveRecordShards.app_env
+      ActiveRecordShards.rails_env
     end
 
-    # Make these few schema related methods available before having switched to
-    # a shard.
-    def with_default_shard(&block)
+    def with_default_shard
       if is_sharded? && current_shard_id.nil? && table_name != ActiveRecord::SchemaMigration.table_name
-        on_first_shard(&block)
+        on_first_shard { yield }
       else
         yield
       end
     end
 
-    if ActiveRecord::VERSION::MAJOR >= 5
-      def load_schema_with_default_shard!
-        with_default_shard { load_schema_without_default_shard! }
-      end
-    else
-      def columns_with_default_shard
-        with_default_shard { columns_without_default_shard }
-      end
+    def load_schema_with_default_shard!
+      with_default_shard { load_schema_without_default_shard! }
     end
 
     def table_exists_with_default_shard?
       with_default_shard { table_exists_without_default_shard? }
     end
 
-    class PrimaryReplicaProxy
+    class MasterSlaveProxy
       def initialize(target, which)
         @target = target
         @which = which
       end
 
-      def method_missing(method, *args, &block) # rubocop:disable Style/MethodMissingSuper, Style/MissingRespondToMissing
-        @target.on_primary_or_replica(@which) { @target.send(method, *args, &block) }
+      def method_missing(method, *args, &block) # rubocop:disable Style/MethodMissing
+        @target.on_master_or_slave(@which) { @target.send(method, *args, &block) }
       end
     end
-
-    MasterSlaveProxy = PrimaryReplicaProxy
   end
 end
 
 case "#{ActiveRecord::VERSION::MAJOR}.#{ActiveRecord::VERSION::MINOR}"
-when '4.2'
-  require 'active_record_shards/connection_switcher-4-2'
 when '5.0'
   require 'active_record_shards/connection_switcher-5-0'
-when '5.1', '5.2'
+when '5.1'
   require 'active_record_shards/connection_switcher-5-1'
-when '6.0'
-  require 'active_record_shards/connection_switcher-6-0'
 else
   raise "ActiveRecordShards is not compatible with #{ActiveRecord::VERSION::STRING}"
 end

data/lib/active_record_shards/default_slave_patches.rb

@@ -1,5 +1,136 @@
-ActiveRecordShards::Deprecation.warn('`DefaultSlavePatches` is deprecated, please use `DefaultReplicaPatches`.')
-
+# frozen_string_literal: true
 module ActiveRecordShards
-  DefaultSlavePatches = DefaultReplicaPatches
+  module DefaultSlavePatches
+    def self.wrap_method_in_on_slave(class_method, base, method)
+      base_methods =
+        if class_method
+          base.methods + base.private_methods
+        else
+          base.instance_methods + base.private_instance_methods
+        end
+
+      return unless base_methods.include?(method)
+      _, method, punctuation = method.to_s.match(/^(.*?)([\?\!]?)$/).to_a
+      base.class_eval <<-RUBY, __FILE__, __LINE__ + 1
+        #{class_method ? 'class << self' : ''}
+          def #{method}_with_default_slave#{punctuation}(*args, &block)
+            on_slave_unless_tx do
+              #{method}_without_default_slave#{punctuation}(*args, &block)
+            end
+          end
+
+          alias_method :#{method}_without_default_slave#{punctuation}, :#{method}#{punctuation}
+          alias_method :#{method}#{punctuation}, :#{method}_with_default_slave#{punctuation}
+        #{class_method ? 'end' : ''}
+      RUBY
+    end
+
+    def columns_with_force_slave(*args, &block)
+      on_cx_switch_block(:slave, construct_ro_scope: false, force: true) do
+        columns_without_force_slave(*args, &block)
+      end
+    end
+
+    def table_exists_with_force_slave?(*args, &block)
+      on_cx_switch_block(:slave, construct_ro_scope: false, force: true) do
+        table_exists_without_force_slave?(*args, &block)
+      end
+    end
+
+    def transaction_with_slave_off(*args, &block)
+      if on_slave_by_default?
+        begin
+          old_val = Thread.current[:_active_record_shards_slave_off]
+          Thread.current[:_active_record_shards_slave_off] = true
+          transaction_without_slave_off(*args, &block)
+        ensure
+          Thread.current[:_active_record_shards_slave_off] = old_val
+        end
+      else
+        transaction_without_slave_off(*args, &block)
+      end
+    end
+
+    module InstanceMethods
+      # fix ActiveRecord to do the right thing, and use our aliased quote_value
+      def quote_value(*args, &block)
+        self.class.quote_value(*args, &block)
+      end
+
+      def reload_with_slave_off(*args, &block)
+        self.class.on_master { reload_without_slave_off(*args, &block) }
+      end
+    end
+
+    CLASS_SLAVE_METHODS = [:find_by_sql, :count_by_sql, :calculate, :find_one, :find_some, :find_every, :exists?].freeze
+
+    def self.extended(base)
+      CLASS_SLAVE_METHODS.each { |m| ActiveRecordShards::DefaultSlavePatches.wrap_method_in_on_slave(true, base, m) }
+
+      base.class_eval do
+        include InstanceMethods
+
+        alias_method :reload_without_slave_off, :reload
+        alias_method :reload, :reload_with_slave_off
+
+        class << self
+          alias_method :columns_without_force_slave, :columns
+          alias_method :columns, :columns_with_force_slave
+
+          alias_method :table_exists_without_force_slave?, :table_exists?
+          alias_method :table_exists?, :table_exists_with_force_slave?
+
+          alias_method :transaction_without_slave_off, :transaction
+          alias_method :transaction, :transaction_with_slave_off
+        end
+      end
+      if ActiveRecord::Associations.const_defined?(:HasAndBelongsToManyAssociation)
+        ActiveRecordShards::DefaultSlavePatches.wrap_method_in_on_slave(false, ActiveRecord::Associations::HasAndBelongsToManyAssociation, :construct_sql)
+        ActiveRecordShards::DefaultSlavePatches.wrap_method_in_on_slave(false, ActiveRecord::Associations::HasAndBelongsToManyAssociation, :construct_find_options!)
+      end
+    end
+
+    def on_slave_unless_tx
+      if on_slave_by_default? && !Thread.current[:_active_record_shards_slave_off]
+        on_slave { yield }
+      else
+        yield
+      end
+    end
+
+    module ActiveRelationPatches
+      def self.included(base)
+        [:calculate, :exists?, :pluck, :find_with_associations].each do |m|
+          ActiveRecordShards::DefaultSlavePatches.wrap_method_in_on_slave(false, base, m)
+        end
+      end
+
+      def on_slave_unless_tx
+        @klass.on_slave_unless_tx { yield }
+      end
+    end
+
+    # in rails 4.1+, they create a join class that's used to pull in records for HABTM.
+    # this simplifies the hell out of our existence, because all we have to do is inerit on-slave-by-default
+    # down from the parent now.
+    module HasAndBelongsToManyBuilderExtension
+      def self.included(base)
+        base.class_eval do
+          alias_method :through_model_without_inherit_default_slave_from_lhs, :through_model
+          alias_method :through_model, :through_model_with_inherit_default_slave_from_lhs
+        end
+      end
+
+      def through_model_with_inherit_default_slave_from_lhs
+        model = through_model_without_inherit_default_slave_from_lhs
+        def model.on_slave_by_default?
+          left_reflection.klass.on_slave_by_default?
+        end
+
+        model.extend(ActiveRecordShards::Ext::ShardedModel) if model.left_reflection.klass.is_sharded?
+
+        model
+      end
+    end
+  end
 end

data/lib/active_record_shards/model.rb

@@ -3,59 +3,44 @@
 module ActiveRecordShards
   module Model
     def not_sharded
-      if self != ActiveRecord::Base && self != base_class
-        raise "You should only call not_sharded on direct descendants of ActiveRecord::Base"
-      end
-
-      self.sharded = false
+      ActiveSupport::Deprecation.warn("Calling not_sharded is deprecated. "\
+                                      "Please ensure to still read from the "\
+                                      "account db slave after removing the "\
+                                      "call.")
     end
 
     def is_sharded? # rubocop:disable Naming/PredicateName
-      if self == ActiveRecord::Base
-        sharded != false && supports_sharding?
-      elsif self == base_class
-        if sharded.nil?
-          ActiveRecord::Base.is_sharded?
-        else
-          sharded != false
-        end
-      else
-        base_class.is_sharded?
-      end
+      false
     end
 
-    def on_replica_by_default?
+    def on_slave_by_default?
       if self == ActiveRecord::Base
         false
       else
         base = base_class
-        if base.instance_variable_defined?(:@on_replica_by_default)
-          base.instance_variable_get(:@on_replica_by_default)
+        if base.instance_variable_defined?(:@on_slave_by_default)
+          base.instance_variable_get(:@on_slave_by_default)
         end
       end
     end
-    alias_method :on_slave_by_default?, :on_replica_by_default?
 
-    def on_replica_by_default=(value)
+    def on_slave_by_default=(value)
       if self == ActiveRecord::Base
-        raise ArgumentError, "Cannot set on_replica_by_default on ActiveRecord::Base"
+        raise ArgumentError, "Cannot set on_slave_by_default on ActiveRecord::Base"
       else
-        base_class.instance_variable_set(:@on_replica_by_default, value)
+        base_class.instance_variable_set(:@on_slave_by_default, value)
       end
     end
-    alias_method :on_slave_by_default=, :on_replica_by_default=
 
     module InstanceMethods
-      def initialize_shard_and_replica
-        @from_replica = !!self.class.current_shard_selection.options[:replica]
+      def initialize_shard_and_slave
+        @from_slave = !!self.class.current_shard_selection.options[:slave]
         @from_shard = self.class.current_shard_selection.options[:shard]
       end
-      alias_method :initialize_shard_and_slave, :initialize_shard_and_replica
 
-      def from_replica?
-        @from_replica
+      def from_slave?
+        @from_slave
       end
-      alias_method :from_slave?, :from_replica?
 
       def from_shard
         @from_shard
@@ -64,11 +49,7 @@ module ActiveRecordShards
 
     def self.extended(base)
       base.send(:include, InstanceMethods)
-      base.after_initialize :initialize_shard_and_replica
+      base.after_initialize :initialize_shard_and_slave
     end
-
-    private
-
-    attr_accessor :sharded
   end
 end

data/lib/active_record_shards/shard_selection.rb

@@ -1,91 +1,52 @@
 # frozen_string_literal: true
-
 module ActiveRecordShards
   class ShardSelection
-    NO_SHARD = :_no_shard
-    cattr_accessor :default_shard
-
-    def initialize
-      @on_replica = false
-      @shard = nil
+    def initialize(shard)
+      @on_slave = false
+      self.shard = shard
     end
 
-    if ActiveRecord::VERSION::MAJOR < 5
-
-      def shard(klass = nil)
-        if (@shard || self.class.default_shard) && (klass.nil? || klass.is_sharded?)
-          if @shard == NO_SHARD
-            nil
-          else
-            @shard || self.class.default_shard
-          end
-        end
-      end
-
-      def shard_name(klass = nil, try_replica = true)
-        the_shard = shard(klass)
-
-        @shard_names ||= {}
-        @shard_names[ActiveRecordShards.app_env] ||= {}
-        @shard_names[ActiveRecordShards.app_env][the_shard] ||= {}
-        @shard_names[ActiveRecordShards.app_env][the_shard][try_replica] ||= {}
-        @shard_names[ActiveRecordShards.app_env][the_shard][try_replica][@on_replica] ||= begin
-          s = ActiveRecordShards.app_env.dup
-          s << "_shard_#{the_shard}" if the_shard
-          s << "_replica"            if @on_replica && try_replica
-          s
-        end
-      end
+    def shard
+      raise "Missing shard information on connection" unless @shard
+      @shard
+    end
 
-    else
+    PRIMARY = "primary".freeze
+    def resolve_connection_name(sharded:, configurations:)
+      resolved_shard = sharded ? shard : nil
 
-      def shard
-        if @shard.nil? || @shard == NO_SHARD
-          nil
-        else
-          @shard || self.class.default_shard
-        end
+      if !resolved_shard && !@on_slave
+        return PRIMARY
       end
 
-      PRIMARY = "primary"
-      def resolve_connection_name(sharded:, configurations:)
-        resolved_shard = sharded ? shard : nil
-        env = ActiveRecordShards.app_env
+      @shard_names ||= {}
+      @shard_names[ActiveRecordShards.rails_env] ||= {}
+      @shard_names[ActiveRecordShards.rails_env][resolved_shard] ||= {}
+      @shard_names[ActiveRecordShards.rails_env][resolved_shard][@on_slave] ||= begin
+        s = ActiveRecordShards.rails_env.dup
+        s << "_shard_#{resolved_shard}" if resolved_shard
 
-        @connection_names ||= {}
-        @connection_names[env] ||= {}
-        @connection_names[env][resolved_shard] ||= {}
-        @connection_names[env][resolved_shard][@on_replica] ||= begin
-          name = env.dup
-          name << "_shard_#{resolved_shard}" if resolved_shard
-          if @on_replica && configurations["#{name}_replica"]
-            "#{name}_replica"
-          else
-            # ActiveRecord always names its default connection pool 'primary'
-            # while everything else is named by the configuration name
-            resolved_shard ? name : PRIMARY
-          end
+        if @on_slave && configurations["#{s}_slave"] # fall back to master connection
+          s << "_slave"
         end
+        s
       end
-
     end
 
     def shard=(new_shard)
-      @shard = (new_shard || NO_SHARD)
+      @shard = Integer(new_shard)
     end
 
-    def on_replica?
-      @on_replica
+    def on_slave?
+      @on_slave
     end
-    alias_method :on_slave?, :on_replica?
 
-    def on_replica=(new_replica)
-      @on_replica = (new_replica == true)
+    def on_slave=(new_slave)
+      @on_slave = (new_slave == true)
     end
-    alias_method :on_slave=, :on_replica=
 
     def options
-      { shard: @shard, replica: @on_replica }
+      { shard: @shard, slave: @on_slave }
     end
   end
 end

data/lib/active_record_shards/shard_support.rb

@@ -1,5 +1,4 @@
 # frozen_string_literal: true
-
 module ActiveRecordShards
   class ShardSupport
     class ShardEnumerator

data/lib/active_record_shards/sharded_model.rb

@@ -0,0 +1,15 @@
+module ActiveRecordShards
+  module Ext
+    module ShardedModel
+      def is_sharded? # rubocop:disable Naming/PredicateName
+        true
+      end
+    end
+  end
+
+  class ShardedModel < ActiveRecord::Base
+    self.abstract_class = true
+
+    extend ActiveRecordShards::Ext::ShardedModel
+  end
+end