Distinct count only for allowed foreign keys

This prevents performance issues

ee/lib/ee/gitlab/usage_data.rb

@@ -340,6 +340,7 @@ module EE
 
         # rubocop:disable CodeReuse/ActiveRecord
         # rubocop: disable UsageData/LargeTable
+        # rubocop: disable UsageData/DistinctCountByLargeForeignKey
         def count_secure_pipelines(time_period)
           return {} if time_period.blank?
 
@@ -357,7 +358,8 @@ module EE
 
           pipelines_with_secure_jobs
         end
-        # rubocop: enabled UsageData/LargeTable
+        # rubocop: enable UsageData/LargeTable
+        # rubocop: enable UsageData/DistinctCountByLargeForeignKey
 
         def approval_merge_request_rule_minimum_id
           strong_memoize(:approval_merge_request_rule_minimum_id) do

lib/gitlab/usage_data.rb

@@ -306,6 +306,7 @@ module Gitlab
         Gitlab::UsageData::Topology.new.topology_usage_data
       end
 
+      # rubocop: disable UsageData/DistinctCountByLargeForeignKey
       def ingress_modsecurity_usage
         ##
         # This method measures usage of the Modsecurity Web Application Firewall across the entire
@@ -326,6 +327,7 @@ module Gitlab
           ingress_modsecurity_not_installed: distinct_count(successful_deployments_with_cluster(::Clusters::Applications::Ingress.modsecurity_not_installed), column)
         }
       end
+      # rubocop: enable UsageData/DistinctCountByLargeForeignKey
 
       # rubocop: disable CodeReuse/ActiveRecord
       def container_expiration_policies_usage
@@ -729,9 +731,11 @@ module Gitlab
       end
 
       # rubocop: disable CodeReuse/ActiveRecord
+      # rubocop: disable UsageData/DistinctCountByLargeForeignKey
       def cluster_applications_user_distinct_count(applications, time_period)
         distinct_count(applications.where(time_period).available.joins(:cluster), 'clusters.user_id')
       end
+      # rubocop: enable UsageData/DistinctCountByLargeForeignKey
 
       def clusters_user_distinct_count(clusters, time_period)
         distinct_count(clusters.where(time_period), :user_id)

rubocop/cop/usage_data/distinct_count_by_large_foreign_key.rb

+# frozen_string_literal: true
+
+module RuboCop
+  module Cop
+    module UsageData
+      # Allows counts only for selected tables' foreign keys for `distinct_count` method.
+      #
+      # Because distinct_counts over large tables' foreign keys will take a long time
+      #
+      # @example
+      #
+      #   # bad because pipeline_id points to a large table
+      #   distinct_count(Ci::Build, :commit_id)
+      #
+      class DistinctCountByLargeForeignKey < RuboCop::Cop::Cop
+        MSG = 'Avoid doing `%s` for large foreign keys.'.freeze
+
+        def_node_matcher :distinct_count?, <<-PATTERN
+          (send _ $:distinct_count $...)
+        PATTERN
+
+        def on_send(node)
+          distinct_count?(node) do |method_name, method_arguments|
+            next unless method_arguments && method_arguments.length >= 2
+            next if allowed_foreign_key?(method_arguments[1])
+
+            add_offense(node, location: :selector, message: format(MSG, method_name))
+          end
+        end
+
+        private
+
+        def allowed_foreign_key?(key)
+          key.type == :sym && allowed_foreign_keys.include?(key.value)
+        end
+
+        def allowed_foreign_keys
+          cop_config['AllowedForeignKeys'] || []
+        end
+      end
+    end
+  end
+end

rubocop/rubocop-usage-data.yml

@@ -30,3 +30,16 @@ UsageData/LargeTable:
     - :arel_table
     - :minimum
     - :maximum
+UsageData/DistinctCountByLargeForeignKey:
+  Enabled: true
+  Include:
+    - 'lib/gitlab/usage_data.rb'
+    - 'ee/lib/ee/gitlab/usage_data.rb'
+  AllowedForeignKeys:
+    - :user_id
+    - :author_id
+    - :creator_id
+    - :owner_id
+    - :project_id
+    - :issue_id
+    - :merge_request_id

spec/rubocop/cop/usage_data/distinct_count_by_large_foreign_key_spec.rb

+# frozen_string_literal: true
+
+require 'fast_spec_helper'
+
+require 'rubocop'
+require 'rubocop/rspec/support'
+
+require_relative '../../../../rubocop/cop/usage_data/distinct_count_by_large_foreign_key'
+
+RSpec.describe RuboCop::Cop::UsageData::DistinctCountByLargeForeignKey, type: :rubocop do
+  include CopHelper
+
+  let(:allowed_foreign_keys) { %i[author_id user_id] }
+
+  let(:config) do
+    RuboCop::Config.new('UsageData/DistinctCountByLargeForeignKey' => {
+                          'AllowedForeignKeys' => allowed_foreign_keys
+                        })
+  end
+
+  subject(:cop) { described_class.new(config) }
+
+  context 'when counting by disallowed key' do
+    it 'register an offence' do
+      inspect_source('distinct_count(Issue, :creator_id)')
+
+      expect(cop.offenses.size).to eq(1)
+    end
+  end
+
+  context 'when calling by allowed key' do
+    it 'does not register an offence' do
+      inspect_source('distinct_count(Issue, :author_id)')
+
+      expect(cop.offenses).to be_empty
+    end
+  end
+end