Add support to generate metrics queries in instrumentation classes

lib/gitlab/usage/metrics/instrumentations/database_metric.rb

@@ -43,12 +43,16 @@ module Gitlab
                     finish: self.class.metric_finish&.call)
           end
 
-          def relation
-            self.class.metric_relation.call.where(time_constraints)
+          def to_sql
+            Gitlab::Usage::Metrics::Query.for(self.class.metric_operation, relation, self.class.column)
           end
 
           private
 
+          def relation
+            self.class.metric_relation.call.where(time_constraints)
+          end
+
           def time_constraints
             case time_frame
             when '28d'

lib/gitlab/usage/metrics/query.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Usage
+    module Metrics
+      class Query
+        class << self
+          def for(operation, relation, column = nil, **extra)
+            case operation
+            when :count
+              count(relation, column)
+            when :distinct_count
+              distinct_count(relation, column)
+            when :sum
+              sum(relation, column)
+            when :estimate_batch_distinct_count
+              estimate_batch_distinct_count(relation, column)
+            when :histogram
+              histogram(relation, column, **extra)
+            else
+              raise ArgumentError, "#{operation} operation not supported"
+            end
+          end
+
+          private
+
+          def count(relation, column = nil)
+            raw_sql(relation, column)
+          end
+
+          def distinct_count(relation, column = nil)
+            raw_sql(relation, column, true)
+          end
+
+          def sum(relation, column)
+            relation.select(relation.all.table[column].sum).to_sql
+          end
+
+          def estimate_batch_distinct_count(relation, column = nil)
+            raw_sql(relation, column, true)
+          end
+
+          # rubocop: disable CodeReuse/ActiveRecord
+          def histogram(relation, column, buckets:, bucket_size: buckets.size)
+            count_grouped = relation.group(column).select(Arel.star.count.as('count_grouped'))
+            cte = Gitlab::SQL::CTE.new(:count_cte, count_grouped)
+
+            bucket_segments = bucket_size - 1
+            width_bucket = Arel::Nodes::NamedFunction
+              .new('WIDTH_BUCKET', [cte.table[:count_grouped], buckets.first, buckets.last, bucket_segments])
+              .as('buckets')
+
+            query = cte
+              .table
+              .project(width_bucket, cte.table[:count])
+              .group('buckets')
+              .order('buckets')
+              .with(cte.to_arel)
+
+            query.to_sql
+          end
+          # rubocop: enable CodeReuse/ActiveRecord
+
+          def raw_sql(relation, column, distinct = false)
+            column ||= relation.primary_key
+            relation.select(relation.all.table[column].count(distinct)).to_sql
+          end
+        end
+      end
+    end
+  end
+end

lib/gitlab/usage_data_queries.rb

@@ -6,43 +6,20 @@ module Gitlab
   class UsageDataQueries < UsageData
     class << self
       def count(relation, column = nil, *args, **kwargs)
-        raw_sql(relation, column)
+        Gitlab::Usage::Metrics::Query.for(:count, relation, column)
       end
 
       def distinct_count(relation, column = nil, *args, **kwargs)
-        raw_sql(relation, column, :distinct)
-      end
-
-      def redis_usage_data(counter = nil, &block)
-        if block_given?
-          { redis_usage_data_block: block.to_s }
-        elsif counter.present?
-          { redis_usage_data_counter: counter }
-        end
+        Gitlab::Usage::Metrics::Query.for(:distinct_count, relation, column)
       end
 
       def sum(relation, column, *args, **kwargs)
-        relation.select(relation.all.table[column].sum).to_sql
+        Gitlab::Usage::Metrics::Query.for(:sum, relation, column)
       end
 
       # rubocop: disable CodeReuse/ActiveRecord
       def histogram(relation, column, buckets:, bucket_size: buckets.size)
-        count_grouped = relation.group(column).select(Arel.star.count.as('count_grouped'))
-        cte = Gitlab::SQL::CTE.new(:count_cte, count_grouped)
-
-        bucket_segments = bucket_size - 1
-        width_bucket = Arel::Nodes::NamedFunction
-          .new('WIDTH_BUCKET', [cte.table[:count_grouped], buckets.first, buckets.last, bucket_segments])
-          .as('buckets')
-
-        query = cte
-          .table
-          .project(width_bucket, cte.table[:count])
-          .group('buckets')
-          .order('buckets')
-          .with(cte.to_arel)
-
-        query.to_sql
+        Gitlab::Usage::Metrics::Query.for(:histogram, relation, column, buckets: buckets, bucket_size: bucket_size)
       end
       # rubocop: enable CodeReuse/ActiveRecord
 
@@ -50,11 +27,11 @@ module Gitlab
       # buckets query, because it can't be used to obtain estimations without
       # supplementary ruby code present in Gitlab::Database::PostgresHll::BatchDistinctCounter
       def estimate_batch_distinct_count(relation, column = nil, *args, **kwargs)
-        raw_sql(relation, column, :distinct)
+        Gitlab::Usage::Metrics::Query.for(:estimate_batch_distinct_count, relation, column)
       end
 
       def add(*args)
-        'SELECT ' + args.map {|arg| "(#{arg})" }.join(' + ')
+        'SELECT ' + args.map { |arg| "(#{arg})" }.join(' + ')
       end
 
       def maximum_id(model, column = nil)
@@ -63,6 +40,14 @@ module Gitlab
       def minimum_id(model, column = nil)
       end
 
+      def redis_usage_data(counter = nil, &block)
+        if block_given?
+          { redis_usage_data_block: block.to_s }
+        elsif counter.present?
+          { redis_usage_data_counter: counter }
+        end
+      end
+
       def jira_service_data
         {
           projects_jira_server_active: 0,
@@ -73,13 +58,6 @@ module Gitlab
       def epics_deepest_relationship_level
         { epics_deepest_relationship_level: 0 }
       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

spec/lib/gitlab/usage/metrics/instrumentations/count_boards_metric_spec.rb

@@ -5,5 +5,8 @@ require 'spec_helper'
 RSpec.describe Gitlab::Usage::Metrics::Instrumentations::CountBoardsMetric do
   let_it_be(:board) { create(:board) }
 
-  it_behaves_like 'a correct instrumented metric value', { time_frame: 'all', data_source: 'database' }, 1
+  let(:expected_value) { 1 }
+  let(:expected_query) { 'SELECT COUNT("boards"."id") FROM "boards"' }
+
+  it_behaves_like 'a correct instrumented metric value and query', { time_frame: 'all' }
 end

spec/lib/gitlab/usage/metrics/instrumentations/count_issues_metric_spec.rb

@@ -5,5 +5,8 @@ require 'spec_helper'
 RSpec.describe Gitlab::Usage::Metrics::Instrumentations::CountIssuesMetric do
   let_it_be(:issue) { create(:issue) }
 
-  it_behaves_like 'a correct instrumented metric value', { time_frame: 'all', data_source: 'database' }, 1
+  let(:expected_value) { 1 }
+  let(:expected_query) { 'SELECT COUNT("issues"."id") FROM "issues"' }
+
+  it_behaves_like 'a correct instrumented metric value and query', { time_frame: 'all' }
 end

spec/lib/gitlab/usage/metrics/instrumentations/count_users_creating_issues_metric_spec.rb

@@ -8,10 +8,18 @@ RSpec.describe Gitlab::Usage::Metrics::Instrumentations::CountUsersCreatingIssue
   let_it_be(:old_issue) { create(:issue, author: author, created_at: 2.months.ago) }
 
   context 'with all time frame' do
-    it_behaves_like 'a correct instrumented metric value', { time_frame: 'all', data_source: 'database' }, 1
+    let(:expected_value) { 1 }
+    let(:expected_query) { 'SELECT COUNT(DISTINCT "issues"."author_id") FROM "issues"' }
+
+    it_behaves_like 'a correct instrumented metric value and query', { time_frame: 'all' }
   end
 
   context 'for 28d time frame' do
-    it_behaves_like 'a correct instrumented metric value', { time_frame: '28d', data_source: 'database' }, 1
+    let(:expected_value) { 1 }
+    let(:start) { 30.days.ago.to_s(:db) }
+    let(:finish) { 2.days.ago.to_s(:db) }
+    let(:expected_query) { "SELECT COUNT(DISTINCT \"issues\".\"author_id\") FROM \"issues\" WHERE \"issues\".\"created_at\" BETWEEN '#{start}' AND '#{finish}'" }
+
+    it_behaves_like 'a correct instrumented metric value and query', { time_frame: '28d' }
   end
 end

spec/lib/gitlab/usage/metrics/instrumentations/hostname_metric_spec.rb

@@ -3,5 +3,7 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::Usage::Metrics::Instrumentations::HostnameMetric do
-  it_behaves_like 'a correct instrumented metric value', { time_frame: 'none', data_source: 'ruby' }, Gitlab.config.gitlab.host
+  let(:expected_value) { Gitlab.config.gitlab.host }
+
+  it_behaves_like 'a correct instrumented metric value', { time_frame: 'none' }
 end

spec/lib/gitlab/usage/metrics/instrumentations/redis_hll_metric_spec.rb

@@ -10,10 +10,19 @@ RSpec.describe Gitlab::Usage::Metrics::Instrumentations::RedisHLLMetric, :clean_
     Gitlab::UsageDataCounters::HLLRedisCounter.track_event(:i_quickactions_approve, values: 2, time: 2.months.ago)
   end
 
-  it_behaves_like 'a correct instrumented metric value', { time_frame: '28d', options: { events: ['i_quickactions_approve'] } }, 2
-  it_behaves_like 'a correct instrumented metric value', { time_frame: '7d', options: { events: ['i_quickactions_approve'] } }, 1
+  context 'for 28d' do
+    let(:expected_value) { 2 }
 
-  it 'raise exception if vents options is not present' do
+    it_behaves_like 'a correct instrumented metric value', { time_frame: '28d', options: { events: ['i_quickactions_approve'] } }
+  end
+
+  context 'for 7d' do
+    let(:expected_value) { 1 }
+
+    it_behaves_like 'a correct instrumented metric value', { time_frame: '7d', options: { events: ['i_quickactions_approve'] } }
+  end
+
+  it 'raise exception if events options is not present' do
     expect { described_class.new(time_frame: '28d') }.to raise_error(ArgumentError)
   end
 end

spec/lib/gitlab/usage/metrics/instrumentations/uuid_metric_spec.rb

@@ -3,5 +3,7 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::Usage::Metrics::Instrumentations::UuidMetric do
-  it_behaves_like 'a correct instrumented metric value', { time_frame: 'none' }, Gitlab::CurrentSettings.uuid
+  let(:expected_value) { Gitlab::CurrentSettings.uuid }
+
+  it_behaves_like 'a correct instrumented metric value', { time_frame: 'none' }
 end

spec/lib/gitlab/usage/metrics/query_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Usage::Metrics::Query do
+  describe '.count' do
+    it 'returns the raw SQL' do
+      expect(described_class.for(:count, User)).to eq('SELECT COUNT("users"."id") FROM "users"')
+    end
+
+    it 'does not mix a nil column with keyword arguments' do
+      expect(described_class.for(:count, User, nil)).to eq('SELECT COUNT("users"."id") FROM "users"')
+    end
+  end
+
+  describe '.distinct_count' do
+    it 'returns the raw SQL' do
+      expect(described_class.for(:distinct_count, Issue, :author_id)).to eq('SELECT COUNT(DISTINCT "issues"."author_id") FROM "issues"')
+    end
+
+    it 'does not mix a nil column with keyword arguments' do
+      expect(described_class.for(:distinct_count, Issue, nil)).to eq('SELECT COUNT(DISTINCT "issues"."id") FROM "issues"')
+    end
+  end
+
+  describe '.sum' do
+    it 'returns the raw SQL' do
+      expect(described_class.for(:sum, Issue, :weight)).to eq('SELECT SUM("issues"."weight") FROM "issues"')
+    end
+  end
+
+  describe 'estimate_batch_distinct_count' do
+    it 'returns the raw SQL' do
+      expect(described_class.for(:estimate_batch_distinct_count, Issue, :author_id)).to eq('SELECT COUNT(DISTINCT "issues"."author_id") FROM "issues"')
+    end
+  end
+
+  describe '.histogram' do
+    it 'returns the histogram sql' do
+      expect(described_class.for(:histogram, AlertManagement::HttpIntegration.active,
+            :project_id, buckets: 1..2, bucket_size: 101))
+        .to match(/^WITH "count_cte" AS #{Gitlab::Database::AsWithMaterialized.materialized_if_supported}/)
+    end
+  end
+
+  describe 'other' do
+    it 'raise ArgumentError error' do
+      expect { described_class.for(:other, nil) }.to raise_error(ArgumentError, 'other operation not supported')
+    end
+  end
+end

spec/lib/gitlab/usage_data_queries_spec.rb

@@ -13,9 +13,7 @@ RSpec.describe Gitlab::UsageDataQueries do
     end
 
     it 'does not mix a nil column with keyword arguments' do
-      expect(described_class).to receive(:raw_sql).with(User, nil)
-
-      described_class.count(User, start: 1, finish: 2)
+      expect(described_class.count(User, nil)).to eq('SELECT COUNT("users"."id") FROM "users"')
     end
   end
 
@@ -25,9 +23,7 @@ RSpec.describe Gitlab::UsageDataQueries do
     end
 
     it 'does not mix a nil column with keyword arguments' do
-      expect(described_class).to receive(:raw_sql).with(Issue, nil, :distinct)
-
-      described_class.distinct_count(Issue, nil, start: 1, finish: 2)
+      expect(described_class.distinct_count(Issue, nil, start: 1, finish: 2)).to eq('SELECT COUNT(DISTINCT "issues"."id") FROM "issues"')
     end
   end
 

spec/support/gitlab/usage/metrics_instrumentation_shared_examples.rb

 # frozen_string_literal: true
 
-RSpec.shared_examples 'a correct instrumented metric value' do |params, expected_value|
+RSpec.shared_examples 'a correct instrumented metric value' do |params|
   let(:time_frame) { params[:time_frame] }
   let(:options) { params[:options] }
+  let(:metric) { described_class.new(time_frame: time_frame, options: options) }
 
   before do
     allow(ActiveRecord::Base.connection).to receive(:transaction_open?).and_return(false)
   end
 
   it 'has correct value' do
-    expect(described_class.new(time_frame: time_frame, options: options).value).to eq(expected_value)
+    expect(metric.value).to eq(expected_value)
   end
 end
+
+RSpec.shared_examples 'a correct instrumented metric query' do |params|
+  let(:time_frame) { params[:time_frame] }
+  let(:options) { params[:options] }
+  let(:metric) { described_class.new(time_frame: time_frame, options: options) }
+
+  around do |example|
+    freeze_time { example.run }
+  end
+
+  before do
+    allow(ActiveRecord::Base.connection).to receive(:transaction_open?).and_return(false)
+  end
+
+  it 'has correct generate query' do
+    expect(metric.to_sql).to eq(expected_query)
+  end
+end
+
+RSpec.shared_examples 'a correct instrumented metric value and query' do |params|
+  it_behaves_like 'a correct instrumented metric value', params
+  it_behaves_like 'a correct instrumented metric query', params
+end