8 files changed, 438 insertions, 137 deletions

diff --git a/app/finders/snippets_finder.rb b/app/finders/snippets_finder.rb
index f90971bb9f6..d3774746cb8 100644
--- a/app/finders/snippets_finder.rb
+++ b/app/finders/snippets_finder.rb
@@ -1,138 +1,149 @@
 # frozen_string_literal: true
 
-# Snippets Finder
+# Finder for retrieving snippets that a user can see, optionally scoped to a
+# project or snippets author.
 #
-# Used to filter Snippets collections by a set of params
+# Basic usage:
 #
-# Arguments.
+#     user = User.find(1)
 #
-# current_user - The current user, nil also can be used.
-# params:
-#   visibility (integer) - Individual snippet visibility: Public(20), internal(10) or private(0).
-#   project (Project) - Project related.
-#   author (User) - Author related.
+#     SnippetsFinder.new(user).execute
 #
-# params are optional
+# To limit the snippets to a specific project, supply the `project:` option:
+#
+#     user = User.find(1)
+#     project = Project.find(1)
+#
+#     SnippetsFinder.new(user, project: project).execute
+#
+# Limiting snippets to an author can be done by supplying the `author:` option:
+#
+#     user = User.find(1)
+#     project = Project.find(1)
+#
+#     SnippetsFinder.new(user, author: user).execute
+#
+# To filter snippets using a specific visibility level, you can provide the
+# `scope:` option:
+#
+#     user = User.find(1)
+#     project = Project.find(1)
+#
+#     SnippetsFinder.new(user, author: user, scope: :are_public).execute
+#
+# Valid `scope:` values are:
+#
+# * `:are_private`
+# * `:are_internal`
+# * `:are_public`
+#
+# Any other value will be ignored.
 class SnippetsFinder < UnionFinder
-  include Gitlab::Allowable
   include FinderMethods
 
-  attr_accessor :current_user, :project, :params
+  attr_accessor :current_user, :project, :author, :scope
 
-  def initialize(current_user, params = {})
+  def initialize(current_user = nil, params = {})
     @current_user = current_user
-    @params = params
     @project = params[:project]
-  end
-
-  def execute
-    items = init_collection
-    items = by_author(items)
-    items = by_visibility(items)
-
-    items.fresh
-  end
-
-  private
-
-  def init_collection
-    if project.present?
-      authorized_snippets_from_project
-    else
-      authorized_snippets
+    @author = params[:author]
+    @scope = params[:scope].to_s
+
+    if project && author
+      raise(
+        ArgumentError,
+        'Filtering by both an author and a project is not supported, ' \
+          'as this finder is not optimised for this use case'
+      )
     end
   end
 
-  def authorized_snippets_from_project
-    if can?(current_user, :read_project_snippet, project)
-      if project.team.member?(current_user)
-        project.snippets
+  def execute
+    base =
+      if project
+        snippets_for_a_single_project
       else
-        project.snippets.public_to_user(current_user)
+        snippets_for_multiple_projects
       end
-    else
-      Snippet.none
-    end
-  end
 
-  # rubocop: disable CodeReuse/ActiveRecord
-  def authorized_snippets
-    # This query was intentionally converted to a raw one to get it work in Rails 5.0.
-    # In Rails 5.0 and 5.1 there's a bug: https://github.com/rails/arel/issues/531
-    # Please convert it back when on rails 5.2 as it works again as expected since 5.2.
-    Snippet.where("#{feature_available_projects} OR #{not_project_related}")
-      .public_or_visible_to_user(current_user)
+    base.with_optional_visibility(visibility_from_scope).fresh
   end
-  # rubocop: enable CodeReuse/ActiveRecord
 
-  # Returns a collection of projects that is either public or visible to the
-  # logged in user.
+  # Produces a query that retrieves snippets from multiple projects.
   #
-  # A caller must pass in a block to modify individual parts of
-  # the query, e.g. to apply .with_feature_available_for_user on top of it.
-  # This is useful for performance as we can stick those additional filters
-  # at the bottom of e.g. the UNION.
-  # rubocop: disable CodeReuse/ActiveRecord
-  def projects_for_user
-    return yield(Project.public_to_user) unless current_user
-
-    # If the current_user is allowed to see all projects,
-    # we can shortcut and just return.
-    return yield(Project.all) if current_user.full_private_access?
-
-    authorized_projects = yield(Project.where('EXISTS (?)', current_user.authorizations_for_projects))
-
-    levels = Gitlab::VisibilityLevel.levels_for_user(current_user)
-    visible_projects = yield(Project.where(visibility_level: levels))
-
-    # We use a UNION here instead of OR clauses since this results in better
-    # performance.
-    Project.from_union([authorized_projects, visible_projects])
-  end
-  # rubocop: enable CodeReuse/ActiveRecord
-
-  def feature_available_projects
-    # Don't return any project related snippets if the user cannot read cross project
-    return table[:id].eq(nil).to_sql unless Ability.allowed?(current_user, :read_cross_project)
-
-    projects = projects_for_user do |part|
-      part.with_feature_available_for_user(:snippets, current_user)
-    end.select(:id)
+  # The resulting query will, depending on the user's permissions, include the
+  # following collections of snippets:
+  #
+  # 1. Snippets that don't belong to any project.
+  # 2. Snippets of projects that are visible to the current user (e.g. snippets
+  #    in public projects).
+  # 3. Snippets of projects that the current user is a member of.
+  #
+  # Each collection is constructed in isolation, allowing for greater control
+  # over the resulting SQL query.
+  def snippets_for_multiple_projects
+    queries = [global_snippets]
+
+    if Ability.allowed?(current_user, :read_cross_project)
+      queries << snippets_of_visible_projects
+      queries << snippets_of_authorized_projects if current_user
+    end
 
-    # This query was intentionally converted to a raw one to get it work in Rails 5.0.
-    # In Rails 5.0 and 5.1 there's a bug: https://github.com/rails/arel/issues/531
-    # Please convert it back when on rails 5.2 as it works again as expected since 5.2.
-    "snippets.project_id IN (#{projects.to_sql})"
+    find_union(queries, Snippet)
   end
 
-  def not_project_related
-    table[:project_id].eq(nil).to_sql
+  def snippets_for_a_single_project
+    Snippet.for_project_with_user(project, current_user)
   end
 
-  def table
-    Snippet.arel_table
+  def global_snippets
+    snippets_for_author_or_visible_to_user.only_global_snippets
   end
 
-  # rubocop: disable CodeReuse/ActiveRecord
-  def by_visibility(items)
-    visibility = params[:visibility] || visibility_from_scope
+  # Returns the snippets that the current user (logged in or not) can view.
+  def snippets_of_visible_projects
+    snippets_for_author_or_visible_to_user
+      .only_include_projects_visible_to(current_user)
+      .only_include_projects_with_snippets_enabled
+  end
 
-    return items unless visibility
+  # Returns the snippets that the currently logged in user has access to by
+  # being a member of the project the snippets belong to.
+  #
+  # This method requires that `current_user` returns a `User` instead of `nil`,
+  # and is optimised for this specific scenario.
+  def snippets_of_authorized_projects
+    base = author ? snippets_for_author : Snippet.all
+
+    base
+      .only_include_projects_with_snippets_enabled(include_private: true)
+      .only_include_authorized_projects(current_user)
+  end
 
-    items.where(visibility_level: visibility)
+  def snippets_for_author_or_visible_to_user
+    if author
+      snippets_for_author
+    elsif current_user
+      Snippet.visible_to_or_authored_by(current_user)
+    else
+      Snippet.public_to_user
+    end
   end
-  # rubocop: enable CodeReuse/ActiveRecord
 
-  # rubocop: disable CodeReuse/ActiveRecord
-  def by_author(items)
-    return items unless params[:author]
+  def snippets_for_author
+    base = author.snippets
 
-    items.where(author_id: params[:author].id)
+    if author == current_user
+      # If the current user is also the author of all snippets, then we can
+      # include private snippets.
+      base
+    else
+      base.public_to_user(current_user)
+    end
   end
-  # rubocop: enable CodeReuse/ActiveRecord
 
   def visibility_from_scope
-    case params[:scope].to_s
+    case scope
     when 'are_private'
       Snippet::PRIVATE
     when 'are_internal'
diff --git a/app/models/project.rb b/app/models/project.rb
index e2e309e8496..4273863d529 100644
--- a/app/models/project.rb
+++ b/app/models/project.rb
@@ -2073,6 +2073,10 @@ class Project < ActiveRecord::Base
     storage_version != LATEST_STORAGE_VERSION
   end
 
+  def snippets_visible?(user = nil)
+    Ability.allowed?(user, :read_project_snippet, self)
+  end
+
   private
 
   def use_hashed_storage
diff --git a/app/models/snippet.rb b/app/models/snippet.rb
index 1c5846b4023..11856b55902 100644
--- a/app/models/snippet.rb
+++ b/app/models/snippet.rb
@@ -63,6 +63,62 @@ class Snippet < ActiveRecord::Base
   attr_spammable :title, spam_title: true
   attr_spammable :content, spam_description: true
 
+  def self.with_optional_visibility(value = nil)
+    if value
+      where(visibility_level: value)
+    else
+      all
+    end
+  end
+
+  def self.only_global_snippets
+    where(project_id: nil)
+  end
+
+  def self.only_include_projects_visible_to(current_user = nil)
+    levels = Gitlab::VisibilityLevel.levels_for_user(current_user)
+
+    joins(:project).where('projects.visibility_level IN (?)', levels)
+  end
+
+  def self.only_include_projects_with_snippets_enabled(include_private: false)
+    column = ProjectFeature.access_level_attribute(:snippets)
+    levels = [ProjectFeature::ENABLED, ProjectFeature::PUBLIC]
+
+    levels << ProjectFeature::PRIVATE if include_private
+
+    joins(project: :project_feature)
+      .where(project_features: { column => levels })
+  end
+
+  def self.only_include_authorized_projects(current_user)
+    where(
+      'EXISTS (?)',
+      ProjectAuthorization
+        .select(1)
+        .where('project_id = snippets.project_id')
+        .where(user_id: current_user.id)
+    )
+  end
+
+  def self.for_project_with_user(project, user = nil)
+    return none unless project.snippets_visible?(user)
+
+    if user && project.team.member?(user)
+      project.snippets
+    else
+      project.snippets.public_to_user(user)
+    end
+  end
+
+  def self.visible_to_or_authored_by(user)
+    where(
+      'snippets.visibility_level IN (?) OR snippets.author_id = ?',
+      Gitlab::VisibilityLevel.levels_for_user(user),
+      user.id
+    )
+  end
+
   def self.reference_prefix
     '$'
   end
@@ -81,27 +137,6 @@ class Snippet < ActiveRecord::Base
     @link_reference_pattern ||= super("snippets", /(?<snippet>\d+)/)
   end
 
-  # Returns a collection of snippets that are either public or visible to the
-  # logged in user.
-  #
-  # This method does not verify the user actually has the access to the project
-  # the snippet is in, so it should be only used on a relation that's already scoped
-  # for project access
-  def self.public_or_visible_to_user(user = nil)
-    if user
-      authorized = user
-        .project_authorizations
-        .select(1)
-        .where('project_authorizations.project_id = snippets.project_id')
-
-      levels = Gitlab::VisibilityLevel.levels_for_user(user)
-
-      where('EXISTS (?) OR snippets.visibility_level IN (?) or snippets.author_id = (?)', authorized, levels, user.id)
-    else
-      public_to_user
-    end
-  end
-
   def to_reference(from = nil, full: false)
     reference = "#{self.class.reference_prefix}#{id}"
 
diff --git a/changelogs/unreleased/refactor-snippets-finder.yml b/changelogs/unreleased/refactor-snippets-finder.yml
new file mode 100644
index 00000000000..37cacf71c14
--- /dev/null
+++ b/changelogs/unreleased/refactor-snippets-finder.yml
@@ -0,0 +1,5 @@
+---
+title: Rewrite SnippetsFinder to improve performance by a factor of 1500
+merge_request:
+author:
+type: performance
diff --git a/lib/api/snippets.rb b/lib/api/snippets.rb
index f1786c15f4f..1ae144ca9c1 100644
--- a/lib/api/snippets.rb
+++ b/lib/api/snippets.rb
@@ -14,7 +14,7 @@ module API
         end
 
         def public_snippets
-          SnippetsFinder.new(current_user, visibility: Snippet::PUBLIC).execute
+          SnippetsFinder.new(current_user, scope: :are_public).execute
         end
       end
 
diff --git a/spec/finders/snippets_finder_spec.rb b/spec/finders/snippets_finder_spec.rb
index 1ae0bd988f2..dfeeb3040c6 100644
--- a/spec/finders/snippets_finder_spec.rb
+++ b/spec/finders/snippets_finder_spec.rb
@@ -4,16 +4,13 @@ describe SnippetsFinder do
   include Gitlab::Allowable
   using RSpec::Parameterized::TableSyntax
 
-  context 'filter by visibility' do
-    let!(:snippet1) { create(:personal_snippet, :private) }
-    let!(:snippet2) { create(:personal_snippet, :internal) }
-    let!(:snippet3) { create(:personal_snippet, :public) }
+  describe '#initialize' do
+    it 'raises ArgumentError when a project and author are given' do
+      user = build(:user)
+      project = build(:project)
 
-    it "returns public snippets when visibility is PUBLIC" do
-      snippets = described_class.new(nil, visibility: Snippet::PUBLIC).execute
-
-      expect(snippets).to include(snippet3)
-      expect(snippets).not_to include(snippet1, snippet2)
+      expect { described_class.new(user, author: user, project: project) }
+        .to raise_error(ArgumentError)
     end
   end
 
@@ -66,21 +63,21 @@ describe SnippetsFinder do
     end
 
     it "returns internal snippets" do
-      snippets = described_class.new(user, author: user, visibility: Snippet::INTERNAL).execute
+      snippets = described_class.new(user, author: user, scope: :are_internal).execute
 
       expect(snippets).to include(snippet2)
       expect(snippets).not_to include(snippet1, snippet3)
     end
 
     it "returns private snippets" do
-      snippets = described_class.new(user, author: user, visibility: Snippet::PRIVATE).execute
+      snippets = described_class.new(user, author: user, scope: :are_private).execute
 
       expect(snippets).to include(snippet1)
       expect(snippets).not_to include(snippet2, snippet3)
     end
 
     it "returns public snippets" do
-      snippets = described_class.new(user, author: user, visibility: Snippet::PUBLIC).execute
+      snippets = described_class.new(user, author: user, scope: :are_public).execute
 
       expect(snippets).to include(snippet3)
       expect(snippets).not_to include(snippet1, snippet2)
@@ -98,6 +95,13 @@ describe SnippetsFinder do
       expect(snippets).to include(snippet3)
       expect(snippets).not_to include(snippet2, snippet1)
     end
+
+    it 'returns all snippets for an admin' do
+      admin = create(:user, :admin)
+      snippets = described_class.new(admin, author: user).execute
+
+      expect(snippets).to include(snippet1, snippet2, snippet3)
+    end
   end
 
   context 'filter by project' do
@@ -126,21 +130,21 @@ describe SnippetsFinder do
     end
 
     it "returns public snippets for non project members" do
-      snippets = described_class.new(user, project: project1, visibility: Snippet::PUBLIC).execute
+      snippets = described_class.new(user, project: project1, scope: :are_public).execute
 
       expect(snippets).to include(@snippet3)
       expect(snippets).not_to include(@snippet1, @snippet2)
     end
 
     it "returns internal snippets for non project members" do
-      snippets = described_class.new(user, project: project1, visibility: Snippet::INTERNAL).execute
+      snippets = described_class.new(user, project: project1, scope: :are_internal).execute
 
       expect(snippets).to include(@snippet2)
       expect(snippets).not_to include(@snippet1, @snippet3)
     end
 
     it "does not return private snippets for non project members" do
-      snippets = described_class.new(user, project: project1, visibility: Snippet::PRIVATE).execute
+      snippets = described_class.new(user, project: project1, scope: :are_private).execute
 
       expect(snippets).not_to include(@snippet1, @snippet2, @snippet3)
     end
@@ -156,10 +160,17 @@ describe SnippetsFinder do
     it "returns private snippets for project members" do
       project1.add_developer(user)
 
-      snippets = described_class.new(user, project: project1, visibility: Snippet::PRIVATE).execute
+      snippets = described_class.new(user, project: project1, scope: :are_private).execute
 
       expect(snippets).to include(@snippet1)
     end
+
+    it 'returns all snippets for an admin' do
+      admin = create(:user, :admin)
+      snippets = described_class.new(admin, project: project1).execute
+
+      expect(snippets).to include(@snippet1, @snippet2, @snippet3)
+    end
   end
 
   describe '#execute' do
@@ -184,4 +195,6 @@ describe SnippetsFinder do
       end
     end
   end
+
+  it_behaves_like 'snippet visibility'
 end
diff --git a/spec/models/project_spec.rb b/spec/models/project_spec.rb
index d4b9a4c8cd6..be2aff73c0a 100644
--- a/spec/models/project_spec.rb
+++ b/spec/models/project_spec.rb
@@ -3963,6 +3963,28 @@ describe Project do
     end
   end
 
+  describe '#snippets_visible?' do
+    it 'returns true when a logged in user can read snippets' do
+      project = create(:project, :public)
+      user = create(:user)
+
+      expect(project.snippets_visible?(user)).to eq(true)
+    end
+
+    it 'returns true when an anonymous user can read snippets' do
+      project = create(:project, :public)
+
+      expect(project.snippets_visible?).to eq(true)
+    end
+
+    it 'returns false when a user can not read snippets' do
+      project = create(:project, :private)
+      user = create(:user)
+
+      expect(project.snippets_visible?(user)).to eq(false)
+    end
+  end
+
   def rugged_config
     rugged_repo(project.repository).config
   end
diff --git a/spec/models/snippet_spec.rb b/spec/models/snippet_spec.rb
index e09d89d235d..7a7272ccb60 100644
--- a/spec/models/snippet_spec.rb
+++ b/spec/models/snippet_spec.rb
@@ -131,6 +131,217 @@ describe Snippet do
     end
   end
 
+  describe '.with_optional_visibility' do
+    context 'when a visibility level is provided' do
+      it 'returns snippets with the given visibility' do
+        create(:snippet, :private)
+
+        snippet = create(:snippet, :public)
+        snippets = described_class
+          .with_optional_visibility(Gitlab::VisibilityLevel::PUBLIC)
+
+        expect(snippets).to eq([snippet])
+      end
+    end
+
+    context 'when a visibility level is not provided' do
+      it 'returns all snippets' do
+        snippet1 = create(:snippet, :public)
+        snippet2 = create(:snippet, :private)
+        snippets = described_class.with_optional_visibility
+
+        expect(snippets).to include(snippet1, snippet2)
+      end
+    end
+  end
+
+  describe '.only_global_snippets' do
+    it 'returns snippets not associated with any projects' do
+      create(:project_snippet)
+
+      snippet = create(:snippet)
+      snippets = described_class.only_global_snippets
+
+      expect(snippets).to eq([snippet])
+    end
+  end
+
+  describe '.only_include_projects_visible_to' do
+    let!(:project1) { create(:project, :public) }
+    let!(:project2) { create(:project, :internal) }
+    let!(:project3) { create(:project, :private) }
+    let!(:snippet1) { create(:project_snippet, project: project1) }
+    let!(:snippet2) { create(:project_snippet, project: project2) }
+    let!(:snippet3) { create(:project_snippet, project: project3) }
+
+    context 'when a user is provided' do
+      it 'returns snippets visible to the user' do
+        user = create(:user)
+
+        snippets = described_class.only_include_projects_visible_to(user)
+
+        expect(snippets).to include(snippet1, snippet2)
+        expect(snippets).not_to include(snippet3)
+      end
+    end
+
+    context 'when a user is not provided' do
+      it 'returns snippets visible to anonymous users' do
+        snippets = described_class.only_include_projects_visible_to
+
+        expect(snippets).to include(snippet1)
+        expect(snippets).not_to include(snippet2, snippet3)
+      end
+    end
+  end
+
+  describe 'only_include_projects_with_snippets_enabled' do
+    context 'when the include_private option is enabled' do
+      it 'includes snippets for projects with snippets set to private' do
+        project = create(:project)
+
+        project.project_feature
+          .update(snippets_access_level: ProjectFeature::PRIVATE)
+
+        snippet = create(:project_snippet, project: project)
+
+        snippets = described_class
+          .only_include_projects_with_snippets_enabled(include_private: true)
+
+        expect(snippets).to eq([snippet])
+      end
+    end
+
+    context 'when the include_private option is not enabled' do
+      it 'does not include snippets for projects that have snippets set to private' do
+        project = create(:project)
+
+        project.project_feature
+          .update(snippets_access_level: ProjectFeature::PRIVATE)
+
+        create(:project_snippet, project: project)
+
+        snippets = described_class.only_include_projects_with_snippets_enabled
+
+        expect(snippets).to be_empty
+      end
+    end
+
+    it 'includes snippets for projects with snippets enabled' do
+      project = create(:project)
+
+      project.project_feature
+        .update(snippets_access_level: ProjectFeature::ENABLED)
+
+      snippet = create(:project_snippet, project: project)
+      snippets = described_class.only_include_projects_with_snippets_enabled
+
+      expect(snippets).to eq([snippet])
+    end
+  end
+
+  describe '.only_include_authorized_projects' do
+    it 'only includes snippets for projects the user is authorized to see' do
+      user = create(:user)
+      project1 = create(:project, :private)
+      project2 = create(:project, :private)
+
+      project1.team.add_developer(user)
+
+      create(:project_snippet, project: project2)
+
+      snippet = create(:project_snippet, project: project1)
+      snippets = described_class.only_include_authorized_projects(user)
+
+      expect(snippets).to eq([snippet])
+    end
+  end
+
+  describe '.for_project_with_user' do
+    context 'when a user is provided' do
+      it 'returns an empty collection if the user can not view the snippets' do
+        project = create(:project, :private)
+        user = create(:user)
+
+        project.project_feature
+          .update(snippets_access_level: ProjectFeature::ENABLED)
+
+        create(:project_snippet, :public, project: project)
+
+        expect(described_class.for_project_with_user(project, user)).to be_empty
+      end
+
+      it 'returns the snippets if the user is a member of the project' do
+        project = create(:project, :private)
+        user = create(:user)
+        snippet = create(:project_snippet, project: project)
+
+        project.team.add_developer(user)
+
+        snippets = described_class.for_project_with_user(project, user)
+
+        expect(snippets).to eq([snippet])
+      end
+
+      it 'returns public snippets for a public project the user is not a member of' do
+        project = create(:project, :public)
+
+        project.project_feature
+          .update(snippets_access_level: ProjectFeature::ENABLED)
+
+        user = create(:user)
+        snippet = create(:project_snippet, :public, project: project)
+
+        create(:project_snippet, :private, project: project)
+
+        snippets = described_class.for_project_with_user(project, user)
+
+        expect(snippets).to eq([snippet])
+      end
+    end
+
+    context 'when a user is not provided' do
+      it 'returns an empty collection for a private project' do
+        project = create(:project, :private)
+
+        project.project_feature
+          .update(snippets_access_level: ProjectFeature::ENABLED)
+
+        create(:project_snippet, :public, project: project)
+
+        expect(described_class.for_project_with_user(project)).to be_empty
+      end
+
+      it 'returns public snippets for a public project' do
+        project = create(:project, :public)
+        snippet = create(:project_snippet, :public, project: project)
+
+        project.project_feature
+          .update(snippets_access_level: ProjectFeature::PUBLIC)
+
+        create(:project_snippet, :private, project: project)
+
+        snippets = described_class.for_project_with_user(project)
+
+        expect(snippets).to eq([snippet])
+      end
+    end
+  end
+
+  describe '.visible_to_or_authored_by' do
+    it 'returns snippets visible to the user' do
+      user = create(:user)
+      snippet1 = create(:snippet, :public)
+      snippet2 = create(:snippet, :private, author: user)
+      snippet3 = create(:snippet, :private)
+
+      snippets = described_class.visible_to_or_authored_by(user)
+
+      expect(snippets).to include(snippet1, snippet2)
+      expect(snippets).not_to include(snippet3)
+    end
+  end
+
   describe '#participants' do
     let(:project) { create(:project, :public) }
     let(:snippet) { create(:snippet, content: 'foo', project: project) }