active_record_shards 3.11.0 → 3.19.0

data/lib/active_record_shards/connection_switcher.rb

@@ -1,13 +1,19 @@
 # frozen_string_literal: true
+
 require 'active_record_shards/shard_support'
 
 module ActiveRecordShards
   module ConnectionSwitcher
-    SHARD_NAMES_CONFIG_KEY = 'shard_names'.freeze
+    SHARD_NAMES_CONFIG_KEY = 'shard_names'
 
     def self.extended(base)
-      base.singleton_class.send(:alias_method, :columns_without_default_shard, :columns)
-      base.singleton_class.send(:alias_method, :columns, :columns_with_default_shard)
+      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, :table_exists_without_default_shard?, :table_exists?)
       base.singleton_class.send(:alias_method, :table_exists?, :table_exists_with_default_shard?)
@@ -18,6 +24,10 @@ module ActiveRecordShards
       switch_connection(shard: new_default_shard)
     end
 
+    def on_primary_db(&block)
+      on_shard(nil, &block)
+    end
+
     def on_shard(shard)
       old_options = current_shard_selection.options
       switch_connection(shard: shard) if supports_sharding?
@@ -26,9 +36,9 @@ module ActiveRecordShards
       switch_connection(old_options)
     end
 
-    def on_first_shard
+    def on_first_shard(&block)
       shard_name = shard_names.first
-      on_shard(shard_name) { yield }
+      on_shard(shard_name, &block)
     end
 
     def shards
@@ -49,70 +59,79 @@ module ActiveRecordShards
       switch_connection(old_options)
     end
 
-    def on_slave_if(condition, &block)
-      condition ? on_slave(&block) : yield
+    def on_replica_if(condition, &block)
+      condition ? on_replica(&block) : yield
     end
+    alias_method :on_slave_if, :on_replica_if
 
-    def on_slave_unless(condition, &block)
-      on_slave_if(!condition, &block)
+    def on_replica_unless(condition, &block)
+      on_replica_if(!condition, &block)
     end
+    alias_method :on_slave_unless, :on_replica_unless
 
-    def on_master_if(condition, &block)
-      condition ? on_master(&block) : yield
+    def on_primary_if(condition, &block)
+      condition ? on_primary(&block) : yield
     end
+    alias_method :on_master_if, :on_primary_if
 
-    def on_master_unless(condition, &block)
-      on_master_if(!condition, &block)
+    def on_primary_unless(condition, &block)
+      on_primary_if(!condition, &block)
     end
+    alias_method :on_master_unless, :on_primary_unless
 
-    def on_master_or_slave(which, &block)
+    def on_primary_or_replica(which, &block)
       if block_given?
         on_cx_switch_block(which, &block)
       else
-        MasterSlaveProxy.new(self, which)
+        PrimaryReplicaProxy.new(self, which)
       end
     end
+    alias_method :on_master_or_slave, :on_primary_or_replica
 
-    # 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
+    # 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
     #     Account.first
     #   end
-    # the first account will be found on the slave DB
+    # the first account will be found on the replica DB
     #
     # For one-liners you can simply do
-    #   Account.on_slave.first
-    def on_slave(&block)
-      on_master_or_slave(:slave, &block)
+    #   Account.on_replica.first
+    def on_replica(&block)
+      on_primary_or_replica(:replica, &block)
     end
+    alias_method :on_slave, :on_replica
 
-    def on_master(&block)
-      on_master_or_slave(:master, &block)
+    def on_primary(&block)
+      on_primary_or_replica(:primary, &block)
     end
+    alias_method :on_master, :on_primary
 
     # just to ease the transition from replica to active_record_shards
-    alias_method :with_slave, :on_slave
-    alias_method :with_slave_if, :on_slave_if
-    alias_method :with_slave_unless, :on_slave_unless
+    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_slave ||= 0
-      @disallow_slave += 1 if which == :master
+      @disallow_replica ||= 0
+      @disallow_replica += 1 if [:primary, :master].include?(which)
 
-      switch_to_slave = force || @disallow_slave.zero?
+      ActiveRecordShards::Deprecation.warn('the `:master` option should be replaced with `:primary`!') if which == :master
+
+      switch_to_replica = force || @disallow_replica.zero?
       old_options = current_shard_selection.options
 
-      switch_connection(slave: switch_to_slave)
+      switch_connection(replica: switch_to_replica)
 
       # 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_slave || construct_ro_scope == false
+      if self == ActiveRecord::Base || !switch_to_replica || construct_ro_scope == false
         yield
       else
         readonly.scoping(&block)
       end
     ensure
-      @disallow_slave -= 1 if which == :master
+      @disallow_replica -= 1 if [:primary, :master].include?(which)
       switch_connection(old_options) if old_options
     end
 
@@ -120,9 +139,10 @@ module ActiveRecordShards
       shard_names.any?
     end
 
-    def on_slave?
-      current_shard_selection.on_slave?
+    def on_replica?
+      current_shard_selection.on_replica?
     end
+    alias_method :on_slave?, :on_replica?
 
     def current_shard_selection
       Thread.current[:shard_selection] ||= ShardSelection.new
@@ -134,8 +154,12 @@ module ActiveRecordShards
 
     def shard_names
       unless config = configurations[shard_env]
-        raise "Did not find #{shard_env} in configurations, did you forget to add it to your database.yml ? (configurations: #{configurations.inspect})"
+        raise "Did not find #{shard_env} in configurations, did you forget to add it to your database config? (configurations: #{configurations.keys.inspect})"
+      end
+      unless config.fetch(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
 
@@ -143,11 +167,15 @@ module ActiveRecordShards
 
     def switch_connection(options)
       if options.any?
-        if options.key?(:slave)
-          current_shard_selection.on_slave = options[:slave]
+        if options.key?(:replica)
+          current_shard_selection.on_replica = options[:replica]
         end
 
         if options.key?(:shard)
+          unless configurations[shard_env]
+            raise "Did not find #{shard_env} in configurations, did you forget to add it to your database config? (configurations: #{configurations.keys.inspect})"
+          end
+
           current_shard_selection.shard = options[:shard]
         end
 
@@ -159,51 +187,51 @@ module ActiveRecordShards
       ActiveRecordShards.rails_env
     end
 
-    if ActiveRecord::VERSION::MAJOR >= 4
-      def with_default_shard
-        if is_sharded? && current_shard_id.nil? && table_name != ActiveRecord::SchemaMigration.table_name
-          on_first_shard { yield }
-        else
-          yield
-        end
-      end
-    else
-      def with_default_shard
-        if is_sharded? && current_shard_id.nil? && table_name != ActiveRecord::Migrator.schema_migrations_table_name
-          on_first_shard { yield }
-        else
-          yield
-        end
+    # Make these few schema related methods available before having switched to
+    # a shard.
+    def with_default_shard(&block)
+      if is_sharded? && current_shard_id.nil? && table_name != ActiveRecord::SchemaMigration.table_name
+        on_first_shard(&block)
+      else
+        yield
       end
     end
 
-    def columns_with_default_shard
-      with_default_shard { columns_without_default_shard }
+    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
     end
 
     def table_exists_with_default_shard?
       with_default_shard { table_exists_without_default_shard? }
     end
 
-    class MasterSlaveProxy
+    class PrimaryReplicaProxy
       def initialize(target, which)
         @target = target
         @which = which
       end
 
-      def method_missing(method, *args, &block) # rubocop:disable Style/MethodMissing
-        @target.on_master_or_slave(@which) { @target.send(method, *args, &block) }
+      def method_missing(method, *args, &block) # rubocop:disable Style/MethodMissingSuper, Style/MissingRespondToMissing
+        @target.on_primary_or_replica(@which) { @target.send(method, *args, &block) }
       end
     end
+
+    MasterSlaveProxy = PrimaryReplicaProxy
   end
 end
 
 case "#{ActiveRecord::VERSION::MAJOR}.#{ActiveRecord::VERSION::MINOR}"
-when '3.2', '4.2'
-  require 'active_record_shards/connection_switcher-4-0'
+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'
+when '5.1', '5.2', '6.0'
   require 'active_record_shards/connection_switcher-5-1'
 else
   raise "ActiveRecordShards is not compatible with #{ActiveRecord::VERSION::STRING}"

data/lib/active_record_shards/default_replica_patches.rb

@@ -0,0 +1,278 @@
+# frozen_string_literal: true
+
+module ActiveRecordShards
+  module DefaultReplicaPatches
+    def self.wrap_method_in_on_replica(class_method, base, method, force_on_replica: false)
+      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
+      # _ALWAYS_ on replica, or only for on `on_replica_by_default = true` models?
+      wrapper = force_on_replica ? 'force_on_replica' : 'on_replica_unless_tx'
+      base.class_eval <<-RUBY, __FILE__, __LINE__ + 1
+        #{class_method ? 'class << self' : ''}
+          def #{method}_with_default_replica#{punctuation}(*args, &block)
+            #{wrapper} do
+              #{method}_without_default_replica#{punctuation}(*args, &block)
+            end
+          end
+
+          alias_method :#{method}_without_default_replica#{punctuation}, :#{method}#{punctuation}
+          alias_method :#{method}#{punctuation}, :#{method}_with_default_replica#{punctuation}
+        #{class_method ? 'end' : ''}
+      RUBY
+    end
+
+    def self.wrap_method_in_on_slave(*args)
+      ActiveRecordShards::Deprecation.deprecation_warning(
+        :'self.wrap_method_in_on_slave',
+        :'self.wrap_method_in_on_replica'
+      )
+      wrap_method_in_on_replica(*args)
+    end
+
+    def transaction_with_replica_off(*args, &block)
+      if on_replica_by_default?
+        begin
+          old_val = Thread.current[:_active_record_shards_in_tx]
+          Thread.current[:_active_record_shards_in_tx] = true
+          transaction_without_replica_off(*args, &block)
+        ensure
+          Thread.current[:_active_record_shards_in_tx] = old_val
+        end
+      else
+        transaction_without_replica_off(*args, &block)
+      end
+    end
+    alias_method :transaction_with_slave_off, :transaction_with_replica_off
+
+    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 on_replica_unless_tx
+        self.class.on_replica_unless_tx { yield }
+      end
+    end
+
+    CLASS_REPLICA_METHODS = [
+      :calculate,
+      :count_by_sql,
+      :exists?,
+      :find,
+      :find_by,
+      :find_by_sql,
+      :find_every,
+      :find_one,
+      :find_some,
+      :get_primary_key
+    ].freeze
+
+    CLASS_FORCE_REPLICA_METHODS = [
+      :replace_bind_variable,
+      :replace_bind_variables,
+      :sanitize_sql_array,
+      :sanitize_sql_hash_for_assignment,
+      :table_exists?
+    ].freeze
+
+    CLASS_SLAVE_METHODS = CLASS_REPLICA_METHODS
+    CLASS_FORCE_SLAVE_METHODS = CLASS_FORCE_REPLICA_METHODS
+
+    def self.extended(base)
+      CLASS_REPLICA_METHODS.each { |m| ActiveRecordShards::DefaultReplicaPatches.wrap_method_in_on_replica(true, base, m) }
+      CLASS_FORCE_REPLICA_METHODS.each { |m| ActiveRecordShards::DefaultReplicaPatches.wrap_method_in_on_replica(true, base, m, force_on_replica: true) }
+
+      if ActiveRecord::VERSION::MAJOR >= 5
+        ActiveRecordShards::DefaultReplicaPatches.wrap_method_in_on_replica(true, base, :load_schema!, force_on_replica: true)
+      else
+        ActiveRecordShards::DefaultReplicaPatches.wrap_method_in_on_replica(true, base, :columns, force_on_replica: true)
+      end
+
+      ActiveRecordShards::DefaultReplicaPatches.wrap_method_in_on_replica(false, base, :reload)
+
+      base.class_eval do
+        include InstanceMethods
+
+        class << self
+          alias_method :transaction_without_replica_off, :transaction
+          alias_method :transaction, :transaction_with_replica_off
+        end
+      end
+    end
+
+    def on_replica_unless_tx(&block)
+      return yield if Thread.current[:_active_record_shards_in_migration]
+      return yield if Thread.current[:_active_record_shards_in_tx]
+
+      if on_replica_by_default?
+        on_replica(&block)
+      else
+        yield
+      end
+    end
+    alias_method :on_slave_unless_tx, :on_replica_unless_tx
+
+    def force_on_replica(&block)
+      return yield if Thread.current[:_active_record_shards_in_migration]
+
+      on_cx_switch_block(:replica, construct_ro_scope: false, force: true, &block)
+    end
+
+    module ActiveRelationPatches
+      def self.included(base)
+        [:calculate, :exists?, :pluck, :load].each do |m|
+          ActiveRecordShards::DefaultReplicaPatches.wrap_method_in_on_replica(false, base, m)
+        end
+
+        if ActiveRecord::VERSION::MAJOR == 4
+          # `where` and `having` clauses call `create_binds`, which will use the primary connection
+          ActiveRecordShards::DefaultReplicaPatches.wrap_method_in_on_replica(false, base, :create_binds, force_on_replica: true)
+        end
+
+        ActiveRecordShards::DefaultReplicaPatches.wrap_method_in_on_replica(false, base, :to_sql, force_on_replica: true)
+      end
+
+      def on_replica_unless_tx
+        @klass.on_replica_unless_tx { yield }
+      end
+    end
+
+    module Rails52RelationPatches
+      def connection
+        return super if Thread.current[:_active_record_shards_in_migration]
+        return super if Thread.current[:_active_record_shards_in_tx]
+
+        if @klass.on_replica_by_default?
+          @klass.on_replica.connection
+        else
+          super
+        end
+      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-replica-by-default
+    # down from the parent now.
+    module Rails41HasAndBelongsToManyBuilderExtension
+      def self.included(base)
+        base.class_eval do
+          alias_method :through_model_without_inherit_default_replica_from_lhs, :through_model
+          alias_method :through_model, :through_model_with_inherit_default_replica_from_lhs
+        end
+      end
+
+      def through_model_with_inherit_default_replica_from_lhs
+        model = through_model_without_inherit_default_replica_from_lhs
+        def model.on_replica_by_default?
+          left_reflection.klass.on_replica_by_default?
+        end
+
+        # also transfer the sharded-ness of the left table to the join model
+        model.not_sharded unless model.left_reflection.klass.is_sharded?
+        model
+      end
+    end
+
+    module AssociationsAssociationAssociationScopePatch
+      def association_scope
+        if klass
+          on_replica_unless_tx { super }
+        else
+          super
+        end
+      end
+
+      def on_replica_unless_tx
+        klass.on_replica_unless_tx { yield }
+      end
+    end
+
+    module AssociationsAssociationFindTargetPatch
+      def find_target
+        if klass
+          on_replica_unless_tx { super }
+        else
+          super
+        end
+      end
+
+      def on_replica_unless_tx
+        klass.on_replica_unless_tx { yield }
+      end
+    end
+
+    module AssociationsAssociationGetRecordsPatch
+      def get_records # rubocop:disable Naming/AccessorMethodName
+        if klass
+          on_replica_unless_tx { super }
+        else
+          super
+        end
+      end
+
+      def on_replica_unless_tx
+        klass.on_replica_unless_tx { yield }
+      end
+    end
+
+    module AssociationsPreloaderAssociationAssociatedRecordsByOwnerPatch
+      def associated_records_by_owner(preloader)
+        if klass
+          on_replica_unless_tx { super }
+        else
+          super
+        end
+      end
+
+      def on_replica_unless_tx
+        klass.on_replica_unless_tx { yield }
+      end
+    end
+
+    module AssociationsPreloaderAssociationLoadRecordsPatch
+      def load_records
+        if klass
+          on_replica_unless_tx { super }
+        else
+          super
+        end
+      end
+
+      def on_replica_unless_tx
+        klass.on_replica_unless_tx { yield }
+      end
+    end
+
+    module TypeCasterConnectionConnectionPatch
+      def connection
+        return super if Thread.current[:_active_record_shards_in_migration]
+        return super if Thread.current[:_active_record_shards_in_tx]
+
+        if @klass.on_replica_by_default?
+          @klass.on_replica.connection
+        else
+          super
+        end
+      end
+    end
+
+    module SchemaDefinePatch
+      def define(info, &block)
+        old_val = Thread.current[:_active_record_shards_in_migration]
+        Thread.current[:_active_record_shards_in_migration] = true
+        super
+      ensure
+        Thread.current[:_active_record_shards_in_migration] = old_val
+      end
+    end
+  end
+end