Incorporate Gitaly's WikiService.WikiGetAllPages RPCgitaly-wiki-get-all-pages

7 files changed, 111 insertions, 23 deletions

diff --git a/GITALY_SERVER_VERSION b/GITALY_SERVER_VERSION
index 4f9b378b40f..7f422a161ae 100644
--- a/GITALY_SERVER_VERSION
+++ b/GITALY_SERVER_VERSION
@@ -1 +1 @@
-0.52.0
+0.53.0
diff --git a/Gemfile b/Gemfile
index 63d3d214a5a..c495e4eceba 100644
--- a/Gemfile
+++ b/Gemfile
@@ -398,7 +398,7 @@ group :ed25519 do
 end
 
 # Gitaly GRPC client
-gem 'gitaly-proto', '~> 0.51.0', require: 'gitaly'
+gem 'gitaly-proto', '~> 0.52.0', require: 'gitaly'
 
 gem 'toml-rb', '~> 0.3.15', require: false
 
diff --git a/Gemfile.lock b/Gemfile.lock
index ae145ca5f69..43555a01037 100644
--- a/Gemfile.lock
+++ b/Gemfile.lock
@@ -274,7 +274,7 @@ GEM
       po_to_json (>= 1.0.0)
       rails (>= 3.2.0)
     gherkin-ruby (0.3.2)
-    gitaly-proto (0.51.0)
+    gitaly-proto (0.52.0)
       google-protobuf (~> 3.1)
       grpc (~> 1.0)
     github-linguist (4.7.6)
@@ -1026,7 +1026,7 @@ DEPENDENCIES
   gettext (~> 3.2.2)
   gettext_i18n_rails (~> 1.8.0)
   gettext_i18n_rails_js (~> 1.2.0)
-  gitaly-proto (~> 0.51.0)
+  gitaly-proto (~> 0.52.0)
   github-linguist (~> 4.7.0)
   gitlab-flowdock-git-hook (~> 1.0.1)
   gitlab-markup (~> 1.6.2)
@@ -1185,4 +1185,4 @@ DEPENDENCIES
   wikicloth (= 0.8.1)
 
 BUNDLED WITH
-   1.15.4
+   1.16.0
diff --git a/lib/gitlab/git/wiki.rb b/lib/gitlab/git/wiki.rb
index 5d36f16abd4..022d1f249a9 100644
--- a/lib/gitlab/git/wiki.rb
+++ b/lib/gitlab/git/wiki.rb
@@ -59,7 +59,13 @@ module Gitlab
       end
 
       def pages
-        gollum_wiki.pages.map { |gollum_page| new_page(gollum_page) }
+        @repository.gitaly_migrate(:wiki_get_all_pages) do |is_enabled|
+          if is_enabled
+            gitaly_get_all_pages
+          else
+            gollum_get_all_pages
+          end
+        end
       end
 
       def page(title:, version: nil, dir: nil)
@@ -179,6 +185,10 @@ module Gitlab
         Gitlab::Git::WikiFile.new(gollum_file)
       end
 
+      def gollum_get_all_pages
+        gollum_wiki.pages.map { |gollum_page| new_page(gollum_page) }
+      end
+
       def gitaly_write_page(name, format, content, commit_details)
         gitaly_wiki_client.write_page(name, format, content, commit_details)
       end
@@ -204,6 +214,12 @@ module Gitlab
 
         Gitlab::Git::WikiFile.new(wiki_file)
       end
+
+      def gitaly_get_all_pages
+        gitaly_wiki_client.get_all_pages.map do |wiki_page, version|
+          Gitlab::Git::WikiPage.new(wiki_page, version)
+        end
+      end
     end
   end
 end
diff --git a/lib/gitlab/gitaly_client/wiki_page.rb b/lib/gitlab/gitaly_client/wiki_page.rb
index 8226278d5f6..98d96fe6211 100644
--- a/lib/gitlab/gitaly_client/wiki_page.rb
+++ b/lib/gitlab/gitaly_client/wiki_page.rb
@@ -11,6 +11,10 @@ module Gitlab
         FIELDS.each do |field|
           instance_variable_set("@#{field}", params[field])
         end
+
+        # All gRPC strings in a response are frozen, so we get an unfrozen
+        # version here so appending to `raw_data` doesn't blow up.
+        @raw_data = @raw_data.dup
       end
 
       def historical?
diff --git a/lib/gitlab/gitaly_client/wiki_service.rb b/lib/gitlab/gitaly_client/wiki_service.rb
index 1a668338f57..8f05f40365e 100644
--- a/lib/gitlab/gitaly_client/wiki_service.rb
+++ b/lib/gitlab/gitaly_client/wiki_service.rb
@@ -81,28 +81,23 @@ module Gitlab
         )
 
         response = GitalyClient.call(@repository.storage, :wiki_service, :wiki_find_page, request)
-        wiki_page = version = nil
 
-        response.each do |message|
-          page = message.page
-          next unless page
+        wiki_page_from_iterator(response)
+      end
 
-          if wiki_page
-            wiki_page.raw_data << page.raw_data
-          else
-            wiki_page = GitalyClient::WikiPage.new(page.to_h)
-            # All gRPC strings in a response are frozen, so we get
-            # an unfrozen version here so appending in the else clause below doesn't blow up.
-            wiki_page.raw_data = wiki_page.raw_data.dup
+      def get_all_pages
+        request = Gitaly::WikiGetAllPagesRequest.new(repository: @gitaly_repo)
+        response = GitalyClient.call(@repository.storage, :wiki_service, :wiki_get_all_pages, request)
+        pages = []
 
-            version = Gitlab::Git::WikiPageVersion.new(
-              Gitlab::Git::Commit.decorate(@repository, page.version.commit),
-              page.version.format
-            )
-          end
+        loop do
+          page, version = wiki_page_from_iterator(response) { |message| message.end_of_page }
+
+          break unless page && version
+          pages << [page, version]
         end
 
-        [wiki_page, version]
+        pages
       end
 
       def find_file(name, revision)
@@ -133,6 +128,35 @@ module Gitlab
 
       private
 
+      # If a block is given and the yielded value is true, iteration will be
+      # stopped early at that point; else the iterator is consumed entirely.
+      # The iterator is traversed with `next` to allow resuming the iteration.
+      def wiki_page_from_iterator(iterator)
+        wiki_page = version = nil
+
+        while message = iterator.next
+          break if block_given? && yield(message)
+
+          page = message.page
+          next unless page
+
+          if wiki_page
+            wiki_page.raw_data << page.raw_data
+          else
+            wiki_page = GitalyClient::WikiPage.new(page.to_h)
+
+            version = Gitlab::Git::WikiPageVersion.new(
+              Gitlab::Git::Commit.decorate(@repository, page.version.commit),
+              page.version.format
+            )
+          end
+        end
+
+        [wiki_page, version]
+      rescue StopIteration
+        [wiki_page, version]
+      end
+
       def gitaly_commit_details(commit_details)
         Gitaly::WikiCommitDetails.new(
           name: GitalyClient.encode(commit_details.name),
diff --git a/spec/lib/gitlab/gitaly_client/wiki_service_spec.rb b/spec/lib/gitlab/gitaly_client/wiki_service_spec.rb
index a9653027587..6ad9f5ef766 100644
--- a/spec/lib/gitlab/gitaly_client/wiki_service_spec.rb
+++ b/spec/lib/gitlab/gitaly_client/wiki_service_spec.rb
@@ -41,4 +41,48 @@ describe Gitlab::GitalyClient::WikiService do
       expect(wiki_page_version.format).to eq('markdown')
     end
   end
+
+  describe '#get_all_pages' do
+    let(:page_2_info) { { title: 'My Page 2', raw_data: 'c', version: page_version } }
+    let(:response) do
+      [
+        Gitaly::WikiGetAllPagesResponse.new(page: Gitaly::WikiPage.new(page_info)),
+        Gitaly::WikiGetAllPagesResponse.new(page: Gitaly::WikiPage.new(raw_data: 'b')),
+        Gitaly::WikiGetAllPagesResponse.new(end_of_page: true),
+        Gitaly::WikiGetAllPagesResponse.new(page: Gitaly::WikiPage.new(page_2_info)),
+        Gitaly::WikiGetAllPagesResponse.new(page: Gitaly::WikiPage.new(raw_data: 'd')),
+        Gitaly::WikiGetAllPagesResponse.new(end_of_page: true)
+      ]
+    end
+    let(:wiki_page_1) { subject[0].first }
+    let(:wiki_page_1_version) { subject[0].last }
+    let(:wiki_page_2) { subject[1].first }
+    let(:wiki_page_2_version) { subject[1].last }
+
+    subject { client.get_all_pages }
+
+    it 'sends a wiki_get_all_pages message' do
+      expect_any_instance_of(Gitaly::WikiService::Stub)
+        .to receive(:wiki_get_all_pages)
+        .with(gitaly_request_with_path(storage_name, relative_path), kind_of(Hash))
+        .and_return([].each)
+
+      subject
+    end
+
+    it 'concatenates the raw data and returns a pair of WikiPage and WikiPageVersion for each page' do
+      expect_any_instance_of(Gitaly::WikiService::Stub)
+        .to receive(:wiki_get_all_pages)
+        .with(gitaly_request_with_path(storage_name, relative_path), kind_of(Hash))
+        .and_return(response.each)
+
+      expect(subject.size).to be(2)
+      expect(wiki_page_1.title).to eq('My Page')
+      expect(wiki_page_1.raw_data).to eq('ab')
+      expect(wiki_page_1_version.format).to eq('markdown')
+      expect(wiki_page_2.title).to eq('My Page 2')
+      expect(wiki_page_2.raw_data).to eq('cd')
+      expect(wiki_page_2_version.format).to eq('markdown')
+    end
+  end
 end