Add pagination support for get_tree_entries RPC

Contributes to https://gitlab.com/gitlab-org/gitlab/-/issues/334140

MR with related Gitaly changes:
https://gitlab.com/gitlab-org/gitaly/-/merge_requests/3611

* Allow to pass pagination parameters to Gitaly get_tree_entries RPC.
* Expose pagination cursor in Tree object.

Changelog: added

app/models/repository.rb

@@ -656,7 +656,7 @@ class Repository
     end
   end
 
-  def tree(sha = :head, path = nil, recursive: false)
+  def tree(sha = :head, path = nil, recursive: false, pagination_params: nil)
     if sha == :head
       return unless head_commit
 
@@ -667,7 +667,7 @@ class Repository
       end
     end
 
-    Tree.new(self, sha, path, recursive: recursive)
+    Tree.new(self, sha, path, recursive: recursive, pagination_params: pagination_params)
   end
 
   def blob_at_branch(branch_name, path)

app/models/tree.rb

@@ -4,9 +4,9 @@ class Tree
   include Gitlab::MarkupHelper
   include Gitlab::Utils::StrongMemoize
 
-  attr_accessor :repository, :sha, :path, :entries
+  attr_accessor :repository, :sha, :path, :entries, :cursor
 
-  def initialize(repository, sha, path = '/', recursive: false)
+  def initialize(repository, sha, path = '/', recursive: false, pagination_params: nil)
     path = '/' if path.blank?
 
     @repository = repository
@@ -14,7 +14,7 @@ class Tree
     @path = path
 
     git_repo = @repository.raw_repository
-    @entries = Gitlab::Git::Tree.where(git_repo, @sha, @path, recursive)
+    @entries, @cursor = Gitlab::Git::Tree.where(git_repo, @sha, @path, recursive, pagination_params)
   end
 
   def readme_path

lib/gitlab/git/rugged_impl/tree.rb

@@ -14,9 +14,12 @@ module Gitlab
           include Gitlab::Git::RuggedImpl::UseRugged
 
           override :tree_entries
-          def tree_entries(repository, sha, path, recursive)
+          def tree_entries(repository, sha, path, recursive, pagination_params = nil)
             if use_rugged?(repository, :rugged_tree_entries)
-              execute_rugged_call(:tree_entries_with_flat_path_from_rugged, repository, sha, path, recursive)
+              [
+                execute_rugged_call(:tree_entries_with_flat_path_from_rugged, repository, sha, path, recursive),
+                nil
+              ]
             else
               super
             end

lib/gitlab/git/tree.rb

@@ -15,15 +15,15 @@ module Gitlab
         # Uses rugged for raw objects
         #
         # Gitaly migration: https://gitlab.com/gitlab-org/gitaly/issues/320
-        def where(repository, sha, path = nil, recursive = false)
+        def where(repository, sha, path = nil, recursive = false, pagination_params = nil)
           path = nil if path == '' || path == '/'
 
-          tree_entries(repository, sha, path, recursive)
+          tree_entries(repository, sha, path, recursive, pagination_params)
         end
 
-        def tree_entries(repository, sha, path, recursive)
+        def tree_entries(repository, sha, path, recursive, pagination_params = nil)
           wrapped_gitaly_errors do
-            repository.gitaly_commit_client.tree_entries(repository, sha, path, recursive)
+            repository.gitaly_commit_client.tree_entries(repository, sha, path, recursive, pagination_params)
           end
         end
 

lib/gitlab/gitaly_client/commit_service.rb

@@ -111,17 +111,22 @@ module Gitlab
         nil
       end
 
-      def tree_entries(repository, revision, path, recursive)
+      def tree_entries(repository, revision, path, recursive, pagination_params)
         request = Gitaly::GetTreeEntriesRequest.new(
           repository: @gitaly_repo,
           revision: encode_binary(revision),
           path: path.present? ? encode_binary(path) : '.',
-          recursive: recursive
+          recursive: recursive,
+          pagination_params: pagination_params
         )
+        request.sort = Gitaly::GetTreeEntriesRequest::SortBy::TREES_FIRST if pagination_params
 
         response = GitalyClient.call(@repository.storage, :commit_service, :get_tree_entries, request, timeout: GitalyClient.medium_timeout)
 
-        response.flat_map do |message|
+        cursor = nil
+
+        entries = response.flat_map do |message|
+          cursor = message.pagination_cursor if message.pagination_cursor
           message.entries.map do |gitaly_tree_entry|
             Gitlab::Git::Tree.new(
               id: gitaly_tree_entry.oid,
@@ -135,6 +140,8 @@ module Gitlab
             )
           end
         end
+
+        [entries, cursor]
       end
 
       def commit_count(ref, options = {})

spec/lib/gitlab/git/tree_spec.rb

@@ -6,29 +6,44 @@ RSpec.describe Gitlab::Git::Tree, :seed_helper do
   let(:repository) { Gitlab::Git::Repository.new('default', TEST_REPO_PATH, '', 'group/project') }
 
   shared_examples :repo do
-    let(:tree) { Gitlab::Git::Tree.where(repository, SeedRepo::Commit::ID) }
+    subject(:tree) { Gitlab::Git::Tree.where(repository, sha, path, recursive, pagination_params) }
 
-    it { expect(tree).to be_kind_of Array }
-    it { expect(tree.empty?).to be_falsey }
-    it { expect(tree.count(&:dir?)).to eq(2) }
-    it { expect(tree.count(&:file?)).to eq(10) }
-    it { expect(tree.count(&:submodule?)).to eq(2) }
+    let(:sha) { SeedRepo::Commit::ID }
+    let(:path) { nil }
+    let(:recursive) { false }
+    let(:pagination_params) { nil }
 
-    it 'returns an empty array when called with an invalid ref' do
-      expect(described_class.where(repository, 'foobar-does-not-exist')).to eq([])
+    let(:entries) { tree.first }
+    let(:cursor) { tree.second }
+
+    it { expect(entries).to be_kind_of Array }
+    it { expect(entries.empty?).to be_falsey }
+    it { expect(entries.count(&:dir?)).to eq(2) }
+    it { expect(entries.count(&:file?)).to eq(10) }
+    it { expect(entries.count(&:submodule?)).to eq(2) }
+    it { expect(cursor&.next_cursor).to be_blank }
+
+    context 'with an invalid ref' do
+      let(:sha) { 'foobar-does-not-exist' }
+
+      it { expect(entries).to eq([]) }
+      it { expect(cursor).to be_nil }
     end
 
-    it 'returns a list of tree objects' do
-      entries = described_class.where(repository, SeedRepo::Commit::ID, 'files', true)
+    context 'when path is provided' do
+      let(:path) { 'files' }
+      let(:recursive) { true }
 
-      expect(entries.map(&:path)).to include('files/html',
-                                             'files/markdown/ruby-style-guide.md')
-      expect(entries.count).to be >= 10
-      expect(entries).to all(be_a(Gitlab::Git::Tree))
+      it 'returns a list of tree objects' do
+        expect(entries.map(&:path)).to include('files/html',
+                                               'files/markdown/ruby-style-guide.md')
+        expect(entries.count).to be >= 10
+        expect(entries).to all(be_a(Gitlab::Git::Tree))
+      end
     end
 
     describe '#dir?' do
-      let(:dir) { tree.select(&:dir?).first }
+      let(:dir) { entries.select(&:dir?).first }
 
       it { expect(dir).to be_kind_of Gitlab::Git::Tree }
       it { expect(dir.id).to eq('3c122d2b7830eca25235131070602575cf8b41a1') }
@@ -41,7 +56,8 @@ RSpec.describe Gitlab::Git::Tree, :seed_helper do
       context :subdir do
         # rubocop: disable Rails/FindBy
         # This is not ActiveRecord where..first
-        let(:subdir) { Gitlab::Git::Tree.where(repository, SeedRepo::Commit::ID, 'files').first }
+        let(:path) { 'files' }
+        let(:subdir) { entries.first }
         # rubocop: enable Rails/FindBy
 
         it { expect(subdir).to be_kind_of Gitlab::Git::Tree }
@@ -55,7 +71,8 @@ RSpec.describe Gitlab::Git::Tree, :seed_helper do
       context :subdir_file do
         # rubocop: disable Rails/FindBy
         # This is not ActiveRecord where..first
-        let(:subdir_file) { Gitlab::Git::Tree.where(repository, SeedRepo::Commit::ID, 'files/ruby').first }
+        let(:path) { 'files/ruby' }
+        let(:subdir_file) { entries.first }
         # rubocop: enable Rails/FindBy
 
         it { expect(subdir_file).to be_kind_of Gitlab::Git::Tree }
@@ -68,10 +85,11 @@ RSpec.describe Gitlab::Git::Tree, :seed_helper do
 
       context :flat_path do
         let(:filename) { 'files/flat/path/correct/content.txt' }
-        let(:oid) { create_file(filename) }
+        let(:sha) { create_file(filename) }
+        let(:path) { 'files/flat' }
         # rubocop: disable Rails/FindBy
         # This is not ActiveRecord where..first
-        let(:subdir_file) { Gitlab::Git::Tree.where(repository, oid, 'files/flat').first }
+        let(:subdir_file) { entries.first }
         # rubocop: enable Rails/FindBy
         let(:repository_rugged) { Rugged::Repository.new(File.join(SEED_STORAGE_PATH, TEST_REPO_PATH)) }
 
@@ -116,7 +134,7 @@ RSpec.describe Gitlab::Git::Tree, :seed_helper do
     end
 
     describe '#file?' do
-      let(:file) { tree.select(&:file?).first }
+      let(:file) { entries.select(&:file?).first }
 
       it { expect(file).to be_kind_of Gitlab::Git::Tree }
       it { expect(file.id).to eq('dfaa3f97ca337e20154a98ac9d0be76ddd1fcc82') }
@@ -125,21 +143,21 @@ RSpec.describe Gitlab::Git::Tree, :seed_helper do
     end
 
     describe '#readme?' do
-      let(:file) { tree.select(&:readme?).first }
+      let(:file) { entries.select(&:readme?).first }
 
       it { expect(file).to be_kind_of Gitlab::Git::Tree }
       it { expect(file.name).to eq('README.md') }
     end
 
     describe '#contributing?' do
-      let(:file) { tree.select(&:contributing?).first }
+      let(:file) { entries.select(&:contributing?).first }
 
       it { expect(file).to be_kind_of Gitlab::Git::Tree }
       it { expect(file.name).to eq('CONTRIBUTING.md') }
     end
 
     describe '#submodule?' do
-      let(:submodule) { tree.select(&:submodule?).first }
+      let(:submodule) { entries.select(&:submodule?).first }
 
       it { expect(submodule).to be_kind_of Gitlab::Git::Tree }
       it { expect(submodule.id).to eq('79bceae69cb5750d6567b223597999bfa91cb3b9') }
@@ -149,7 +167,16 @@ RSpec.describe Gitlab::Git::Tree, :seed_helper do
   end
 
   describe '.where with Gitaly enabled' do
-    it_behaves_like :repo
+    it_behaves_like :repo do
+      context 'with pagination parameters' do
+        let(:pagination_params) { { limit: 3, page_token: nil } }
+
+        it 'returns paginated list of tree objects' do
+          expect(entries.count).to eq(3)
+          expect(cursor.next_cursor).to be_present
+        end
+      end
+    end
   end
 
   describe '.where with Rugged enabled', :enable_rugged do
@@ -161,6 +188,15 @@ RSpec.describe Gitlab::Git::Tree, :seed_helper do
       described_class.where(repository, SeedRepo::Commit::ID, 'files', false)
     end
 
-    it_behaves_like :repo
+    it_behaves_like :repo do
+      context 'with pagination parameters' do
+        let(:pagination_params) { { limit: 3, page_token: nil } }
+
+        it 'does not support pagination' do
+          expect(entries.count).to be >= 10
+          expect(cursor).to be_nil
+        end
+      end
+    end
   end
 end

spec/lib/gitlab/gitaly_client/commit_service_spec.rb

@@ -169,7 +169,11 @@ RSpec.describe Gitlab::GitalyClient::CommitService do
   end
 
   describe '#tree_entries' do
+    subject { client.tree_entries(repository, revision, path, recursive, pagination_params) }
+
     let(:path) { '/' }
+    let(:recursive) { false }
+    let(:pagination_params) { nil }
 
     it 'sends a get_tree_entries message' do
       expect_any_instance_of(Gitaly::CommitService::Stub)
@@ -177,7 +181,7 @@ RSpec.describe Gitlab::GitalyClient::CommitService do
         .with(gitaly_request_with_path(storage_name, relative_path), kind_of(Hash))
         .and_return([])
 
-      client.tree_entries(repository, revision, path, false)
+      is_expected.to eq([[], nil])
     end
 
     context 'with UTF-8 params strings' do
@@ -190,7 +194,26 @@ RSpec.describe Gitlab::GitalyClient::CommitService do
           .with(gitaly_request_with_path(storage_name, relative_path), kind_of(Hash))
           .and_return([])
 
-        client.tree_entries(repository, revision, path, false)
+        is_expected.to eq([[], nil])
+      end
+    end
+
+    context 'with pagination parameters' do
+      let(:pagination_params) { { limit: 3, page_token: nil } }
+
+      it 'responds with a pagination cursor' do
+        pagination_cursor = Gitaly::PaginationCursor.new(next_cursor: 'aabbccdd')
+        response = Gitaly::GetTreeEntriesResponse.new(
+          entries: [],
+          pagination_cursor: pagination_cursor
+        )
+
+        expect_any_instance_of(Gitaly::CommitService::Stub)
+          .to receive(:get_tree_entries)
+          .with(gitaly_request_with_path(storage_name, relative_path), kind_of(Hash))
+          .and_return([response])
+
+        is_expected.to eq([[], pagination_cursor])
       end
     end
   end

spec/models/repository_spec.rb

@@ -2523,24 +2523,46 @@ RSpec.describe Repository do
   end
 
   shared_examples '#tree' do
+    subject { repository.tree(sha, path, recursive: recursive, pagination_params: pagination_params) }
+
+    let(:sha) { :head }
+    let(:path) { nil }
+    let(:recursive) { false }
+    let(:pagination_params) { nil }
+
     context 'using a non-existing repository' do
       before do
         allow(repository).to receive(:head_commit).and_return(nil)
       end
 
-      it 'returns nil' do
-        expect(repository.tree(:head)).to be_nil
-      end
+      it { is_expected.to be_nil }
+
+      context 'when path is defined' do
+        let(:path) { 'README.md' }
 
-      it 'returns nil when using a path' do
-        expect(repository.tree(:head, 'README.md')).to be_nil
+        it { is_expected.to be_nil }
       end
     end
 
     context 'using an existing repository' do
-      it 'returns a Tree' do
-        expect(repository.tree(:head)).to be_an_instance_of(Tree)
-        expect(repository.tree('v1.1.1')).to be_an_instance_of(Tree)
+      it { is_expected.to be_an_instance_of(Tree) }
+
+      context 'when different sha is set' do
+        let(:sha) { 'v1.1.1' }
+
+        it { is_expected.to be_an_instance_of(Tree) }
+      end
+
+      context 'when recursive is true' do
+        let(:recursive) { true }
+
+        it { is_expected.to be_an_instance_of(Tree) }
+      end
+
+      context 'with pagination parameters' do
+        let(:pagination_params) { { limit: 10, page_token: nil } }
+
+        it { is_expected.to be_an_instance_of(Tree) }
       end
     end
   end

spec/models/tree_spec.rb

@@ -6,7 +6,7 @@ RSpec.describe Tree do
   let(:repository) { create(:project, :repository).repository }
   let(:sha) { repository.root_ref }
 
-  subject { described_class.new(repository, '54fcc214') }
+  subject(:tree) { described_class.new(repository, '54fcc214') }
 
   describe '#readme' do
     before do
@@ -66,4 +66,10 @@ RSpec.describe Tree do
       expect(subject.readme.name).to eq 'README.md'
     end
   end
+
+  describe '#cursor' do
+    subject { tree.cursor }
+
+    it { is_expected.to be_an_instance_of(Gitaly::PaginationCursor) }
+  end
 end