Add latest changes from gitlab-org/gitlab@master

1 files changed, 58 insertions, 12 deletions

diff --git a/spec/workers/database/batched_background_migration/execution_worker_spec.rb b/spec/workers/database/batched_background_migration/execution_worker_spec.rb
index ae879feced9..d6add0c68de 100644
--- a/spec/workers/database/batched_background_migration/execution_worker_spec.rb
+++ b/spec/workers/database/batched_background_migration/execution_worker_spec.rb
@@ -3,20 +3,32 @@
 require 'spec_helper'
 
 RSpec.describe Database::BatchedBackgroundMigration::ExecutionWorker, :clean_gitlab_redis_shared_state do
+  include ExclusiveLeaseHelpers
+
   describe '#perform' do
     let(:database_name) { Gitlab::Database::MAIN_DATABASE_NAME.to_sym }
+    let(:base_model) { Gitlab::Database.database_base_models[database_name] }
+    let(:table_name) { :events }
+    let(:job_interval) { 5.minutes }
+    let(:lease_timeout) { job_interval * described_class::LEASE_TIMEOUT_MULTIPLIER }
+    let(:interval_variance) { described_class::INTERVAL_VARIANCE }
 
     subject(:worker) { described_class.new }
 
     context 'when the feature flag is disabled' do
+      let(:migration) do
+        create(:batched_background_migration, :active, interval: job_interval, table_name: table_name)
+      end
+
       before do
         stub_feature_flags(execute_batched_migrations_on_schedule: false)
       end
 
       it 'does nothing' do
-        expect(worker).not_to receive(:run)
+        expect(Gitlab::Database::BackgroundMigration::BatchedMigration).not_to receive(:find_executable)
+        expect(worker).not_to receive(:run_migration_job)
 
-        worker.perform(database_name, 123)
+        worker.perform(database_name, migration.id)
       end
     end
 
@@ -25,19 +37,28 @@ RSpec.describe Database::BatchedBackgroundMigration::ExecutionWorker, :clean_git
         stub_feature_flags(execute_batched_migrations_on_schedule: true)
       end
 
+      context 'when the provided database is sharing config' do
+        it 'does nothing' do
+          ci_model = Gitlab::Database.database_base_models['ci']
+          expect(Gitlab::Database).to receive(:db_config_share_with)
+            .with(ci_model.connection_db_config).and_return('main')
+
+          expect(Gitlab::Database::BackgroundMigration::BatchedMigration).not_to receive(:find_executable)
+          expect(worker).not_to receive(:run_migration_job)
+
+          worker.perform(:ci, 123)
+        end
+      end
+
       context 'when migration does not exist' do
         it 'does nothing' do
-          expect(worker).not_to receive(:run)
+          expect(worker).not_to receive(:run_migration_job)
 
           worker.perform(database_name, non_existing_record_id)
         end
       end
 
       context 'when migration exist' do
-        let(:base_model) { Gitlab::Database.database_base_models[database_name] }
-        let(:table_name) { :events }
-        let(:job_interval) { 5.minutes }
-        let(:interval_variance) { described_class::INTERVAL_VARIANCE }
         let(:migration) do
           create(:batched_background_migration, :active, interval: job_interval, table_name: table_name)
         end
@@ -54,7 +75,7 @@ RSpec.describe Database::BatchedBackgroundMigration::ExecutionWorker, :clean_git
 
             expect(migration).to receive(:active?).and_return(false)
 
-            expect(worker).not_to receive(:run)
+            expect(worker).not_to receive(:run_migration_job)
 
             worker.perform(database_name, migration.id)
           end
@@ -63,16 +84,38 @@ RSpec.describe Database::BatchedBackgroundMigration::ExecutionWorker, :clean_git
         context 'when the interval has not elapsed' do
           it 'does not run the migration' do
             expect(Gitlab::Database::SharedModel).to receive(:using_connection).with(base_model.connection).and_yield
-
             expect(migration).to receive(:interval_elapsed?).with(variance: interval_variance).and_return(false)
-
-            expect(worker).not_to receive(:run)
+            expect(worker).not_to receive(:run_migration_job)
 
             worker.perform(database_name, migration.id)
           end
         end
 
         context 'when the migration is still active and the interval has elapsed' do
+          let(:table_name_lease_key) do
+            "#{described_class.name.underscore}:database_name:#{database_name}:" \
+              "table_name:#{table_name}"
+          end
+
+          context 'when can not obtain lease on the table name' do
+            it 'does nothing' do
+              stub_exclusive_lease_taken(table_name_lease_key, timeout: lease_timeout)
+
+              expect(worker).not_to receive(:run_migration_job)
+
+              worker.perform(database_name, migration.id)
+            end
+          end
+
+          it 'always cleans up the exclusive lease' do
+            expect_to_obtain_exclusive_lease(table_name_lease_key, 'uuid-table-name', timeout: lease_timeout)
+            expect_to_cancel_exclusive_lease(table_name_lease_key, 'uuid-table-name')
+
+            expect(worker).to receive(:run_migration_job).and_raise(RuntimeError, 'I broke')
+
+            expect { worker.perform(database_name, migration.id) }.to raise_error(RuntimeError, 'I broke')
+          end
+
           it 'runs the migration' do
             expect(Gitlab::Database::SharedModel).to receive(:using_connection).with(base_model.connection).and_yield
 
@@ -80,7 +123,10 @@ RSpec.describe Database::BatchedBackgroundMigration::ExecutionWorker, :clean_git
               expect(instance).to receive(:run_migration_job).with(migration)
             end
 
-            expect(worker).to receive(:run).and_call_original
+            expect_to_obtain_exclusive_lease(table_name_lease_key, 'uuid-table-name', timeout: lease_timeout)
+            expect_to_cancel_exclusive_lease(table_name_lease_key, 'uuid-table-name')
+
+            expect(worker).to receive(:run_migration_job).and_call_original
 
             worker.perform(database_name, migration.id)
           end