1 files changed, 5 insertions, 1 deletions
diff --git a/lib/gitlab/database/batch_count.rb b/lib/gitlab/database/batch_count.rb
index 92a41bb36ee..7a064fb4005 100644
--- a/lib/gitlab/database/batch_count.rb
+++ b/lib/gitlab/database/batch_count.rb
@@ -35,6 +35,10 @@ module Gitlab
         BatchCounter.new(relation, column: column).count(batch_size: batch_size, start: start, finish: finish)
       end
 
+      def batch_count_with_timeout(relation, column = nil, batch_size: nil, start: nil, finish: nil, timeout: nil, partial_results: nil)
+        BatchCounter.new(relation, column: column).count_with_timeout(batch_size: batch_size, start: start, finish: finish, timeout: timeout, partial_results: partial_results)
+      end
+
       def batch_distinct_count(relation, column = nil, batch_size: nil, start: nil, finish: nil)
         BatchCounter.new(relation, column: column).count(mode: :distinct, batch_size: batch_size, start: start, finish: finish)
       end
@@ -44,7 +48,7 @@ module Gitlab
       end
 
       def batch_average(relation, column, batch_size: nil, start: nil, finish: nil)
-        BatchCounter.new(relation, column: nil, operation: :average, operation_args: [column]).count(batch_size: batch_size, start: start, finish: finish)
+        BatchAverageCounter.new(relation, column).count(batch_size: batch_size)
       end
 
       class << self