Denormalize notes permission data & update when permissions change

ee/app/models/ee/project.rb

@@ -105,6 +105,7 @@ module EE
       has_many :incident_management_oncall_schedules, class_name: 'IncidentManagement::OncallSchedule', inverse_of: :project
 
       elastic_index_dependant_association :issues, on_change: :visibility_level
+      elastic_index_dependant_association :notes, on_change: :visibility_level
 
       scope :with_shared_runners_limit_enabled, -> do
         if ::Ci::Runner.has_shared_runners_with_non_zero_public_cost?

ee/app/models/ee/project_feature.rb

@@ -5,6 +5,7 @@ module EE
     extend ActiveSupport::Concern
 
     EE_FEATURES = %i(requirements security_and_compliance).freeze
+    NOTES_PERMISSION_TRACKED_FIELDS = %w(issues_access_level repository_access_level merge_requests_access_level snippets_access_level).freeze
 
     prepended do
       set_available_features(EE_FEATURES)
@@ -14,7 +15,11 @@ module EE
         if project.maintaining_elasticsearch?
           project.maintain_elasticsearch_update
 
-          ElasticAssociationIndexerWorker.perform_async(self.project.class.name, project_id, ['issues']) if elasticsearch_project_associations_need_updating?
+          associations_to_update = []
+          associations_to_update << 'issues' if elasticsearch_project_issues_need_updating?
+          associations_to_update << 'notes' if elasticsearch_project_notes_need_updating?
+
+          ElasticAssociationIndexerWorker.perform_async(self.project.class.name, project_id, associations_to_update) if associations_to_update.any?
         end
       end
 
@@ -23,7 +28,11 @@ module EE
 
       private
 
-      def elasticsearch_project_associations_need_updating?
+      def elasticsearch_project_notes_need_updating?
+        self.previous_changes.keys.any? { |key| NOTES_PERMISSION_TRACKED_FIELDS.include?(key) }
+      end
+
+      def elasticsearch_project_issues_need_updating?
         self.previous_changes.key?(:issues_access_level)
       end
     end

ee/elastic/migrate/20210128163600_add_permissions_data_to_notes_documents.rb

+# frozen_string_literal: true
+
+class AddPermissionsDataToNotesDocuments < Elastic::Migration
+  batched!
+  throttle_delay 1.minute
+
+  QUERY_BATCH_SIZE = 5000
+  UPDATE_BATCH_SIZE = 100
+
+  def migrate
+    if completed?
+      log "Skipping adding permission data to notes documents migration since it is already applied"
+      return
+    end
+
+    log "Adding permission data to notes documents for batch of #{QUERY_BATCH_SIZE} documents"
+
+    # use filter query to prevent scores from being calculated
+    query = {
+      size: QUERY_BATCH_SIZE,
+      query: {
+        bool: {
+          filter: {
+            bool: {
+              must: note_type_filter,
+              should: [
+                field_does_not_exist('visibility_level'),
+                field_does_not_exist_for_type('issues_access_level', 'Issue'),
+                field_does_not_exist_for_type('repository_access_level', 'Commit'),
+                field_does_not_exist_for_type('merge_requests_access_level', 'MergeRequest'),
+                field_does_not_exist_for_type('snippets_access_level', 'Snippet')
+              ],
+              minimum_should_match: 1
+            }
+          }
+        }
+      }
+    }
+
+    results = client.search(index: helper.target_index_name, body: query)
+    hits = results.dig('hits', 'hits') || []
+
+    document_references = hits.map do |hit|
+      id = hit.dig('_source', 'id')
+      es_id = hit.dig('_id')
+      es_parent = hit.dig('_source', 'join_field', 'parent')
+
+      # ensure that any notes missing from the database will be removed from Elasticsearch
+      # as the data is back-filled
+      Gitlab::Elastic::DocumentReference.new(Note, id, es_id, es_parent)
+    end
+
+    document_references.each_slice(UPDATE_BATCH_SIZE) do |refs|
+      Elastic::ProcessBookkeepingService.track!(*refs)
+    end
+
+    log "Adding permission data to notes documents is completed for batch of #{document_references.size} documents"
+  end
+
+  def completed?
+    log "completed check: Refreshing #{helper.target_index_name}"
+    helper.refresh_index(index_name: helper.target_index_name)
+
+    query = {
+      size: 0,
+      query: note_type_filter,
+      aggs: {
+        notes: {
+          filter: {
+            bool: {
+              should: [
+                field_does_not_exist('visibility_level'),
+                field_does_not_exist_for_type('issues_access_level', 'Issue'),
+                field_does_not_exist_for_type('repository_access_level', 'Commit'),
+                field_does_not_exist_for_type('merge_requests_access_level', 'MergeRequest'),
+                field_does_not_exist_for_type('snippets_access_level', 'Snippet')
+              ],
+              minimum_should_match: 1
+            }
+          }
+        }
+      }
+    }
+
+    results = client.search(index: helper.target_index_name, body: query)
+    doc_count = results.dig('aggregations', 'notes', 'doc_count')
+    log "Migration has #{doc_count} documents remaining" if doc_count
+
+    doc_count && doc_count == 0
+  end
+
+  private
+
+  def note_type_filter
+    {
+      term: {
+        type: {
+          value: 'note'
+        }
+      }
+    }
+  end
+
+  def field_does_not_exist(field)
+    {
+      bool: {
+        must_not: {
+          exists: {
+            field: field
+          }
+        }
+      }
+    }
+  end
+
+  def field_does_not_exist_for_type(field, type)
+    query = field_does_not_exist(field)
+    query[:bool][:must] = {
+      term: {
+        noteable_type: {
+          value: type
+        }
+      }
+    }
+    query
+  end
+end

ee/lib/elastic/instance_proxy_util.rb

@@ -54,6 +54,15 @@ module Elastic
       nil
     end
 
+    # protect against missing project_feature and set visibility to PRIVATE
+    # if the project_feature is missing on a project
+    def safely_read_project_feature_for_elasticsearch(feature)
+      target.project.project_feature.issues_access_level
+    rescue NoMethodError => e
+      Gitlab::ErrorTracking.track_exception(e, project_id: target.project_id, id: target.id, class: target.class)
+      ProjectFeature::PRIVATE
+    end
+
     def apply_field_limit(result)
       return result unless result.is_a? String
 

ee/lib/elastic/latest/issue_instance_proxy.rb

@@ -14,15 +14,7 @@ module Elastic
 
         data['assignee_id'] = safely_read_attribute_for_elasticsearch(:assignee_ids)
         data['visibility_level'] = target.project.visibility_level
-
-        # protect against missing project_feature and set visibility to PRIVATE
-        # if the project_feature is missing on a project
-        begin
-          data['issues_access_level'] = target.project.project_feature.issues_access_level
-        rescue NoMethodError => e
-          Gitlab::ErrorTracking.track_exception(e, project_id: target.project_id, issue_id: target.id)
-          data['issues_access_level'] = ProjectFeature::PRIVATE
-        end
+        data['issues_access_level'] = safely_read_project_feature_for_elasticsearch(:issues_access_level)
 
         data.merge(generic_attributes)
       end

ee/lib/elastic/latest/note_instance_proxy.rb

@@ -14,14 +14,23 @@ module Elastic
           data[attr.to_s] = safely_read_attribute_for_elasticsearch(attr)
         end
 
-        if noteable.is_a?(Issue)
+        if for_issue?
           data['issue'] = {
             'assignee_id' => noteable.assignee_ids,
             'author_id' => noteable.author_id,
             'confidential' => noteable.confidential
           }
+          data['issues_access_level'] = safely_read_project_feature_for_elasticsearch(:issues_access_level)
+        elsif for_snippet?
+          data['snippets_access_level'] = safely_read_project_feature_for_elasticsearch(:snippets_access_level)
+        elsif for_merge_request?
+          data['merge_requests_access_level'] = safely_read_project_feature_for_elasticsearch(:merge_requests_access_level)
+        elsif for_commit?
+          data['repository_access_level'] = safely_read_project_feature_for_elasticsearch(:repository_access_level)
         end
 
+        data['visibility_level'] = target.project.visibility_level
+
         data.merge(generic_attributes)
       end
     end

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

+# frozen_string_literal: true
+
+require 'spec_helper'
+require File.expand_path('ee/elastic/migrate/20210128163600_add_permissions_data_to_notes_documents.rb')
+
+RSpec.describe AddPermissionsDataToNotesDocuments, :elastic, :sidekiq_inline do
+  let(:version) { 20210128163600 }
+  let(:migration) { described_class.new(version) }
+  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)
+    end
+  end
+
+  describe '#migrate' do
+    let!(:note_on_commit) { create(:note_on_commit) }
+    let!(:note_on_issue) { create(:note_on_issue) }
+    let!(:note_on_merge_request) { create(:note_on_merge_request) }
+    let!(:note_on_snippet) { create(:note_on_project_snippet) }
+
+    before do
+      ensure_elasticsearch_index!
+    end
+
+    context 'when migration is completed' do
+      it 'does not execute delete_by_query' do
+        expect(migration.completed?).to be_truthy
+        expect(::Elastic::ProcessBookkeepingService).not_to receive(:track!)
+
+        migration.migrate
+      end
+    end
+
+    context 'migration process' do
+      before do
+        remove_permission_data_for_notes([note_on_commit, note_on_issue, note_on_merge_request, note_on_snippet])
+      end
+
+      it 'adds permission data to the index' do
+        # track calls are batched in groups of 100
+        expect(::Elastic::ProcessBookkeepingService).to receive(:track!).once do |*tracked_refs|
+          expect(tracked_refs.count).to eq(4)
+        end
+
+        migration.migrate
+      end
+
+      it 'only updates note documents missing data', :aggregate_failures do
+        add_permission_data_for_notes([note_on_issue, note_on_snippet, note_on_merge_request])
+
+        expected = [Gitlab::Elastic::DocumentReference.new(Note, note_on_commit.id, note_on_commit.es_id, note_on_commit.es_parent)]
+        expect(::Elastic::ProcessBookkeepingService).to receive(:track!).with(*expected).once
+
+        migration.migrate
+      end
+
+      it 'processes in batches', :aggregate_failures do
+        stub_const("#{described_class}::QUERY_BATCH_SIZE", 2)
+        stub_const("#{described_class}::UPDATE_BATCH_SIZE", 1)
+
+        expect(::Elastic::ProcessBookkeepingService).to receive(:track!).exactly(4).times.and_call_original
+
+        migration.migrate
+
+        ensure_elasticsearch_index!
+
+        migration.migrate
+      end
+    end
+  end
+
+  describe '#completed?' do
+    let!(:note_on_commit) { create(:note_on_commit) }
+
+    before do
+      ensure_elasticsearch_index!
+    end
+
+    subject { migration.completed? }
+
+    context 'when documents are missing permissions data' do
+      before do
+        remove_permission_data_for_notes([note_on_commit])
+      end
+
+      it { is_expected.to be_falsey }
+    end
+
+    context 'when no documents are missing permissions data' do
+      it { is_expected.to be_truthy }
+    end
+
+    it 'refreshes the index' do
+      expect(helper).to receive(:refresh_index)
+
+      subject
+    end
+  end
+
+  private
+
+  def add_permission_data_for_notes(notes)
+    script =  {
+      source: "ctx._source['visibility_level'] = params.visibility_level; ctx._source['issues_access_level'] = params.visibility_level; ctx._source['merge_requests_access_level'] = params.visibility_level; ctx._source['snippets_access_level'] = params.visibility_level; ctx._source['repository_access_level'] = params.visibility_level;",
+      lang: "painless",
+      params: {
+        visibility_level: Gitlab::VisibilityLevel::PRIVATE
+      }
+    }
+
+    update_by_query(notes, script)
+  end
+
+  def remove_permission_data_for_notes(notes)
+    script =  {
+      source: "ctx._source.remove('visibility_level'); ctx._source.remove('repository_access_level'); ctx._source.remove('snippets_access_level'); ctx._source.remove('merge_requests_access_level'); ctx._source.remove('issues_access_level');"
+    }
+
+    update_by_query(notes, script)
+  end
+
+  def update_by_query(notes, script)
+    note_ids = notes.map { |i| i.id }
+
+    client = Note.__elasticsearch__.client
+    client.update_by_query({
+                             index: Note.__elasticsearch__.index_name,
+                             wait_for_completion: true, # run synchronously
+                             refresh: true, # make operation visible to search
+                             body: {
+                               script: script,
+                               query: {
+                                 bool: {
+                                   must: [
+                                     {
+                                       terms: {
+                                         id: note_ids
+                                       }
+                                     },
+                                     {
+                                       term: {
+                                         type: {
+                                           value: 'note'
+                                         }
+                                       }
+                                     }
+                                   ]
+                                 }
+                               }
+                             }
+                           })
+  end
+end

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

@@ -88,9 +88,12 @@ RSpec.describe Note, :elastic do
   it "returns json with all needed elements" do
     assignee = create(:user)
     issue = create(:issue, assignees: [assignee])
-    note = create(:note, noteable: issue, project: issue.project)
+    issue_note = create(:note, noteable: issue, project: issue.project)
+    commit_note = create(:note_on_commit, project: issue.project)
+    merge_request_note = create(:note_on_merge_request, project: issue.project)
+    snippet_note = create(:note_on_project_snippet, project: issue.project)
 
-    expected_hash = note.attributes.extract!(
+    expected_hash = issue_note.attributes.extract!(
       'id',
       'note',
       'project_id',
@@ -105,14 +108,19 @@ RSpec.describe Note, :elastic do
         'author_id' => issue.author_id,
         'confidential' => issue.confidential
       },
-      'type' => note.es_type,
+      'type' => issue_note.es_type,
       'join_field' => {
-        'name' => note.es_type,
-        'parent' => note.es_parent
-      }
+        'name' => issue_note.es_type,
+        'parent' => issue_note.es_parent
+      },
+      'visibility_level' => issue.project.visibility_level,
+      'issues_access_level' => issue.project.issues_access_level
     })
 
-    expect(note.__elasticsearch__.as_indexed_json).to eq(expected_hash)
+    expect(issue_note.__elasticsearch__.as_indexed_json).to eq(expected_hash)
+    expect(commit_note.__elasticsearch__.as_indexed_json).to have_key('repository_access_level')
+    expect(merge_request_note.__elasticsearch__.as_indexed_json).to have_key('merge_requests_access_level')
+    expect(snippet_note.__elasticsearch__.as_indexed_json).to have_key('snippets_access_level')
   end
 
   it 'does not track system note updates' do

ee/spec/models/project_feature_spec.rb

@@ -28,13 +28,14 @@ RSpec.describe ProjectFeature do
       allow(project).to receive(:maintaining_elasticsearch?).and_return(true)
     end
 
-    where(:feature, :worker_expected) do
-      'issues'          | true
-      'wiki'            | false
-      'builds'          | false
-      'merge_requests'  | false
-      'repository'      | false
-      'pages'           | false
+    where(:feature, :worker_expected, :associations) do
+      'issues'          | true        | %w[issues notes]
+      'wiki'            | false       | nil
+      'builds'          | false       | nil
+      'merge_requests'  | true        | %w[notes]
+      'repository'      | true        | %w[notes]
+      'snippets'        | true        | %w[notes]
+      'pages'           | false       | nil
     end
 
     with_them do
@@ -42,7 +43,7 @@ RSpec.describe ProjectFeature do
         expect(project).to receive(:maintain_elasticsearch_update)
 
         if worker_expected
-          expect(ElasticAssociationIndexerWorker).to receive(:perform_async).with('Project', project.id, ['issues'])
+          expect(ElasticAssociationIndexerWorker).to receive(:perform_async).with('Project', project.id, associations)
         else
           expect(ElasticAssociationIndexerWorker).not_to receive(:perform_async)
         end

ee/spec/models/project_spec.rb

@@ -2638,8 +2638,8 @@ RSpec.describe Project do
       let!(:issue) { create(:issue, project: project) }
 
       context 'when updating the visibility_level' do
-        it 'triggers ElasticAssociationIndexerWorker to update issues' do
-          expect(ElasticAssociationIndexerWorker).to receive(:perform_async).with('Project', project.id, ['issues'])
+        it 'triggers ElasticAssociationIndexerWorker to update issues and notes' do
+          expect(ElasticAssociationIndexerWorker).to receive(:perform_async).with('Project', project.id, %w[issues notes])
 
           project.update!(visibility_level: Gitlab::VisibilityLevel::PRIVATE)
         end

ee/spec/services/groups/transfer_service_spec.rb

@@ -67,11 +67,11 @@ RSpec.describe Groups::TransferService, '#execute' do
 
           expect(::Gitlab::CurrentSettings).not_to receive(:invalidate_elasticsearch_indexes_cache_for_project!)
           expect(Elastic::ProcessBookkeepingService).to receive(:track!).with(project1)
-          expect(ElasticAssociationIndexerWorker).to receive(:perform_async).with('Project', project1.id, ['issues'])
+          expect(ElasticAssociationIndexerWorker).to receive(:perform_async).with('Project', project1.id, %w[issues notes])
           expect(Elastic::ProcessBookkeepingService).to receive(:track!).with(project2)
-          expect(ElasticAssociationIndexerWorker).to receive(:perform_async).with('Project', project2.id, ['issues'])
+          expect(ElasticAssociationIndexerWorker).to receive(:perform_async).with('Project', project2.id, %w[issues notes])
           expect(Elastic::ProcessBookkeepingService).not_to receive(:track!).with(project3)
-          expect(ElasticAssociationIndexerWorker).not_to receive(:perform_async).with('Project', project3.id, ['issues'])
+          expect(ElasticAssociationIndexerWorker).not_to receive(:perform_async).with('Project', project3.id, %w[issues notes])
 
           transfer_service.execute(new_group)
 

ee/spec/services/projects/transfer_service_spec.rb

@@ -72,9 +72,9 @@ RSpec.describe Projects::TransferService do
     context 'when visibility level changes' do
       let_it_be(:group) { create(:group, :private) }
 
-      it 'reindexes the project and associated issues' do
+      it 'reindexes the project and associated issues and notes' do
         expect(Elastic::ProcessBookkeepingService).to receive(:track!).with(project)
-        expect(ElasticAssociationIndexerWorker).to receive(:perform_async).with('Project', project.id, ['issues'])
+        expect(ElasticAssociationIndexerWorker).to receive(:perform_async).with('Project', project.id, %w[issues notes])
 
         subject.execute(group)
       end