Add rubocop rule for count large table in usage data

  - Add test
  - Add configuration for CountLargeTable cop
  - Add test for allowed methods
  - Add more tables
  - Add rules exceptions
  - Use Include usage_data.rb files
  - Rename cop to LargeTable
  - Rename test file
  - Add cop comments

.rubocop.yml

@@ -9,6 +9,7 @@ require:
 inherit_from:
   - .rubocop_todo.yml
   - ./rubocop/rubocop-migrations.yml
+  - ./rubocop/rubocop-usage-data.yml
 
 inherit_mode:
   merge:

ee/lib/ee/gitlab/usage_data.rb

@@ -76,7 +76,9 @@ module EE
           if license
             usage_data[:license_md5] = license.md5
             usage_data[:license_id] = license.license_id
+            # rubocop: disable UsageData/LargeTable
             usage_data[:historical_max_users] = ::HistoricalData.max_historical_user_count
+            # rubocop: enable UsageData/LargeTable
             usage_data[:licensee] = license.licensee
             usage_data[:license_user_count] = license.restricted_user_count
             usage_data[:license_starts_at] = license.starts_at
@@ -98,6 +100,9 @@ module EE
         end
 
         def approval_rules_counts
+          # rubocop: disable UsageData/LargeTable:
+          projects_with_service_desk = ::Project.where(service_desk_enabled: true)
+          # rubocop: enable UsageData/LargeTable:
           {
             approval_project_rules: count(ApprovalProjectRule),
             approval_project_rules_with_target_branch: count(ApprovalProjectRulesProtectedBranch, :approval_project_rule_id)
@@ -192,7 +197,9 @@ module EE
         end
 
         def epics_deepest_relationship_level
+          # rubocop: disable UsageData/LargeTable
           { epics_deepest_relationship_level: ::Epic.deepest_relationship_level.to_i }
+          # rubocop: enable UsageData/LargeTable
         end
 
         # Omitted because no user, creator or author associated: `auto_devops_disabled`, `auto_devops_enabled`

lib/gitlab/usage_data.rb

@@ -10,7 +10,6 @@
 #   alt_usage_data { Gitlab::VERSION }
 #   redis_usage_data(Gitlab::UsageDataCounters::WikiPageCounter)
 #   redis_usage_data { ::Gitlab::UsageCounters::PodLogs.usage_totals[:total] }
-
 module Gitlab
   class UsageData
     BATCH_SIZE = 100
@@ -332,14 +331,18 @@ module Gitlab
         finish = ::Project.maximum(:id)
 
         results[:projects_with_expiration_policy_disabled] = distinct_count(::ContainerExpirationPolicy.where(enabled: false), :project_id, start: start, finish: finish)
+        # rubocop: disable UsageData/LargeTable
         base = ::ContainerExpirationPolicy.active
+        # rubocop: enable UsageData/LargeTable
         results[:projects_with_expiration_policy_enabled] = distinct_count(base, :project_id, start: start, finish: finish)
 
+        # rubocop: disable UsageData/LargeTable
         %i[keep_n cadence older_than].each do |option|
           ::ContainerExpirationPolicy.public_send("#{option}_options").keys.each do |value| # rubocop: disable GitlabSecurity/PublicSend
             results["projects_with_expiration_policy_enabled_with_#{option}_set_to_#{value}".to_sym] = distinct_count(base.where(option => value), :project_id, start: start, finish: finish)
           end
         end
+        # rubocop: enable UsageData/LargeTable
 
         results[:projects_with_expiration_policy_enabled_with_keep_n_unset] = distinct_count(base.where(keep_n: nil), :project_id, start: start, finish: finish)
         results[:projects_with_expiration_policy_enabled_with_older_than_unset] = distinct_count(base.where(older_than: nil), :project_id, start: start, finish: finish)
@@ -350,9 +353,11 @@ module Gitlab
 
       # rubocop: disable CodeReuse/ActiveRecord
       def services_usage
+        # rubocop: disable UsageData/LargeTable:
         Service.available_services_names.without('jira').each_with_object({}) do |service_name, response|
           response["projects_#{service_name}_active".to_sym] = count(Service.active.where(template: false, type: "#{service_name}_service".camelize))
         end.merge(jira_usage).merge(jira_import_usage)
+        # rubocop: enable UsageData/LargeTable:
       end
 
       def jira_usage
@@ -365,6 +370,7 @@ module Gitlab
           projects_jira_active: 0
         }
 
+        # rubocop: disable UsageData/LargeTable:
         JiraService.active.includes(:jira_tracker_data).find_in_batches(batch_size: BATCH_SIZE) do |services|
           counts = services.group_by do |service|
             # TODO: Simplify as part of https://gitlab.com/gitlab-org/gitlab/issues/29404
@@ -376,21 +382,24 @@ module Gitlab
           results[:projects_jira_cloud_active] += counts[:cloud].size if counts[:cloud]
           results[:projects_jira_active] += services.size
         end
-
+        # rubocop: enable UsageData/LargeTable:
         results
       rescue ActiveRecord::StatementInvalid
         { projects_jira_server_active: FALLBACK, projects_jira_cloud_active: FALLBACK, projects_jira_active: FALLBACK }
       end
 
+      # rubocop: disable UsageData/LargeTable
       def successful_deployments_with_cluster(scope)
         scope
           .joins(cluster: :deployments)
           .merge(Clusters::Cluster.enabled)
           .merge(Deployment.success)
       end
+      # rubocop: enable UsageData/LargeTable
       # rubocop: enable CodeReuse/ActiveRecord
 
       def jira_import_usage
+        # rubocop: disable UsageData/LargeTable
         finished_jira_imports = JiraImportState.finished
 
         {
@@ -398,6 +407,7 @@ module Gitlab
           jira_imports_projects_count: distinct_count(finished_jira_imports, :project_id),
           jira_imports_total_imported_issues_count: alt_usage_data { JiraImportState.finished_imports_count }
         }
+        # rubocop: enable UsageData/LargeTable
       end
 
       def user_preferences_usage
@@ -445,6 +455,7 @@ module Gitlab
       end
 
       # rubocop: disable CodeReuse/ActiveRecord
+      # rubocop: disable UsageData/LargeTable
       def usage_activity_by_stage_configure(time_period)
         {
           clusters_applications_cert_managers: cluster_applications_user_distinct_count(::Clusters::Applications::CertManager, time_period),
@@ -465,6 +476,7 @@ module Gitlab
           project_clusters_enabled: clusters_user_distinct_count(::Clusters::Cluster.enabled.project_type, time_period)
         }
       end
+      # rubocop: enable UsageData/LargeTable
       # rubocop: enable CodeReuse/ActiveRecord
 
       # rubocop: disable CodeReuse/ActiveRecord

rubocop/cop/usage_data/event_counters.rb

-# frozen_string_literal: true
-
-module RuboCop
-  module Cop
-    module Gitlab
-      class EventCounters < RuboCop::Cop::Cop
-        MSG = 'Use the `count` or `distinct_count`'
-
-        def_node_matcher :events_table, <<~PATTERN
-          $(send (const {nil cbase} :Event ) ...)
-        PATTERN
-
-        def on_send(node)
-          return unless usage_data_files?(node)
-
-          matched_node = events_table(node)
-
-          return if matched_node.nil?
-
-          method_name = matched_node.parent.children[1]
-
-          add_offense(node, location: :expression) unless [:count, :distinct_count].include?(method_name)
-        end
-
-        private
-
-        def usage_data_files?(node)
-          path = node.location.expression.source_buffer.name
-          path.include?('usage_data.rb')
-        end
-      end
-    end
-  end
-end

rubocop/cop/usage_data/large_table.rb

+# frozen_string_literal: true
+
+module RuboCop
+  module Cop
+    module UsageData
+      class LargeTable < RuboCop::Cop::Cop
+        # This cop checks that batch count and distinct_count are used in usage_data.rb files in metrics based on ActiveRecord models.
+        #
+        # @example
+        #
+        # # bad
+        # Issue.count
+        # List.assignee.count
+        # ::Ci::Pipeline.auto_devops_source.count
+        # ZoomMeeting.distinct.count(:issue_id)
+        #
+        # # Good
+        # count(Issue)
+        # count(List.assignee)
+        # count(::Ci::Pipeline.auto_devops_source)
+        # distinct_count(ZoomMeeting, :issue_id)
+        MSG = 'Use one of the %{count_methods} methods for counting on %{class_name}'
+
+        # Match one level const as Issue, Gitlab
+        def_node_matcher :one_level_node, <<~PATTERN
+          (send
+            (const {nil? cbase} $...)
+          $...)
+        PATTERN
+
+        # Match two level const as ::Clusters::Cluster, ::Ci::Pipeline
+        def_node_matcher :two_level_node, <<~PATTERN
+          (send
+            (const
+              (const {nil? cbase} $...)
+            $...)
+          $...)
+        PATTERN
+
+        def_node_matcher :non_related_class?, <<~PATTERN
+          (send
+            (const {nil? cbase} ${#non_related?})
+          $...)
+        PATTERN
+
+        def on_send(node)
+          one_level_matches = one_level_node(node)
+          two_level_matches = two_level_node(node)
+
+          return unless Array(one_level_matches).any? || Array(two_level_matches).any?
+
+          if one_level_matches
+            class_name = one_level_matches[0].first
+            method_used = one_level_matches[1]&.first
+          else
+            class_name = "#{two_level_matches[0].first}::#{two_level_matches[1].first}".to_sym
+            method_used = two_level_matches[2]&.first
+          end
+
+          return if non_related?(class_name) || allowed_methods.include?(method_used)
+
+          counters_used = node.ancestors.any? { |ancestor| allowed_method?(ancestor) }
+
+          unless counters_used
+            add_offense(node, location: :expression, message: format(MSG, count_methods: count_methods.join(', '), class_name: class_name))
+          end
+        end
+
+        private
+
+        def count_methods
+          cop_config['CountMethods'] || []
+        end
+
+        def allowed_methods
+          cop_config['AllowedMethods'] || []
+        end
+
+        def non_related_classes
+          cop_config['NonRelatedClasses'] || []
+        end
+
+        def non_related?(class_name)
+          non_related_classes.include?(class_name)
+        end
+
+        def allowed_method?(ancestor)
+          ancestor.send_type? && !ancestor.dot? && count_methods.include?(ancestor.method_name)
+        end
+      end
+    end
+  end
+end

rubocop/rubocop-usage-data.yml

+UsageData/LargeTable:
+  Enabled: true
+  Include:
+    - 'lib/gitlab/usage_data.rb'
+    - 'ee/lib/ee/gitlab/usage_data.rb'
+  NonRelatedClasses:
+    - :EE::Gitlab::UsageData
+    - :Feature
+    - :Gitlab
+    - :Gitlab::AppLogger
+    - :Gitlab::Auth
+    - :Gitlab::CurrentSettings
+    - :Gitlab::Database
+    - :Gitlab::ErrorTracking
+    - :Gitlab::Geo
+    - :Gitlab::Git
+    - :Gitlab::IncomingEmail
+    - :Gitlab::Metrics
+    - :Gitlab::Runtime
+    - :Gitaly::Server
+    - :Gitlab::UsageData
+    - :License
+    - :Rails
+    - :Time
+    - :SECURE_PRODUCT_TYPES
+    - :Settings
+  CountMethods:
+    - :count
+    - :distinct_count
+  AllowedMethods:
+    - :arel_table
+    - :minimum
+    - :maximum

spec/rubocop/cop/usage_data/large_table_spec.rb

+# frozen_string_literal: true
+
+require 'fast_spec_helper'
+
+require 'rubocop'
+require 'rubocop/rspec/support'
+
+require_relative '../../../../rubocop/cop/usage_data/large_table'
+
+RSpec.describe RuboCop::Cop::UsageData::LargeTable, type: :rubocop do
+  include CopHelper
+
+  let(:large_tables) { %i[Rails Time] }
+  let(:count_methods) { %i[count distinct_count] }
+  let(:allowed_methods) { %i[minimum maximum] }
+
+  let(:config) do
+    RuboCop::Config.new('UsageData/LargeTable' => {
+                          'NonRelatedClasses' => large_tables,
+                          'CountMethods' => count_methods,
+                          'AllowedMethods' => allowed_methods
+                        })
+  end
+
+  subject(:cop) { described_class.new(config) }
+
+  context 'when in usage_data files' do
+    before do
+      allow(cop).to receive(:usage_data_files?).and_return(true)
+    end
+
+    context 'with large tables' do
+      context 'when counting' do
+        let(:source) do
+          <<~SRC
+            Issue.count
+          SRC
+        end
+
+        let(:source_with_count_ancestor) do
+          <<~SRC
+            Issue.active.count
+          SRC
+        end
+
+        let(:correct_source) do
+          <<~SRC
+            count(Issue)
+          SRC
+        end
+
+        let(:correct_source_with_module) do
+          <<~SRC
+            count(Ci::Build.active)
+          SRC
+        end
+
+        let(:incorrect_source_with_module) do
+          <<~SRC
+            Ci::Build.active.count
+          SRC
+        end
+
+        it 'registers an offence' do
+          inspect_source(source)
+
+          expect(cop.offenses.size).to eq(1)
+        end
+
+        it 'registers an offence with .count' do
+          inspect_source(source_with_count_ancestor)
+
+          expect(cop.offenses.size).to eq(1)
+        end
+
+        it 'does not register an offence' do
+          inspect_source(correct_source)
+
+          expect(cop.offenses).to be_empty
+        end
+      end
+
+      context 'when using allowed methods' do
+        let(:source) do
+          <<~SRC
+            Issue.minimum
+          SRC
+        end
+
+        it 'does not register an offence' do
+          inspect_source(source)
+
+          expect(cop.offenses).to be_empty
+        end
+      end
+    end
+
+    context 'with non related class' do
+      let(:source) do
+        <<~SRC
+          Rails.count
+        SRC
+      end
+
+      it 'does not registers an offence' do
+        inspect_source(source)
+
+        expect(cop.offenses).to be_empty
+      end
+    end
+  end
+end