Merge branch 'sy-resolve-oncall-n-plus-one' into 'master'

Improve performance of on-call schedules SQL queries

See merge request gitlab-org/gitlab!58209

app/graphql/resolvers/concerns/looks_ahead.rb

@@ -15,12 +15,7 @@ module LooksAhead
   end
 
   def apply_lookahead(query)
-    selection = node_selection
-
-    includes = preloads.each.flat_map do |name, requirements|
-      selection&.selects?(name) ? requirements : []
-    end
-    all_preloads = (unconditional_includes + includes).uniq
+    all_preloads = (unconditional_includes + filtered_preloads).uniq
 
     return query if all_preloads.empty?
 
@@ -37,6 +32,14 @@ module LooksAhead
     {}
   end
 
+  def filtered_preloads
+    selection = node_selection
+
+    preloads.each.flat_map do |name, requirements|
+      selection&.selects?(name) ? requirements : []
+    end
+  end
+
   def node_selection
     return unless lookahead
 

app/graphql/resolvers/group_members_resolver.rb

@@ -13,12 +13,6 @@ module Resolvers
 
     private
 
-    def preloads
-      {
-      user: [:user, :source]
-      }
-    end
-
     def finder_class
       GroupMembersFinder
     end

app/graphql/resolvers/members_resolver.rb

@@ -21,6 +21,12 @@ module Resolvers
 
     private
 
+    def preloads
+      {
+        user: [:user, :source]
+      }
+    end
+
     def finder_class
       # override in subclass
     end

doc/development/api_graphql_styleguide.md

@@ -1014,6 +1014,26 @@ class MyThingResolver < BaseResolver
 end
 ```
 
+By default, fields defined in `#preloads` will be preloaded if that field
+is selected in the query. Occasionally, finer control may be
+needed to avoid preloading too much or incorrect content.
+
+Extending the above example, we might want to preload a different
+association if certain fields are requested together. This can
+be done by overriding `#filtered_preloads`:
+
+```ruby
+class MyThingResolver < BaseResolver
+  # ...
+
+  def filtered_preloads
+    return [:alternate_attribute] if lookahead.selects?(:field_one) && lookahead.selects?(:field_two)
+
+    super
+  end
+end
+```
+
 The final thing that is needed is that every field that uses this resolver needs
 to advertise the need for lookahead:
 

ee/app/graphql/ee/types/project_type.rb

@@ -128,6 +128,7 @@ module EE
               ::Types::IncidentManagement::OncallScheduleType.connection_type,
               null: true,
               description: 'Incident Management On-call schedules of the project.',
+              extras: [:lookahead],
               resolver: ::Resolvers::IncidentManagement::OncallScheduleResolver
 
         field :api_fuzzing_ci_configuration,

ee/app/graphql/resolvers/incident_management/oncall_schedule_resolver.rb

@@ -3,12 +3,48 @@
 module Resolvers
   module IncidentManagement
     class OncallScheduleResolver < BaseResolver
+      extend ::Gitlab::Utils::Override
+      include LooksAhead
+
       alias_method :project, :object
 
       type Types::IncidentManagement::OncallScheduleType.connection_type, null: true
 
-      def resolve(**args)
-        ::IncidentManagement::OncallSchedulesFinder.new(context[:current_user], project).execute
+      def resolve_with_lookahead(**args)
+        apply_lookahead(::IncidentManagement::OncallSchedulesFinder.new(context[:current_user], project).execute)
+      end
+
+      private
+
+      # Tailor preloads to requested rotation fields instead of
+      # using LooksAhead#preloads to bulk-load all rotation associations
+      override :filtered_preloads
+      def filtered_preloads
+        rotation = rotation_selection
+
+        return [] unless rotation
+        return [{ rotations: { active_participants: :user } }] if rotation.selects?(:participants)
+        return [{ rotations: :active_participants }] if will_generate_shifts?(rotation)
+
+        [:rotations]
+      end
+
+      # @param rotation [GraphQL::Execution::Lookahead]
+      def will_generate_shifts?(rotation)
+        return false unless rotation.selects?(:shifts)
+
+        rotation.selection(:shifts).arguments[:end_time] > Time.current
+      end
+
+      def rotation_selection
+        rotations = node_selection&.selection(:rotations)
+        return unless rotations&.selected?
+
+        if rotations.selects?(:nodes)
+          rotations.selection(:nodes)
+        elsif rotations.selects?(:edges)
+          rotations.selection(:edges).selection(:node)
+        end
       end
     end
   end

ee/app/graphql/types/incident_management/oncall_rotation_type.rb

@@ -58,7 +58,7 @@ module Types
             resolver: ::Resolvers::IncidentManagement::OncallShiftsResolver
 
       def participants
-        object.participants.not_removed
+        object.active_participants
       end
     end
   end

ee/app/graphql/types/incident_management/oncall_shift_type.rb

@@ -2,12 +2,11 @@
 
 module Types
   module IncidentManagement
+    # rubocop: disable Graphql/AuthorizeTypes
     class OncallShiftType < BaseObject
       graphql_name 'IncidentManagementOncallShift'
       description 'A block of time for which a participant is on-call.'
 
-      authorize :read_incident_management_oncall_schedule
-
       field :participant,
             ::Types::IncidentManagement::OncallParticipantType,
             null: true,
@@ -22,6 +21,11 @@ module Types
             Types::TimeType,
             null: true,
             description: 'End time of the on-call shift.'
+
+      def participant
+        Gitlab::Graphql::Loaders::OncallParticipantLoader.new(object.participant_id).find
+      end
     end
+    # rubocop: enable Graphql/AuthorizeTypes
   end
 end

ee/changelogs/unreleased/sy-resolve-oncall-n-plus-one.yml

+---
+title: Improve performance of on-call schedules SQL queries
+merge_request: 58209
+author:
+type: performance

ee/lib/gitlab/graphql/loaders/oncall_participant_loader.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Graphql
+    module Loaders
+      class OncallParticipantLoader
+        attr_reader :participant_id
+
+        def initialize(participant_id)
+          @participant_id = participant_id
+        end
+
+        # rubocop: disable CodeReuse/ActiveRecord
+        def find
+          BatchLoader::GraphQL.for(participant_id.to_i).batch do |ids, loader|
+            results = ::IncidentManagement::OncallParticipant.includes(:user).id_in(ids)
+
+            results.each { |participant| loader.call(participant.id, participant) }
+          end
+        end
+        # rubocop: enable CodeReuse/ActiveRecord
+      end
+    end
+  end
+end

ee/spec/graphql/types/incident_management/oncall_shift_type_spec.rb

@@ -5,8 +5,6 @@ require 'spec_helper'
 RSpec.describe GitlabSchema.types['IncidentManagementOncallShift'] do
   specify { expect(described_class.graphql_name).to eq('IncidentManagementOncallShift') }
 
-  specify { expect(described_class).to require_graphql_authorizations(:read_incident_management_oncall_schedule) }
-
   it 'exposes the expected fields' do
     expected_fields = %i[
       participant

ee/spec/lib/gitlab/graphql/loaders/oncall_participant_loader_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Graphql::Loaders::OncallParticipantLoader do
+  describe '#find' do
+    let_it_be(:participant1) { create(:incident_management_oncall_participant) }
+    let_it_be(:participant2) { create(:incident_management_oncall_participant) }
+    let_it_be(:participant3) { create(:incident_management_oncall_participant) }
+
+    it 'finds a participant by id' do
+      first_result = described_class.new(participant1.id).find
+      second_result = described_class.new(participant2.id).find
+
+      expect(first_result.sync).to eq(participant1)
+      expect(second_result.sync).to eq(participant2)
+    end
+
+    it 'includes the user association' do
+      expect do
+        [described_class.new(participant3.id).find,
+         described_class.new(participant2.id).find,
+         described_class.new(participant1.id).find].map(&:sync).map(&:user)
+      end.not_to exceed_query_limit(2)
+    end
+  end
+end

ee/spec/requests/api/graphql/project/incident_management/oncall_participants_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe 'getting Incident Management on-call shifts' do
+  include GraphqlHelpers
+
+  let_it_be(:participant) { create(:incident_management_oncall_participant, :utc, :with_developer_access) }
+  let_it_be(:rotation) { participant.rotation }
+  let_it_be(:project) { rotation.project }
+  let_it_be(:current_user) { participant.user }
+
+  let(:fields) do
+    <<~QUERY
+      nodes {
+        rotations {
+          nodes {
+            participants {
+              nodes {
+                id
+                colorPalette
+                colorWeight
+                user { id }
+              }
+            }
+          }
+        }
+      }
+    QUERY
+  end
+
+  let(:query) do
+    graphql_query_for(
+      'project',
+      { 'fullPath' => project.full_path },
+      query_graphql_field('incidentManagementOncallSchedules', {}, fields)
+    )
+  end
+
+  let(:participants) do
+    graphql_data
+      .dig('project', 'incidentManagementOncallSchedules', 'nodes').first
+      .dig('rotations', 'nodes').first
+      .dig('participants', 'nodes')
+  end
+
+  before do
+    stub_licensed_features(oncall_schedules: true)
+    post_graphql(query, current_user: current_user)
+  end
+
+  it_behaves_like 'a working graphql query'
+
+  it 'returns the correct properties of the on-call shifts' do
+    expect(participants.first).to include(
+      'id' => participant.to_global_id.to_s,
+      'user' => { 'id' => participant.user.to_global_id.to_s },
+      'colorWeight' => '50',
+      'colorPalette' => 'blue'
+    )
+  end
+
+  context 'performance' do
+    shared_examples 'avoids N+1 queries' do
+      specify do
+        base_count = ActiveRecord::QueryRecorder.new do
+          post_graphql(query, current_user: current_user)
+        end
+
+        action
+
+        expect { post_graphql(query, current_user: current_user) }.not_to exceed_query_limit(base_count)
+      end
+    end
+
+    context 'for additional participant' do
+      let(:action) { create(:incident_management_oncall_participant, rotation: rotation) }
+
+      it_behaves_like 'avoids N+1 queries'
+    end
+
+    context 'for additional rotation with participants' do
+      let(:action) { create(:incident_management_oncall_rotation, :with_participants, schedule: rotation.schedule) }
+
+      it_behaves_like 'avoids N+1 queries'
+    end
+  end
+end

ee/spec/requests/api/graphql/project/incident_management/oncall_shifts_spec.rb

@@ -17,7 +17,7 @@ RSpec.describe 'getting Incident Management on-call shifts' do
   let(:shift_fields) do
     <<~QUERY
       nodes {
-        participant { id }
+        participant { id user { id } }
         endsAt
         startsAt
       }
@@ -60,12 +60,98 @@ RSpec.describe 'getting Incident Management on-call shifts' do
 
   it 'returns the correct properties of the on-call shifts' do
     expect(shifts.first).to include(
-      'participant' => { 'id' => participant.to_global_id.to_s },
+      'participant' => {
+        'id' => participant.to_global_id.to_s,
+        'user' => { 'id' => participant.user.to_global_id.to_s }
+      },
       'startsAt' => params[:start_time],
       'endsAt' => params[:end_time]
     )
   end
 
+  context 'performance' do
+    shared_examples 'avoids N+1 queries' do
+      specify do
+        base_count = ActiveRecord::QueryRecorder.new do
+          post_graphql(query, current_user: current_user)
+        end
+
+        action
+
+        expect { post_graphql(query, current_user: current_user) }.not_to exceed_query_limit(base_count)
+      end
+    end
+
+    shared_examples 'avoids N+1 queries for additional generated shift' do
+      include_examples 'avoids N+1 queries' do
+        let(:action) { params[:end_time] = ends_at.next_day.iso8601 }
+      end
+    end
+
+    shared_examples 'avoids N+1 queries for additional historical shift' do
+      include_examples 'avoids N+1 queries' do
+        let(:action) { create(:incident_management_oncall_shift, participant: participant, starts_at: last_shift.ends_at) }
+      end
+    end
+
+    shared_examples 'avoids N+1 queries for additional participant' do
+      include_examples 'avoids N+1 queries' do
+        let(:action) { create(:incident_management_oncall_participant, rotation: rotation) }
+      end
+    end
+
+    shared_examples 'avoids N+1 queries for additional rotation with participants' do
+      include_examples 'avoids N+1 queries' do
+        let(:action) { create(:incident_management_oncall_rotation, :with_participants, schedule: rotation.schedule) }
+      end
+    end
+
+    shared_examples 'adds only one query for each additional rotation with participants' do
+      specify do
+        base_count = ActiveRecord::QueryRecorder.new do
+          post_graphql(query, current_user: current_user)
+        end
+
+        create(:incident_management_oncall_rotation, :with_participants, schedule: rotation.schedule)
+        create(:incident_management_oncall_rotation, :with_participants, schedule: rotation.schedule)
+
+        expect { post_graphql(query, current_user: current_user) }.not_to exceed_query_limit(base_count).with_threshold(2)
+      end
+    end
+
+    context 'for past and future shifts' do
+      let_it_be(:last_shift) { create(:incident_management_oncall_shift, participant: participant) }
+
+      let(:ends_at) { rotation.starts_at + 2 * rotation.shift_cycle_duration }
+
+      it_behaves_like 'avoids N+1 queries for additional generated shift'
+      it_behaves_like 'avoids N+1 queries for additional historical shift'
+      it_behaves_like 'avoids N+1 queries for additional participant'
+      it_behaves_like 'adds only one query for each additional rotation with participants'
+    end
+
+    context 'for future shifts only' do
+      let(:starts_at) { rotation.starts_at + rotation.shift_cycle_duration }
+      let(:ends_at) { rotation.starts_at + 2 * rotation.shift_cycle_duration }
+
+      it_behaves_like 'avoids N+1 queries for additional generated shift'
+      it_behaves_like 'avoids N+1 queries for additional participant'
+      it_behaves_like 'avoids N+1 queries for additional rotation with participants'
+    end
+
+    context 'for past shifts only' do
+      let_it_be(:last_shift) { create(:incident_management_oncall_shift, participant: participant) }
+
+      around do |example|
+        travel_to(starts_at + 1.5 * rotation.shift_cycle_duration) { example.run }
+      end
+
+      it_behaves_like 'avoids N+1 queries for additional historical shift'
+      it_behaves_like 'avoids N+1 queries for additional participant'
+      it_behaves_like 'adds only one query for each additional rotation with participants'
+    end
+  end
+
   context "without required argument starts_at" do
     let(:params) { { end_time: ends_at.iso8601 } }