Add elastic_reference and es_parent methods in epic model

ee/app/models/concerns/search/elastic/epics_search.rb

-# frozen_string_literal: true
-
-module Search
-  module Elastic
-    module EpicsSearch
-      extend ActiveSupport::Concern
-
-      include ::Elastic::ApplicationVersionedSearch
-
-      included do
-        extend ::Gitlab::Utils::Override
-
-        override :maintain_elasticsearch_create
-        def maintain_elasticsearch_create
-          ::Elastic::ProcessBookkeepingService.track!(*get_indexing_data)
-        end
-
-        override :maintain_elasticsearch_update
-        def maintain_elasticsearch_update(updated_attributes: previous_changes.keys)
-          ::Elastic::ProcessBookkeepingService.track!(*get_indexing_data)
-          associations_to_update = associations_needing_elasticsearch_update(updated_attributes)
-          return unless associations_to_update.present?
-
-          ElasticAssociationIndexerWorker.perform_async(self.class.name, id, associations_to_update)
-        end
-
-        override :maintain_elasticsearch_destroy
-        def maintain_elasticsearch_destroy
-          ::Elastic::ProcessBookkeepingService.track!(*get_indexing_data)
-        end
-      end
-
-      private
-
-      def get_indexing_data
-        [Search::Elastic::References::WorkItem.new(
-          issue_id, "group_#{group.root_ancestor.id}")]
-      end
-    end
-  end
-end

ee/app/models/ee/epic.rb

@@ -3,6 +3,7 @@
 module EE
   module Epic
     extend ActiveSupport::Concern
+    extend ::Gitlab::Utils::Override
 
     prepended do
       include AtomicInternalId
@@ -23,7 +24,7 @@ module Epic
       include EachBatch
       include ::Exportable
       include Epics::MetadataCacheUpdate
-      include ::Search::Elastic::EpicsSearch
+      include ::Elastic::ApplicationVersionedSearch
       include Elastic::UpdateAssociatedEpicsOnDateChange
 
       # we'd need to make sure these override the existing associations so we prepend this.
@@ -562,6 +563,15 @@ def use_elasticsearch?
       group.use_elasticsearch?
     end
 
+    def es_parent
+      "group_#{group.root_ancestor.id}"
+    end
+
+    override :elastic_reference
+    def elastic_reference
+      ::Search::Elastic::References::WorkItem.new(issue_id, es_parent).serialize
+    end
+
     def issues_readable_by(current_user, preload: nil)
       related_issues = self.class.related_issues(ids: id, preload: preload)
 

ee/app/models/ee/work_item.rb

@@ -44,6 +44,7 @@ module WorkItem
           project: :project_feature
         )
 
+        preloaded_data.each(&:lazy_labels)
         ::Namespaces::Preloaders::NamespaceRootAncestorPreloader.new(preloaded_data.map(&:namespace)).execute
 
         if ::Feature.enabled?(:search_work_items_index_notes, ::Feature.current_request)

ee/spec/elastic_integration/epic_index_spec.rb

@@ -28,11 +28,11 @@
     context 'when an epic is created' do
       let(:epic) { build(:epic, group: group) }
 
-      it 'tracks the work item' do
+      it 'tracks the epic and work item' do
         expect(::Elastic::ProcessBookkeepingService).to receive(:track!).once do |*tracked_refs|
           expect(tracked_refs.count).to eq(1)
-          expect(tracked_refs[0].class).to eq(Search::Elastic::References::WorkItem)
-          expect(tracked_refs[0].identifier).to eq(epic.issue_id)
+          expect(tracked_refs[0].class).to eq(Epic)
+          expect(tracked_refs[0].id).to eq(epic.id)
         end
         expect(::Elastic::ProcessBookkeepingService).to receive(:track!).once do |*tracked_refs|
           expect(tracked_refs.count).to eq(1)
@@ -48,23 +48,23 @@
     end
 
     context 'when an epic is updated' do
-      it 'tracks the work item' do
+      it 'tracks the epic' do
         expect(::Elastic::ProcessBookkeepingService).to receive(:track!).once do |*tracked_refs|
           expect(tracked_refs.count).to eq(1)
-          expect(tracked_refs[0]).to be_a_kind_of(Search::Elastic::References::WorkItem)
-          expect(tracked_refs[0].identifier).to eq(epic.issue_id)
+          expect(tracked_refs[0].class).to eq(Epic)
+          expect(tracked_refs[0].id).to eq(epic.id)
         end
         epic.update!(title: 'A new title')
       end
     end
 
     context 'when an epic is deleted' do
-      it 'tracks the work item' do
+      it 'tracks the work item and epic' do
         expect(::Elastic::ProcessBookkeepingService).to receive(:track!).with(*[WorkItem.find(epic.issue_id)]).once
         expect(::Elastic::ProcessBookkeepingService).to receive(:track!).once do |*tracked_refs|
           expect(tracked_refs.count).to eq(1)
-          expect(tracked_refs[0]).to be_a_kind_of(Search::Elastic::References::WorkItem)
-          expect(tracked_refs[0].identifier).to eq(epic.issue_id)
+          expect(tracked_refs[0].class).to eq(Epic)
+          expect(tracked_refs[0].id).to eq(epic.id)
         end
         epic.destroy!
       end
@@ -145,8 +145,8 @@
       it 'tracks the epic' do
         expect(::Elastic::ProcessBookkeepingService).to receive(:track!).once do |*tracked_refs|
           expect(tracked_refs.count).to eq(1)
-          expect(tracked_refs[0]).to be_a_kind_of(Search::Elastic::References::WorkItem)
-          expect(tracked_refs[0].identifier).to eq(epic.issue_id)
+          expect(tracked_refs[0]).to be_a_kind_of(Epic)
+          expect(tracked_refs[0].id).to eq(epic.id)
         end
         epic.update!(group: parent_group)
       end

ee/spec/models/concerns/search/elastic/epics_search_spec.rb

-# frozen_string_literal: true
-
-require 'spec_helper'
-
-RSpec.describe ::Search::Elastic::EpicsSearch, :elastic, feature_category: :global_search do
-  let_it_be(:epic) { create(:epic) }
-
-  describe '#maintain_elasticsearch_create' do
-    it 'calls track! for epic' do
-      expect(::Elastic::ProcessBookkeepingService).to receive(:track!).once do |*tracked_refs|
-        expect(tracked_refs.count).to eq(1)
-        expect(tracked_refs[0]).to be_a_kind_of(Search::Elastic::References::WorkItem)
-        expect(tracked_refs[0].identifier).to eq(epic.issue_id)
-      end
-      epic.maintain_elasticsearch_create
-    end
-  end
-
-  describe '#maintain_elasticsearch_destory' do
-    it 'calls track! for epic' do
-      expect(::Elastic::ProcessBookkeepingService).to receive(:track!).once do |*tracked_refs|
-        expect(tracked_refs.count).to eq(1)
-        expect(tracked_refs[0]).to be_a_kind_of(Search::Elastic::References::WorkItem)
-        expect(tracked_refs[0].identifier).to eq(epic.issue_id)
-      end
-
-      epic.maintain_elasticsearch_destroy
-    end
-  end
-
-  describe '#maintain_elasticsearch_update' do
-    it 'calls track! for epic' do
-      expect(::Elastic::ProcessBookkeepingService).to receive(:track!).once do |*tracked_refs|
-        expect(tracked_refs.count).to eq(1)
-        expect(tracked_refs[0]).to be_a_kind_of(Search::Elastic::References::WorkItem)
-        expect(tracked_refs[0].identifier).to eq(epic.issue_id)
-      end
-
-      epic.maintain_elasticsearch_update
-    end
-
-    context 'when we have associations to update' do
-      before do
-        allow(Epic).to receive(:elastic_index_dependants).and_return([{ association_name: :issues,
-                                                                        on_change: 'title' }])
-      end
-
-      it 'calls track! for epic and updates the associations' do
-        expect(::Elastic::ProcessBookkeepingService).to receive(:track!).once do |*tracked_refs|
-          expect(tracked_refs.count).to eq(1)
-          expect(tracked_refs[0]).to be_a_kind_of(Search::Elastic::References::WorkItem)
-          expect(tracked_refs[0].identifier).to eq(epic.issue_id)
-        end
-
-        expect(::ElasticAssociationIndexerWorker).to receive(:perform_async).with('Epic', epic.id, ['issues']).once
-        epic.maintain_elasticsearch_update
-      end
-    end
-  end
-end

ee/spec/models/epic_spec.rb

@@ -1364,8 +1364,8 @@ def as_item(item)
         it 'calls ::Elastic::ProcessBookkeepingService.track! when the epic is updated' do
           expect(::Elastic::ProcessBookkeepingService).to receive(:track!).once do |*tracked_refs|
             expect(tracked_refs.count).to eq(1)
-            expect(tracked_refs[0]).to be_a_kind_of(Search::Elastic::References::WorkItem)
-            expect(tracked_refs[0].identifier).to eq(epic.issue_id)
+            expect(tracked_refs[0]).to be_a_kind_of(Epic)
+            expect(tracked_refs[0].id).to eq(epic.id)
           end
           epic.update!(title: 'A new title')
         end
@@ -1833,4 +1833,20 @@ def as_item(item)
       end
     end
   end
+
+  describe '#elastic_reference' do
+    let_it_be(:epic) { create(:epic, group: group) }
+
+    it 'returns the string representation for the elasticsearch' do
+      expect(epic.elastic_reference).to eq("WorkItem|#{epic.issue_id}|#{epic.es_parent}")
+    end
+  end
+
+  describe '#es_parent' do
+    let_it_be(:epic) { create(:epic, group: group) }
+
+    it 'returns to correct routing id' do
+      expect(epic.es_parent).to eq("group_#{group.root_ancestor.id}")
+    end
+  end
 end

ee/spec/services/elastic/process_bookkeeping_service_spec.rb

@@ -493,6 +493,70 @@
 
           expect { described_class.new.execute }.not_to exceed_all_query_limit(control)
         end
+
+        it 'does not have N+1 queries for epics' do
+          epics = create_list(:epic, 2, :use_fixed_dates)
+
+          described_class.track!(*epics)
+
+          control = ActiveRecord::QueryRecorder.new(skip_cached: false) { described_class.new.execute }
+
+          epics += create_list(:epic, 3, :use_fixed_dates)
+
+          described_class.track!(*epics)
+
+          expect do
+            described_class.new.execute
+          end.not_to exceed_all_query_limit(control)
+        end
+
+        it 'does not have N+1 queries for epics in a group with multiple parents' do
+          parent_group = create(:group)
+          group = create(:group, parent: parent_group)
+
+          epics = create_list(:epic, 2, group: group)
+
+          described_class.track!(*epics)
+
+          control = ActiveRecord::QueryRecorder.new(skip_cached: false) { described_class.new.execute }
+
+          epics += create_list(:epic, 3, group: create(:group, parent: parent_group))
+
+          described_class.track!(*epics)
+
+          expect do
+            described_class.new.execute
+          end.not_to exceed_all_query_limit(control)
+        end
+
+        it 'does not have N+1 queries for epics with inherited dates' do
+          child_epic = create(:epic, :use_fixed_dates)
+          milestone = create(:milestone, :with_dates)
+
+          epics = create_list(:epic, 2)
+          epics.each do |epic|
+            epic.start_date_sourcing_epic = child_epic
+            epic.due_date_sourcing_milestone = milestone
+            epic.save!
+          end
+
+          described_class.track!(*epics)
+
+          control = ActiveRecord::QueryRecorder.new(skip_cached: false) { described_class.new.execute }
+
+          epics += create_list(:epic, 3)
+          epics.each do |epic|
+            epic.start_date_sourcing_epic = child_epic
+            epic.due_date_sourcing_milestone = milestone
+            epic.save!
+          end
+
+          described_class.track!(*epics)
+
+          expect do
+            described_class.new.execute
+          end.not_to exceed_all_query_limit(control)
+        end
       end
 
       it 'does not have N+1 queries for notes' do
@@ -611,64 +675,6 @@
           expect { described_class.new.execute }.not_to exceed_all_query_limit(control)
         end
       end
-
-      it 'does not have N+1 queries for epics' do
-        epics = create_list(:epic, 2, :use_fixed_dates)
-
-        described_class.track!(*epics)
-
-        control = ActiveRecord::QueryRecorder.new(skip_cached: false) { described_class.new.execute }
-
-        epics += create_list(:epic, 3, :use_fixed_dates)
-
-        described_class.track!(*epics)
-
-        expect { described_class.new.execute }.not_to exceed_all_query_limit(control)
-      end
-
-      it 'does not have N+1 queries for epics with inherited dates' do
-        child_epic = create(:epic, :use_fixed_dates)
-        milestone = create(:milestone, :with_dates)
-
-        epics = create_list(:epic, 2)
-        epics.each do |epic|
-          epic.start_date_sourcing_epic = child_epic
-          epic.due_date_sourcing_milestone = milestone
-          epic.save!
-        end
-
-        described_class.track!(*epics)
-
-        control = ActiveRecord::QueryRecorder.new(skip_cached: false) { described_class.new.execute }
-
-        epics += create_list(:epic, 3)
-        epics.each do |epic|
-          epic.start_date_sourcing_epic = child_epic
-          epic.due_date_sourcing_milestone = milestone
-          epic.save!
-        end
-
-        described_class.track!(*epics)
-
-        expect { described_class.new.execute }.not_to exceed_all_query_limit(control)
-      end
-
-      it 'does not have N+1 queries for epics in a group with multiple parents' do
-        parent_group = create(:group)
-        group = create(:group, parent: parent_group)
-
-        epics = create_list(:epic, 2, group: group)
-
-        described_class.track!(*epics)
-
-        control = ActiveRecord::QueryRecorder.new(skip_cached: false) { described_class.new.execute }
-
-        epics += create_list(:epic, 3, group: create(:group, parent: parent_group))
-
-        described_class.track!(*epics)
-
-        expect { described_class.new.execute }.not_to exceed_all_query_limit(control)
-      end
     end
 
     def expect_processing(*refs, failures: [])