Rake task to generate raw SQLs for usage ping

Adds a rake task and a raw SQL mode for usage data counters

changelogs/unreleased/244298-rake-task-to-dump-usage-ping-sql-yaml-json.yml

+---
+title: Rake task to generate raw SQLs for usage ping
+merge_request: 41091
+author:
+type: added

ee/lib/ee/gitlab/usage_data.rb

@@ -116,7 +116,7 @@ module EE
 
           # handle license rename https://gitlab.com/gitlab-org/gitlab/issues/8911
           license_scan_count = results.delete(:license_scanning_jobs)
-          results[:license_management_jobs] += license_scan_count > 0 ? license_scan_count : 0
+          results[:license_management_jobs] += license_scan_count > 0 ? license_scan_count : 0 if license_scan_count.is_a?(Integer)
 
           results
         end
@@ -324,7 +324,7 @@ module EE
           # handle license rename https://gitlab.com/gitlab-org/gitlab/issues/8911
           combined_license_key = "#{prefix}license_management_jobs".to_sym
           license_scan_count = results.delete("#{prefix}license_scanning_jobs".to_sym)
-          results[combined_license_key] += license_scan_count > 0 ? license_scan_count : 0
+          results[combined_license_key] += license_scan_count > 0 ? license_scan_count : 0 if license_scan_count.is_a?(Integer)
 
           super.merge(results)
         end

lib/gitlab/usage_data_queries.rb

+# frozen_string_literal: true
+
+module Gitlab
+  class UsageDataQueries < UsageData
+    class << self
+      def count(relation, column = nil, *rest)
+        raw_sql(relation, column)
+      end
+
+      def distinct_count(relation, column = nil, *rest)
+        raw_sql(relation, column, :distinct)
+      end
+
+      private
+
+      def raw_sql(relation, column, distinct = nil)
+        column ||= relation.primary_key
+        relation.select(relation.all.table[column].count(distinct)).to_sql
+      end
+    end
+  end
+end

lib/tasks/gitlab/usage_data.rake

+namespace :gitlab do
+  namespace :usage_data do
+    desc 'GitLab | UsageData | Generate raw SQLs for usage ping in YAML'
+    task dump_sql_in_yaml: :environment do
+      puts Gitlab::UsageDataQueries.uncached_data.to_yaml
+    end
+
+    desc 'GitLab | UsageData | Generate raw SQLs for usage ping in JSON'
+    task dump_sql_in_json: :environment do
+      puts Gitlab::Json.pretty_generate(Gitlab::UsageDataQueries.uncached_data)
+    end
+  end
+end

spec/lib/gitlab/usage_data_queries_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::UsageDataQueries do
+  before do
+    allow(ActiveRecord::Base.connection).to receive(:transaction_open?).and_return(false)
+  end
+
+  describe '.count' do
+    it 'returns the raw SQL' do
+      expect(described_class.count(User)).to start_with('SELECT COUNT("users"."id") FROM "users"')
+    end
+  end
+
+  describe '.distinct_count' do
+    it 'returns the raw SQL' do
+      expect(described_class.distinct_count(Issue, :author_id)).to eq('SELECT COUNT(DISTINCT "issues"."author_id") FROM "issues"')
+    end
+  end
+end

spec/tasks/gitlab/usage_data_rake_spec.rb

+# frozen_string_literal: true
+
+require 'rake_helper'
+
+RSpec.describe 'gitlab:usage data take tasks' do
+  before do
+    Rake.application.rake_require 'tasks/gitlab/usage_data'
+    # stub prometheus external http calls https://gitlab.com/gitlab-org/gitlab/-/issues/245277
+    stub_request(:get, %r{^http://::1:9090/api/v1/query\?query=.*})
+      .to_return(
+        status: 200,
+        body: [{}].to_json,
+        headers: { 'Content-Type' => 'application/json' }
+      )
+  end
+
+  describe 'dump_sql_in_yaml' do
+    it 'dumps SQL queries in yaml format' do
+      expect { run_rake_task('gitlab:usage_data:dump_sql_in_yaml') }.to output(/.*recorded_at:.*/).to_stdout
+    end
+  end
+
+  describe 'dump_sql_in_json' do
+    it 'dumps SQL queries in json format' do
+      expect { run_rake_task('gitlab:usage_data:dump_sql_in_json') }.to output(/.*"recorded_at":.*/).to_stdout
+    end
+  end
+end