Add latest changes from gitlab-org/gitlab@master

4 files changed, 118 insertions, 2 deletions

diff --git a/spec/workers/bulk_imports/pipeline_batch_worker_spec.rb b/spec/workers/bulk_imports/pipeline_batch_worker_spec.rb
index 9ac297ae757..c459c17b1bc 100644
--- a/spec/workers/bulk_imports/pipeline_batch_worker_spec.rb
+++ b/spec/workers/bulk_imports/pipeline_batch_worker_spec.rb
@@ -13,6 +13,10 @@ RSpec.describe BulkImports::PipelineBatchWorker, feature_category: :importers do
         @context = context
       end
 
+      def self.relation
+        'labels'
+      end
+
       def run
         @context.tracker.finish!
       end
@@ -202,4 +206,55 @@ RSpec.describe BulkImports::PipelineBatchWorker, feature_category: :importers do
       expect(batch.reload).to be_failed
     end
   end
+
+  context 'with stop signal from database health check' do
+    around do |example|
+      with_sidekiq_server_middleware do |chain|
+        chain.add Gitlab::SidekiqMiddleware::SkipJobs
+        Sidekiq::Testing.inline! { example.run }
+      end
+    end
+
+    before do
+      stub_feature_flags("drop_sidekiq_jobs_#{described_class.name}": false)
+
+      stop_signal = instance_double("Gitlab::Database::HealthStatus::Signals::Stop", stop?: true)
+      allow(Gitlab::Database::HealthStatus).to receive(:evaluate).and_return([stop_signal])
+    end
+
+    it 'defers the job by set time' do
+      expect_next_instance_of(described_class) do |worker|
+        expect(worker).not_to receive(:perform).with(batch.id)
+      end
+
+      expect(described_class).to receive(:perform_in).with(described_class::DEFER_ON_HEALTH_DELAY, batch.id)
+
+      described_class.perform_async(batch.id)
+    end
+
+    it 'lazy evaluates schema and tables', :aggregate_failures do
+      block = described_class.database_health_check_attrs[:block]
+
+      job_args = [batch.id]
+
+      schema, table = block.call([job_args])
+
+      expect(schema).to eq(:gitlab_main_cell)
+      expect(table).to eq(['labels'])
+    end
+
+    context 'when `bulk_import_deferred_workers` feature flag is disabled' do
+      it 'does not defer job execution' do
+        stub_feature_flags(bulk_import_deferred_workers: false)
+
+        expect_next_instance_of(described_class) do |worker|
+          expect(worker).to receive(:perform).with(batch.id)
+        end
+
+        expect(described_class).not_to receive(:perform_in)
+
+        described_class.perform_async(batch.id)
+      end
+    end
+  end
 end
diff --git a/spec/workers/bulk_imports/pipeline_worker_spec.rb b/spec/workers/bulk_imports/pipeline_worker_spec.rb
index 6ea7334f6a6..d99b3e9de73 100644
--- a/spec/workers/bulk_imports/pipeline_worker_spec.rb
+++ b/spec/workers/bulk_imports/pipeline_worker_spec.rb
@@ -9,6 +9,10 @@ RSpec.describe BulkImports::PipelineWorker, feature_category: :importers do
 
       def run; end
 
+      def self.relation
+        'labels'
+      end
+
       def self.file_extraction_pipeline?
         false
       end
@@ -155,6 +159,62 @@ RSpec.describe BulkImports::PipelineWorker, feature_category: :importers do
     end
   end
 
+  context 'with stop signal from database health check' do
+    around do |example|
+      with_sidekiq_server_middleware do |chain|
+        chain.add Gitlab::SidekiqMiddleware::SkipJobs
+        Sidekiq::Testing.inline! { example.run }
+      end
+    end
+
+    before do
+      stub_feature_flags("drop_sidekiq_jobs_#{described_class.name}": false)
+
+      stop_signal = instance_double("Gitlab::Database::HealthStatus::Signals::Stop", stop?: true)
+      allow(Gitlab::Database::HealthStatus).to receive(:evaluate).and_return([stop_signal])
+    end
+
+    it 'defers the job by set time' do
+      expect_next_instance_of(described_class) do |worker|
+        expect(worker).not_to receive(:perform).with(pipeline_tracker.id, pipeline_tracker.stage, entity.id)
+      end
+
+      expect(described_class).to receive(:perform_in).with(
+        described_class::DEFER_ON_HEALTH_DELAY,
+        pipeline_tracker.id,
+        pipeline_tracker.stage,
+        entity.id
+      )
+
+      described_class.perform_async(pipeline_tracker.id, pipeline_tracker.stage, entity.id)
+    end
+
+    it 'lazy evaluates schema and tables', :aggregate_failures do
+      block = described_class.database_health_check_attrs[:block]
+
+      job_args = [pipeline_tracker.id, pipeline_tracker.stage, entity.id]
+
+      schema, table = block.call([job_args])
+
+      expect(schema).to eq(:gitlab_main_cell)
+      expect(table).to eq(['labels'])
+    end
+
+    context 'when `bulk_import_deferred_workers` feature flag is disabled' do
+      it 'does not defer job execution' do
+        stub_feature_flags(bulk_import_deferred_workers: false)
+
+        expect_next_instance_of(described_class) do |worker|
+          expect(worker).to receive(:perform).with(pipeline_tracker.id, pipeline_tracker.stage, entity.id)
+        end
+
+        expect(described_class).not_to receive(:perform_in)
+
+        described_class.perform_async(pipeline_tracker.id, pipeline_tracker.stage, entity.id)
+      end
+    end
+  end
+
   context 'when pipeline is finished' do
     let(:pipeline_tracker) do
       create(
diff --git a/spec/workers/ci/pipeline_success_unlock_artifacts_worker_spec.rb b/spec/workers/ci/pipeline_success_unlock_artifacts_worker_spec.rb
index 60a34fdab53..d5f3c2b92fb 100644
--- a/spec/workers/ci/pipeline_success_unlock_artifacts_worker_spec.rb
+++ b/spec/workers/ci/pipeline_success_unlock_artifacts_worker_spec.rb
@@ -75,7 +75,8 @@ RSpec.describe Ci::PipelineSuccessUnlockArtifactsWorker, feature_category: :buil
       expect(described_class.database_health_check_attrs).to eq(
         gitlab_schema: :gitlab_ci,
         delay_by: described_class::DEFAULT_DEFER_DELAY,
-        tables: [:ci_job_artifacts]
+        tables: [:ci_job_artifacts],
+        block: nil
       )
     end
   end
diff --git a/spec/workers/concerns/worker_attributes_spec.rb b/spec/workers/concerns/worker_attributes_spec.rb
index 90c07a9c959..767a55162fb 100644
--- a/spec/workers/concerns/worker_attributes_spec.rb
+++ b/spec/workers/concerns/worker_attributes_spec.rb
@@ -37,7 +37,7 @@ RSpec.describe WorkerAttributes, feature_category: :shared do
       :worker_has_external_dependencies? | :worker_has_external_dependencies! | false            | []                                 | true
       :idempotent?                       | :idempotent!                       | false            | []                                 | true
       :big_payload?                      | :big_payload!                      | false            | []                                 | true
-      :database_health_check_attrs       | :defer_on_database_health_signal   | nil              | [:gitlab_main, [:users], 1.minute] | { gitlab_schema: :gitlab_main, tables: [:users], delay_by: 1.minute }
+      :database_health_check_attrs       | :defer_on_database_health_signal   | nil              | [:gitlab_main, [:users], 1.minute] | { gitlab_schema: :gitlab_main, tables: [:users], delay_by: 1.minute, block: nil }
     end
     # rubocop: enable Layout/LineLength