Migrate notes (comments) to a separate index

This is the Advanced Search migration of notes into a separate indices.
It pauses indexing and migrates all notes into a new index. After the
migration is complete it will unpause indexing and switch all queries to
using this new index.

ee/app/models/concerns/elastic/application_versioned_search.rb

@@ -93,7 +93,7 @@ module Elastic
       end
 
       def use_separate_indices?
-        Gitlab::Elastic::Helper::ES_SEPARATE_CLASSES.include?(self) && Elastic::DataMigrationService.migration_has_finished?(:migrate_issues_to_separate_index)
+        false
       end
 
       # Mark a dependant association as needing to be updated when a specific

ee/app/models/ee/issue.rb

@@ -79,6 +79,11 @@ module EE
       def with_api_entity_associations
         super.preload(:epic)
       end
+
+      # override
+      def use_separate_indices?
+        Elastic::DataMigrationService.migration_has_finished?(:migrate_issues_to_separate_index)
+      end
     end
 
     # override

ee/app/models/ee/note.rb

@@ -20,6 +20,13 @@ module EE
       end
     end
 
+    class_methods do
+      # override
+      def use_separate_indices?
+        Elastic::DataMigrationService.migration_has_finished?(:migrate_notes_to_separate_index)
+      end
+    end
+
     # Original method in Elastic::ApplicationSearch
     def searchable?
       !system && super

ee/app/services/elastic/data_migration_service.rb

@@ -25,6 +25,10 @@ module Elastic
         migrations.find { |m| m.version == version }
       end
 
+      def find_by_name(name)
+        migrations.find { |migration| migration.name_for_key == name.to_s.underscore }
+      end
+
       def drop_migration_has_finished_cache!(migration)
         Rails.cache.delete cache_key(:migration_has_finished, migration.name_for_key)
       end
@@ -36,7 +40,7 @@ module Elastic
       end
 
       def migration_has_finished_uncached?(name)
-        migration = migrations.find { |migration| migration.name_for_key == name.to_s.underscore }
+        migration = find_by_name(name)
 
         !!migration&.load_from_index&.dig('_source', 'completed')
       end

ee/app/workers/concerns/elastic/migration_helper.rb

+# frozen_string_literal: true
+
+module Elastic
+  module MigrationHelper
+    private
+
+    def document_type
+      raise NotImplementedError
+    end
+
+    def document_type_fields
+      raise NotImplementedError
+    end
+
+    def document_type_plural
+      document_type.pluralize
+    end
+
+    def get_number_of_shards
+      helper.get_settings(index_name: new_index_name).dig('number_of_shards').to_i
+    end
+
+    def default_index_name
+      helper.target_name
+    end
+
+    def new_index_name
+      "#{default_index_name}-#{document_type_plural}"
+    end
+
+    def original_documents_count
+      query = {
+        size: 0,
+        aggs: {
+          documents: {
+            filter: {
+              term: {
+                type: {
+                  value: document_type
+                }
+              }
+            }
+          }
+        }
+      }
+
+      results = client.search(index: default_index_name, body: query)
+      results.dig('aggregations', 'documents', 'doc_count')
+    end
+
+    def new_documents_count
+      helper.documents_count(index_name: new_index_name)
+    end
+
+    def reindexing_cleanup!
+      helper.delete_index(index_name: new_index_name) if helper.index_exists?(index_name: new_index_name)
+    end
+
+    def reindex(slice:, max_slices:)
+      body = reindex_query(slice: slice, max_slices: max_slices)
+
+      response = client.reindex(body: body, wait_for_completion: false)
+
+      response['task']
+    end
+
+    def reindexing_completed?(task_id:)
+      response = helper.task_status(task_id: task_id)
+      completed = response['completed']
+
+      return false unless completed
+
+      stats = response['response']
+      if stats['failures'].present?
+        log_raise "Reindexing failed with #{stats['failures']}"
+      end
+
+      if stats['total'] != (stats['updated'] + stats['created'] + stats['deleted'])
+        log_raise "Slice reindexing seems to have failed, total is not equal to updated + created + deleted"
+      end
+
+      true
+    end
+
+    def reindex_query(slice:, max_slices:)
+      {
+        source: {
+          index: default_index_name,
+          _source: document_type_fields,
+          query: {
+            match: {
+              type: document_type
+            }
+          },
+          slice: {
+            id: slice,
+            max: max_slices
+          }
+        },
+        dest: {
+          index: new_index_name
+        }
+      }
+    end
+  end
+end

ee/app/workers/elastic_delete_project_worker.rb

@@ -19,11 +19,17 @@ class ElasticDeleteProjectWorker
   def indices
     helper = Gitlab::Elastic::Helper.default
 
+    index_names = [helper.target_name]
+
     if Elastic::DataMigrationService.migration_has_finished?(:migrate_issues_to_separate_index)
-      [helper.target_name] + helper.standalone_indices_proxies.map(&:index_name)
-    else
-      [helper.target_name]
+      index_names << helper.standalone_indices_proxies(target_classes: [Issue]).map(&:index_name)
+    end
+
+    if Elastic::DataMigrationService.migration_has_finished?(:migrate_notes_to_separate_index)
+      index_names << helper.standalone_indices_proxies(target_classes: [Note]).map(&:index_name)
     end
+
+    index_names
   end
 
   def remove_project_and_children_documents(project_id, es_id)

ee/changelogs/unreleased/297646-move-notes-to-separate-index.yml

+---
+title: Move Notes (Comments) to dedicated ES index
+merge_request: 53013
+author:
+type: changed

ee/elastic/migrate/20210302104500_migrate_notes_to_separate_index.rb

+# frozen_string_literal: true
+
+class MigrateNotesToSeparateIndex < Elastic::Migration
+  include Elastic::MigrationHelper
+
+  pause_indexing!
+  batched!
+  throttle_delay 1.minute
+
+  MAX_ATTEMPTS_PER_SLICE = 30
+
+  def migrate
+    # On initial batch we only create index
+    if migration_state[:slice].blank?
+      reindexing_cleanup! # support retries
+
+      log "Create standalone #{document_type_plural} index under #{new_index_name}"
+      helper.create_standalone_indices(target_classes: [Note])
+
+      options = {
+        slice: 0,
+        retry_attempt: 0,
+        max_slices: get_number_of_shards
+      }
+      set_migration_state(options)
+
+      return
+    end
+
+    retry_attempt = migration_state[:retry_attempt].to_i
+    slice = migration_state[:slice]
+    task_id = migration_state[:task_id]
+    max_slices = migration_state[:max_slices]
+
+    if retry_attempt >= MAX_ATTEMPTS_PER_SLICE
+      fail_migration_halt_error!(retry_attempt: retry_attempt)
+      return
+    end
+
+    return unless slice < max_slices
+
+    if task_id
+      log "Checking reindexing status for slice:#{slice} | task_id:#{task_id}"
+
+      if reindexing_completed?(task_id: task_id)
+        log "Reindexing is completed for slice:#{slice} | task_id:#{task_id}"
+
+        set_migration_state(
+          slice: slice + 1,
+          task_id: nil,
+          retry_attempt: 0, # We reset retry_attempt for a next slice
+          max_slices: max_slices
+        )
+      else
+        log "Reindexing is still in progress for slice:#{slice} | task_id:#{task_id}"
+      end
+
+      return
+    end
+
+    log "Launching reindexing for slice:#{slice} | max_slices:#{max_slices}"
+
+    task_id = reindex(slice: slice, max_slices: max_slices)
+
+    log "Reindexing for slice:#{slice} | max_slices:#{max_slices} is started with task_id:#{task_id}"
+
+    set_migration_state(
+      slice: slice,
+      task_id: task_id,
+      max_slices: max_slices
+    )
+  rescue StandardError => e
+    log "migrate failed, increasing migration_state for slice:#{slice} retry_attempt:#{retry_attempt} error:#{e.message}"
+
+    set_migration_state(
+      slice: slice,
+      task_id: task_id,
+      retry_attempt: retry_attempt + 1,
+      max_slices: max_slices
+    )
+
+    raise e
+  end
+
+  def completed?
+    helper.refresh_index(index_name: new_index_name)
+
+    original_count = original_documents_count
+    new_count = new_documents_count
+    log "Checking to see if migration is completed based on index counts: original_count:#{original_count}, new_count:#{new_count}"
+
+    original_count == new_count
+  end
+
+  private
+
+  def document_type
+    'note'
+  end
+
+  def document_type_fields
+    %w(
+      type
+      id
+      note
+      project_id
+      noteable_type
+      noteable_id
+      created_at
+      updated_at
+      confidential
+      issue
+      visibility_level
+      issues_access_level
+      repository_access_level
+      merge_requests_access_level
+      snippets_access_level
+    )
+  end
+end

ee/lib/elastic/class_proxy_util.rb

@@ -12,7 +12,11 @@ module Elastic
       super(target)
 
       const_name = if use_separate_indices
-                     "#{target.name}Config"
+                     if target.superclass.abstract_class?
+                       "#{target.name}Config"
+                     else
+                       "#{target.superclass.name}Config"
+                     end
                    else
                      'Config'
                    end

ee/lib/elastic/instance_proxy_util.rb

@@ -10,7 +10,11 @@ module Elastic
       super(target)
 
       const_name = if use_separate_indices
-                     "#{target.class.name}Config"
+                     if target.class.superclass.abstract_class?
+                       "#{target.class.name}Config"
+                     else
+                       "#{target.class.superclass.name}Config"
+                     end
                    else
                      'Config'
                    end

ee/lib/elastic/latest/note_config.rb

+# frozen_string_literal: true
+
+module Elastic
+  module Latest
+    module NoteConfig
+      # To obtain settings and mappings methods
+      extend Elasticsearch::Model::Indexing::ClassMethods
+      extend Elasticsearch::Model::Naming::ClassMethods
+
+      self.document_type = 'doc'
+      self.index_name = [Rails.application.class.module_parent_name.downcase, Rails.env, 'notes'].join('-')
+
+      settings Elastic::Latest::Config.settings.to_hash
+
+      mappings dynamic: 'strict' do
+        indexes :type, type: :keyword
+
+        indexes :id, type: :integer
+
+        indexes :note, type: :text, index_options: 'positions'
+        indexes :project_id, type: :integer
+
+        indexes :noteable_type, type: :keyword
+        indexes :noteable_id, type: :keyword
+
+        indexes :created_at, type: :date
+        indexes :updated_at, type: :date
+
+        indexes :confidential, type: :boolean
+
+        indexes :visibility_level, type: :integer
+        indexes :issues_access_level, type: :integer
+        indexes :repository_access_level, type: :integer
+        indexes :merge_requests_access_level, type: :integer
+        indexes :snippets_access_level, type: :integer
+
+        indexes :issue do
+          indexes :assignee_id, type: :integer
+          indexes :author_id, type: :integer
+          indexes :confidential, type: :boolean
+        end
+      end
+    end
+  end
+end

ee/lib/elastic/latest/note_instance_proxy.rb

@@ -38,6 +38,14 @@ module Elastic
         data.merge(generic_attributes)
       end
 
+      def generic_attributes
+        if Elastic::DataMigrationService.migration_has_finished?(:migrate_notes_to_separate_index)
+          super.except('join_field')
+        else
+          super
+        end
+      end
+
       private
 
       def merge_project_feature_access_level(data)

ee/lib/gitlab/elastic/helper.rb

@@ -14,7 +14,8 @@ module Gitlab
       ].freeze
 
       ES_SEPARATE_CLASSES = [
-        Issue
+        Issue,
+        Note
       ].freeze
 
       attr_reader :version, :client
@@ -260,7 +261,7 @@ module Gitlab
       end
 
       def get_settings(index_name: nil)
-        index = index_name || target_index_name
+        index = target_index_name(target: index_name)
         settings = client.indices.get_settings(index: index)
         settings.dig(index, 'settings', 'index')
       end

ee/spec/elastic/migrate/20210127154600_remove_permissions_data_from_notes_documents_spec.rb

@@ -11,6 +11,7 @@ RSpec.describe RemovePermissionsDataFromNotesDocuments, :elastic, :sidekiq_inlin
   before do
     stub_ee_application_setting(elasticsearch_search: true, elasticsearch_indexing: true)
     allow(migration).to receive(:helper).and_return(helper)
+    set_elasticsearch_migration_to :remove_permissions_data_from_notes_documents, including: false
   end
 
   describe 'migration_options' do
@@ -47,12 +48,9 @@ RSpec.describe RemovePermissionsDataFromNotesDocuments, :elastic, :sidekiq_inlin
       before do
         add_permission_data_for_notes([note_on_commit, note_on_issue, note_on_merge_request, note_on_snippet])
 
-        allow(Elastic::DataMigrationService).to receive(:migration_has_finished?).and_call_original
         # migrations are completed by default in test environments
         # required to prevent the `as_indexed_json` method from populating the permissions fields
-        allow(Elastic::DataMigrationService).to receive(:migration_has_finished?)
-                                                  .with(:remove_permissions_data_from_notes_documents)
-                                                  .and_return(false)
+        set_elasticsearch_migration_to version, including: false
       end
 
       it 'queues documents for indexing' do

ee/spec/elastic/migrate/202101281636000_add_permissions_data_to_notes_documents_spec.rb

@@ -9,6 +9,7 @@ RSpec.describe AddPermissionsDataToNotesDocuments, :elastic, :sidekiq_inline do
 
   before do
     stub_ee_application_setting(elasticsearch_search: true, elasticsearch_indexing: true)
+    set_elasticsearch_migration_to :add_permissions_data_to_notes_documents, including: false
   end
 
   describe 'migration_options' do

ee/spec/elastic/migrate/20210201104800_migrate_notes_to_separate_index_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+require File.expand_path('ee/elastic/migrate/20210302104500_migrate_notes_to_separate_index.rb')
+
+RSpec.describe MigrateNotesToSeparateIndex do
+  let(:version) { 20210302104500 }
+  let(:migration) { described_class.new(version) }
+  let(:index_name) { "#{es_helper.target_name}-notes" }
+  let(:helper) { Gitlab::Elastic::Helper.new }
+
+  before do
+    stub_ee_application_setting(elasticsearch_search: true, elasticsearch_indexing: true)
+    allow(migration).to receive(:helper).and_return(helper)
+  end
+
+  describe 'migration_options' do
+    it 'has migration options set', :aggregate_failures do
+      expect(migration.batched?).to be_truthy
+      expect(migration.throttle_delay).to eq(1.minute)
+      expect(migration.pause_indexing?).to be_truthy
+    end
+  end
+
+  describe '.migrate', :elastic, :clean_gitlab_redis_shared_state do
+    before do
+      set_elasticsearch_migration_to :migrate_notes_to_separate_index, including: false
+    end
+
+    context 'initial launch' do
+      before do
+        es_helper.delete_index(index_name: es_helper.target_index_name(target: index_name))
+      end
+
+      it 'creates an index and sets migration_state' do
+        expect { migration.migrate }.to change { es_helper.alias_exists?(name: index_name) }.from(false).to(true)
+
+        expect(migration.migration_state).to include(slice: 0, max_slices: 5)
+      end
+    end
+
+    context 'batch run' do
+      it 'sets migration_state task_id' do
+        allow(migration).to receive(:reindex).and_return('task_id')
+
+        migration.set_migration_state(slice: 0, max_slices: 5)
+
+        migration.migrate
+
+        expect(migration.migration_state).to include(slice: 0, max_slices: 5, task_id: 'task_id')
+      end
+
+      it 'sets next slice after task check' do
+        allow(migration).to receive(:reindexing_completed?).and_return(true)
+
+        migration.set_migration_state(slice: 0, max_slices: 5, task_id: 'task_id')
+
+        migration.migrate
+
+        expect(migration.migration_state).to include(slice: 1, max_slices: 5, task_id: nil)
+      end
+
+      it 'resets retry_attempt for a next slice' do
+        allow(migration).to receive(:reindexing_completed?).and_return(true)
+
+        migration.set_migration_state(slice: 0, max_slices: 5, retry_attempt: 5, task_id: 'task_id')
+
+        migration.migrate
+
+        expect(migration.migration_state).to match(slice: 1, max_slices: 5, retry_attempt: 0, task_id: nil)
+      end
+
+      context 'reindexing is still in progress' do
+        before do
+          allow(migration).to receive(:reindexing_completed?).and_return(false)
+        end
+
+        it 'does nothing' do
+          migration.set_migration_state(slice: 0, max_slices: 5, task_id: 'task_id')
+
+          migration.migrate
+
+          expect(migration).not_to receive(:reindex)
+        end
+      end
+
+      context 'with notes in elastic' do
+        # Create notes on different projects to ensure they are spread across
+        # all shards. If they all end up in 1 ES shard then they'll be migrated
+        # in a single slice.
+        let!(:notes) { Array.new(10).map { create(:note, project: create(:project)) } }
+
+        before do
+          ensure_elasticsearch_index!
+        end
+
+        it 'migrates all notes' do
+          slices = 2
+
+          migration.set_migration_state(slice: 0, max_slices: slices)
+
+          migration.migrate
+
+          50.times do |i| # Max 0.5s waiting
+            break if migration.completed?
+
+            sleep 0.01
+            migration.migrate
+          end
+
+          expect(migration.completed?).to be_truthy
+          expect(es_helper.documents_count(index_name: "#{es_helper.target_name}-notes")).to eq(notes.count)
+        end
+      end
+    end
+
+    context 'failed run' do
+      context 'exception is raised' do
+        before do
+          allow(migration).to receive(:reindex).and_raise(StandardError)
+        end
+
+        it 'increases retry_attempt' do
+          migration.set_migration_state(slice: 0, max_slices: 2, retry_attempt: 1)
+
+          expect { migration.migrate }.to raise_error(StandardError)
+          expect(migration.migration_state).to match(slice: 0, max_slices: 2, retry_attempt: 2, task_id: nil)
+        end
+
+        it 'fails the migration after too many attempts' do
+          migration.set_migration_state(slice: 0, max_slices: 2, retry_attempt: 30)
+
+          migration.migrate
+
+          expect(migration.migration_state).to match(slice: 0, max_slices: 2, retry_attempt: 30, halted: true, halted_indexing_unpaused: false)
+          expect(migration).not_to receive(:reindex)
+        end
+      end
+
+      context 'elasticsearch failures' do
+        context 'total is not equal' do
+          before do
+            allow(helper).to receive(:task_status).and_return({ "completed" => true, "response" => { "total" => 60, "updated" => 0, "created" => 45, "deleted" => 0, "failures" => [] } })
+          end
+
+          it 'raises an error' do
+            migration.set_migration_state(slice: 0, max_slices: 2, task_id: 'task_id')
+
+            expect { migration.migrate }.to raise_error(/total is not equal/)
+          end
+        end
+
+        context 'reindexing failues' do
+          before do
+            allow(helper).to receive(:task_status).with(task_id: 'task_id').and_return({ "completed" => true, "response" => { "total" => 60, "updated" => 0, "created" => 0, "deleted" => 0, "failures" => [{ "type": "es_rejected_execution_exception" }] } })
+          end
+
+          it 'raises an error' do
+            migration.set_migration_state(slice: 0, max_slices: 2, task_id: 'task_id')
+
+            expect { migration.migrate }.to raise_error(/failed with/)
+          end
+        end
+      end
+    end
+  end
+
+  describe '.completed?' do
+    subject { migration.completed? }
+
+    let(:original_count) { 5 }
+
+    before do
+      allow(helper).to receive(:refresh_index).and_return(true)
+      allow(migration).to receive(:original_documents_count).and_return(original_count)
+      allow(migration).to receive(:new_documents_count).and_return(new_count)
+    end
+
+    context 'counts are equal' do
+      let(:new_count) { original_count }
+
+      it 'returns true' do
+        is_expected.to be_truthy
+      end
+    end
+
+    context 'counts are not equal' do
+      let(:new_count) { original_count - 1 }
+
+      it 'returns true' do
+        is_expected.to be_falsey
+      end
+    end
+  end
+end

ee/spec/lib/elastic/latest/note_config_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Elastic::Latest::NoteConfig do
+  describe '.document_type' do
+    it 'returns config' do
+      expect(described_class.document_type).to eq('doc')
+    end
+  end
+
+  describe '.settings' do
+    it 'returns config' do
+      expect(described_class.settings).to be_a(Elasticsearch::Model::Indexing::Settings)
+    end
+  end
+
+  describe '.mappings' do
+    it 'returns config' do
+      expect(described_class.mapping).to be_a(Elasticsearch::Model::Indexing::Mappings)
+    end
+  end
+end

ee/spec/models/concerns/elastic/note_spec.rb

@@ -110,10 +110,6 @@ RSpec.describe Note, :elastic do
                   'confidential' => issue.confidential
                 },
                 'type' => note.es_type,
-                'join_field' => {
-                  'name' => note.es_type,
-                  'parent' => note.es_parent
-                },
                 'visibility_level' => project.visibility_level,
                 'issues_access_level' => project.issues_access_level
               })
@@ -166,6 +162,9 @@ RSpec.describe Note, :elastic do
         allow(Elastic::DataMigrationService).to receive(:migration_has_finished?)
                                                   .with(:remove_permissions_data_from_notes_documents)
                                                   .and_return(false)
+        allow(Elastic::DataMigrationService).to receive(:migration_has_finished?)
+                                                  .with(:migrate_notes_to_separate_index)
+                                                  .and_return(false)
 
         expect(note_json).not_to have_key(access_level) if access_level.present?
         expect(note_json).not_to have_key('visibility_level')

ee/spec/services/search/global_service_spec.rb

@@ -21,23 +21,12 @@ RSpec.describe Search::GlobalService do
   context 'issue search' do
     let(:results) { described_class.new(nil, search: '*').execute.objects('issues') }
 
-    before do
-      allow(Elastic::DataMigrationService).to receive(:migration_has_finished?).and_call_original
-      allow(Elastic::DataMigrationService).to receive(:migration_has_finished?)
-                                                .with(:migrate_issues_to_separate_index)
-                                                .and_return(false)
-    end
-
     it_behaves_like 'search query applies joins based on migrations shared examples', :add_new_data_to_issues_documents
   end
 
   context 'notes search' do
     let(:results) { described_class.new(nil, search: '*').execute.objects('notes') }
 
-    before do
-      allow(Elastic::DataMigrationService).to receive(:migration_has_finished?).and_call_original
-    end
-
     it_behaves_like 'search query applies joins based on migrations shared examples', :add_permissions_data_to_notes_documents
   end
 
@@ -110,17 +99,13 @@ RSpec.describe Search::GlobalService do
         end
 
         with_them do
-          before do
-            allow(Elastic::DataMigrationService).to receive(:migration_has_finished?).and_call_original
-          end
-
           context 'when add_permissions_data_to_notes_documents migration is finished' do
             it_behaves_like 'search respects visibility'
           end
 
           context 'when add_permissions_data_to_notes_documents migration is not finished' do
             before do
-              allow(Elastic::DataMigrationService).to receive(:migration_has_finished?).with(:add_permissions_data_to_notes_documents).and_return(false)
+              set_elasticsearch_migration_to :add_permissions_data_to_notes_documents, including: false
             end
 
             it_behaves_like 'search respects visibility'
@@ -136,17 +121,13 @@ RSpec.describe Search::GlobalService do
         end
 
         with_them do
-          before do
-            allow(Elastic::DataMigrationService).to receive(:migration_has_finished?).and_call_original
-          end
-
           context 'when add_permissions_data_to_notes_documents migration is finished' do
             it_behaves_like 'search respects visibility'
           end
 
           context 'when add_permissions_data_to_notes_documents migration is not finished' do
             before do
-              allow(Elastic::DataMigrationService).to receive(:migration_has_finished?).with(:add_permissions_data_to_notes_documents).and_return(false)
+              set_elasticsearch_migration_to :add_permissions_data_to_notes_documents, including: false
             end
 
             it_behaves_like 'search respects visibility'
@@ -164,7 +145,6 @@ RSpec.describe Search::GlobalService do
 
         with_them do
           before do
-            allow(Elastic::DataMigrationService).to receive(:migration_has_finished?).and_call_original
             project.repository.index_commits_and_blobs
           end
 
@@ -174,7 +154,7 @@ RSpec.describe Search::GlobalService do
 
           context 'when add_permissions_data_to_notes_documents migration is not finished' do
             before do
-              allow(Elastic::DataMigrationService).to receive(:migration_has_finished?).with(:add_permissions_data_to_notes_documents).and_return(false)
+              set_elasticsearch_migration_to :add_permissions_data_to_notes_documents, including: false
             end
 
             it_behaves_like 'search respects visibility'
@@ -190,17 +170,13 @@ RSpec.describe Search::GlobalService do
         end
 
         with_them do
-          before do
-            allow(Elastic::DataMigrationService).to receive(:migration_has_finished?).and_call_original
-          end
-
           context 'when add_permissions_data_to_notes_documents migration is finished' do
             it_behaves_like 'search respects visibility'
           end
 
           context 'when add_permissions_data_to_notes_documents migration is not finished' do
             before do
-              allow(Elastic::DataMigrationService).to receive(:migration_has_finished?).with(:add_permissions_data_to_notes_documents).and_return(false)
+              set_elasticsearch_migration_to :add_permissions_data_to_notes_documents, including: false
             end
 
             it_behaves_like 'search respects visibility'
@@ -230,13 +206,7 @@ RSpec.describe Search::GlobalService do
       # instances which haven't finished the migration yet
       context 'when add_new_data_to_issues_documents migration is not finished' do
         before do
-          allow(Elastic::DataMigrationService).to receive(:migration_has_finished?).and_call_original
-          allow(Elastic::DataMigrationService).to receive(:migration_has_finished?)
-            .with(:add_new_data_to_issues_documents)
-            .and_return(false)
-          allow(Elastic::DataMigrationService).to receive(:migration_has_finished?)
-            .with(:migrate_issues_to_separate_index)
-            .and_return(false)
+          set_elasticsearch_migration_to :add_new_data_to_issues_documents, including: false
         end
 
         # issue cannot be defined prior to the migration mocks because it
@@ -455,18 +425,12 @@ RSpec.describe Search::GlobalService do
   context 'confidential notes' do
     let(:project) { create(:project, :public, :repository) }
 
-    before do
-      allow(Elastic::DataMigrationService).to receive(:migration_has_finished?).and_call_original
-    end
-
     context 'with notes on issues' do
       let(:noteable) { create :issue, project: project }
 
       context 'when add_permissions_data_to_notes_documents migration has not finished' do
         before do
-          allow(Elastic::DataMigrationService).to receive(:migration_has_finished?)
-                                                    .with(:add_permissions_data_to_notes_documents)
-                                                    .and_return(false)
+          set_elasticsearch_migration_to :add_permissions_data_to_notes_documents, including: false
         end
 
         it_behaves_like 'search notes shared examples', :note_on_issue
@@ -474,9 +438,7 @@ RSpec.describe Search::GlobalService do
 
       context 'when add_permissions_data_to_notes_documents migration has finished' do
         before do
-          allow(Elastic::DataMigrationService).to receive(:migration_has_finished?)
-                                                    .with(:add_permissions_data_to_notes_documents)
-                                                    .and_return(true)
+          set_elasticsearch_migration_to :add_permissions_data_to_notes_documents, including: true
         end
 
         it_behaves_like 'search notes shared examples', :note_on_issue
@@ -488,9 +450,7 @@ RSpec.describe Search::GlobalService do
 
       context 'when add_permissions_data_to_notes_documents migration has not finished' do
         before do
-          allow(Elastic::DataMigrationService).to receive(:migration_has_finished?)
-                                                    .with(:add_permissions_data_to_notes_documents)
-                                                    .and_return(false)
+          set_elasticsearch_migration_to :add_permissions_data_to_notes_documents, including: false
         end
 
         it_behaves_like 'search notes shared examples', :note_on_merge_request
@@ -498,9 +458,7 @@ RSpec.describe Search::GlobalService do
 
       context 'when add_permissions_data_to_notes_documents migration has finished' do
         before do
-          allow(Elastic::DataMigrationService).to receive(:migration_has_finished?)
-                                                    .with(:add_permissions_data_to_notes_documents)
-                                                    .and_return(true)
+          set_elasticsearch_migration_to :add_permissions_data_to_notes_documents, including: true
         end
 
         it_behaves_like 'search notes shared examples', :note_on_merge_request
@@ -512,9 +470,7 @@ RSpec.describe Search::GlobalService do
 
       context 'when add_permissions_data_to_notes_documents migration has not finished' do
         before do
-          allow(Elastic::DataMigrationService).to receive(:migration_has_finished?)
-                                                    .with(:add_permissions_data_to_notes_documents)
-                                                    .and_return(false)
+          set_elasticsearch_migration_to :add_permissions_data_to_notes_documents, including: false
         end
 
         it_behaves_like 'search notes shared examples', :note_on_commit
@@ -522,9 +478,7 @@ RSpec.describe Search::GlobalService do
 
       context 'when add_permissions_data_to_notes_documents migration has finished' do
         before do
-          allow(Elastic::DataMigrationService).to receive(:migration_has_finished?)
-                                                    .with(:add_permissions_data_to_notes_documents)
-                                                    .and_return(true)
+          set_elasticsearch_migration_to :add_permissions_data_to_notes_documents, including: true
         end
 
         it_behaves_like 'search notes shared examples', :note_on_commit

ee/spec/services/search/group_service_spec.rb

@@ -72,10 +72,6 @@ RSpec.describe Search::GroupService, :elastic do
     let_it_be(:project) { create(:project, namespace: group) }
     let(:results) { described_class.new(nil, group, search: 'test').execute.objects('notes') }
 
-    before do
-      allow(Elastic::DataMigrationService).to receive(:migration_has_finished?).and_call_original
-    end
-
     it_behaves_like 'search query applies joins based on migrations shared examples', :add_permissions_data_to_notes_documents
   end
 
@@ -156,17 +152,13 @@ RSpec.describe Search::GroupService, :elastic do
         end
 
         with_them do
-          before do
-            allow(Elastic::DataMigrationService).to receive(:migration_has_finished?).and_call_original
-          end
-
           context 'when add_permissions_data_to_notes_documents migration is finished' do
             it_behaves_like 'search respects visibility'
           end
 
           context 'when add_permissions_data_to_notes_documents migration is not finished' do
             before do
-              allow(Elastic::DataMigrationService).to receive(:migration_has_finished?).with(:add_permissions_data_to_notes_documents).and_return(false)
+              set_elasticsearch_migration_to :add_permissions_data_to_notes_documents, including: false
             end
 
             it_behaves_like 'search respects visibility'
@@ -183,17 +175,13 @@ RSpec.describe Search::GroupService, :elastic do
         end
 
         with_them do
-          before do
-            allow(Elastic::DataMigrationService).to receive(:migration_has_finished?).and_call_original
-          end
-
           context 'when add_permissions_data_to_notes_documents migration is finished' do
             it_behaves_like 'search respects visibility'
           end
 
           context 'when add_permissions_data_to_notes_documents migration is not finished' do
             before do
-              allow(Elastic::DataMigrationService).to receive(:migration_has_finished?).with(:add_permissions_data_to_notes_documents).and_return(false)
+              set_elasticsearch_migration_to :add_permissions_data_to_notes_documents, including: false
             end
 
             it_behaves_like 'search respects visibility'
@@ -213,8 +201,6 @@ RSpec.describe Search::GroupService, :elastic do
 
         with_them do
           before do
-            allow(Elastic::DataMigrationService).to receive(:migration_has_finished?).and_call_original
-
             project.repository.index_commits_and_blobs
             project2.repository.index_commits_and_blobs
           end
@@ -225,7 +211,7 @@ RSpec.describe Search::GroupService, :elastic do
 
           context 'when add_permissions_data_to_notes_documents migration is not finished' do
             before do
-              allow(Elastic::DataMigrationService).to receive(:migration_has_finished?).with(:add_permissions_data_to_notes_documents).and_return(false)
+              set_elasticsearch_migration_to :add_permissions_data_to_notes_documents, including: false
             end
 
             it_behaves_like 'search respects visibility'
@@ -242,17 +228,13 @@ RSpec.describe Search::GroupService, :elastic do
         end
 
         with_them do
-          before do
-            allow(Elastic::DataMigrationService).to receive(:migration_has_finished?).and_call_original
-          end
-
           context 'when add_permissions_data_to_notes_documents migration is finished' do
             it_behaves_like 'search respects visibility'
           end
 
           context 'when add_permissions_data_to_notes_documents migration is not finished' do
             before do
-              allow(Elastic::DataMigrationService).to receive(:migration_has_finished?).with(:add_permissions_data_to_notes_documents).and_return(false)
+              set_elasticsearch_migration_to :add_permissions_data_to_notes_documents, including: false
             end
 
             it_behaves_like 'search respects visibility'

ee/spec/services/search/project_service_spec.rb

@@ -38,10 +38,6 @@ RSpec.describe Search::ProjectService do
     let_it_be(:project) { create(:project) }
     let(:results) { described_class.new(project, nil, search: 'test').execute.objects('notes') }
 
-    before do
-      allow(Elastic::DataMigrationService).to receive(:migration_has_finished?).and_call_original
-    end
-
     it_behaves_like 'search query applies joins based on migrations shared examples', :add_permissions_data_to_notes_documents
   end
 
@@ -122,17 +118,13 @@ RSpec.describe Search::ProjectService do
         end
 
         with_them do
-          before do
-            allow(Elastic::DataMigrationService).to receive(:migration_has_finished?).and_call_original
-          end
-
           context 'when add_permissions_data_to_notes_documents migration is finished' do
             it_behaves_like 'search respects visibility'
           end
 
           context 'when add_permissions_data_to_notes_documents migration is not finished' do
             before do
-              allow(Elastic::DataMigrationService).to receive(:migration_has_finished?).with(:add_permissions_data_to_notes_documents).and_return(false)
+              set_elasticsearch_migration_to :add_permissions_data_to_notes_documents, including: false
             end
 
             it_behaves_like 'search respects visibility'
@@ -149,17 +141,13 @@ RSpec.describe Search::ProjectService do
         end
 
         with_them do
-          before do
-            allow(Elastic::DataMigrationService).to receive(:migration_has_finished?).and_call_original
-          end
-
           context 'when add_permissions_data_to_notes_documents migration is finished' do
             it_behaves_like 'search respects visibility'
           end
 
           context 'when add_permissions_data_to_notes_documents migration is not finished' do
             before do
-              allow(Elastic::DataMigrationService).to receive(:migration_has_finished?).with(:add_permissions_data_to_notes_documents).and_return(false)
+              set_elasticsearch_migration_to :add_permissions_data_to_notes_documents, including: false
             end
 
             it_behaves_like 'search respects visibility'
@@ -181,8 +169,6 @@ RSpec.describe Search::ProjectService do
           before do
             project.repository.index_commits_and_blobs
             project2.repository.index_commits_and_blobs
-
-            allow(Elastic::DataMigrationService).to receive(:migration_has_finished?).and_call_original
           end
 
           context 'when add_permissions_data_to_notes_documents migration is finished' do
@@ -191,7 +177,7 @@ RSpec.describe Search::ProjectService do
 
           context 'when add_permissions_data_to_notes_documents migration is not finished' do
             before do
-              allow(Elastic::DataMigrationService).to receive(:migration_has_finished?).with(:add_permissions_data_to_notes_documents).and_return(false)
+              set_elasticsearch_migration_to :add_permissions_data_to_notes_documents, including: false
             end
 
             it_behaves_like 'search respects visibility'
@@ -208,17 +194,13 @@ RSpec.describe Search::ProjectService do
         end
 
         with_them do
-          before do
-            allow(Elastic::DataMigrationService).to receive(:migration_has_finished?).and_call_original
-          end
-
           context 'when add_permissions_data_to_notes_documents migration is finished' do
             it_behaves_like 'search respects visibility'
           end
 
           context 'when add_permissions_data_to_notes_documents migration is not finished' do
             before do
-              allow(Elastic::DataMigrationService).to receive(:migration_has_finished?).with(:add_permissions_data_to_notes_documents).and_return(false)
+              set_elasticsearch_migration_to :add_permissions_data_to_notes_documents, including: false
             end
 
             it_behaves_like 'search respects visibility'

ee/spec/support/helpers/elasticsearch_helpers.rb

@@ -39,6 +39,26 @@ module ElasticsearchHelpers
     es_helper.refresh_index(index_name: es_helper.migrations_index_name)
   end
 
+  def set_elasticsearch_migration_to(name_or_version, including: true)
+    version = if name_or_version.is_a?(Numeric)
+                name_or_version
+              else
+                Elastic::DataMigrationService.find_by_name(name_or_version).version
+              end
+
+    Elastic::DataMigrationService.migrations.each do |migration|
+      return_value = if including
+                       migration.version <= version
+                     else
+                       migration.version < version
+                     end
+
+      allow(Elastic::DataMigrationService).to receive(:migration_has_finished?)
+        .with(migration.name_for_key.to_sym)
+        .and_return(return_value)
+    end
+  end
+
   def warm_elasticsearch_migrations_cache!
     ::Elastic::DataMigrationService.migrations.each do |migration|
       ::Elastic::DataMigrationService.migration_has_finished?(migration.name.underscore.to_sym)

ee/spec/support/shared_examples/services/search_query_applies_joins_based_on_migrations_shared_examples.rb

@@ -7,9 +7,7 @@ RSpec.shared_examples 'search query applies joins based on migrations shared exa
 
     context "when #{migration_name} migration is finished" do
       before do
-        allow(Elastic::DataMigrationService).to receive(:migration_has_finished?)
-                                                  .with(migration_name)
-                                                  .and_return(true)
+        set_elasticsearch_migration_to migration_name, including: true
       end
 
       it 'does not use joins to apply permissions' do
@@ -25,9 +23,7 @@ RSpec.shared_examples 'search query applies joins based on migrations shared exa
 
     context "when #{migration_name} migration is not finished" do
       before do
-        allow(Elastic::DataMigrationService).to receive(:migration_has_finished?)
-                                                  .with(migration_name)
-                                                  .and_return(false)
+        set_elasticsearch_migration_to migration_name, including: false
       end
 
       it 'uses joins to apply permissions' do