From 6f0447dba04ade76f8ea5fba59bbb48625afb9b6 Mon Sep 17 00:00:00 2001
From: GitLab Bot <gitlab-bot@gitlab.com>
Date: Wed, 3 Nov 2021 03:13:15 +0000
Subject: [PATCH] Add latest changes from gitlab-org/gitlab@master

---
 .gitlab/ci/rails.gitlab-ci.yml                  |  1 +
 Gemfile                                         |  2 +-
 Gemfile.lock                                    |  4 ++--
 app/models/ci/application_record.rb             |  4 ++++
 app/models/ci/ci_database_record.rb             | 17 -----------------
 app/models/ci/instance_variable.rb              |  2 +-
 doc/administration/repository_storage_paths.md  |  2 +-
 doc/development/database/multiple_databases.md  | 10 ----------
 lib/gitlab/blob_helper.rb                       |  2 +-
 lib/gitlab/database.rb                          |  2 +-
 spec/lib/api/ci/helpers/runner_spec.rb          |  8 ++++----
 .../bare_repository_import/importer_spec.rb     |  2 --
 spec/lib/gitlab/blob_helper_spec.rb             | 12 ++++++++++--
 .../load_balancing/configuration_spec.rb        |  8 +++++---
 .../load_balancing/load_balancer_spec.rb        |  6 +++---
 spec/lib/gitlab/database/load_balancing_spec.rb |  2 +-
 .../database/schema_migrations/context_spec.rb  |  2 +-
 spec/lib/marginalia_spec.rb                     |  6 +++---
 spec/models/ci/build_spec.rb                    |  2 +-
 spec/models/ci/runner_spec.rb                   |  2 +-
 spec/support/helpers/usage_data_helpers.rb      |  2 ++
 .../lib/gitlab/ci/ci_trace_shared_examples.rb   | 16 ++++++++--------
 .../database/multiple_databases_spec.rb         |  6 +++---
 .../usage_trends/counter_job_worker_spec.rb     |  3 ++-
 24 files changed, 56 insertions(+), 67 deletions(-)
 delete mode 100644 app/models/ci/ci_database_record.rb

diff --git a/.gitlab/ci/rails.gitlab-ci.yml b/.gitlab/ci/rails.gitlab-ci.yml
index 5ff387ee3a7..2d6ddf1c997 100644
--- a/.gitlab/ci/rails.gitlab-ci.yml
+++ b/.gitlab/ci/rails.gitlab-ci.yml
@@ -26,6 +26,7 @@
 .decomposed-database-rspec:
   variables:
     DECOMPOSED_DB: "true"
+    GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci: "main"
 
 .rspec-base:
   extends: .rails-job-base
diff --git a/Gemfile b/Gemfile
index 536060d2b99..6146d6e011f 100644
--- a/Gemfile
+++ b/Gemfile
@@ -399,7 +399,7 @@ group :development, :test do
 end
 
 group :development, :test, :danger do
-  gem 'gitlab-dangerfiles', '~> 2.3.2', require: false
+  gem 'gitlab-dangerfiles', '~> 2.4.0', require: false
 end
 
 group :development, :test, :coverage do
diff --git a/Gemfile.lock b/Gemfile.lock
index 4cb585a75f5..b82fdfc81c7 100644
--- a/Gemfile.lock
+++ b/Gemfile.lock
@@ -458,7 +458,7 @@ GEM
       terminal-table (~> 1.5, >= 1.5.1)
     gitlab-chronic (0.10.5)
       numerizer (~> 0.2)
-    gitlab-dangerfiles (2.3.2)
+    gitlab-dangerfiles (2.4.0)
       danger (>= 8.3.1)
       danger-gitlab (>= 8.0.0)
     gitlab-experiment (0.6.4)
@@ -1464,7 +1464,7 @@ DEPENDENCIES
   gitaly (~> 14.4.0.pre.rc43)
   github-markup (~> 1.7.0)
   gitlab-chronic (~> 0.10.5)
-  gitlab-dangerfiles (~> 2.3.2)
+  gitlab-dangerfiles (~> 2.4.0)
   gitlab-experiment (~> 0.6.4)
   gitlab-fog-azure-rm (~> 1.2.0)
   gitlab-labkit (~> 0.21.1)
diff --git a/app/models/ci/application_record.rb b/app/models/ci/application_record.rb
index 9d4a8f0648e..ea7b1104e36 100644
--- a/app/models/ci/application_record.rb
+++ b/app/models/ci/application_record.rb
@@ -4,6 +4,10 @@ module Ci
   class ApplicationRecord < ::ApplicationRecord
     self.abstract_class = true
 
+    if Gitlab::Database.has_config?(:ci)
+      connects_to database: { writing: :ci, reading: :ci }
+    end
+
     def self.table_name_prefix
       'ci_'
     end
diff --git a/app/models/ci/ci_database_record.rb b/app/models/ci/ci_database_record.rb
deleted file mode 100644
index e2b832a28e7..00000000000
--- a/app/models/ci/ci_database_record.rb
+++ /dev/null
@@ -1,17 +0,0 @@
-# frozen_string_literal: true
-
-module Ci
-  # TODO: https://gitlab.com/groups/gitlab-org/-/epics/6168
-  #
-  # Do not use this yet outside of `ci_instance_variables`.
-  # This class is part of a migration to move all CI classes to a new separate database.
-  # Initially we are only going to be moving the `Ci::InstanceVariable` model and it will be duplicated in the main and CI tables
-  # Do not extend this class in any other models.
-  class CiDatabaseRecord < Ci::ApplicationRecord
-    self.abstract_class = true
-
-    if Gitlab::Database.has_config?(:ci)
-      connects_to database: { writing: :ci, reading: :ci }
-    end
-  end
-end
diff --git a/app/models/ci/instance_variable.rb b/app/models/ci/instance_variable.rb
index f4aa935b983..da9d4dea537 100644
--- a/app/models/ci/instance_variable.rb
+++ b/app/models/ci/instance_variable.rb
@@ -1,7 +1,7 @@
 # frozen_string_literal: true
 
 module Ci
-  class InstanceVariable < Ci::CiDatabaseRecord
+  class InstanceVariable < Ci::ApplicationRecord
     extend Gitlab::ProcessMemoryCache::Helper
     include Ci::NewHasVariable
     include Ci::Maskable
diff --git a/doc/administration/repository_storage_paths.md b/doc/administration/repository_storage_paths.md
index ed50d0e7263..00aa8c214c5 100644
--- a/doc/administration/repository_storage_paths.md
+++ b/doc/administration/repository_storage_paths.md
@@ -14,7 +14,7 @@ storage is either:
 - A `path`, which points directly to the directory where the repositories are stored. GitLab
   directly accessing a directory containing repositories
   [is deprecated](https://gitlab.com/gitlab-org/gitaly/-/issues/1690).
-  GitLab should be configured to access GitLab repositories though a `gitaly_address`.
+  GitLab should be configured to access GitLab repositories through a `gitaly_address`.
 
 GitLab allows you to define multiple repository storages to distribute the storage load between
 several mount points. For example:
diff --git a/doc/development/database/multiple_databases.md b/doc/development/database/multiple_databases.md
index e21606bd3e7..a17ad798305 100644
--- a/doc/development/database/multiple_databases.md
+++ b/doc/development/database/multiple_databases.md
@@ -88,16 +88,6 @@ test: &test
       statement_timeout: 120s
 ```
 
-### Migrations
-
-Place any migrations that affect `Ci::CiDatabaseRecord` models
-and their tables in two directories:
-
-- `db/migrate`
-- `db/ci_migrate`
-
-We aim to keep the schema for both tables the same across both databases.
-
 <!--
 NOTE: The `validate_cross_joins!` method in `spec/support/database/prevent_cross_joins.rb` references
       the following heading in the code, so if you make a change to this heading, make sure to update
diff --git a/lib/gitlab/blob_helper.rb b/lib/gitlab/blob_helper.rb
index c5b183d113d..9e4ea934edb 100644
--- a/lib/gitlab/blob_helper.rb
+++ b/lib/gitlab/blob_helper.rb
@@ -47,7 +47,7 @@ module Gitlab
     end
 
     def image?
-      ['.png', '.jpg', '.jpeg', '.gif', '.svg'].include?(extname.downcase)
+      ['.png', '.jpg', '.jpeg', '.gif', '.svg', '.webp'].include?(extname.downcase)
     end
 
     # Internal: Lookup mime type for extension.
diff --git a/lib/gitlab/database.rb b/lib/gitlab/database.rb
index b560d4cbca8..e5ce3984bee 100644
--- a/lib/gitlab/database.rb
+++ b/lib/gitlab/database.rb
@@ -59,7 +59,7 @@ module Gitlab
         # that inher from ActiveRecord::Base; not just our own models that
         # inherit from ApplicationRecord.
         main: ::ActiveRecord::Base,
-        ci: ::Ci::CiDatabaseRecord.connection_class? ? ::Ci::CiDatabaseRecord : nil
+        ci: ::Ci::ApplicationRecord.connection_class? ? ::Ci::ApplicationRecord : nil
       }.compact.freeze
     end
 
diff --git a/spec/lib/api/ci/helpers/runner_spec.rb b/spec/lib/api/ci/helpers/runner_spec.rb
index cc871d66d40..37277e7dcbd 100644
--- a/spec/lib/api/ci/helpers/runner_spec.rb
+++ b/spec/lib/api/ci/helpers/runner_spec.rb
@@ -15,7 +15,7 @@ RSpec.describe API::Ci::Helpers::Runner do
     it 'handles sticking of a build when a build ID is specified' do
       allow(helper).to receive(:params).and_return(id: build.id)
 
-      expect(ApplicationRecord.sticking)
+      expect(Ci::Build.sticking)
         .to receive(:stick_or_unstick_request)
         .with({}, :build, build.id)
 
@@ -25,7 +25,7 @@ RSpec.describe API::Ci::Helpers::Runner do
     it 'does not handle sticking if no build ID was specified' do
       allow(helper).to receive(:params).and_return({})
 
-      expect(ApplicationRecord.sticking)
+      expect(Ci::Build.sticking)
         .not_to receive(:stick_or_unstick_request)
 
       helper.current_job
@@ -44,7 +44,7 @@ RSpec.describe API::Ci::Helpers::Runner do
     it 'handles sticking of a runner if a token is specified' do
       allow(helper).to receive(:params).and_return(token: runner.token)
 
-      expect(ApplicationRecord.sticking)
+      expect(Ci::Runner.sticking)
         .to receive(:stick_or_unstick_request)
         .with({}, :runner, runner.token)
 
@@ -54,7 +54,7 @@ RSpec.describe API::Ci::Helpers::Runner do
     it 'does not handle sticking if no token was specified' do
       allow(helper).to receive(:params).and_return({})
 
-      expect(ApplicationRecord.sticking)
+      expect(Ci::Runner.sticking)
         .not_to receive(:stick_or_unstick_request)
 
       helper.current_runner
diff --git a/spec/lib/gitlab/bare_repository_import/importer_spec.rb b/spec/lib/gitlab/bare_repository_import/importer_spec.rb
index e09430a858c..b0d721a74ce 100644
--- a/spec/lib/gitlab/bare_repository_import/importer_spec.rb
+++ b/spec/lib/gitlab/bare_repository_import/importer_spec.rb
@@ -89,10 +89,8 @@ RSpec.describe Gitlab::BareRepositoryImport::Importer, :seed_helper do
 
         project = Project.find_by_full_path(project_path)
         repo_path = "#{project.disk_path}.git"
-        hook_path = File.join(repo_path, 'hooks')
 
         expect(gitlab_shell.repository_exists?(project.repository_storage, repo_path)).to be(true)
-        expect(TestEnv.storage_dir_exists?(project.repository_storage, hook_path)).to be(true)
       end
 
       context 'hashed storage enabled' do
diff --git a/spec/lib/gitlab/blob_helper_spec.rb b/spec/lib/gitlab/blob_helper_spec.rb
index 65fa5bf0120..a2f20dcd4fc 100644
--- a/spec/lib/gitlab/blob_helper_spec.rb
+++ b/spec/lib/gitlab/blob_helper_spec.rb
@@ -7,6 +7,7 @@ RSpec.describe Gitlab::BlobHelper do
 
   let(:project) { create(:project) }
   let(:blob) { fake_blob(path: 'file.txt') }
+  let(:webp_blob) { fake_blob(path: 'file.webp') }
   let(:large_blob) { fake_blob(path: 'test.pdf', size: 2.megabytes, binary: true) }
 
   describe '#extname' do
@@ -62,8 +63,15 @@ RSpec.describe Gitlab::BlobHelper do
   end
 
   describe '#image?' do
-    it 'returns false' do
-      expect(blob.image?).to be_falsey
+    context 'with a .txt file' do
+      it 'returns false' do
+        expect(blob.image?).to be_falsey
+      end
+    end
+    context 'with a .webp file' do
+      it 'returns true' do
+        expect(webp_blob.image?).to be_truthy
+      end
     end
   end
 
diff --git a/spec/lib/gitlab/database/load_balancing/configuration_spec.rb b/spec/lib/gitlab/database/load_balancing/configuration_spec.rb
index 435da8b6ae2..eef248afdf2 100644
--- a/spec/lib/gitlab/database/load_balancing/configuration_spec.rb
+++ b/spec/lib/gitlab/database/load_balancing/configuration_spec.rb
@@ -195,7 +195,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::Configuration do
   end
 
   describe '#replica_db_config' do
-    let(:model) { double(:model, connection_db_config: db_config, connection_specification_name: 'Ci::CiDatabaseRecord') }
+    let(:model) { double(:model, connection_db_config: db_config, connection_specification_name: 'Ci::ApplicationRecord') }
     let(:config) { described_class.for_model(model) }
 
     it 'returns exactly db_config' do
@@ -212,12 +212,14 @@ RSpec.describe Gitlab::Database::LoadBalancing::Configuration do
   end
 
   describe 'reuse_primary_connection!' do
-    let(:model) { double(:model, connection_db_config: db_config, connection_specification_name: 'Ci::CiDatabaseRecord') }
+    let(:model) { double(:model, connection_db_config: db_config, connection_specification_name: 'Ci::ApplicationRecord') }
     let(:config) { described_class.for_model(model) }
 
     context 'when GITLAB_LOAD_BALANCING_REUSE_PRIMARY_* not configured' do
       it 'the primary connection uses default specification' do
-        expect(config.primary_connection_specification_name).to eq('Ci::CiDatabaseRecord')
+        stub_env('GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci', nil)
+
+        expect(config.primary_connection_specification_name).to eq('Ci::ApplicationRecord')
       end
     end
 
diff --git a/spec/lib/gitlab/database/load_balancing/load_balancer_spec.rb b/spec/lib/gitlab/database/load_balancing/load_balancer_spec.rb
index 269a1f872dc..883c8473704 100644
--- a/spec/lib/gitlab/database/load_balancing/load_balancer_spec.rb
+++ b/spec/lib/gitlab/database/load_balancing/load_balancer_spec.rb
@@ -199,7 +199,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::LoadBalancer, :request_store do
 
     it 'does not create conflicts with other load balancers when caching hosts' do
       ci_config = Gitlab::Database::LoadBalancing::Configuration
-        .new(Ci::CiDatabaseRecord, [db_host, db_host])
+        .new(Ci::ApplicationRecord, [db_host, db_host])
 
       lb1 = described_class.new(config)
       lb2 = described_class.new(ci_config)
@@ -455,7 +455,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::LoadBalancer, :request_store do
   end
 
   describe 'primary connection re-use', :reestablished_active_record_base do
-    let(:model) { Ci::CiDatabaseRecord }
+    let(:model) { Ci::ApplicationRecord }
 
     before do
       # fake additional Database
@@ -483,7 +483,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::LoadBalancer, :request_store do
     end
 
     describe '#read_write' do
-      it 'returns Ci::CiDatabaseRecord connection' do
+      it 'returns Ci::ApplicationRecord connection' do
         expect { |b| lb.read_write(&b) }.to yield_with_args do |args|
           expect(args.pool.db_config.name).to eq('ci')
         end
diff --git a/spec/lib/gitlab/database/load_balancing_spec.rb b/spec/lib/gitlab/database/load_balancing_spec.rb
index bf5314e2c34..bbadcf73b10 100644
--- a/spec/lib/gitlab/database/load_balancing_spec.rb
+++ b/spec/lib/gitlab/database/load_balancing_spec.rb
@@ -10,7 +10,7 @@ RSpec.describe Gitlab::Database::LoadBalancing do
       expect(models).to include(ActiveRecord::Base)
 
       if Gitlab::Database.has_config?(:ci)
-        expect(models).to include(Ci::CiDatabaseRecord)
+        expect(models).to include(Ci::ApplicationRecord)
       end
     end
 
diff --git a/spec/lib/gitlab/database/schema_migrations/context_spec.rb b/spec/lib/gitlab/database/schema_migrations/context_spec.rb
index 0323fa22b78..07c97ea0ec3 100644
--- a/spec/lib/gitlab/database/schema_migrations/context_spec.rb
+++ b/spec/lib/gitlab/database/schema_migrations/context_spec.rb
@@ -14,7 +14,7 @@ RSpec.describe Gitlab::Database::SchemaMigrations::Context do
     end
 
     context 'CI database' do
-      let(:connection_class) { Ci::CiDatabaseRecord }
+      let(:connection_class) { Ci::ApplicationRecord }
 
       it 'returns a directory path that is database specific' do
         skip_if_multiple_databases_not_setup
diff --git a/spec/lib/marginalia_spec.rb b/spec/lib/marginalia_spec.rb
index 3f39d969dbd..53048ae2e6b 100644
--- a/spec/lib/marginalia_spec.rb
+++ b/spec/lib/marginalia_spec.rb
@@ -59,14 +59,14 @@ RSpec.describe 'Marginalia spec' do
           "application"    => "test",
           "endpoint_id"    => "MarginaliaTestController#first_user",
           "correlation_id" => correlation_id,
-          "db_config_name" => "ci"
+          "db_config_name" => ENV['GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci'] == 'main' ? 'main' : 'ci'
         }
       end
 
-      before do |example|
+      before do
         skip_if_multiple_databases_not_setup
 
-        allow(User).to receive(:connection) { Ci::CiDatabaseRecord.connection }
+        allow(User).to receive(:connection) { Ci::ApplicationRecord.connection }
       end
 
       it 'generates a query that includes the component and value' do
diff --git a/spec/models/ci/build_spec.rb b/spec/models/ci/build_spec.rb
index 50978ac017c..1d7f4475f2c 100644
--- a/spec/models/ci/build_spec.rb
+++ b/spec/models/ci/build_spec.rb
@@ -351,7 +351,7 @@ RSpec.describe Ci::Build do
     it 'sticks the build if the status changed' do
       job = create(:ci_build, :pending)
 
-      expect(ApplicationRecord.sticking).to receive(:stick)
+      expect(described_class.sticking).to receive(:stick)
         .with(:build, job.id)
 
       job.update!(status: :running)
diff --git a/spec/models/ci/runner_spec.rb b/spec/models/ci/runner_spec.rb
index 1b679496437..08c471c6bfb 100644
--- a/spec/models/ci/runner_spec.rb
+++ b/spec/models/ci/runner_spec.rb
@@ -405,7 +405,7 @@ RSpec.describe Ci::Runner do
     it 'sticks the runner to the primary and calls the original method' do
       runner = create(:ci_runner)
 
-      expect(ApplicationRecord.sticking).to receive(:stick)
+      expect(described_class.sticking).to receive(:stick)
         .with(:runner, runner.id)
 
       expect(Gitlab::Workhorse).to receive(:set_key_and_notify)
diff --git a/spec/support/helpers/usage_data_helpers.rb b/spec/support/helpers/usage_data_helpers.rb
index 5ead1813439..5865bafd382 100644
--- a/spec/support/helpers/usage_data_helpers.rb
+++ b/spec/support/helpers/usage_data_helpers.rb
@@ -162,6 +162,8 @@ module UsageDataHelpers
 
   def stub_usage_data_connections
     allow(ActiveRecord::Base.connection).to receive(:transaction_open?).and_return(false)
+    allow(::Ci::ApplicationRecord.connection).to receive(:transaction_open?).and_return(false) if ::Ci::ApplicationRecord.connection_class?
+
     allow(Gitlab::Prometheus::Internal).to receive(:prometheus_enabled?).and_return(false)
   end
 
diff --git a/spec/support/shared_examples/lib/gitlab/ci/ci_trace_shared_examples.rb b/spec/support/shared_examples/lib/gitlab/ci/ci_trace_shared_examples.rb
index 26adc94bd12..a3c67210a4a 100644
--- a/spec/support/shared_examples/lib/gitlab/ci/ci_trace_shared_examples.rb
+++ b/spec/support/shared_examples/lib/gitlab/ci/ci_trace_shared_examples.rb
@@ -35,8 +35,8 @@ RSpec.shared_examples 'common trace features' do
         stub_feature_flags(gitlab_ci_archived_trace_consistent_reads: trace.job.project)
       end
 
-      it 'calls ::ApplicationRecord.sticking.unstick_or_continue_sticking' do
-        expect(::ApplicationRecord.sticking).to receive(:unstick_or_continue_sticking)
+      it 'calls ::Ci::Build.sticking.unstick_or_continue_sticking' do
+        expect(::Ci::Build.sticking).to receive(:unstick_or_continue_sticking)
           .with(described_class::LOAD_BALANCING_STICKING_NAMESPACE, trace.job.id)
           .and_call_original
 
@@ -49,8 +49,8 @@ RSpec.shared_examples 'common trace features' do
         stub_feature_flags(gitlab_ci_archived_trace_consistent_reads: false)
       end
 
-      it 'does not call ::ApplicationRecord.sticking.unstick_or_continue_sticking' do
-        expect(::ApplicationRecord.sticking).not_to receive(:unstick_or_continue_sticking)
+      it 'does not call ::Ci::Build.sticking.unstick_or_continue_sticking' do
+        expect(::Ci::Build.sticking).not_to receive(:unstick_or_continue_sticking)
 
         trace.read { |stream| stream }
       end
@@ -305,8 +305,8 @@ RSpec.shared_examples 'common trace features' do
           stub_feature_flags(gitlab_ci_archived_trace_consistent_reads: trace.job.project)
         end
 
-        it 'calls ::ApplicationRecord.sticking.stick' do
-          expect(::ApplicationRecord.sticking).to receive(:stick)
+        it 'calls ::Ci::Build.sticking.stick' do
+          expect(::Ci::Build.sticking).to receive(:stick)
             .with(described_class::LOAD_BALANCING_STICKING_NAMESPACE, trace.job.id)
             .and_call_original
 
@@ -319,8 +319,8 @@ RSpec.shared_examples 'common trace features' do
           stub_feature_flags(gitlab_ci_archived_trace_consistent_reads: false)
         end
 
-        it 'does not call ::ApplicationRecord.sticking.stick' do
-          expect(::ApplicationRecord.sticking).not_to receive(:stick)
+        it 'does not call ::Ci::Build.sticking.stick' do
+          expect(::Ci::Build.sticking).not_to receive(:stick)
 
           subject
         end
diff --git a/spec/support_specs/database/multiple_databases_spec.rb b/spec/support_specs/database/multiple_databases_spec.rb
index 6ad15fd6594..10d1a8277c6 100644
--- a/spec/support_specs/database/multiple_databases_spec.rb
+++ b/spec/support_specs/database/multiple_databases_spec.rb
@@ -19,19 +19,19 @@ RSpec.describe 'Database::MultipleDatabases' do
         end
       end
 
-      context 'on Ci::CiDatabaseRecord' do
+      context 'on Ci::ApplicationRecord' do
         before do
           skip_if_multiple_databases_not_setup
         end
 
         it 'raises exception' do
-          expect { Ci::CiDatabaseRecord.establish_connection(:ci) }.to raise_error /Cannot re-establish/
+          expect { Ci::ApplicationRecord.establish_connection(:ci) }.to raise_error /Cannot re-establish/
         end
 
         context 'when using with_reestablished_active_record_base' do
           it 'does not raise exception' do
             with_reestablished_active_record_base do
-              expect { Ci::CiDatabaseRecord.establish_connection(:main) }.not_to raise_error
+              expect { Ci::ApplicationRecord.establish_connection(:main) }.not_to raise_error
             end
           end
         end
diff --git a/spec/workers/analytics/usage_trends/counter_job_worker_spec.rb b/spec/workers/analytics/usage_trends/counter_job_worker_spec.rb
index dd180229d12..c45ec20fe5a 100644
--- a/spec/workers/analytics/usage_trends/counter_job_worker_spec.rb
+++ b/spec/workers/analytics/usage_trends/counter_job_worker_spec.rb
@@ -11,7 +11,8 @@ RSpec.describe Analytics::UsageTrends::CounterJobWorker do
   let(:job_args) { [users_measurement_identifier, user_1.id, user_2.id, recorded_at] }
 
   before do
-    allow(::Analytics::UsageTrends::Measurement.connection).to receive(:transaction_open?).and_return(false)
+    allow(::ApplicationRecord.connection).to receive(:transaction_open?).and_return(false)
+    allow(::Ci::ApplicationRecord.connection).to receive(:transaction_open?).and_return(false) if ::Ci::ApplicationRecord.connection_class?
   end
 
   include_examples 'an idempotent worker' do