Incorporate feedback fixes

Moved all query methods to a model concern

ee/app/models/concerns/ee/usage_statistics.rb

+# frozen_string_literal: true
+
+module EE
+  module UsageStatistics
+    extend ActiveSupport::Concern
+
+    class_methods do
+      def distinct_count_by(column = nil, fallback = -1)
+        distinct.count(column)
+      rescue ActiveRecord::StatementInvalid
+        fallback
+      end
+    end
+  end
+end

ee/app/models/ee/group_member.rb

@@ -6,6 +6,7 @@ module EE
 
     prepended do
       extend ::Gitlab::Utils::Override
+      include UsageStatistics
 
       validate :sso_enforcement, if: :group
 

ee/app/models/ldap_key.rb

 # frozen_string_literal: true
 
 class LDAPKey < Key
+  include EE::UsageStatistics # rubocop: disable Cop/InjectEnterpriseEditionModule
+
   def can_delete?
     false
   end

ee/lib/ee/gitlab/usage_data.rb

@@ -174,20 +174,21 @@ module EE
         end
 
         # Source: https://gitlab.com/gitlab-data/analytics/blob/master/transform/snowflake-dbt/data/ping_metrics_to_stage_mapping_data.csv
-        # rubocop: disable CodeReuse/ActiveRecord
         def usage_activity_by_stage
           {
             usage_activity_by_stage: {
-              manage: {
-                groups: count(::GroupMember.distinct, count_by: :user_id),
-                ldap_group_links: -1, # no creator / user ID
-                ldap_keys: count(::LDAPKey.distinct, count_by: :user_id),
-                ldap_users: count(::GroupMember.distinct.of_ldap_type, count_by: :user_id)
-              }
+              manage: usage_activity_by_stage_manage
             }
           }
         end
-        # rubocop: enable CodeReuse/ActiveRecord
+
+        def usage_activity_by_stage_manage
+          {
+            groups: ::GroupMember.distinct_count_by(:user_id),
+            ldap_keys: ::LDAPKey.distinct_count_by(:user_id),
+            ldap_users: ::GroupMember.of_ldap_type.distinct_count_by(:user_id)
+          }
+        end
       end
     end
   end

ee/spec/lib/ee/gitlab/usage_data_spec.rb

+# frozen_string_literal: true
+
 require 'spec_helper'
 
 describe Gitlab::UsageData do
   describe '.uncached_data' do
     context 'when the :usage_activity_by_stage feaure is not enabled' do
-      it 'does not include usage_activity_by_stage data' do
+      before do
         stub_feature_flags(usage_activity_by_stage: false)
+      end
 
+      it 'does not include usage_activity_by_stage data' do
         expect(described_class.uncached_data).not_to include(:usage_activity_by_stage)
       end
     end
 
-    context 'when the :usage_activity_by_stage feaure is enabled' do
-      before do
-        stub_feature_flags(usage_activity_by_stage: true)
-      end
-
+    context 'when the :usage_activity_by_stage feature is enabled' do
       it 'includes usage_activity_by_stage data' do
         expect(described_class.uncached_data).to include(:usage_activity_by_stage)
       end
@@ -26,7 +26,7 @@ describe Gitlab::UsageData do
           create(:key, type: 'LDAPKey', user: user)
           create(:group_member, ldap: true, user: user)
 
-          expect(described_class.uncached_data[:usage_activity_by_stage][:manage]).to eq(groups: 1, ldap_group_links: -1, ldap_keys: 1, ldap_users: 1)
+          expect(described_class.uncached_data[:usage_activity_by_stage][:manage]).to eq(groups: 1, ldap_keys: 1, ldap_users: 1)
         end
       end
     end

ee/spec/models/concerns/ee/usage_statistics_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+describe EE::UsageStatistics do
+  describe '.distinct_count_by' do
+    let(:user_1) { create(:user) }
+    let(:user_2) { create(:user) }
+
+    context 'two records created by the same user' do
+      let!(:models_created_by_user_1) { create_list(:group_member, 2, user: user_1)}
+
+      it 'returns a count of 1' do
+        expect(::GroupMember.distinct_count_by(:user_id)).to eq(1)
+      end
+    end
+
+    context 'one record created by each user' do
+      let!(:model_created_by_user_1) { create(:group_member, user: user_1)}
+      let!(:model_created_by_user_2) { create(:group_member, user: user_2)}
+
+      it 'returns a count of 2' do
+        expect(::GroupMember.distinct_count_by(:user_id)).to eq(2)
+      end
+    end
+
+    context 'the count query times out' do
+      before do
+        allow_any_instance_of(ActiveRecord::Relation)
+          .to receive(:count).and_raise(ActiveRecord::StatementInvalid.new(''))
+      end
+
+      it 'does not raise an error' do
+        expect { ::GroupMember.distinct_count_by(:user_id) }.not_to raise_error
+      end
+
+      it 'returns -1' do
+        expect(::GroupMember.distinct_count_by(:user_id)).to eq(-1)
+      end
+    end
+  end
+end