Allow UsageData.count to use count_by:

2 files changed, 8 insertions, 2 deletions

diff --git a/lib/gitlab/usage_data.rb b/lib/gitlab/usage_data.rb
index 7572c0bdbfd..f02d9f7d7e7 100644
--- a/lib/gitlab/usage_data.rb
+++ b/lib/gitlab/usage_data.rb
@@ -176,8 +176,8 @@ module Gitlab
         {} # augmented in EE
       end
 
-      def count(relation, fallback: -1)
-        relation.count
+      def count(relation, count_by: nil, fallback: -1)
+        count_by ? relation.count(count_by) : relation.count
       rescue ActiveRecord::StatementInvalid
         fallback
       end
diff --git a/spec/lib/gitlab/usage_data_spec.rb b/spec/lib/gitlab/usage_data_spec.rb
index 90a534de202..7fe60fd214c 100644
--- a/spec/lib/gitlab/usage_data_spec.rb
+++ b/spec/lib/gitlab/usage_data_spec.rb
@@ -234,6 +234,12 @@ describe Gitlab::UsageData do
       expect(described_class.count(relation)).to eq(1)
     end
 
+    it 'returns the count for count_by when provided' do
+      allow(relation).to receive(:count).with(:creator_id).and_return(2)
+
+      expect(described_class.count(relation, count_by: :creator_id)).to eq(2)
+    end
+
     it 'returns the fallback value when counting fails' do
       allow(relation).to receive(:count).and_raise(ActiveRecord::StatementInvalid.new(''))