active_record_shards 4.0.0.beta8 → 5.1.0

data/lib/active_record_shards/default_replica_patches.rb

@@ -0,0 +1,237 @@
+# 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
+          ruby2_keywords(:#{method}_with_default_replica#{punctuation}) if respond_to?(:ruby2_keywords, true)
+          alias_method :#{method}_without_default_replica#{punctuation}, :#{method}#{punctuation}
+          alias_method :#{method}#{punctuation}, :#{method}_with_default_replica#{punctuation}
+        #{class_method ? 'end' : ''}
+      RUBY
+    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
+    ruby2_keywords(:transaction_with_replica_off) if respond_to?(:ruby2_keywords, true)
+
+    module InstanceMethods
+      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
+
+    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) }
+
+      ActiveRecordShards::DefaultReplicaPatches.wrap_method_in_on_replica(true, base, :load_schema!, force_on_replica: true)
+      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
+
+    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
+
+        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 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

data/lib/active_record_shards/default_shard.rb

@@ -0,0 +1,27 @@
+# frozen_string_literal: true
+
+module ActiveRecordShards
+  module DefaultShard
+    def default_shard=(new_default_shard)
+      if ars_shard_type?(new_default_shard)
+        ActiveRecordShards::ShardSelection.ars_default_shard = new_default_shard
+        switch_connection(shard: new_default_shard)
+      else
+        super
+      end
+    end
+
+    private
+
+    def ars_shard_type?(shard)
+      return true if ActiveRecord.version < Gem::Version.new('6.1')
+      return true if shard.nil?
+      return true if shard == :_no_shard
+      return true if shard.is_a?(Integer)
+
+      false
+    end
+  end
+end
+
+ActiveRecord::Base.singleton_class.prepend(ActiveRecordShards::DefaultShard)

data/lib/active_record_shards/migration.rb

@@ -0,0 +1,124 @@
+# frozen_string_literal: true
+
+module ActiveRecord
+  class Migrator
+    def self.shards_migration_context
+      if ActiveRecord::VERSION::MAJOR >= 6
+        ActiveRecord::MigrationContext.new(ActiveRecord::Migrator.migrations_paths, ActiveRecord::SchemaMigration)
+      elsif ActiveRecord::VERSION::STRING >= '5.2.0'
+        ActiveRecord::MigrationContext.new(ActiveRecord::Migrator.migrations_paths)
+      else
+        self
+      end
+    end
+
+    def initialize_with_sharding(*args)
+      initialize_without_sharding(*args)
+
+      # Rails creates the internal tables on the unsharded DB. We make them
+      # manually on the sharded DBs.
+      ActiveRecord::Base.on_all_shards do
+        ActiveRecord::SchemaMigration.create_table
+        ActiveRecord::InternalMetadata.create_table
+      end
+    end
+    ruby2_keywords(:initialize_with_sharding) if respond_to?(:ruby2_keywords, true)
+    alias_method :initialize_without_sharding, :initialize
+    alias_method :initialize, :initialize_with_sharding
+
+    def run_with_sharding
+      ActiveRecord::Base.on_shard(nil) { run_without_sharding }
+      ActiveRecord::Base.on_all_shards { run_without_sharding }
+    end
+    alias_method :run_without_sharding, :run
+    alias_method :run, :run_with_sharding
+
+    def migrate_with_sharding
+      ActiveRecord::Base.on_shard(nil) { migrate_without_sharding }
+      ActiveRecord::Base.on_all_shards { migrate_without_sharding }
+    end
+    alias_method :migrate_without_sharding, :migrate
+    alias_method :migrate, :migrate_with_sharding
+
+    # don't allow Migrator class to cache versions
+    undef migrated
+    def migrated
+      self.class.shards_migration_context.get_all_versions
+    end
+
+    # list of pending migrations is any migrations that haven't run on all shards.
+    undef pending_migrations
+    def pending_migrations
+      pending, _missing = self.class.shard_status(migrations.map(&:version))
+      pending = pending.values.flatten
+      migrations.select { |m| pending.include?(m.version) }
+    end
+
+    # public
+    # list of pending and missing versions per shard
+    # [{1 => [1234567]}, {1 => [2345678]}]
+    def self.shard_status(versions)
+      pending = {}
+      missing = {}
+
+      collect = lambda do |shard|
+        migrated = shards_migration_context.get_all_versions
+
+        p = versions - migrated
+        pending[shard] = p if p.any?
+
+        m = migrated - versions
+        missing[shard] = m if m.any?
+      end
+
+      ActiveRecord::Base.on_shard(nil) { collect.call(nil) }
+      ActiveRecord::Base.on_all_shards { |shard| collect.call(shard) }
+
+      [pending, missing]
+    end
+  end
+end
+
+module ActiveRecordShards
+  module MigrationClassExtension
+    attr_accessor :migration_shard
+
+    def shard(arg = nil)
+      self.migration_shard = arg
+    end
+  end
+
+  module ActualMigrationExtension
+    def migrate_with_forced_shard(direction)
+      if migration_shard.blank?
+        raise "#{name}: Can't run migrations without a shard spec: this may be :all, :none,
+                 or a specific shard (for data-fixups).  please call shard(arg) in your migration."
+      end
+
+      shard = ActiveRecord::Base.current_shard_selection.shard
+
+      if shard.nil?
+        return if migration_shard != :none
+      else
+        return if migration_shard == :none
+        return if migration_shard != :all && migration_shard.to_s != shard.to_s
+      end
+
+      migrate_without_forced_shard(direction)
+    end
+
+    def migration_shard
+      self.class.migration_shard
+    end
+  end
+end
+
+ActiveRecord::Migration.class_eval do
+  extend ActiveRecordShards::MigrationClassExtension
+  include ActiveRecordShards::ActualMigrationExtension
+
+  alias_method :migrate_without_forced_shard, :migrate
+  alias_method :migrate, :migrate_with_forced_shard
+end
+
+ActiveRecord::MigrationProxy.delegate :migration_shard, to: :migration

data/lib/active_record_shards/model.rb

@@ -3,53 +3,68 @@
 module ActiveRecordShards
   module Model
     def not_sharded
-      message = "Calling not_sharded is deprecated. "\
-        "Please ensure to still read from the "\
-        "account db slave after removing the "\
-        "call."
-      if ActiveRecord::Base.logger
-        ActiveRecord::Base.logger.warn(message)
-      else
-        Kernel.warn(message)
+      if self != ActiveRecord::Base && self != base_class
+        raise "You should only call not_sharded on direct descendants of ActiveRecord::Base"
       end
+
+      self.sharded = false
     end
 
     def is_sharded? # rubocop:disable Naming/PredicateName
-      false
+      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
     end
 
-    def on_slave_by_default?
+    def on_replica_by_default?
       if self == ActiveRecord::Base
         false
       else
         base = base_class
-        if base.instance_variable_defined?(:@on_slave_by_default)
-          base.instance_variable_get(:@on_slave_by_default)
+        if base.instance_variable_defined?(:@on_replica_by_default)
+          base.instance_variable_get(:@on_replica_by_default)
         end
       end
     end
 
-    def on_slave_by_default=(value)
+    def on_replica_by_default=(value)
       if self == ActiveRecord::Base
-        raise ArgumentError, "Cannot set on_slave_by_default on ActiveRecord::Base"
+        raise ArgumentError, "Cannot set on_replica_by_default on ActiveRecord::Base"
       else
-        base_class.instance_variable_set(:@on_slave_by_default, value)
+        base_class.instance_variable_set(:@on_replica_by_default, value)
       end
     end
 
     module InstanceMethods
-      def initialize_slave
-        @from_slave = !!self.class.current_slave_selection
+      def initialize_shard_and_replica
+        @from_replica = !!self.class.current_shard_selection.options[:replica]
+        @from_shard = self.class.current_shard_selection.options[:shard]
+      end
+
+      def from_replica?
+        @from_replica
       end
 
-      def from_slave?
-        @from_slave
+      def from_shard
+        @from_shard
       end
     end
 
     def self.extended(base)
-      base.include(InstanceMethods)
-      base.after_initialize :initialize_slave
+      base.send(:include, InstanceMethods)
+      base.after_initialize :initialize_shard_and_replica
     end
+
+    private
+
+    attr_accessor :sharded
   end
 end

data/lib/active_record_shards/schema_dumper_extension.rb

@@ -0,0 +1,42 @@
+# frozen_string_literal: true
+
+module ActiveRecordShards
+  module SchemaDumperExtension
+    def dump(stream)
+      stream = super(stream)
+      original_connection = @connection
+
+      if ActiveRecord::Base.supports_sharding?
+        ActiveRecord::Base.on_first_shard do
+          @connection = ActiveRecord::Base.connection
+          shard_header(stream)
+          extensions(stream)
+          tables(stream)
+          shard_trailer(stream)
+        end
+      end
+
+      stream
+    ensure
+      @connection = original_connection
+    end
+
+    def shard_header(stream)
+      define_params = @version ? "version: #{@version}" : ""
+
+      stream.puts <<~HEADER
+
+
+        # This section generated by active_record_shards
+
+        ActiveRecord::Base.on_all_shards do
+        ActiveRecord::Schema.define(#{define_params}) do
+
+      HEADER
+    end
+
+    def shard_trailer(stream)
+      stream.puts "end\nend"
+    end
+  end
+end

data/lib/active_record_shards/shard_selection.rb

@@ -1,18 +1,58 @@
 # frozen_string_literal: true
+
 module ActiveRecordShards
   class ShardSelection
-    attr_reader :shard
+    NO_SHARD = :_no_shard
+    cattr_accessor :ars_default_shard
+
+    def initialize
+      @on_replica = false
+      @shard = nil
+    end
+
+    def shard
+      if @shard.nil? || @shard == NO_SHARD
+        nil
+      else
+        @shard || self.class.ars_default_shard
+      end
+    end
+
+    PRIMARY = "primary"
+    def resolve_connection_name(sharded:, configurations:)
+      resolved_shard = sharded ? shard : nil
+      env = ActiveRecordShards.app_env
+
+      @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
+      end
+    end
+
+    def shard=(new_shard)
+      @shard = (new_shard || NO_SHARD)
+    end
 
-    def initialize(shard)
-      @shard = Integer(shard)
+    def on_replica?
+      @on_replica
     end
 
-    def connection_config
-      { shard: shard }
+    def on_replica=(new_replica)
+      @on_replica = (new_replica == true)
     end
 
-    def on_shard?
-      true
+    def options
+      { shard: @shard, replica: @on_replica }
     end
   end
 end

data/lib/active_record_shards/shard_support.rb

@@ -1,11 +1,12 @@
 # frozen_string_literal: true
+
 module ActiveRecordShards
   class ShardSupport
     class ShardEnumerator
       include Enumerable
 
       def each(&block)
-        ActiveRecordShards::ShardedModel.on_all_shards(&block)
+        ActiveRecord::Base.on_all_shards(&block)
       end
     end
 
@@ -22,15 +23,14 @@ module ActiveRecordShards
 
       exception = nil
       enum.each do
-        begin
-          record = @scope.find(*find_args)
-          return record if record
-        rescue ActiveRecord::RecordNotFound => e
-          exception = e
-        end
+        record = @scope.find(*find_args)
+        return record if record
+      rescue ActiveRecord::RecordNotFound => e
+        exception = e
       end
       raise exception
     end
+    ruby2_keywords(:find) if respond_to?(:ruby2_keywords, true)
 
     def count
       enum.inject(0) { |accum, _shard| @scope.clone.count + accum }

data/lib/active_record_shards/sql_comments.rb

@@ -1,16 +1,23 @@
-# show which connection was picked to debug master/slave slowness when both servers are the same
+# show which connection was picked to debug primary/replica slowness when both servers are the same
 module ActiveRecordShards
   module SqlComments
     module Methods
       def execute(query, name = nil)
-        slave = ActiveRecord::Base.current_slave_selection
-        query += " /* #{slave ? 'slave' : 'master'} */"
+        shard = ActiveRecord::Base.current_shard_selection.shard
+        shard_text = shard ? "shard #{shard}" : 'unsharded'
+        replica = ActiveRecord::Base.current_shard_selection.on_replica?
+        replica_text = replica ? 'replica' : 'primary'
+        query = "/* #{shard_text} #{replica_text} */ " + query
         super(query, name)
       end
     end
 
     def self.enable
-      ActiveRecord::Base.connection.class.prepend(Methods)
+      ActiveRecord::Base.on_replica do
+        ActiveRecord::Base.on_shard(nil) do
+          ActiveRecord::Base.connection.class.prepend(Methods)
+        end
+      end
     end
   end
 end