Add a new total_count_with_limit to Kaminari's methods52674-api-v4-projects-project_id-jobs-endpoint-hits-statement-timeout

Signed-off-by: Rémy Coutable <remy@rymai.me>

3 files changed, 48 insertions, 28 deletions

diff --git a/config/initializers/kaminari_active_record_relation_methods_with_limit.rb b/config/initializers/kaminari_active_record_relation_methods_with_limit.rb
new file mode 100644
index 00000000000..cc20b83b234
--- /dev/null
+++ b/config/initializers/kaminari_active_record_relation_methods_with_limit.rb
@@ -0,0 +1,41 @@
+module Kaminari
+  # Active Record specific page scope methods implementations
+  module ActiveRecordRelationMethodsWithLimit
+    MAX_COUNT_LIMIT = 10_000
+
+    # This is a modified version of
+    # https://github.com/kaminari/kaminari/blob/c5186f5d9b7f23299d115408e62047447fd3189d/kaminari-activerecord/lib/kaminari/activerecord/active_record_relation_methods.rb#L17-L41
+    # that limit the COUNT query to 10,000 to avoid query timeouts.
+    # rubocop: disable Gitlab/ModuleWithInstanceVariables
+    def total_count_with_limit(column_name = :all, _options = nil) #:nodoc:
+      return @total_count if defined?(@total_count) && @total_count
+
+      # There are some cases that total count can be deduced from loaded records
+      if loaded?
+        # Total count has to be 0 if loaded records are 0
+        return @total_count = 0 if (current_page == 1) && @records.empty?
+        # Total count is calculable at the last page
+        return @total_count = (current_page - 1) * limit_value + @records.length if @records.any? && (@records.length < limit_value)
+      end
+
+      # #count overrides the #select which could include generated columns referenced in #order, so skip #order here, where it's irrelevant to the result anyway
+      c = except(:offset, :limit, :order)
+      # Remove includes only if they are irrelevant
+      c = c.except(:includes) unless references_eager_loaded_tables?
+      # .group returns an OrderedHash that responds to #count
+      # The following line was modified from `c = c.count(:all)`
+      c = c.limit(MAX_COUNT_LIMIT + 1).count(column_name)
+      @total_count =
+        if c.is_a?(Hash) || c.is_a?(ActiveSupport::OrderedHash)
+          c.count
+        elsif c.respond_to? :count
+          c.count(column_name)
+        else
+          c
+        end
+    end
+    # rubocop: enable Gitlab/ModuleWithInstanceVariables
+
+    Kaminari::ActiveRecordRelationMethods.prepend(self)
+  end
+end
diff --git a/lib/api/helpers/pagination.rb b/lib/api/helpers/pagination.rb
index cda2557328e..8c01e1276e8 100644
--- a/lib/api/helpers/pagination.rb
+++ b/lib/api/helpers/pagination.rb
@@ -3,8 +3,6 @@
 module API
   module Helpers
     module Pagination
-      TOTAL_COUNT_LIMIT = 1000
-
       def paginate(relation)
         strategy = if params[:pagination] == 'keyset' && Feature.enabled?('api_keyset_pagination')
                      KeysetPaginationStrategy
@@ -191,30 +189,11 @@ module API
           pagination_data = relation.page(params[:page]).per(params[:per_page])
           return pagination_data unless pagination_data.is_a?(ActiveRecord::Relation)
 
-          limited_total_count = total_count_with_limit(relation)
-          if limited_total_count > TOTAL_COUNT_LIMIT
+          limited_total_count = pagination_data.total_count_with_limit
+          if limited_total_count > Kaminari::ActiveRecordRelationMethods::MAX_COUNT_LIMIT
             pagination_data.without_count
           else
-            # There's currently no way to manually set @total_count in a clean way:
-            # https://github.com/kaminari/kaminari/blob/c5186f5d9b7f23299d115408e62047447fd3189d/kaminari-activerecord/lib/kaminari/activerecord/active_record_relation_methods.rb#L17-L41
-            pagination_data.tap { |data| data.instance_variable_set(:@total_count, limited_total_count) }
-          end
-        end
-
-        # This is a modified version of https://github.com/kaminari/kaminari/blob/c5186f5d9b7f23299d115408e62047447fd3189d/kaminari-activerecord/lib/kaminari/activerecord/active_record_relation_methods.rb#L17-L41
-        def total_count_with_limit(relation)
-          # #count overrides the #select which could include generated columns referenced in #order, so skip #order here, where it's irrelevant to the result anyway
-          relation = relation.except(:offset, :limit, :order)
-          # Remove includes only if they are irrelevant
-          relation = relation.except(:includes) unless relation.__send__(:references_eager_loaded_tables?)
-          # .group returns an OrderedHash that responds to #count
-          count = relation.limit(TOTAL_COUNT_LIMIT + 1).count(:all)
-          if count.is_a?(Hash) || count.is_a?(ActiveSupport::OrderedHash)
-            count.count
-          elsif count.respond_to?(:count)
-            count.count(:all)
-          else
-            count
+            pagination_data
           end
         end
 
diff --git a/spec/lib/api/helpers/pagination_spec.rb b/spec/lib/api/helpers/pagination_spec.rb
index 752931469fa..72c6c55315b 100644
--- a/spec/lib/api/helpers/pagination_spec.rb
+++ b/spec/lib/api/helpers/pagination_spec.rb
@@ -259,9 +259,9 @@ describe API::Helpers::Pagination do
           subject.paginate(resource)
         end
 
-        context 'when resources count is less than TOTAL_COUNT_LIMIT' do
+        context 'when resources count is less than MAX_COUNT_LIMIT' do
           before do
-            stub_const("#{described_class}::TOTAL_COUNT_LIMIT", 4)
+            stub_const("::Kaminari::ActiveRecordRelationMethods::MAX_COUNT_LIMIT", 4)
           end
 
           it 'executes only one SELECT COUNT query' do
@@ -269,9 +269,9 @@ describe API::Helpers::Pagination do
           end
         end
 
-        context 'when resources count is more than TOTAL_COUNT_LIMIT' do
+        context 'when resources count is more than MAX_COUNT_LIMIT' do
           before do
-            stub_const("#{described_class}::TOTAL_COUNT_LIMIT", 2)
+            stub_const("::Kaminari::ActiveRecordRelationMethods::MAX_COUNT_LIMIT", 2)
           end
 
           it 'executes only one SELECT COUNT query' do