Fix broken use cases with disable_joins

Backport changes from a few additional Rails PRs that fix certain use
cases for `disable_joins`.

Also add an additional fix to address polymorphic relationships not
properly filtering by the polymorphic type.

config/initializers/00_rails_disable_joins.rb

 # frozen_string_literal: true
 
+# Backported from Rails 7.0
+# Initial support for has_many :through was implemented in https://github.com/rails/rails/pull/41937
+# Support for has_one :through was implemented in https://github.com/rails/rails/pull/42079
+raise 'DisableJoins patch is only to be used with versions of Rails < 7.0' unless Rails::VERSION::MAJOR < 7
+
 module DisableJoins
   module ConfigurableDisableJoins
     extend ActiveSupport::Concern
@@ -110,29 +115,50 @@ module DisableJoins
       unscoped = association.klass.unscoped
       reverse_chain = get_chain(source_reflection, association, unscoped.alias_tracker).reverse
 
-      last_reflection, last_ordered, last_join_ids = last_scope_chain(reverse_chain, owner)
+      previous_reflection, last_reflection, last_ordered, last_join_ids = last_scope_chain(reverse_chain, owner)
 
-      add_constraints(last_reflection, last_reflection.join_primary_key, last_join_ids, owner, last_ordered)
+      add_constraints(last_reflection, last_reflection.join_primary_key, last_join_ids, owner, last_ordered,
+          previous_reflection: previous_reflection)
     end
 
     private
 
     def last_scope_chain(reverse_chain, owner)
-      first_scope = [reverse_chain.shift, false, [owner.id]]
+      # Pulled from https://github.com/rails/rails/pull/42448
+      # Fixes cases where the foreign key is not id
+      first_item = reverse_chain.shift
+      first_scope = [nil, first_item, false, [owner._read_attribute(first_item.join_foreign_key)]]
 
-      reverse_chain.inject(first_scope) do |(reflection, ordered, join_ids), next_reflection|
+      reverse_chain.inject(first_scope) do |(previous_reflection, reflection, ordered, join_ids), next_reflection|
         key = reflection.join_primary_key
-        records = add_constraints(reflection, key, join_ids, owner, ordered)
+        records = add_constraints(reflection, key, join_ids, owner, ordered, previous_reflection: previous_reflection)
         foreign_key = next_reflection.join_foreign_key
         record_ids = records.pluck(foreign_key)
         records_ordered = records && records.order_values.any?
 
-        [next_reflection, records_ordered, record_ids]
+        [reflection, next_reflection, records_ordered, record_ids]
       end
     end
 
-    def add_constraints(reflection, key, join_ids, owner, ordered)
+    def add_constraints(reflection, key, join_ids, owner, ordered, previous_reflection: nil)
       scope = reflection.build_scope(reflection.aliased_table).where(key => join_ids)
+
+      # Pulled from https://github.com/rails/rails/pull/42590
+      # Fixes cases where used with an STI type
+      relation = reflection.klass.scope_for_association
+      scope.merge!(
+        relation.except(:select, :create_with, :includes, :preload, :eager_load, :joins, :left_outer_joins)
+      )
+
+      # Attempt to fix use case where we have a polymorphic relationship
+      # Build on an additional scope to filter by the polymorphic type
+      if reflection.type
+        polymorphic_class = previous_reflection.try(:klass) || owner.class
+
+        polymorphic_type = transform_value(polymorphic_class.polymorphic_name)
+        scope = apply_scope(scope, reflection.aliased_table, reflection.type, polymorphic_type)
+      end
+
       scope = reflection.constraints.inject(scope) do |memo, scope_chain_item|
         item = eval_scope(reflection, scope_chain_item, owner)
         scope.unscope!(*item.unscope_values)

spec/initializers/00_rails_disable_joins_spec.rb

@@ -83,7 +83,7 @@ RSpec.describe 'DisableJoins' do
 
   context 'querying has_one :through when disable_joins is set' do
     before do
-      ApplicationRecord.connection.execute(<<~SQL)
+      create_tables(<<~SQL)
         CREATE TABLE _test_primary_records (
           id serial NOT NULL PRIMARY KEY);
 
@@ -96,12 +96,8 @@ RSpec.describe 'DisableJoins' do
           id serial NOT NULL PRIMARY KEY);
       SQL
 
-      bridge_model.reset_column_information
-      secondary_model.reset_column_information
-
       primary_model.has_one :test_bridge, anonymous_class: bridge_model, foreign_key: :primary_record_id
       bridge_model.belongs_to :test_secondary, anonymous_class: secondary_model, foreign_key: :secondary_record_id
-
       primary_model.has_one :test_secondary, through: :test_bridge, anonymous_class: secondary_model,
         disable_joins: -> { joins_disabled_flag }
 
@@ -137,7 +133,7 @@ RSpec.describe 'DisableJoins' do
 
   context 'querying has_many :through when disable_joins is set' do
     before do
-      ApplicationRecord.connection.execute(<<~SQL)
+      create_tables(<<~SQL)
         CREATE TABLE _test_primary_records (
           id serial NOT NULL PRIMARY KEY);
 
@@ -150,12 +146,8 @@ RSpec.describe 'DisableJoins' do
           bridge_record_id int NOT NULL);
       SQL
 
-      bridge_model.reset_column_information
-      secondary_model.reset_column_information
-
       primary_model.has_many :test_bridges, anonymous_class: bridge_model, foreign_key: :primary_record_id
       bridge_model.has_many :test_secondaries, anonymous_class: secondary_model, foreign_key: :bridge_record_id
-
       primary_model.has_many :test_secondaries, through: :test_bridges, anonymous_class: secondary_model,
         disable_joins: -> { disabled_join_flag }
 
@@ -188,4 +180,109 @@ RSpec.describe 'DisableJoins' do
       end
     end
   end
+
+  context 'querying STI relationships' do
+    let(:child_bridge_model) do
+      Class.new(bridge_model) do
+        def self.name
+          'ChildBridge'
+        end
+      end
+    end
+
+    let(:child_secondary_model) do
+      Class.new(secondary_model) do
+        def self.name
+          'ChildSecondary'
+        end
+      end
+    end
+
+    before do
+      create_tables(<<~SQL)
+        CREATE TABLE _test_primary_records (
+          id serial NOT NULL PRIMARY KEY);
+
+        CREATE TABLE _test_bridge_records (
+          id serial NOT NULL PRIMARY KEY,
+          primary_record_id int NOT NULL,
+          type text);
+
+        CREATE TABLE _test_secondary_records (
+          id serial NOT NULL PRIMARY KEY,
+          bridge_record_id int NOT NULL,
+          type text);
+      SQL
+
+      primary_model.has_many :child_bridges, anonymous_class: child_bridge_model, foreign_key: :primary_record_id
+      child_bridge_model.has_one :child_secondary, anonymous_class: child_secondary_model, foreign_key: :bridge_record_id
+      primary_model.has_many :child_secondaries, through: :child_bridges, anonymous_class: child_secondary_model, disable_joins: true
+
+      primary_record = primary_model.create!
+      parent_bridge_record = bridge_model.create!(primary_record_id: primary_record.id)
+      child_bridge_record = child_bridge_model.create!(primary_record_id: primary_record.id)
+
+      secondary_model.create!(bridge_record_id: child_bridge_record.id)
+      child_secondary_model.create!(bridge_record_id: parent_bridge_record.id)
+      child_secondary_model.create!(bridge_record_id: child_bridge_record.id)
+    end
+
+    it 'filters correctly by the STI type across multiple queries' do
+      primary_record = primary_model.first
+
+      query_recorder = ActiveRecord::QueryRecorder.new do
+        expect(primary_record.child_secondaries.count).to eq(1)
+      end
+
+      expect(query_recorder.count).to eq(2)
+    end
+  end
+
+  context 'querying polymorphic relationships' do
+    before do
+      create_tables(<<~SQL)
+        CREATE TABLE _test_primary_records (
+          id serial NOT NULL PRIMARY KEY);
+
+        CREATE TABLE _test_bridge_records (
+          id serial NOT NULL PRIMARY KEY,
+          primaryable_id int NOT NULL,
+          primaryable_type text NOT NULL);
+
+        CREATE TABLE _test_secondary_records (
+          id serial NOT NULL PRIMARY KEY,
+          bridgeable_id int NOT NULL,
+          bridgeable_type text NOT NULL);
+      SQL
+
+      primary_model.has_many :test_bridges, anonymous_class: bridge_model, foreign_key: :primaryable_id, as: :primaryable
+      bridge_model.has_one :test_secondaries, anonymous_class: secondary_model, foreign_key: :bridgeable_id, as: :bridgeable
+      primary_model.has_many :test_secondaries, through: :test_bridges, anonymous_class: secondary_model, disable_joins: true
+
+      primary_record = primary_model.create!
+      primary_bridge_record = bridge_model.create!(primaryable_id: primary_record.id, primaryable_type: 'TestPrimary')
+      nonprimary_bridge_record = bridge_model.create!(primaryable_id: primary_record.id, primaryable_type: 'NonPrimary')
+
+      secondary_model.create!(bridgeable_id: primary_bridge_record.id, bridgeable_type: 'TestBridge')
+      secondary_model.create!(bridgeable_id: nonprimary_bridge_record.id, bridgeable_type: 'TestBridge')
+      secondary_model.create!(bridgeable_id: primary_bridge_record.id, bridgeable_type: 'NonBridge')
+    end
+
+    it 'filters correctly by the polymorphic type across multiple queries' do
+      primary_record = primary_model.first
+
+      query_recorder = ActiveRecord::QueryRecorder.new do
+        expect(primary_record.test_secondaries.count).to eq(1)
+      end
+
+      expect(query_recorder.count).to eq(2)
+    end
+  end
+
+  def create_tables(table_sql)
+    ApplicationRecord.connection.execute(table_sql)
+
+    bridge_model.reset_column_information
+    secondary_model.reset_column_information
+  end
 end