1 files changed, 14 insertions, 36 deletions

diff --git a/lib/gitlab/usage_data_queries.rb b/lib/gitlab/usage_data_queries.rb
index 1c776501fdb..da01b68e8fc 100644
--- a/lib/gitlab/usage_data_queries.rb
+++ b/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