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,35 +69,55 @@ RSpec.describe MergeRequest, :elastic do
     expect(results.first.title).to eq('bla-bla merge request')
   end
 
-  it "returns json with all needed elements" do
-    merge_request = create :merge_request
-    merge_request.project.update!(visibility_level: Gitlab::VisibilityLevel::INTERNAL)
-    create(:award_emoji, :upvote, awardable: merge_request)
-
-    expected_hash = merge_request.attributes.extract!(
-      'id',
-      'iid',
-      'target_branch',
-      'source_branch',
-      'title',
-      'description',
-      'created_at',
-      'updated_at',
-      'state',
-      'merge_status',
-      'source_project_id',
-      'target_project_id',
-      'author_id'
-    ).merge({
-      'state' => merge_request.state,
-      'type' => merge_request.es_type,
-      'merge_requests_access_level' => ProjectFeature::ENABLED,
-      'visibility_level' => Gitlab::VisibilityLevel::INTERNAL,
-      'project_id' => merge_request.target_project.id,
-      'upvotes' => 1
-    })
-
-    expect(merge_request.__elasticsearch__.as_indexed_json).to eq(expected_hash)
+  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)
+      create(:award_emoji, :upvote, awardable: merge_request)
+
+      expected_hash = merge_request.attributes.extract!(
+        'id',
+        'iid',
+        'target_branch',
+        'source_branch',
+        'title',
+        'description',
+        'created_at',
+        'updated_at',
+        'state',
+        'merge_status',
+        'source_project_id',
+        'target_project_id',
+        'author_id'
+      ).merge({
+        'state' => merge_request.state,
+        'type' => merge_request.es_type,
+        'merge_requests_access_level' => ProjectFeature::ENABLED,
+        'visibility_level' => Gitlab::VisibilityLevel::INTERNAL,
+        'project_id' => merge_request.target_project.id,
+        '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