Fix AddUpvotesToMergeRequests migration

EE: true
Changelog: fixed

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

@@ -5,14 +5,6 @@ module Elastic
     UPDATE_BATCH_SIZE = 100
 
     def migrate
-      # only create mappings on the first run
-      if respond_to?(:new_mappings) && migration_state[:update_mappings].blank?
-        update_mapping!(index_name, { properties: new_mappings })
-
-        options = { update_mappings: true }
-        set_migration_state(options)
-      end
-
       if completed?
         log "Skipping adding #{field_name} field to #{index_name} documents migration since it is already applied"
         return
@@ -109,9 +101,5 @@ module Elastic
 
       UPDATE_BATCH_SIZE
     end
-
-    def update_mapping!(index_name, mappings)
-      helper.update_mapping(index_name: index_name, mappings: mappings)
-    end
   end
 end

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

+# frozen_string_literal: true
+
+module Elastic
+  module MigrationUpdateMappingsHelper
+    def migrate
+      if completed?
+        log "Skipping updating #{index_name} mapping migration since it is already applied"
+        return
+      end
+
+      log "Adding #{new_mappings.inspect} to #{index_name} mapping"
+      update_mapping!(index_name, { properties: new_mappings })
+    end
+
+    def completed?
+      helper.refresh_index(index_name: index_name)
+
+      mappings = helper.get_mapping(index_name: index_name)
+
+      # Check if mappings include all new_mappings
+      new_mappings.keys.map(&:to_s).to_set.subset?(mappings.keys.to_set)
+    end
+
+    private
+
+    def index_name
+      raise NotImplementedError
+    end
+
+    def new_mappings
+      raise NotImplementedError
+    end
+
+    def update_mapping!(index_name, mappings)
+      helper.update_mapping(index_name: index_name, mappings: mappings)
+    end
+  end
+end

ee/elastic/migrate/20210722112500_add_upvotes_mappings_to_merge_requests.rb

+# frozen_string_literal: true
+
+class AddUpvotesMappingsToMergeRequests < Elastic::Migration
+  include Elastic::MigrationUpdateMappingsHelper
+
+  DOCUMENT_TYPE = MergeRequest
+
+  private
+
+  def index_name
+    DOCUMENT_TYPE.__elasticsearch__.index_name
+  end
+
+  def new_mappings
+    { upvotes: { type: 'integer' } }
+  end
+end

ee/elastic/migrate/20210722112600_add_upvotes_to_merge_requests.rb

@@ -9,10 +9,6 @@ class AddUpvotesToMergeRequests < Elastic::Migration
 
   DOCUMENT_TYPE = MergeRequest
 
-  def new_mappings
-    { upvotes: { type: 'integer' } }
-  end
-
   private
 
   def index_name

ee/elastic/migrate/20210813134600_add_namespace_ancestry_to_issues_mapping.rb

 # frozen_string_literal: true
 
 class AddNamespaceAncestryToIssuesMapping < Elastic::Migration
-  include Elastic::MigrationBackfillHelper
+  include Elastic::MigrationUpdateMappingsHelper
 
   DOCUMENT_KLASS = Issue
 

ee/lib/elastic/latest/merge_request_instance_proxy.rb

@@ -29,7 +29,7 @@ module Elastic
 
         data['visibility_level'] = target.project.visibility_level
         data['merge_requests_access_level'] = safely_read_project_feature_for_elasticsearch(:merge_requests)
-        data['upvotes'] = target.lazy_upvotes_count
+        data['upvotes'] = target.lazy_upvotes_count if Elastic::DataMigrationService.migration_has_finished?(:add_upvotes_mappings_to_merge_requests)
 
         data.merge(generic_attributes)
       end

ee/spec/elastic/migrate/20210722112500_add_upvotes_mappings_to_merge_requests_spec.rb

+# frozen_string_literal: true
+require 'spec_helper'
+require File.expand_path('ee/elastic/migrate/20210722112500_add_upvotes_mappings_to_merge_requests.rb')
+
+RSpec.describe AddUpvotesMappingsToMergeRequests, :elastic, :sidekiq_inline do
+  let(:version) { 20210722112500 }
+  let(:migration) { described_class.new(version) }
+  let(:helper) { Gitlab::Elastic::Helper.new }
+
+  before do
+    allow(migration).to receive(:helper).and_return(helper)
+  end
+
+  describe '.migrate' do
+    subject { migration.migrate }
+
+    context 'when migration is already completed' do
+      it 'does not modify data' do
+        expect(helper).not_to receive(:update_mapping)
+
+        subject
+      end
+    end
+
+    context 'migration process' do
+      before do
+        allow(helper).to receive(:get_mapping).and_return({})
+      end
+
+      it 'updates the issues index mappings' do
+        expect(helper).to receive(:update_mapping)
+
+        subject
+      end
+    end
+  end
+
+  describe '.completed?' do
+    context 'mapping has been updated' do
+      specify { expect(migration).to be_completed }
+    end
+
+    context 'mapping has not been updated' do
+      before do
+        allow(helper).to receive(:get_mapping).and_return({})
+      end
+
+      specify { expect(migration).not_to be_completed }
+    end
+  end
+end

ee/spec/elastic/migrate/20210722112600_add_upvotes_to_merge_requests_spec.rb

@@ -73,25 +73,6 @@ RSpec.describe AddUpvotesToMergeRequests, :elastic, :sidekiq_inline do
 
         migration.migrate
       end
-
-      it 'only updates mappings on the first run' do
-        helper = Gitlab::Elastic::Helper.new
-        allow(Gitlab::Elastic::Helper).to receive(:default).and_return(helper)
-        allow(migration).to receive(:batch_size).and_return(2)
-
-        expect(helper).to receive(:update_mapping).and_call_original
-
-        # cannot use subject in spec because it is memoized
-        migration.migrate
-
-        ensure_elasticsearch_index!
-
-        expect(migration.migration_state).to include(update_mappings: true)
-
-        expect(helper).not_to receive(:update_mapping)
-
-        migration.migrate
-      end
     end
   end
 

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

@@ -69,6 +69,7 @@ RSpec.describe MergeRequest, :elastic do
     expect(results.first.title).to eq('bla-bla merge request')
   end
 
+  shared_examples 'merge request json fields' do |remove_fields|
     it "returns json with all needed elements" do
       merge_request = create :merge_request
       merge_request.project.update!(visibility_level: Gitlab::VisibilityLevel::INTERNAL)
@@ -97,8 +98,27 @@ RSpec.describe MergeRequest, :elastic do
         'upvotes' => 1
       })
 
+      expected_hash = expected_hash.except(*remove_fields.map(&:to_s))
+
       expect(merge_request.__elasticsearch__.as_indexed_json).to eq(expected_hash)
     end
+  end
+
+  context 'add_upvotes_mappings_to_merge_requests is not completed' do
+    before do
+      set_elasticsearch_migration_to(:add_upvotes_mappings_to_merge_requests, including: false)
+    end
+
+    include_context 'merge request json fields', [:upvotes]
+  end
+
+  context 'add_upvotes_mappings_to_merge_requests is completed' do
+    before do
+      set_elasticsearch_migration_to(:add_upvotes_mappings_to_merge_requests, including: true)
+    end
+
+    include_context 'merge request json fields', []
+  end
 
   it 'handles when a project is missing project_feature' do
     merge_request = create :merge_request