1 files changed, 18 insertions, 3 deletions
diff --git a/lib/gitlab/analytics/cycle_analytics/data_collector.rb b/lib/gitlab/analytics/cycle_analytics/data_collector.rb
index 10a008a76d5..56179533ffb 100644
--- a/lib/gitlab/analytics/cycle_analytics/data_collector.rb
+++ b/lib/gitlab/analytics/cycle_analytics/data_collector.rb
@@ -12,6 +12,8 @@ module Gitlab
       class DataCollector
         include Gitlab::Utils::StrongMemoize
 
+        MAX_COUNT = 1001
+
         delegate :serialized_records, to: :records_fetcher
 
         def initialize(stage:, params: {})
@@ -27,13 +29,19 @@ module Gitlab
 
         def median
           strong_memoize(:median) do
-            Median.new(stage: stage, query: query)
+            Median.new(stage: stage, query: query, params: params)
           end
         end
 
         def average
           strong_memoize(:average) do
-            Average.new(stage: stage, query: query)
+            Average.new(stage: stage, query: query, params: params)
+          end
+        end
+
+        def count
+          strong_memoize(:count) do
+            limit_count
           end
         end
 
@@ -44,9 +52,16 @@ module Gitlab
         def query
           BaseQueryBuilder.new(stage: stage, params: params).build
         end
+
+        # Limiting the maximum number of records so the COUNT(*) query stays efficient for large groups.
+        # COUNT = 1001, show 1000+ on the UI
+        # COUNT < 1001, show the actual number on the UI
+        def limit_count
+          query.limit(MAX_COUNT).count
+        end
       end
     end
   end
 end
 
-Gitlab::Analytics::CycleAnalytics::DataCollector.prepend_if_ee('EE::Gitlab::Analytics::CycleAnalytics::DataCollector')
+Gitlab::Analytics::CycleAnalytics::DataCollector.prepend_mod_with('Gitlab::Analytics::CycleAnalytics::DataCollector')