Use elasticsearch routing for querying

When using advanced search in the scope of a project we're querying
all of the elasticsearch shards instead of one. To fix that we started
using routing parameter when we search inside of a project

ee/lib/elastic/latest/application_class_proxy.rb

@@ -4,6 +4,14 @@ module Elastic
   module Latest
     class ApplicationClassProxy < Elasticsearch::Model::Proxy::ClassMethodsProxy
       include ClassProxyUtil
+      include Elastic::Latest::Routing
+
+      def search(query, search_options = {})
+        es_options = routing_options(search_options)
+
+        # Calling elasticsearch-ruby method
+        super(query, es_options)
+      end
 
       def es_type
         target.name.underscore

ee/lib/elastic/latest/git_class_proxy.rb

@@ -22,6 +22,10 @@ module Elastic
 
       private
 
+      def extract_repository_ids(options)
+        [options[:repository_id]].flatten
+      end
+
       def search_commit(query, page: 1, per: 20, options: {})
         page ||= 1
 
@@ -49,10 +53,11 @@ module Elastic
           query_hash[:track_scores] = true
         end
 
-        if options[:repository_id]
+        repository_ids = extract_repository_ids(options)
+        if repository_ids.any?
           query_hash[:query][:bool][:filter] << {
             terms: {
-              'commit.rid' => [options[:repository_id]].flatten
+              'commit.rid' => repository_ids
             }
           }
         end
@@ -77,7 +82,7 @@ module Elastic
 
         query_hash[:sort] = [:_score]
 
-        res = search(query_hash)
+        res = search(query_hash, options)
         {
           results: res.results,
           total_count: res.size
@@ -114,10 +119,11 @@ module Elastic
 
         query_hash[:query][:bool][:filter] += query.elasticsearch_filters(:blob)
 
-        if options[:repository_id]
+        repository_ids = extract_repository_ids(options)
+        if repository_ids.any?
           query_hash[:query][:bool][:filter] << {
             terms: {
-              'blob.rid' => [options[:repository_id]].flatten
+              'blob.rid' => repository_ids
             }
           }
         end
@@ -150,7 +156,9 @@ module Elastic
           }
         end
 
-        res = search(query_hash)
+        options[:project_ids] = repository_ids.map { |id| id.to_s[/\d+/].to_i } if type == :wiki_blob && repository_ids.any?
+
+        res = search(query_hash, options)
 
         {
           results: res.results,

ee/lib/elastic/latest/issue_class_proxy.rb

@@ -15,7 +15,7 @@ module Elastic
         query_hash = project_ids_filter(query_hash, options)
         query_hash = confidentiality_filter(query_hash, options[:current_user])
 
-        search(query_hash)
+        search(query_hash, options)
       end
 
       private

ee/lib/elastic/latest/merge_request_class_proxy.rb

@@ -14,7 +14,7 @@ module Elastic
         options[:features] = 'merge_requests'
         query_hash = project_ids_filter(query_hash, options)
 
-        search(query_hash)
+        search(query_hash, options)
       end
     end
   end

ee/lib/elastic/latest/milestone_class_proxy.rb

@@ -10,7 +10,7 @@ module Elastic
 
         query_hash = project_ids_filter(query_hash, options)
 
-        search(query_hash)
+        search(query_hash, options)
       end
     end
   end

ee/lib/elastic/latest/note_class_proxy.rb

@@ -23,7 +23,7 @@ module Elastic
 
         query_hash[:highlight] = highlight_options(options[:in])
 
-        search(query_hash)
+        search(query_hash, options)
       end
 
       private

ee/lib/elastic/latest/project_class_proxy.rb

@@ -44,7 +44,7 @@ module Elastic
 
         query_hash[:sort] = [:_score]
 
-        search(query_hash)
+        search(query_hash, options)
       end
     end
   end

ee/lib/elastic/latest/routing.rb

+# frozen_string_literal: true
+
+module Elastic
+  module Latest
+    module Routing
+      extend ActiveSupport::Concern
+
+      def routing_options(options)
+        return {} if Feature.disabled?(:elasticsearch_use_routing)
+        return {} if options[:public_and_internal_projects]
+
+        ids = if options[:project_id]
+                [options[:project_id]]
+              elsif options[:project_ids]
+                options[:project_ids]
+              elsif options[:repository_id]
+                [options[:repository_id]]
+              else
+                []
+              end
+
+        return {} if ids == :any
+
+        routing = build_routing(ids)
+
+        return {} if routing.blank?
+
+        { routing: routing }
+      end
+
+      private
+
+      def build_routing(ids)
+        ids.map { |id| "project_#{id}" }.join(',')
+      end
+    end
+  end
+end

ee/lib/elastic/latest/snippet_class_proxy.rb

@@ -7,14 +7,14 @@ module Elastic
         query_hash = basic_query_hash(%w(title file_name), query)
         query_hash = filter(query_hash, options)
 
-        search(query_hash)
+        search(query_hash, options)
       end
 
       def elastic_search_code(query, options: {})
         query_hash = basic_query_hash(%w(content), query)
         query_hash = filter(query_hash, options)
 
-        search(query_hash)
+        search(query_hash, options)
       end
 
       def es_type

ee/lib/elastic/v12p1/routing.rb

+# frozen_string_literal: true
+
+module Elastic
+  module V12p1
+    Routing = Elastic::Latest::Routing
+  end
+end

ee/lib/gitlab/elastic/search_results.rb

@@ -226,9 +226,12 @@ module Gitlab
         return Kaminari.paginate_array([]) if query.blank?
 
         strong_memoize(:blobs) do
-          opt = {
-            additional_filter: repository_filter
-          }
+          project_ids = visible_project_ids
+
+          opt = base_options.merge(
+            additional_filter: repository_filter(project_ids),
+            project_ids: project_ids
+          )
 
           Repository.elastic_search(
             query,
@@ -242,9 +245,12 @@ module Gitlab
         return Kaminari.paginate_array([]) if query.blank?
 
         strong_memoize(:wiki_blobs) do
-          opt = {
-            additional_filter: wiki_filter
-          }
+          project_ids = visible_project_ids(visible_for_guests: true)
+
+          opt = base_options.merge(
+            additional_filter: wiki_filter(project_ids),
+            project_ids: project_ids
+          )
 
           ProjectWiki.elastic_search(
             query,
@@ -264,9 +270,12 @@ module Gitlab
         return Kaminari.paginate_array([]) if query.blank?
 
         strong_memoize(:commits) do
-          options = {
-            additional_filter: repository_filter
-          }
+          project_ids = visible_project_ids
+
+          options = base_options.merge(
+            additional_filter: repository_filter(project_ids),
+            project_ids: project_ids
+          )
 
           Repository.find_commits_by_message_with_elastic(
             query,
@@ -277,16 +286,19 @@ module Gitlab
         end
       end
 
-      def wiki_filter
-        blob_filter(:wiki, visible_for_guests: true)
+      def wiki_filter(project_ids)
+        blob_filter(:wiki, project_ids)
+      end
+
+      def repository_filter(project_ids)
+        blob_filter(:repository, project_ids)
       end
 
-      def repository_filter
-        blob_filter(:repository)
+      def visible_project_ids(visible_for_guests: false)
+        visible_for_guests ? limit_project_ids : non_guest_project_ids
       end
 
-      def blob_filter(feature, visible_for_guests: false)
-        project_ids = visible_for_guests ? limit_project_ids : non_guest_project_ids
+      def blob_filter(feature, project_ids)
         key_name = "#{feature}_access_level"
 
         conditions =

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

+# frozen_string_literal: true
+
+require 'fast_spec_helper'
+require 'support/helpers/stub_feature_flags'
+
+describe Elastic::Latest::Routing do
+  let(:proxified_class) { Issue }
+  let(:included_class) { Elastic::Latest::ApplicationClassProxy }
+
+  subject { included_class.new(proxified_class) }
+
+  let(:project_ids) { [1, 2, 3] }
+  let(:project_routing) { 'project_1,project_2,project_3' }
+
+  describe '#search' do
+    it 'calls routing_options with empty hash' do
+      expect(subject).to receive(:routing_options).and_return({})
+
+      subject.search('q')
+    end
+
+    it 'calls routing_options with correct routing' do
+      expect(subject).to receive(:routing_options).and_return({ routing: project_routing })
+
+      subject.search('q', project_ids: project_ids)
+    end
+  end
+
+  describe '#routing_options' do
+    include StubFeatureFlags
+
+    context 'when feature flag is enabled' do
+      before do
+        stub_feature_flags(elasticsearch_use_routing: true)
+      end
+
+      it 'returns correct options for project_id' do
+        expect(subject.routing_options({ project_id: 1 })).to eq({ routing: 'project_1' })
+      end
+
+      it 'returns correct options for repository_id' do
+        expect(subject.routing_options({ repository_id: 1 })).to eq({ routing: 'project_1' })
+      end
+
+      it 'returns correct options for project_ids' do
+        expect(subject.routing_options({ project_ids: project_ids })).to eq({ routing: project_routing })
+      end
+
+      it 'returns empty hash when provided an empty array' do
+        expect(subject.routing_options({ project_ids: [] })).to eq({})
+      end
+
+      it 'returns empty hash when provided :any to project_ids' do
+        expect(subject.routing_options({ project_ids: :any })).to eq({})
+      end
+
+      it 'returns empty hash when public projects flag is passed' do
+        expect(subject.routing_options({ project_ids: project_ids, public_and_internal_projects: true })).to eq({})
+      end
+
+      it 'uses project_ids rather than repository_id when both are supplied' do
+        options = { project_ids: project_ids, repository_id: 'wiki_5' }
+
+        expect(subject.routing_options(options)).to eq({ routing: project_routing })
+      end
+    end
+
+    context 'when feature flag is disabled' do
+      before do
+        stub_feature_flags(elasticsearch_use_routing: false)
+      end
+
+      it 'returns empty hash for project_ids' do
+        expect(subject.routing_options({ project_ids: project_ids })).to eq({})
+      end
+
+      it 'returns empty hash for empty options' do
+        expect(subject.routing_options({})).to eq({})
+      end
+    end
+  end
+end