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,8 +66,12 @@ module Timebox
       groups = groups.compact if groups.is_a? Array
       groups = [] if groups.nil?
 
+      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
       where('start_date is not NULL or due_date is not NULL')

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,13 +225,14 @@ describe Milestone do
     end
   end
 
+  shared_examples '#for_projects_and_groups' do
     describe '#for_projects_and_groups' do
-    let(:project) { create(:project) }
-    let(:project_other) { create(:project) }
-    let(:group) { create(:group) }
-    let(:group_other) { create(:group) }
+      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 do
+      before(:all) do
         create(:milestone, project: project)
         create(:milestone, project: project_other)
         create(:milestone, group: group)
@@ -290,6 +291,23 @@ describe Milestone do
         expect(milestones).to be_empty
       end
     end
+  end
+
+  context 'when `optimized_timebox_queries` feature flag is enabled' do
+    before do
+      stub_feature_flags(optimized_timebox_queries: true)
+    end
+
+    it_behaves_like '#for_projects_and_groups'
+  end
+
+  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
     let(:group_1) { create(:group) }