Merge branch 'ali/use-batch-sum-counter' into 'master'

Use batch sum counter to sum Jira imported issues

See merge request gitlab-org/gitlab!40654

app/models/jira_import_state.rb

@@ -110,10 +110,6 @@ class JiraImportState < ApplicationRecord
     )
   end
 
-  def self.finished_imports_count
-    finished.sum(:imported_issues_count)
-  end
-
   def mark_as_failed(error_message)
     sanitized_message = Gitlab::UrlSanitizer.sanitize(error_message)
 

lib/gitlab/usage_data.rb

@@ -428,7 +428,7 @@ module Gitlab
         {
           jira_imports_total_imported_count: count(finished_jira_imports),
           jira_imports_projects_count: distinct_count(finished_jira_imports, :project_id),
-          jira_imports_total_imported_issues_count: alt_usage_data { JiraImportState.finished_imports_count }
+          jira_imports_total_imported_issues_count: sum(JiraImportState.finished, :imported_issues_count)
         }
         # rubocop: enable UsageData/LargeTable
       end

lib/gitlab/usage_data_queries.rb

 # frozen_string_literal: true
 
 module Gitlab
+  # This class is used by the `gitlab:usage_data:dump_sql` rake tasks to output SQL instead of running it.
+  # See https://gitlab.com/gitlab-org/gitlab/-/merge_requests/41091
   class UsageDataQueries < UsageData
     class << self
       def count(relation, column = nil, *rest)
@@ -19,6 +21,10 @@ module Gitlab
         end
       end
 
+      def sum(relation, column, *rest)
+        relation.select(relation.all.table[column].sum).to_sql # rubocop:disable CodeReuse/ActiveRecord
+      end
+
       private
 
       def raw_sql(relation, column, distinct = nil)

lib/gitlab/utils/usage_data.rb

@@ -59,6 +59,12 @@ module Gitlab
         FALLBACK
       end
 
+      def sum(relation, column, batch_size: nil, start: nil, finish: nil)
+        Gitlab::Database::BatchCount.batch_sum(relation, column, batch_size: batch_size, start: start, finish: finish)
+      rescue ActiveRecord::StatementInvalid
+        FALLBACK
+      end
+
       def alt_usage_data(value = nil, fallback: FALLBACK, &block)
         if block_given?
           yield

spec/lib/gitlab/usage_data_queries_spec.rb

@@ -32,4 +32,10 @@ RSpec.describe Gitlab::UsageDataQueries do
       expect(redis_usage_data[:redis_usage_data_block]).to start_with('#<Proc:')
     end
   end
+
+  describe '.sum' do
+    it 'returns the raw SQL' do
+      expect(described_class.sum(Issue, :weight)).to eq('SELECT SUM("issues"."weight") FROM "issues"')
+    end
+  end
 end

spec/lib/gitlab/utils/usage_data_spec.rb

@@ -37,6 +37,28 @@ RSpec.describe Gitlab::Utils::UsageData do
     end
   end
 
+  describe '#sum' do
+    let(:relation) { double(:relation) }
+
+    it 'returns the count when counting succeeds' do
+      allow(Gitlab::Database::BatchCount)
+        .to receive(:batch_sum)
+        .with(relation, :column, batch_size: 100, start: 2, finish: 3)
+        .and_return(1)
+
+      expect(described_class.sum(relation, :column, batch_size: 100, start: 2, finish: 3)).to eq(1)
+    end
+
+    it 'returns the fallback value when counting fails' do
+      stub_const("Gitlab::Utils::UsageData::FALLBACK", 15)
+      allow(Gitlab::Database::BatchCount)
+        .to receive(:batch_sum)
+        .and_raise(ActiveRecord::StatementInvalid.new(''))
+
+      expect(described_class.sum(relation, :column)).to eq(15)
+    end
+  end
+
   describe '#alt_usage_data' do
     it 'returns the fallback when it gets an error' do
       expect(described_class.alt_usage_data { raise StandardError } ).to eq(-1)