Expose blob search aggregations to SearchController

app/controllers/search_controller.rb

@@ -46,6 +46,7 @@ class SearchController < ApplicationController
     @search_results = @search_service.search_results
     @search_objects = @search_service.search_objects
     @search_highlight = @search_service.search_highlight
+    @aggregations = @search_service.search_aggregations
 
     increment_search_counters
   end

app/services/search_service.rb

@@ -75,6 +75,10 @@ class SearchService
     search_results.highlight_map(scope)
   end
 
+  def search_aggregations
+    search_results.aggregations(scope)
+  end
+
   private
 
   def page

ee/lib/elastic/latest/git_class_proxy.rb

@@ -36,6 +36,12 @@ module Elastic
         end
       end
 
+      def blob_aggregations
+        return @aggregations if defined?(@aggregations) # rubocop:disable Gitlab/ModuleWithInstanceVariables
+
+        []
+      end
+
       private
 
       def options_filter_context(type, options)
@@ -201,7 +207,7 @@ module Elastic
           }
         end
 
-        if type == 'blob' && !options[:count_only] && ::Feature.enabled?(:search_blobs_language_aggregation, options[:current_user], default_enabled: :yaml)
+        if include_aggregations?(type, options[:count_only], options[:current_user])
           query_hash[:aggs] = {
             language: {
               composite: {
@@ -243,6 +249,9 @@ module Elastic
           options: options
         )[type.pluralize.to_sym][:results]
 
+        # Retrieve aggregations for blob type queries
+        @aggregations ||= ::Gitlab::Search::AggregationParser.call(response.response.aggregations) if include_aggregations?(type, options[:count_only], options[:current_user]) # rubocop:disable Gitlab/ModuleWithInstanceVariables
+
         items, total_count = yield_each_search_result(response, type, preload_method, &blk)
 
         # Before "map" we had a paginated array so we need to recover it
@@ -301,6 +310,10 @@ module Elastic
           end
         end
       end
+
+      def include_aggregations?(type, count_only, current_user)
+        type == 'blob' && !count_only && ::Feature.enabled?(:search_blobs_language_aggregation, current_user, default_enabled: :yaml)
+      end
     end
   end
 end

ee/lib/elastic/latest/git_instance_proxy.rb

@@ -28,6 +28,10 @@ module Elastic
         self.class.elastic_search_as_found_blob(query, page: page, per: per, options: options, preload_method: preload_method)
       end
 
+      def blob_aggregations
+        self.class.blob_aggregations
+      end
+
       def delete_index_for_commits_and_blobs(wiki: false)
         types =
           if wiki

ee/lib/gitlab/elastic/project_search_results.rb

@@ -18,8 +18,8 @@ module Gitlab
       private
 
       def blobs(page: 1, per_page: DEFAULT_PER_PAGE, count_only: false, preload_method: nil)
-        return Kaminari.paginate_array([]) unless Ability.allowed?(@current_user, :download_code, project)
         return Kaminari.paginate_array([]) if project.empty_repo? || query.blank?
+        return Kaminari.paginate_array([]) unless Ability.allowed?(@current_user, :download_code, project)
 
         strong_memoize(memoize_key(:blobs, count_only: count_only)) do
           project.repository.__elasticsearch__.elastic_search_as_found_blob(
@@ -33,8 +33,8 @@ module Gitlab
       end
 
       def wiki_blobs(page: 1, per_page: DEFAULT_PER_PAGE, count_only: false)
-        return Kaminari.paginate_array([]) unless Ability.allowed?(@current_user, :read_wiki, project)
         return Kaminari.paginate_array([]) unless project.wiki_enabled? && !project.wiki.empty? && query.present?
+        return Kaminari.paginate_array([]) unless Ability.allowed?(@current_user, :read_wiki, project)
 
         strong_memoize(memoize_key(:wiki_blobs, count_only: count_only)) do
           project.wiki.__elasticsearch__.elastic_search_as_wiki_page(
@@ -60,8 +60,8 @@ module Gitlab
       end
 
       def commits(page: 1, per_page: DEFAULT_PER_PAGE, preload_method: nil, count_only: false)
-        return Kaminari.paginate_array([]) unless Ability.allowed?(@current_user, :download_code, project)
         return Kaminari.paginate_array([]) if project.empty_repo? || query.blank?
+        return Kaminari.paginate_array([]) unless Ability.allowed?(@current_user, :download_code, project)
 
         strong_memoize(memoize_key(:commits, count_only: count_only)) do
           project.repository.find_commits_by_message_with_elastic(
@@ -73,6 +73,15 @@ module Gitlab
           )
         end
       end
+
+      def blob_aggregations
+        return [] if project.empty_repo? || query.blank?
+        return [] unless Ability.allowed?(@current_user, :download_code, project)
+
+        strong_memoize(:blob_aggregations) do
+          project.repository.__elasticsearch__.blob_aggregations
+        end
+      end
     end
   end
 end

ee/lib/gitlab/elastic/search_results.rb

@@ -214,6 +214,15 @@ module Gitlab
         )
       end
 
+      def aggregations(scope)
+        case scope
+        when 'blobs'
+          blob_aggregations
+        else
+          []
+        end
+      end
+
       private
 
       # Apply some eager loading to the `records` of an ES result object without
@@ -351,6 +360,14 @@ module Gitlab
           number_with_delimiter(count)
         end
       end
+
+      def blob_aggregations
+        return [] if query.blank?
+
+        strong_memoize(:blob_aggregations) do
+          Repository.__elasticsearch__.blob_aggregations
+        end
+      end
     end
   end
 end

ee/lib/gitlab/search/aggregation.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Search
+    class Aggregation
+      attr_reader :name, :buckets
+
+      def initialize(name, elastic_aggregation_buckets)
+        @name = name
+        @buckets = parse_buckets(elastic_aggregation_buckets)
+      end
+
+      private
+
+      def parse_buckets(buckets)
+        return [] unless buckets
+
+        buckets.map do |b|
+          { key: b['key'], count: b['doc_count'] }
+        end
+      end
+    end
+  end
+end

ee/lib/gitlab/search/aggregation_parser.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Search
+    class AggregationParser
+      def self.call(aggregations)
+        return [] unless aggregations
+
+        aggregations.keys.map do |key|
+          ::Gitlab::Search::Aggregation.new(key, aggregations[key].buckets)
+        end
+      end
+    end
+  end
+end

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

@@ -13,7 +13,7 @@ RSpec.describe Elastic::Latest::GitClassProxy, :elastic do
     before do
       stub_ee_application_setting(elasticsearch_search: true, elasticsearch_indexing: true)
 
-      Gitlab::Elastic::Indexer.new(project).run
+      project.repository.index_commits_and_blobs
       ensure_elasticsearch_index!
     end
 
@@ -33,6 +33,57 @@ RSpec.describe Elastic::Latest::GitClassProxy, :elastic do
     end
   end
 
+  describe '#blob_aggregations', :sidekiq_inline do
+    before do
+      stub_ee_application_setting(elasticsearch_search: true, elasticsearch_indexing: true)
+
+      project.repository.index_commits_and_blobs
+      ensure_elasticsearch_index!
+    end
+
+    it 'returns aggregations' do
+      subject.elastic_search_as_found_blob('This guide details how contribute to GitLab')
+      result = subject.blob_aggregations
+
+      expect(result.first.name).to eq('language')
+      expect(result.first.buckets.first[:key]).to eq({ 'language' => 'Markdown' })
+      expect(result.first.buckets.first[:count]).to eq(2)
+    end
+
+    context 'when search_blobs_language_aggregation feature flag is disabled' do
+      before do
+        stub_feature_flags(search_blobs_language_aggregation: false)
+      end
+
+      it 'returns empty array' do
+        subject.elastic_search_as_found_blob('This guide details how contribute to GitLab')
+        result = subject.blob_aggregations
+
+        expect(result).to match_array([])
+      end
+    end
+
+    context 'when search type is not blobs' do
+      let(:included_class) { Elastic::Latest::ProjectWikiClassProxy }
+
+      it 'returns empty array' do
+        subject.elastic_search_as_found_blob('This guide details how contribute to GitLab')
+        result = subject.blob_aggregations
+
+        expect(result).to match_array([])
+      end
+    end
+
+    context 'when count_only search' do
+      it 'returns empty array' do
+        subject.elastic_search_as_found_blob('This guide details how contribute to GitLab', options: { count_only: true })
+        result = subject.blob_aggregations
+
+        expect(result).to match_array([])
+      end
+    end
+  end
+
   it "names elasticsearch queries" do
     subject.elastic_search_as_found_blob('*')
 

ee/spec/lib/gitlab/elastic/project_search_results_spec.rb

@@ -124,4 +124,58 @@ RSpec.describe Gitlab::Elastic::ProjectSearchResults, :elastic, :clean_gitlab_re
     include_examples 'does not hit Elasticsearch twice for objects and counts', %w[notes blobs wiki_blobs commits issues merge_requests milestones]
     include_examples 'does not load results for count only queries', %w[notes blobs wiki_blobs commits issues merge_requests milestones]
   end
+
+  describe '#aggregations' do
+    using RSpec::Parameterized::TableSyntax
+
+    subject { described_class.new(user, query, project: project).aggregations(scope) }
+
+    where(:scope, :expected) do
+      'milestones'     | []
+      'notes'          | []
+      'issues'         | []
+      'merge_requests' | []
+      'wiki_blobs'     | []
+      'commits'        | []
+      'users'          | []
+      'unknown'        | []
+      'blobs'          | [::Gitlab::Search::Aggregation.new('language', nil)]
+    end
+
+    with_them do
+      before do
+        allow(project.repository.__elasticsearch__).to receive(:blob_aggregations).and_return(expected) if scope == 'blobs'
+      end
+
+      it_behaves_like 'loads aggregations'
+    end
+
+    context 'project search specific gates for blob scope' do
+      let(:scope) { 'blobs' }
+
+      context 'when query is blank' do
+        let(:query) { nil }
+
+        it 'returns the an empty array' do
+          expect(subject).to match_array([])
+        end
+      end
+
+      context 'when project has an empty repository' do
+        it 'returns an empty array' do
+          allow(project).to receive(:empty_repo?).and_return(true)
+
+          expect(subject).to match_array([])
+        end
+      end
+
+      context 'when user does not have download_code permission on project' do
+        it 'returns an empty array' do
+          allow(Ability).to receive(:allowed?).with(user, :download_code, project).and_return(false)
+
+          expect(subject).to match_array([])
+        end
+      end
+    end
+  end
 end

ee/spec/lib/gitlab/elastic/search_results_spec.rb

@@ -63,6 +63,34 @@ RSpec.describe Gitlab::Elastic::SearchResults, :elastic, :clean_gitlab_redis_sha
     end
   end
 
+  describe '#aggregations' do
+    using RSpec::Parameterized::TableSyntax
+
+    subject { described_class.new(user, query, limit_project_ids).aggregations(scope) }
+
+    where(:scope, :expected) do
+      'projects'       | []
+      'milestones'     | []
+      'notes'          | []
+      'issues'         | []
+      'merge_requests' | []
+      'wiki_blobs'     | []
+      'commits'        | []
+      'users'          | []
+      'epics'          | []
+      'unknown'        | []
+      'blobs'          | [::Gitlab::Search::Aggregation.new('language', nil)]
+    end
+
+    with_them do
+      before do
+        allow(Repository.__elasticsearch__).to receive(:blob_aggregations).and_return(expected) if scope == 'blobs'
+      end
+
+      it_behaves_like 'loads aggregations'
+    end
+  end
+
   shared_examples_for 'a paginated object' do |object_type|
     let(:results) { described_class.new(user, 'hello world', limit_project_ids) }
 

ee/spec/lib/gitlab/search/aggregation_parser_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Search::AggregationParser do
+  let(:search) do
+    Elasticsearch::Model::Searching::SearchRequest.new(Issue, '*').tap do |request|
+      allow(request).to receive(:execute!).and_return(elastic_aggregations)
+    end
+  end
+
+  let(:aggregations) do
+    Elasticsearch::Model::Response::Response.new(Issue, search).aggregations
+  end
+
+  describe '.call' do
+    subject { described_class.call(aggregations) }
+
+    context 'when elasticsearch buckets are provided' do
+      let(:elastic_aggregations) do
+        {
+          'aggregations' =>
+            {
+              'test' =>
+                {
+                  'after_key' => { 'test' => 'HTML', 'rid' => '3' },
+                  'buckets' =>
+                    [
+                      { 'key' => { 'test' => 'C', 'rid' => '3' }, 'doc_count' => 142 },
+                      { 'key' => { 'test' => 'C++', 'rid' => '3' }, 'doc_count' => 6 },
+                      { 'key' => { 'test' => 'CSS', 'rid' => '3' }, 'doc_count' => 1 }
+                    ]
+                },
+              'test2' =>
+                {
+                  'after_key' => { 'test2' => '1' },
+                  'buckets' =>
+                    [
+                      { 'key' => { 'test2' => '1' }, 'doc_count' => 1000 },
+                      { 'key' => { 'test2' => '2' }, 'doc_count' => 3 }
+                    ]
+                }
+            }
+        }
+      end
+
+      it 'parses the results' do
+        expected_buckets_1 = [
+          { key: { 'test': 'C', 'rid': '3' }, count: 142 },
+          { key: { 'test': 'C++', 'rid': '3' }, count: 6 },
+          { key: { 'test': 'CSS', 'rid': '3' }, count: 1 }
+        ]
+        expected_buckets_2 = [
+          { key: { 'test2': '1' }, count: 1000 },
+          { key: { 'test2': '2' }, count: 3 }
+        ]
+
+        expect(subject.length).to eq(2)
+        expect(subject.first.name).to eq('test')
+        expect(subject.first.buckets).to match_array(expected_buckets_1)
+        expect(subject.second.name).to eq('test2')
+        expect(subject.second.buckets).to match_array(expected_buckets_2)
+      end
+    end
+
+    context 'aggregations are not present' do
+      let(:elastic_aggregations) { {} }
+
+      it 'parses the results' do
+        expect(subject).to match_array([])
+      end
+    end
+  end
+end

ee/spec/lib/gitlab/search/aggregation_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Search::Aggregation do
+  describe 'parsing bucket results' do
+    subject { described_class.new('language', aggregation_buckets) }
+
+    context 'when elasticsearch buckets are provided' do
+      let(:aggregation_buckets) { [{ 'key': { 'language': 'ruby' }, 'doc_count': 10 }, { 'key': { 'language': 'java' }, 'doc_count': 20 }].map(&:with_indifferent_access) }
+
+      it 'parses the results' do
+        expected = [{ key: { 'language': 'ruby' }, count: 10 }, { key: { 'language': 'java' }, count: 20 }]
+
+        expect(subject.buckets).to match_array(expected)
+      end
+    end
+
+    context 'when elasticsearch buckets are not provided' do
+      let(:aggregation_buckets) { nil }
+
+      it 'parses the results' do
+        expect(subject.buckets).to match_array([])
+      end
+    end
+  end
+end

ee/spec/support/shared_examples/lib/gitlab/elastic/search_results_shared_examples.rb

@@ -45,3 +45,19 @@ RSpec.shared_examples 'does not load results for count only queries' do |scopes|
     end
   end
 end
+
+RSpec.shared_examples 'loads aggregations' do
+  let(:query) { 'hello world' }
+
+  it 'returns the expected aggregations' do
+    expect(subject).to match_array(expected)
+  end
+
+  context 'when query is blank' do
+    let(:query) { nil }
+
+    it 'returns an empty array' do
+      expect(subject).to match_array([])
+    end
+  end
+end

lib/gitlab/search_results.rb

@@ -115,6 +115,11 @@ module Gitlab
       {}
     end
 
+    # aggregations are only performed by Elasticsearch backed results
+    def aggregations(scope)
+      []
+    end
+
     private
 
     def collection_for(scope)

spec/lib/gitlab/search_results_spec.rb

@@ -96,6 +96,18 @@ RSpec.describe Gitlab::SearchResults do
       end
     end
 
+    describe '#aggregations' do
+      where(:scope) do
+        %w(projects issues merge_requests blobs commits wiki_blobs epics milestones users unknown)
+      end
+
+      with_them do
+        it 'returns an empty array' do
+          expect(results.aggregations(scope)).to match_array([])
+        end
+      end
+    end
+
     context "when count_limit is lower than total amount" do
       before do
         allow(results).to receive(:count_limit).and_return(1)