Move all CI movels to CI database connection if CI db configured

.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

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

app/models/ci/ci_database_record.rb

-# 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

app/models/ci/instance_variable.rb

 # frozen_string_literal: true
 
 module Ci
-  class InstanceVariable < Ci::CiDatabaseRecord
+  class InstanceVariable < Ci::ApplicationRecord
     extend Gitlab::ProcessMemoryCache::Helper
     include Ci::NewHasVariable
     include Ci::Maskable

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

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
 

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

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
 

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

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
 

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

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

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)

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)

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
 

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

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

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