Fix to pass condition for sharding when calling association relation …

…with additional where conditions

lib/active_record/turntable/active_record_ext.rb

@@ -30,7 +30,7 @@ module ActiveRecordExt
       ActiveRecord::Migration.include(ActiveRecord::Turntable::Migration)
       ActiveRecord::ConnectionAdapters::ConnectionHandler.prepend(ConnectionHandlerExtension)
       ActiveRecord::Associations::Preloader::Association.prepend(AssociationPreloader)
-      ActiveRecord::Associations::Association.include(Association)
+      ActiveRecord::Associations::Association.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

@@ -3,49 +3,24 @@
 module ActiveRecord::Turntable
   module ActiveRecordExt
     module Association
-      extend ActiveSupport::Concern
       include ShardingCondition
 
-      included do
-        ActiveRecord::Associations::SingularAssociation.prepend(SingularAssociationExt)
-        ActiveRecord::Associations::CollectionAssociation.prepend(CollectionAssociationExt)
+      def self.prepended(mod)
         ActiveRecord::Associations::Builder::Association::VALID_OPTIONS << :foreign_shard_key
       end
 
-      # @note Inject to add sharding condition for singular association
-      module SingularAssociationExt
-        private
-          def get_records
-            # OPTIMIZE: statement caching
-            if should_use_shard_key?
-              return turntable_scope(scope).limit(1).records
-            end
+      protected
 
-            super
-          end
-      end
-
-      # @note Inject to add sharding condition for collection association
-      module CollectionAssociationExt
-        private
-          def get_records
-            # OPTIMIZE: statement caching
-            if should_use_shard_key?
-              return turntable_scope(scope).to_a
-            end
+      # @note Override to pass shard key conditions
+      def target_scope
+        return super unless should_use_shard_key?
 
-            super
-          end
+        scope = klass.where(
+                  klass.turntable_shard_key =>
+                    owner.send(foreign_shard_key)
+                )
+        super.merge!(scope)
       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
     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