Use UNION instead of OR for milestone queries

This MR optimizes some of the milestones (timebox like model) queries by
using UNION instead of OR.

app/controllers/groups/milestones_controller.rb

@@ -12,7 +12,7 @@ class Groups::MilestonesController < Groups::ApplicationController
   def index
     respond_to do |format|
       format.html do
-        @milestone_states = Milestone.states_count(group_projects_with_access, [group])
+        @milestone_states = Milestone.states_count(group_projects_with_access.without_order, [group])
         @milestones = milestones.page(params[:page])
       end
       format.json do

app/models/application_record.rb

@@ -5,6 +5,10 @@ class ApplicationRecord < ActiveRecord::Base
 
   alias_method :reset, :reload
 
+  def self.without_order
+    reorder(nil)
+  end
+
   def self.id_in(ids)
     where(id: ids)
   end

app/models/concerns/timebox.rb

@@ -9,6 +9,7 @@ module Timebox
   include IidRoutes
   include Referable
   include StripAttribute
+  include FromUnion
 
   TimeboxStruct = Struct.new(:title, :name, :id) do
     # Ensure these models match the interface required for exporting
@@ -65,7 +66,11 @@ module Timebox
       groups = groups.compact if groups.is_a? Array
       groups = [] if groups.nil?
 
-      where(project_id: projects).or(where(group_id: groups))
+      if Feature.enabled?(:optimized_timebox_queries)
+        from_union([where(project_id: projects), where(group_id: groups)], remove_duplicates: false)
+      else
+        where(project_id: projects).or(where(group_id: groups))
+      end
     end
 
     scope :within_timeframe, -> (start_date, end_date) do

changelogs/unreleased/optimize-milestones-page.yml

+---
+title: Optimize SQL queries on Milestone index page
+merge_request: 32953
+author:
+type: performance

spec/models/milestone_spec.rb

@@ -225,70 +225,88 @@ describe Milestone do
     end
   end
 
-  describe '#for_projects_and_groups' do
-    let(:project) { create(:project) }
-    let(:project_other) { create(:project) }
-    let(:group) { create(:group) }
-    let(:group_other) { create(:group) }
+  shared_examples '#for_projects_and_groups' do
+    describe '#for_projects_and_groups' do
+      let_it_be(:project) { create(:project) }
+      let_it_be(:project_other) { create(:project) }
+      let_it_be(:group) { create(:group) }
+      let_it_be(:group_other) { create(:group) }
+
+      before(:all) do
+        create(:milestone, project: project)
+        create(:milestone, project: project_other)
+        create(:milestone, group: group)
+        create(:milestone, group: group_other)
+      end
 
-    before do
-      create(:milestone, project: project)
-      create(:milestone, project: project_other)
-      create(:milestone, group: group)
-      create(:milestone, group: group_other)
-    end
+      subject { described_class.for_projects_and_groups(projects, groups) }
+
+      shared_examples 'filters by projects and groups' do
+        it 'returns milestones filtered by project' do
+          milestones = described_class.for_projects_and_groups(projects, [])
+
+          expect(milestones.count).to eq(1)
+          expect(milestones.first.project_id).to eq(project.id)
+        end
+
+        it 'returns milestones filtered by group' do
+          milestones = described_class.for_projects_and_groups([], groups)
 
-    subject { described_class.for_projects_and_groups(projects, groups) }
+          expect(milestones.count).to eq(1)
+          expect(milestones.first.group_id).to eq(group.id)
+        end
 
-    shared_examples 'filters by projects and groups' do
-      it 'returns milestones filtered by project' do
-        milestones = described_class.for_projects_and_groups(projects, [])
+        it 'returns milestones filtered by both project and group' do
+          milestones = described_class.for_projects_and_groups(projects, groups)
 
-        expect(milestones.count).to eq(1)
-        expect(milestones.first.project_id).to eq(project.id)
+          expect(milestones.count).to eq(2)
+          expect(milestones).to contain_exactly(project.milestones.first, group.milestones.first)
+        end
       end
 
-      it 'returns milestones filtered by group' do
-        milestones = described_class.for_projects_and_groups([], groups)
+      context 'ids as params' do
+        let(:projects) { [project.id] }
+        let(:groups) { [group.id] }
 
-        expect(milestones.count).to eq(1)
-        expect(milestones.first.group_id).to eq(group.id)
+        it_behaves_like 'filters by projects and groups'
       end
 
-      it 'returns milestones filtered by both project and group' do
-        milestones = described_class.for_projects_and_groups(projects, groups)
+      context 'relations as params' do
+        let(:projects) { Project.where(id: project.id).select(:id) }
+        let(:groups) { Group.where(id: group.id).select(:id) }
 
-        expect(milestones.count).to eq(2)
-        expect(milestones).to contain_exactly(project.milestones.first, group.milestones.first)
+        it_behaves_like 'filters by projects and groups'
       end
-    end
 
-    context 'ids as params' do
-      let(:projects) { [project.id] }
-      let(:groups) { [group.id] }
+      context 'objects as params' do
+        let(:projects) { [project] }
+        let(:groups) { [group] }
 
-      it_behaves_like 'filters by projects and groups'
-    end
+        it_behaves_like 'filters by projects and groups'
+      end
 
-    context 'relations as params' do
-      let(:projects) { Project.where(id: project.id).select(:id) }
-      let(:groups) { Group.where(id: group.id).select(:id) }
+      it 'returns no records if projects and groups are nil' do
+        milestones = described_class.for_projects_and_groups(nil, nil)
 
-      it_behaves_like 'filters by projects and groups'
+        expect(milestones).to be_empty
+      end
     end
+  end
 
-    context 'objects as params' do
-      let(:projects) { [project] }
-      let(:groups) { [group] }
-
-      it_behaves_like 'filters by projects and groups'
+  context 'when `optimized_timebox_queries` feature flag is enabled' do
+    before do
+      stub_feature_flags(optimized_timebox_queries: true)
     end
 
-    it 'returns no records if projects and groups are nil' do
-      milestones = described_class.for_projects_and_groups(nil, nil)
+    it_behaves_like '#for_projects_and_groups'
+  end
 
-      expect(milestones).to be_empty
+  context 'when `optimized_timebox_queries` feature flag is disabled' do
+    before do
+      stub_feature_flags(optimized_timebox_queries: false)
     end
+
+    it_behaves_like '#for_projects_and_groups'
   end
 
   describe '.upcoming_ids' do