Merge branch '345279-hll-distinct-counter-uses-relation-connection' into 'master'

HLL BatchDistinctCounter uses relation connection

See merge request gitlab-org/gitlab!74112

.rubocop_manual_todo.yml

@@ -2402,7 +2402,6 @@ Database/MultipleDatabases:
     - 'lib/gitlab/database/migrations/observers/query_log.rb'
     - 'lib/gitlab/database/multi_threaded_migration.rb'
     - 'lib/gitlab/database/partitioning_migration_helpers/backfill_partitioned_table.rb'
-    - 'lib/gitlab/database/postgres_hll/batch_distinct_counter.rb'
     - 'lib/gitlab/database/postgresql_adapter/dump_schema_versions_mixin.rb'
     - 'lib/gitlab/database/postgresql_database_tasks/load_schema_versions_mixin.rb'
     - 'lib/gitlab/database.rb'

ee/spec/lib/gitlab/usage_data_metrics_spec.rb

@@ -13,7 +13,9 @@ RSpec.describe Gitlab::UsageDataMetrics do
     end
 
     before do
-      allow(ActiveRecord::Base.connection).to receive(:transaction_open?).and_return(false) # rubocop:disable Database/MultipleDatabases
+      allow_next_instance_of(Gitlab::Database::BatchCounter) do |batch_counter|
+        allow(batch_counter).to receive(:transaction_open?).and_return(false)
+      end
     end
 
     context 'with instrumentation_class' do

lib/gitlab/database/batch_counter.rb

@@ -31,7 +31,7 @@ module Gitlab
       end
 
       def count(batch_size: nil, mode: :itself, start: nil, finish: nil)
-        raise 'BatchCount can not be run inside a transaction' if @relation.connection.transaction_open?
+        raise 'BatchCount can not be run inside a transaction' if transaction_open?
 
         check_mode!(mode)
 
@@ -87,6 +87,10 @@ module Gitlab
         results
       end
 
+      def transaction_open?
+        @relation.connection.transaction_open?
+      end
+
       def merge_results(results, object)
         return object unless results
 

lib/gitlab/database/postgres_hll/batch_distinct_counter.rb

@@ -57,7 +57,7 @@ module Gitlab
         # @param finish final pkey range
         # @return [Gitlab::Database::PostgresHll::Buckets] HyperLogLog data structure instance that can estimate number of unique elements
         def execute(batch_size: nil, start: nil, finish: nil)
-          raise 'BatchCount can not be run inside a transaction' if ActiveRecord::Base.connection.transaction_open? # rubocop: disable Database/MultipleDatabases
+          raise 'BatchCount can not be run inside a transaction' if transaction_open?
 
           batch_size ||= DEFAULT_BATCH_SIZE
           start = actual_start(start)
@@ -79,6 +79,10 @@ module Gitlab
 
         private
 
+        def transaction_open?
+          @relation.connection.transaction_open?
+        end
+
         def unwanted_configuration?(start, finish, batch_size)
           batch_size <= MIN_REQUIRED_BATCH_SIZE ||
             (finish - start) >= MAX_DATA_VOLUME ||

spec/lib/gitlab/database/batch_count_spec.rb

@@ -19,7 +19,7 @@ RSpec.describe Gitlab::Database::BatchCount do
   end
 
   before do
-    allow(ActiveRecord::Base.connection).to receive(:transaction_open?).and_return(in_transaction)
+    allow(model.connection).to receive(:transaction_open?).and_return(in_transaction)
   end
 
   def calculate_batch_size(batch_size)

spec/lib/gitlab/database/postgres_hll/batch_distinct_counter_spec.rb

@@ -21,7 +21,7 @@ RSpec.describe Gitlab::Database::PostgresHll::BatchDistinctCounter do
   end
 
   before do
-    allow(ActiveRecord::Base.connection).to receive(:transaction_open?).and_return(in_transaction)
+    allow(model.connection).to receive(:transaction_open?).and_return(in_transaction)
   end
 
   context 'unit test for different counting parameters' do

spec/lib/gitlab/usage_data_metrics_spec.rb

@@ -13,7 +13,9 @@ RSpec.describe Gitlab::UsageDataMetrics do
     end
 
     before do
-      allow(ActiveRecord::Base.connection).to receive(:transaction_open?).and_return(false)
+      allow_next_instance_of(Gitlab::Database::BatchCounter) do |batch_counter|
+        allow(batch_counter).to receive(:transaction_open?).and_return(false)
+      end
     end
 
     context 'with instrumentation_class' do

spec/lib/gitlab/utils/usage_data_spec.rb

@@ -70,7 +70,7 @@ RSpec.describe Gitlab::Utils::UsageData do
     let(:relation) { double(:relation, connection: double(:connection)) }
 
     before do
-      allow(ActiveRecord::Base.connection).to receive(:transaction_open?).and_return(false) # rubocop: disable Database/MultipleDatabases
+      allow(relation.connection).to receive(:transaction_open?).and_return(false)
     end
 
     it 'delegates counting to counter class instance' do
@@ -122,7 +122,7 @@ RSpec.describe Gitlab::Utils::UsageData do
       let(:ci_builds_estimated_cardinality) { 2.0809220082170614 }
 
       before do
-        allow(ActiveRecord::Base.connection).to receive(:transaction_open?).and_return(false) # rubocop: disable Database/MultipleDatabases
+        allow(model.connection).to receive(:transaction_open?).and_return(false)
       end
 
       context 'different counting parameters' do

spec/support/gitlab/usage/metrics_instrumentation_shared_examples.rb

@@ -6,7 +6,9 @@ RSpec.shared_examples 'a correct instrumented metric value' do |params|
   let(:metric) { described_class.new(time_frame: time_frame, options: options) }
 
   before do
-    allow(ActiveRecord::Base.connection).to receive(:transaction_open?).and_return(false)
+    if described_class.respond_to?(:relation) && described_class.relation.respond_to?(:connection)
+      allow(described_class.relation.connection).to receive(:transaction_open?).and_return(false)
+    end
   end
 
   it 'has correct value' do