Add estimate_batch_distinct_count operation block passing support

doc/development/service_ping/metrics_instrumentation.md

@@ -110,14 +110,14 @@ end
 
 There is support for:
 
-- `count`, `distinct_count` for [database metrics](#database-metrics).
+- `count`, `distinct_count`, `estimate_batch_distinct_count` for [database metrics](#database-metrics).
 - [Redis metrics](#redis-metrics).
 - [Redis HLL metrics](#redis-hyperloglog-metrics).
 - [Generic metrics](#generic-metrics), which are metrics based on settings or configurations.
 
 Currently, there is no support for:
 
-- `add`, `sum`, `histogram`, `estimate_batch_distinct_count` for database metrics.
+- `add`, `sum`, `histogram` for database metrics.
 
 You can [track the progress to support these](https://gitlab.com/groups/gitlab-org/-/epics/6118).
 
@@ -128,7 +128,7 @@ To create a stub instrumentation for a Service Ping metric, you can use a dedica
 The generator takes the class name as an argument and the following options:
 
 - `--type=TYPE` Required. Indicates the metric type. It must be one of: `database`, `generic`, `redis`.
-- `--operation` Required for `database` type. It must be one of: `count`, `distinct_count`.
+- `--operation` Required for `database` type. It must be one of: `count`, `distinct_count`, `estimate_batch_distinct_count`.
 - `--ee` Indicates if the metric is for EE.
 
 ```shell

lib/generators/gitlab/usage_metric_generator.rb

@@ -15,7 +15,7 @@ module Gitlab
       redis: 'Redis'
     }.freeze
 
-    ALLOWED_OPERATIONS = %w(count distinct_count).freeze
+    ALLOWED_OPERATIONS = %w(count distinct_count estimate_batch_distinct_count).freeze
 
     source_root File.expand_path('usage_metric/templates', __dir__)
 

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

@@ -33,16 +33,17 @@ module Gitlab
               @metric_relation = block
             end
 
-            def operation(symbol, column: nil)
+            def operation(symbol, column: nil, &block)
               @metric_operation = symbol
               @column = column
+              @metric_operation_block = block if block_given?
             end
 
             def cache_start_and_finish_as(cache_key)
               @cache_key = cache_key
             end
 
-            attr_reader :metric_operation, :metric_relation, :metric_start, :metric_finish, :column, :cache_key
+            attr_reader :metric_operation, :metric_relation, :metric_start, :metric_finish, :metric_operation_block, :column, :cache_key
           end
 
           def value
@@ -52,7 +53,8 @@ module Gitlab
               .call(relation,
                     self.class.column,
                     start: start,
-                    finish: finish)
+                    finish: finish,
+                    &self.class.metric_operation_block)
           end
 
           def to_sql

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

@@ -4,11 +4,11 @@ require 'spec_helper'
 
 RSpec.describe Gitlab::Usage::Metrics::Instrumentations::DatabaseMetric do
   subject do
-    described_class.tap do |m|
-      m.relation { Issue }
-      m.operation :count
-      m.start { m.relation.minimum(:id) }
-      m.finish { m.relation.maximum(:id) }
+    described_class.tap do |metric_class|
+      metric_class.relation { Issue }
+      metric_class.operation :count
+      metric_class.start { metric_class.relation.minimum(:id) }
+      metric_class.finish { metric_class.relation.maximum(:id) }
     end.new(time_frame: 'all')
   end
 
@@ -38,9 +38,9 @@ RSpec.describe Gitlab::Usage::Metrics::Instrumentations::DatabaseMetric do
 
     context 'with start and finish not called' do
       subject do
-        described_class.tap do |m|
-          m.relation { Issue }
-          m.operation :count
+        described_class.tap do |metric_class|
+          metric_class.relation { Issue }
+          metric_class.operation :count
         end.new(time_frame: 'all')
       end
 
@@ -51,12 +51,12 @@ RSpec.describe Gitlab::Usage::Metrics::Instrumentations::DatabaseMetric do
 
     context 'with cache_start_and_finish_as called' do
       subject do
-        described_class.tap do |m|
-          m.relation { Issue }
-          m.operation :count
-          m.start { m.relation.minimum(:id) }
-          m.finish { m.relation.maximum(:id) }
-          m.cache_start_and_finish_as :special_issue_count
+        described_class.tap do |metric_class|
+          metric_class.relation { Issue }
+          metric_class.operation :count
+          metric_class.start { metric_class.relation.minimum(:id) }
+          metric_class.finish { metric_class.relation.maximum(:id) }
+          metric_class.cache_start_and_finish_as :special_issue_count
         end.new(time_frame: 'all')
       end
 
@@ -71,5 +71,45 @@ RSpec.describe Gitlab::Usage::Metrics::Instrumentations::DatabaseMetric do
         expect(Rails.cache.read('metric_instrumentation/special_issue_count_maximum_id')).to eq(issues.max_by(&:id).id)
       end
     end
+
+    context 'with estimate_batch_distinct_count' do
+      subject do
+        described_class.tap do |metric_class|
+          metric_class.relation { Issue }
+          metric_class.operation(:estimate_batch_distinct_count)
+          metric_class.start { metric_class.relation.minimum(:id) }
+          metric_class.finish { metric_class.relation.maximum(:id) }
+        end.new(time_frame: 'all')
+      end
+
+      it 'calculates a correct result' do
+        expect(subject.value).to be_within(Gitlab::Database::PostgresHll::BatchDistinctCounter::ERROR_RATE).percent_of(3)
+      end
+
+      context 'with block passed to operation' do
+        let(:buckets) { double('Buckets').as_null_object }
+
+        subject do
+          described_class.tap do |metric_class|
+            metric_class.relation { Issue }
+            metric_class.operation(:estimate_batch_distinct_count) do |result|
+              result.foo
+            end
+            metric_class.start { metric_class.relation.minimum(:id) }
+            metric_class.finish { metric_class.relation.maximum(:id) }
+          end.new(time_frame: 'all')
+        end
+
+        before do
+          allow(Gitlab::Database::PostgresHll::Buckets).to receive(:new).and_return(buckets)
+        end
+
+        it 'calls the block passing HLL buckets as an argument' do
+          expect(buckets).to receive(:foo)
+
+          subject.value
+        end
+      end
+    end
   end
 end