1 files changed, 36 insertions, 0 deletions
diff --git a/lib/gitlab/usage_data_queries.rb b/lib/gitlab/usage_data_queries.rb
new file mode 100644
index 00000000000..bacd63ab282
--- /dev/null
+++ b/lib/gitlab/usage_data_queries.rb
@@ -0,0 +1,36 @@
+# 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)
+        raw_sql(relation, column)
+      end
+
+      def distinct_count(relation, column = nil, *rest)
+        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
+      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)
+        column ||= relation.primary_key
+        relation.select(relation.all.table[column].count(distinct)).to_sql
+      end
+    end
+  end
+end