diff options
| -rw-r--r-- | lib/gitlab/usage_data.rb | 4 | ||||
| -rw-r--r-- | spec/lib/gitlab/usage_data_spec.rb | 6 |
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('')) |