Add GitalyKeysetPager for branch list API

- optional `pagination: 'keyset'` can be passed in to enable keyset
pagination
- adds next page link to Link and Links header

changelogs/unreleased/change_pagination_for_branch_list_api.yml

+---
+title: Allow optional keyset pagination for branch list API
+merge_request: 37524
+author:
+type: added

lib/api/branches.rb

@@ -40,12 +40,8 @@ module API
 
         repository = user_project.repository
 
-        if Feature.enabled?(:branch_list_keyset_pagination, user_project)
-          branches = BranchesFinder.new(repository, declared_params(include_missing: false)).execute(gitaly_pagination: true)
-        else
-          branches = BranchesFinder.new(repository, declared_params(include_missing: false)).execute
-          branches = paginate(::Kaminari.paginate_array(branches))
-        end
+        branches_finder = BranchesFinder.new(repository, declared_params(include_missing: false))
+        branches = Gitlab::Pagination::GitalyKeysetPager.new(self, user_project).paginate(branches_finder)
 
         merged_branch_names = repository.merged_branch_names(branches.map(&:name))
 

lib/gitlab/pagination/gitaly_keyset_pager.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Pagination
+    class GitalyKeysetPager
+      attr_reader :request_context, :project
+      delegate :params, to: :request_context
+
+      def initialize(request_context, project)
+        @request_context = request_context
+        @project = project
+      end
+
+      # It is expected that the given finder will respond to `execute` method with `gitaly_pagination: true` option
+      # and supports pagination via gitaly.
+      def paginate(finder)
+        return paginate_via_gitaly(finder) if keyset_pagination_enabled?
+
+        branches = ::Kaminari.paginate_array(finder.execute)
+        Gitlab::Pagination::OffsetPagination
+          .new(request_context)
+          .paginate(branches)
+      end
+
+      private
+
+      def keyset_pagination_enabled?
+        Feature.enabled?(:branch_list_keyset_pagination, project) && params[:pagination] == 'keyset'
+      end
+
+      def paginate_via_gitaly(finder)
+        finder.execute(gitaly_pagination: true).tap do |records|
+          apply_headers(records)
+        end
+      end
+
+      def apply_headers(records)
+        if records.count == params[:per_page]
+          Gitlab::Pagination::Keyset::HeaderBuilder
+            .new(request_context)
+            .add_next_page_header(
+              query_params_for(records.last)
+            )
+        end
+      end
+
+      def query_params_for(record)
+        # NOTE: page_token is name for now, but it could be dynamic if we have other gitaly finders
+        # that is based on something other than name
+        { page_token: record.name }
+      end
+    end
+  end
+end

lib/gitlab/pagination/keyset/header_builder.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Pagination
+    module Keyset
+      class HeaderBuilder
+        attr_reader :request_context
+        delegate :params, :header, :request, to: :request_context
+
+        def initialize(request_context)
+          @request_context = request_context
+        end
+
+        def add_next_page_header(query_params)
+          link = next_page_link(page_href(query_params))
+          header('Links', link)
+          header('Link', link)
+        end
+
+        private
+
+        def next_page_link(href)
+          %(<#{href}>; rel="next")
+        end
+
+        def page_href(query_params)
+          base_request_uri.tap do |uri|
+            uri.query = updated_params(query_params).to_query
+          end.to_s
+        end
+
+        def base_request_uri
+          @base_request_uri ||= URI.parse(request.url).tap do |uri|
+            uri.host = Gitlab.config.gitlab.host
+            uri.port = Gitlab.config.gitlab.port
+          end
+        end
+
+        def updated_params(query_params)
+          params.merge(query_params)
+        end
+      end
+    end
+  end
+end

lib/gitlab/pagination/keyset/request_context.rb

@@ -24,9 +24,11 @@ module Gitlab
         end
 
         def apply_headers(next_page)
-          link = pagination_links(next_page)
-          request.header('Links', link)
-          request.header('Link', link)
+          Gitlab::Pagination::Keyset::HeaderBuilder
+            .new(request)
+            .add_next_page_header(
+              query_params_for(next_page)
+            )
         end
 
         private
@@ -63,25 +65,8 @@ module Gitlab
           end
         end
 
-        def page_href(page)
-          base_request_uri.tap do |uri|
-            uri.query = query_params_for(page).to_query
-          end.to_s
-        end
-
-        def pagination_links(next_page)
-          %(<#{page_href(next_page)}>; rel="next")
-        end
-
-        def base_request_uri
-          @base_request_uri ||= URI.parse(request.request.url).tap do |uri|
-            uri.host = Gitlab.config.gitlab.host
-            uri.port = Gitlab.config.gitlab.port
-          end
-        end
-
         def query_params_for(page)
-          request.params.merge(lower_bounds_params(page))
+          lower_bounds_params(page)
         end
       end
     end

spec/lib/gitlab/pagination/gitaly_keyset_pager_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Pagination::GitalyKeysetPager do
+  let(:pager) { described_class.new(request_context, project) }
+
+  let_it_be(:project) { create(:project, :repository) }
+
+  let(:request_context) { double("request context") }
+  let(:finder) { double("branch finder") }
+  let(:custom_port) { 8080 }
+  let(:incoming_api_projects_url) { "#{Gitlab.config.gitlab.url}:#{custom_port}/api/v4/projects" }
+
+  before do
+    stub_config_setting(port: custom_port)
+  end
+
+  describe '.paginate' do
+    let(:base_query) { { per_page: 2 } }
+    let(:query) { base_query }
+
+    before do
+      allow(request_context).to receive(:params).and_return(query)
+      allow(request_context).to receive(:header)
+    end
+
+    shared_examples_for 'offset pagination' do
+      let(:paginated_array) { double 'paginated array' }
+      let(:branches) { [] }
+
+      it 'uses offset pagination' do
+        expect(finder).to receive(:execute).and_return(branches)
+        expect(Kaminari).to receive(:paginate_array).with(branches).and_return(paginated_array)
+        expect_next_instance_of(Gitlab::Pagination::OffsetPagination) do |offset_pagination|
+          expect(offset_pagination).to receive(:paginate).with(paginated_array)
+        end
+
+        pager.paginate(finder)
+      end
+    end
+
+    context 'with branch_list_keyset_pagination feature off' do
+      before do
+        stub_feature_flags(branch_list_keyset_pagination: false)
+      end
+
+      context 'without keyset pagination option' do
+        it_behaves_like 'offset pagination'
+      end
+
+      context 'with keyset pagination option' do
+        let(:query) { base_query.merge(pagination: 'keyset') }
+
+        it_behaves_like 'offset pagination'
+      end
+    end
+
+    context 'with branch_list_keyset_pagination feature on' do
+      before do
+        stub_feature_flags(branch_list_keyset_pagination: project)
+      end
+
+      context 'without keyset pagination option' do
+        it_behaves_like 'offset pagination'
+      end
+
+      context 'with keyset pagination option' do
+        let(:query) { base_query.merge(pagination: 'keyset') }
+        let(:fake_request) { double(url: "#{incoming_api_projects_url}?#{query.to_query}") }
+
+        before do
+          allow(request_context).to receive(:request).and_return(fake_request)
+          expect(finder).to receive(:execute).with(gitaly_pagination: true).and_return(branches)
+        end
+
+        context 'when next page could be available' do
+          let(:branch1) { double 'branch', name: 'branch1' }
+          let(:branch2) { double 'branch', name: 'branch2' }
+          let(:branches) { [branch1, branch2] }
+
+          let(:expected_next_page_link) { %Q(<#{incoming_api_projects_url}?#{query.merge(page_token: branch2.name).to_query}>; rel="next") }
+
+          it 'uses keyset pagination and adds link headers' do
+            expect(request_context).to receive(:header).with('Links', expected_next_page_link)
+            expect(request_context).to receive(:header).with('Link', expected_next_page_link)
+
+            pager.paginate(finder)
+          end
+        end
+
+        context 'when the current page is the last page' do
+          let(:branch1) { double 'branch', name: 'branch1' }
+          let(:branches) { [branch1] }
+
+          it 'uses keyset pagination without link headers' do
+            expect(request_context).not_to receive(:header).with('Links', anything)
+            expect(request_context).not_to receive(:header).with('Link', anything)
+
+            pager.paginate(finder)
+          end
+        end
+      end
+    end
+  end
+end

spec/requests/api/branches_spec.rb

@@ -39,9 +39,11 @@ RSpec.describe API::Branches do
       end
 
       context 'with branch_list_keyset_pagination feature off' do
-        context 'with legacy pagination params' do
+        let(:base_params) { {} }
+
+        context 'with offset pagination params' do
           it 'returns the repository branches' do
-            get api(route, current_user), params: { per_page: 100 }
+            get api(route, current_user), params: base_params.merge(per_page: 100)
 
             expect(response).to have_gitlab_http_status(:ok)
             expect(response).to match_response_schema('public_api/v4/branches')
@@ -53,7 +55,7 @@ RSpec.describe API::Branches do
           it 'determines only a limited number of merged branch names' do
             expect(API::Entities::Branch).to receive(:represent).with(anything, has_up_to_merged_branch_names_count(2)).and_call_original
 
-            get api(route, current_user), params: { per_page: 2 }
+            get api(route, current_user), params: base_params.merge(per_page: 2)
 
             expect(response).to have_gitlab_http_status(:ok)
             expect(json_response.count).to eq 2
@@ -64,7 +66,7 @@ RSpec.describe API::Branches do
           it 'merge status matches reality on paginated input' do
             expected_first_branch_name = project.repository.branches_sorted_by('name')[20].name
 
-            get api(route, current_user), params: { per_page: 20, page: 2 }
+            get api(route, current_user), params: base_params.merge(per_page: 20, page: 2)
 
             expect(response).to have_gitlab_http_status(:ok)
             expect(json_response.count).to eq 20
@@ -74,11 +76,11 @@ RSpec.describe API::Branches do
           end
         end
 
-        context 'with gitaly pagination params ' do
+        context 'with gitaly pagination params' do
           it 'merge status matches reality on paginated input' do
             expected_first_branch_name = project.repository.branches_sorted_by('name').first.name
 
-            get api(route, current_user), params: { per_page: 20, page_token: 'feature' }
+            get api(route, current_user), params: base_params.merge(per_page: 20, page_token: 'feature')
 
             expect(response).to have_gitlab_http_status(:ok)
             expect(json_response.count).to eq 20
@@ -91,15 +93,19 @@ RSpec.describe API::Branches do
 
       context 'with branch_list_keyset_pagination feature on' do
         before do
-          stub_feature_flags(branch_list_keyset_pagination: true)
+          stub_feature_flags(branch_list_keyset_pagination: project)
         end
 
+        context 'with keyset pagination option' do
+          let(:base_params) { { pagination: 'keyset' } }
+
           context 'with gitaly pagination params ' do
             it 'returns the repository branches' do
-            get api(route, current_user), params: { per_page: 100 }
+              get api(route, current_user), params: base_params.merge(per_page: 100)
 
               expect(response).to have_gitlab_http_status(:ok)
               expect(response).to match_response_schema('public_api/v4/branches')
+              expect(response.headers).not_to include('Link', 'Links')
               branch_names = json_response.map { |x| x['name'] }
               expect(branch_names).to match_array(project.repository.branch_names)
             end
@@ -107,9 +113,10 @@ RSpec.describe API::Branches do
             it 'determines only a limited number of merged branch names' do
               expect(API::Entities::Branch).to receive(:represent).with(anything, has_up_to_merged_branch_names_count(2)).and_call_original
 
-            get api(route, current_user), params: { per_page: 2 }
+              get api(route, current_user), params: base_params.merge(per_page: 2)
 
               expect(response).to have_gitlab_http_status(:ok)
+              expect(response.headers).to include('Link', 'Links')
               expect(json_response.count).to eq 2
 
               check_merge_status(json_response)
@@ -118,7 +125,7 @@ RSpec.describe API::Branches do
             it 'merge status matches reality on paginated input' do
               expected_first_branch_name = project.repository.branches_sorted_by('name').drop_while { |b| b.name <= 'feature' }.first.name
 
-            get api(route, current_user), params: { per_page: 20, page_token: 'feature' }
+              get api(route, current_user), params: base_params.merge(per_page: 20, page_token: 'feature')
 
               expect(response).to have_gitlab_http_status(:ok)
               expect(json_response.count).to eq 20
@@ -128,10 +135,10 @@ RSpec.describe API::Branches do
             end
           end
 
-        context 'with legacy pagination params' do
+          context 'with offset pagination params' do
             it 'ignores legacy pagination params' do
               expected_first_branch_name = project.repository.branches_sorted_by('name').first.name
-            get api(route, current_user), params: { per_page: 20, page: 2 }
+              get api(route, current_user), params: base_params.merge(per_page: 20, page: 2)
 
               expect(response).to have_gitlab_http_status(:ok)
               expect(json_response.first['name']).to eq(expected_first_branch_name)
@@ -140,6 +147,7 @@ RSpec.describe API::Branches do
             end
           end
         end
+      end
 
       context 'when repository is disabled' do
         include_context 'disabled repository'