Migrate recursive tree entries fetching to Gitaly

2 files changed, 22 insertions, 6 deletions

diff --git a/lib/gitlab/git/tree.rb b/lib/gitlab/git/tree.rb
index ba6058fd3c9..b6ceb542dd1 100644
--- a/lib/gitlab/git/tree.rb
+++ b/lib/gitlab/git/tree.rb
@@ -14,14 +14,14 @@ module Gitlab
         # Uses rugged for raw objects
         #
         # Gitaly migration: https://gitlab.com/gitlab-org/gitaly/issues/320
-        def where(repository, sha, path = nil)
+        def where(repository, sha, path = nil, recursive = false)
           path = nil if path == '' || path == '/'
 
           Gitlab::GitalyClient.migrate(:tree_entries) do |is_enabled|
             if is_enabled
-              repository.gitaly_commit_client.tree_entries(repository, sha, path)
+              repository.gitaly_commit_client.tree_entries(repository, sha, path, recursive)
             else
-              tree_entries_from_rugged(repository, sha, path)
+              tree_entries_from_rugged(repository, sha, path, recursive)
             end
           end
         end
@@ -57,7 +57,22 @@ module Gitlab
           end
         end
 
-        def tree_entries_from_rugged(repository, sha, path)
+        def tree_entries_from_rugged(repository, sha, path, recursive)
+          current_path_entries = get_tree_entries_from_rugged(repository, sha, path)
+          ordered_entries = []
+
+          current_path_entries.each do |entry|
+            ordered_entries << entry
+
+            if recursive && entry.dir?
+              ordered_entries.concat(tree_entries_from_rugged(repository, sha, entry.path, true))
+            end
+          end
+
+          ordered_entries
+        end
+
+        def get_tree_entries_from_rugged(repository, sha, path)
           commit = repository.lookup(sha)
           root_tree = commit.tree
 
diff --git a/lib/gitlab/gitaly_client/commit_service.rb b/lib/gitlab/gitaly_client/commit_service.rb
index 269a048cf5d..d60f57717b5 100644
--- a/lib/gitlab/gitaly_client/commit_service.rb
+++ b/lib/gitlab/gitaly_client/commit_service.rb
@@ -105,11 +105,12 @@ module Gitlab
         entry unless entry.oid.blank?
       end
 
-      def tree_entries(repository, revision, path)
+      def tree_entries(repository, revision, path, recursive)
         request = Gitaly::GetTreeEntriesRequest.new(
           repository: @gitaly_repo,
           revision: encode_binary(revision),
-          path: path.present? ? encode_binary(path) : '.'
+          path: path.present? ? encode_binary(path) : '.',
+          recursive: recursive
         )
 
         response = GitalyClient.call(@repository.storage, :commit_service, :get_tree_entries, request, timeout: GitalyClient.medium_timeout)