Merge branch 'pb-backport-disable-joins-from-rails' into 'master'

Backport :disable_joins option from Rails main

See merge request gitlab-org/gitlab!66400

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
+
+ActiveRecord::Associations::Association.prepend(GemExtensions::ActiveRecord::Association)
+# Temporarily allow :disable_joins to accept a lambda argument, to control rollout with feature flags
+ActiveRecord::Associations::Association.prepend(GemExtensions::ActiveRecord::ConfigurableDisableJoins)
+ActiveRecord::Associations::Builder::HasOne.prepend(GemExtensions::ActiveRecord::Associations::Builder::HasOne)
+ActiveRecord::Associations::Builder::HasMany.prepend(GemExtensions::ActiveRecord::Associations::Builder::HasMany)
+ActiveRecord::Associations::HasOneThroughAssociation.prepend(GemExtensions::ActiveRecord::Associations::HasOneThroughAssociation)
+ActiveRecord::Associations::HasManyThroughAssociation.prepend(GemExtensions::ActiveRecord::Associations::HasManyThroughAssociation)
+ActiveRecord::Associations::Preloader::ThroughAssociation.prepend(GemExtensions::ActiveRecord::Associations::Preloader::ThroughAssociation)
+ActiveRecord::Base.extend(GemExtensions::ActiveRecord::DelegateCache)

lib/gem_extensions/active_record/association.rb

+# frozen_string_literal: true
+
+module GemExtensions
+  module ActiveRecord
+    module Association
+      extend ActiveSupport::Concern
+
+      attr_reader :disable_joins
+
+      def initialize(owner, reflection)
+        super
+
+        @disable_joins = @reflection.options[:disable_joins] || false
+      end
+
+      def scope
+        if disable_joins
+          DisableJoins::Associations::AssociationScope.create.scope(self)
+        else
+          super
+        end
+      end
+
+      def association_scope
+        if klass
+          @association_scope ||= begin # rubocop:disable Gitlab/ModuleWithInstanceVariables
+            if disable_joins
+              DisableJoins::Associations::AssociationScope.scope(self)
+            else
+              super
+            end
+          end
+        end
+      end
+    end
+  end
+end

lib/gem_extensions/active_record/associations/builder/has_many.rb

+# frozen_string_literal: true
+
+module GemExtensions
+  module ActiveRecord
+    module Associations
+      module Builder
+        module HasMany
+          extend ActiveSupport::Concern
+
+          class_methods do
+            def valid_options(options)
+              valid = super
+              valid += [:disable_joins] if options[:disable_joins] && options[:through]
+              valid
+            end
+          end
+        end
+      end
+    end
+  end
+end

lib/gem_extensions/active_record/associations/builder/has_one.rb

+# frozen_string_literal: true
+
+module GemExtensions
+  module ActiveRecord
+    module Associations
+      module Builder
+        module HasOne
+          extend ActiveSupport::Concern
+
+          class_methods do
+            def valid_options(options)
+              valid = super
+              valid += [:disable_joins] if options[:disable_joins] && options[:through]
+              valid
+            end
+          end
+        end
+      end
+    end
+  end
+end

lib/gem_extensions/active_record/associations/has_many_through_association.rb

+# frozen_string_literal: true
+
+module GemExtensions
+  module ActiveRecord
+    module Associations
+      module HasManyThroughAssociation
+        extend ActiveSupport::Concern
+
+        def find_target
+          return [] unless target_reflection_has_associated_record?
+          return scope.to_a if disable_joins
+
+          super
+        end
+      end
+    end
+  end
+end

lib/gem_extensions/active_record/associations/has_one_through_association.rb

+# frozen_string_literal: true
+
+module GemExtensions
+  module ActiveRecord
+    module Associations
+      module HasOneThroughAssociation
+        extend ActiveSupport::Concern
+
+        def find_target
+          return scope.first if disable_joins
+
+          super
+        end
+      end
+    end
+  end
+end

lib/gem_extensions/active_record/associations/preloader/through_association.rb

+# frozen_string_literal: true
+
+module GemExtensions
+  module ActiveRecord
+    module Associations
+      module Preloader
+        module ThroughAssociation
+          extend ActiveSupport::Concern
+
+          def through_scope
+            scope = through_reflection.klass.unscoped
+            options = reflection.options
+
+            return scope if options[:disable_joins]
+
+            super
+          end
+        end
+      end
+    end
+  end
+end

lib/gem_extensions/active_record/configurable_disable_joins.rb

+# frozen_string_literal: true
+
+module GemExtensions
+  module ActiveRecord
+    module ConfigurableDisableJoins
+      extend ActiveSupport::Concern
+
+      def disable_joins
+        # rubocop:disable Gitlab/ModuleWithInstanceVariables
+        return @disable_joins.call if @disable_joins.is_a?(Proc)
+
+        @disable_joins
+        # rubocop:enable Gitlab/ModuleWithInstanceVariables
+      end
+    end
+  end
+end

lib/gem_extensions/active_record/delegate_cache.rb

+# frozen_string_literal: true
+
+module GemExtensions
+  module ActiveRecord
+    module DelegateCache
+      def relation_delegate_class(klass)
+        @relation_delegate_cache2[klass] || super # rubocop:disable Gitlab/ModuleWithInstanceVariables
+      end
+
+      def initialize_relation_delegate_cache_disable_joins
+        @relation_delegate_cache2 = {} # rubocop:disable Gitlab/ModuleWithInstanceVariables
+
+        [
+          DisableJoins::Relation
+        ].each do |klass|
+          delegate = Class.new(klass) do
+            include ::ActiveRecord::Delegation::ClassSpecificRelation
+          end
+          include_relation_methods(delegate)
+          mangled_name = klass.name.gsub("::", "_")
+          const_set mangled_name, delegate
+          private_constant mangled_name
+
+          @relation_delegate_cache2[klass] = delegate # rubocop:disable Gitlab/ModuleWithInstanceVariables
+        end
+      end
+
+      def inherited(child_class)
+        child_class.initialize_relation_delegate_cache_disable_joins
+        super
+      end
+    end
+  end
+end

lib/gem_extensions/active_record/disable_joins/associations/association_scope.rb

+# frozen_string_literal: true
+
+module GemExtensions
+  module ActiveRecord
+    module DisableJoins
+      module Associations
+        class AssociationScope < ::ActiveRecord::Associations::AssociationScope # :nodoc:
+          def scope(association)
+            source_reflection = association.reflection
+            owner = association.owner
+            unscoped = association.klass.unscoped
+            reverse_chain = get_chain(source_reflection, association, unscoped.alias_tracker).reverse
+
+            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,
+                previous_reflection: previous_reflection)
+          end
+
+          private
+
+          def last_scope_chain(reverse_chain, owner)
+            # 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 |(previous_reflection, reflection, ordered, join_ids), next_reflection|
+              key = reflection.join_primary_key
+              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) # rubocop:disable CodeReuse/ActiveRecord
+              records_ordered = records && records.order_values.any?
+
+              [reflection, next_reflection, records_ordered, record_ids]
+            end
+          end
+
+          def add_constraints(reflection, key, join_ids, owner, ordered, previous_reflection: nil)
+            scope = reflection.build_scope(reflection.aliased_table).where(key => join_ids) # rubocop:disable CodeReuse/ActiveRecord
+
+            # 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)
+              scope.where_clause += item.where_clause
+              scope.order_values = item.order_values | scope.order_values
+              scope
+            end
+
+            if scope.order_values.empty? && ordered
+              split_scope = DisableJoins::Relation.create(scope.klass, key, join_ids)
+              split_scope.where_clause += scope.where_clause
+              split_scope
+            else
+              scope
+            end
+          end
+        end
+      end
+    end
+  end
+end

lib/gem_extensions/active_record/disable_joins/relation.rb

+# frozen_string_literal: true
+
+module GemExtensions
+  module ActiveRecord
+    module DisableJoins
+      class Relation < ::ActiveRecord::Relation
+        attr_reader :ids, :key
+
+        def initialize(klass, key, ids)
+          @ids = ids.uniq
+          @key = key
+          super(klass)
+        end
+
+        def limit(value)
+          records.take(value) # rubocop:disable CodeReuse/ActiveRecord
+        end
+
+        def first(limit = nil)
+          if limit
+            records.limit(limit).first
+          else
+            records.first
+          end
+        end
+
+        def load
+          super
+          records = @records
+
+          records_by_id = records.group_by do |record|
+            record[key]
+          end
+
+          records = ids.flat_map { |id| records_by_id[id.to_i] }
+          records.compact!
+
+          @records = records
+        end
+      end
+    end
+  end
+end

spec/initializers/00_rails_disable_joins_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe 'DisableJoins' do
+  let(:primary_model) do
+    Class.new(ApplicationRecord) do
+      self.table_name = '_test_primary_records'
+
+      def self.name
+        'TestPrimary'
+      end
+    end
+  end
+
+  let(:bridge_model) do
+    Class.new(ApplicationRecord) do
+      self.table_name = '_test_bridge_records'
+
+      def self.name
+        'TestBridge'
+      end
+    end
+  end
+
+  let(:secondary_model) do
+    Class.new(ApplicationRecord) do
+      self.table_name = '_test_secondary_records'
+
+      def self.name
+        'TestSecondary'
+      end
+    end
+  end
+
+  context 'passing disable_joins as an association option' do
+    context 'when the association is a bare has_one' do
+      it 'disallows the disable_joins option' do
+        expect do
+          primary_model.has_one :test_bridge, disable_joins: true
+        end.to raise_error(ArgumentError, /Unknown key: :disable_joins/)
+      end
+    end
+
+    context 'when the association is a belongs_to' do
+      it 'disallows the disable_joins option' do
+        expect do
+          bridge_model.belongs_to :test_secondary, disable_joins: true
+        end.to raise_error(ArgumentError, /Unknown key: :disable_joins/)
+      end
+    end
+
+    context 'when the association is has_one :through' do
+      it 'allows the disable_joins option' do
+        primary_model.has_one :test_bridge
+        bridge_model.belongs_to :test_secondary
+
+        expect do
+          primary_model.has_one :test_secondary, through: :test_bridge, disable_joins: true
+        end.not_to raise_error
+      end
+    end
+
+    context 'when the association is a bare has_many' do
+      it 'disallows the disable_joins option' do
+        expect do
+          primary_model.has_many :test_bridges, disable_joins: true
+        end.to raise_error(ArgumentError, /Unknown key: :disable_joins/)
+      end
+    end
+
+    context 'when the association is a has_many :through' do
+      it 'allows the disable_joins option' do
+        primary_model.has_many :test_bridges
+        bridge_model.belongs_to :test_secondary
+
+        expect do
+          primary_model.has_many :test_secondaries, through: :test_bridges, disable_joins: true
+        end.not_to raise_error
+      end
+    end
+  end
+
+  context 'querying has_one :through when disable_joins is set' 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,
+          primary_record_id int NOT NULL,
+          secondary_record_id int NOT NULL);
+
+        CREATE TABLE _test_secondary_records (
+          id serial NOT NULL PRIMARY KEY);
+      SQL
+
+      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 }
+
+      primary_record = primary_model.create!
+      secondary_record = secondary_model.create!
+      bridge_model.create!(primary_record_id: primary_record.id, secondary_record_id: secondary_record.id)
+    end
+
+    context 'when disable_joins evaluates to true' do
+      let(:joins_disabled_flag) { true }
+
+      it 'executes separate queries' do
+        primary_record = primary_model.first
+
+        query_count = ActiveRecord::QueryRecorder.new { primary_record.test_secondary }.count
+
+        expect(query_count).to eq(2)
+      end
+    end
+
+    context 'when disable_joins evalutes to false' do
+      let(:joins_disabled_flag) { false }
+
+      it 'executes a single query' do
+        primary_record = primary_model.first
+
+        query_count = ActiveRecord::QueryRecorder.new { primary_record.test_secondary }.count
+
+        expect(query_count).to eq(1)
+      end
+    end
+  end
+
+  context 'querying has_many :through when disable_joins is set' 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,
+          primary_record_id int NOT NULL);
+
+        CREATE TABLE _test_secondary_records (
+          id serial NOT NULL PRIMARY KEY,
+          bridge_record_id int NOT NULL);
+      SQL
+
+      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 }
+
+      primary_record = primary_model.create!
+      bridge_record = bridge_model.create!(primary_record_id: primary_record.id)
+      secondary_model.create!(bridge_record_id: bridge_record.id)
+    end
+
+    context 'when disable_joins evaluates to true' do
+      let(:disabled_join_flag) { true }
+
+      it 'executes separate queries' do
+        primary_record = primary_model.first
+
+        query_count = ActiveRecord::QueryRecorder.new { primary_record.test_secondaries.first }.count
+
+        expect(query_count).to eq(2)
+      end
+    end
+
+    context 'when disable_joins evalutes to false' do
+      let(:disabled_join_flag) { false }
+
+      it 'executes a single query' do
+        primary_record = primary_model.first
+
+        query_count = ActiveRecord::QueryRecorder.new { primary_record.test_secondaries.first }.count
+
+        expect(query_count).to eq(1)
+      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