Fix to propagate shard conditions to AssociationRelation too

backport from develop-3-0 branch

When calling `@user.tasks.find_by(name: 'xxx')`

before:

  SELECT `tasks`.* FROM `tasks` WHERE `tasks`.`name` = "xxx" # raises CannotSpecifyShardError

after:

  SELECT `tasks`.* FROM `tasks` WHERE `tasks`.`user_id` = 1 AND `tasks`.`name` = "xxx"

lib/active_record/turntable.rb

@@ -38,6 +38,7 @@ module ActiveRecord::Turntable
     autoload :Mixer
     autoload :PoolProxy
     autoload :Shard
+    autoload :ShardingCondition
     autoload :SeqShard
     autoload :Sequencer
   end

lib/active_record/turntable/active_record_ext.rb

@@ -31,7 +31,7 @@ module ActiveRecordExt
         include ConnectionHandlerExtension
       end
       ActiveRecord::Associations::Preloader::Association.send(:include, AssociationPreloader)
-      ActiveRecord::Associations::Association.send(:include, Association)
+      ActiveRecord::Associations::Association.send(:prepend, Association)
       require 'active_record/turntable/active_record_ext/fixtures'
       require 'active_record/turntable/active_record_ext/migration_proxy'
       require 'active_record/turntable/active_record_ext/activerecord_import_ext'

lib/active_record/turntable/active_record_ext/association.rb

@@ -1,139 +1,32 @@
-require 'active_record/associations'
+require "active_record/associations"
 
 module ActiveRecord::Turntable
   module ActiveRecordExt
     module Association
-      extend ActiveSupport::Concern
+      include ShardingCondition
 
-      included do
-        ActiveRecord::Associations::SingularAssociation.send(:include, SingularAssociationExt)
-        ActiveRecord::Associations::CollectionAssociation.send(:include, CollectionAssociationExt)
-        ActiveRecord::Associations::Builder::Association.valid_options += [:foreign_shard_key]
+      def self.prepended(mod)
+        ActiveRecord::Associations::Builder::Association.valid_options << :foreign_shard_key
       end
 
-      private
-
-      def turntable_scope(scope, bind = nil)
-        if should_use_shard_key?
-          scope = scope.where(klass.turntable_shard_key => owner.send(foreign_shard_key))
-        end
-        scope
-      end
-
-      module SingularAssociationExt
-        extend ActiveSupport::Concern
-
-        included do
-          if Util.ar42_or_later?
-            alias_method_chain :get_records, :turntable
-          else
-            alias_method_chain :find_target, :turntable
-          end
-        end
-
-        # @note Override to add sharding condition for singular association
-        if Util.ar42_or_later?
-          def get_records_with_turntable
-            if reflection.scope_chain.any?(&:any?) ||
-                scope.eager_loading? ||
-                klass.current_scope ||
-                klass.default_scopes.any? ||
-                should_use_shard_key? # OPTIMIZE: Use bind values if cachable scope
-
-              return turntable_scope(scope).limit(1).to_a
-            end
-
-            conn = klass.connection
-            sc = reflection.association_scope_cache(conn, owner) do
-              ActiveRecord::StatementCache.create(conn) { |params|
-                as = ActiveRecord::Associations::AssociationScope.create { params.bind }
-                target_scope.merge(as.scope(self, conn)).limit(1)
-              }
-            end
-
-            binds = ActiveRecord::Associations::AssociationScope.get_bind_values(owner, reflection.chain)
-            sc.execute binds, klass, klass.connection
-          end
-        elsif Util.ar41_or_later?
-          def find_target_with_turntable
-            if record = turntable_scope(scope).take
-              set_inverse_instance record
-            end
-          end
-        else
-          def find_target_with_turntable
-            turntable_scope(scope).take.tap { |record| set_inverse_instance(record) }
-          end
-        end
-      end
+      protected
 
-      module CollectionAssociationExt
-        extend ActiveSupport::Concern
+        # @note Override to pass shard key conditions
+        def target_scope
+          return super unless should_use_shard_key?
 
-        included do
-          if Util.ar42_or_later?
-            alias_method_chain :get_records, :turntable
-          else
-            alias_method_chain :find_target, :turntable
-          end
+          scope = klass.where(
+                    klass.turntable_shard_key =>
+                      owner.send(foreign_shard_key)
+                  )
+          super.merge!(scope)
         end
 
-        private
-
-        if Util.ar42_or_later?
-          def get_records_with_turntable
-            if reflection.scope_chain.any?(&:any?) ||
-                scope.eager_loading? ||
-                klass.current_scope ||
-                klass.default_scopes.any? ||
-                should_use_shard_key? # OPTIMIZE: Use bind values if cachable scope
-
-              return turntable_scope(scope).to_a
-            end
-
-            conn = klass.connection
-            sc = reflection.association_scope_cache(conn, owner) do
-              ActiveRecord::StatementCache.create(conn) { |params|
-                as = ActiveRecord::Associations::AssociationScope.create { params.bind }
-                target_scope.merge as.scope(self, conn)
-              }
-            end
-
-            binds = ActiveRecord::Associations::AssociationScope.get_bind_values(owner, reflection.chain)
-            sc.execute binds, klass, klass.connection
-          end
-        else
-          # @note Override to add sharding condition for collection association
-          def find_target_with_turntable
-            records =
-              if options[:finder_sql]
-                reflection.klass.find_by_sql(custom_finder_sql)
-              else
-                current_scope = scope
-                if should_use_shard_key?
-                  current_scope = current_scope.where(klass.turntable_shard_key => owner.send(foreign_shard_key))
-                end
-                current_scope.to_a
-              end
-            records.each { |record| set_inverse_instance(record) }
-            records
-          end
-        end
-      end
-
       private
 
-      def foreign_shard_key
-        options[:foreign_shard_key] || owner.turntable_shard_key
-      end
-
-      def should_use_shard_key?
-        same_association_shard_key? || !!options[:foreign_shard_key]
-      end
-
-      def same_association_shard_key?
-        owner.class.turntable_enabled? && klass.turntable_enabled? && foreign_shard_key == klass.turntable_shard_key
-      end
+        def skip_statement_cache?
+          super || should_use_shard_key?
+        end
     end
   end
 end

lib/active_record/turntable/sharding_condition.rb

@@ -0,0 +1,23 @@
+module ActiveRecord::Turntable
+  module ShardingCondition
+    private
+
+    def foreign_shard_key
+      options[:foreign_shard_key] || foreign_target_model.turntable_shard_key
+    end
+
+    def foreign_target_model
+      respond_to?(:model) ? model : owner
+    end
+
+    def should_use_shard_key?
+      sharded_by_same_key? || !!options[:foreign_shard_key]
+    end
+
+    def sharded_by_same_key?
+      foreign_target_model.turntable_enabled? &&
+        klass.turntable_enabled? &&
+        foreign_shard_key == klass.turntable_shard_key
+    end
+  end
+end

spec/active_record/turntable/active_record_ext/association_spec.rb

@@ -47,11 +47,20 @@
       ActiveRecord::Base.turntable_config.instance_variable_get(:@config)[:raise_on_not_specified_shard_query] = true
     end
     let(:cards_user) { CardsUser.where(user: user).first }
+    let(:cards_users_history) { cards_users_histories.find { |history| history.user_id == user.id } }
 
     context "associated objects has same turntable_key" do
-      subject { cards_user.cards_users_histories.to_a }
-      it { expect { subject }.to_not raise_error }
-      it { is_expected.to include(*cards_users_histories.select { |history| history.cards_user_id == cards_user.id }) }
+      context "AssociationRelation#to_a" do
+        subject { cards_user.cards_users_histories.to_a }
+        it { expect { subject }.to_not raise_error }
+        it { is_expected.to include(*cards_users_histories.select { |history| history.cards_user_id == cards_user.id }) }
+      end
+
+      context "AssociationRelation#where" do
+        subject { cards_user.cards_users_histories.where(id: cards_users_history.id).to_a }
+        it { expect { subject }.to_not raise_error }
+        it { is_expected.to include(cards_users_history) }
+      end
     end
 
     context "associated objects has different turntable_key" do