1 files changed, 147 insertions, 0 deletions
diff --git a/spec/lib/gitlab/database/partitioning_migration_helpers/table_management_helpers_spec.rb b/spec/lib/gitlab/database/partitioning_migration_helpers/table_management_helpers_spec.rb
index 86f79b213ae..44ef0b307fe 100644
--- a/spec/lib/gitlab/database/partitioning_migration_helpers/table_management_helpers_spec.rb
+++ b/spec/lib/gitlab/database/partitioning_migration_helpers/table_management_helpers_spec.rb
@@ -480,6 +480,153 @@ RSpec.describe Gitlab::Database::PartitioningMigrationHelpers::TableManagementHe
     end
   end
 
+  describe '#finalize_backfilling_partitioned_table' do
+    let(:source_table) { 'todos' }
+    let(:source_column) { 'id' }
+
+    context 'when the table is not allowed' do
+      let(:source_table) { :this_table_is_not_allowed }
+
+      it 'raises an error' do
+        expect(migration).to receive(:assert_table_is_allowed).with(source_table).and_call_original
+
+        expect do
+          migration.finalize_backfilling_partitioned_table source_table
+        end.to raise_error(/#{source_table} is not allowed for use/)
+      end
+    end
+
+    context 'when the partitioned table does not exist' do
+      it 'raises an error' do
+        expect(migration).to receive(:table_exists?).with(partitioned_table).and_return(false)
+
+        expect do
+          migration.finalize_backfilling_partitioned_table source_table
+        end.to raise_error(/could not find partitioned table for #{source_table}/)
+      end
+    end
+
+    context 'finishing pending background migration jobs' do
+      let(:source_table_double) { double('table name') }
+      let(:raw_arguments) { [1, 50_000, source_table_double, partitioned_table, source_column] }
+
+      before do
+        allow(migration).to receive(:table_exists?).with(partitioned_table).and_return(true)
+        allow(migration).to receive(:copy_missed_records)
+        allow(migration).to receive(:execute).with(/VACUUM/)
+      end
+
+      it 'finishes remaining jobs for the correct table' do
+        expect_next_instance_of(described_class::JobArguments) do |job_arguments|
+          expect(job_arguments).to receive(:source_table_name).and_call_original
+        end
+
+        expect(Gitlab::BackgroundMigration).to receive(:steal)
+          .with(described_class::MIGRATION_CLASS_NAME)
+          .and_yield(raw_arguments)
+
+        expect(source_table_double).to receive(:==).with(source_table.to_s)
+
+        migration.finalize_backfilling_partitioned_table source_table
+      end
+    end
+
+    context 'when there is missed data' do
+      let(:partitioned_model) { Class.new(ActiveRecord::Base) }
+      let(:timestamp) { Time.utc(2019, 12, 1, 12).round }
+      let!(:todo1) { create(:todo, created_at: timestamp, updated_at: timestamp) }
+      let!(:todo2) { create(:todo, created_at: timestamp, updated_at: timestamp) }
+      let!(:todo3) { create(:todo, created_at: timestamp, updated_at: timestamp) }
+      let!(:todo4) { create(:todo, created_at: timestamp, updated_at: timestamp) }
+
+      let!(:pending_job1) do
+        create(:background_migration_job,
+               class_name: described_class::MIGRATION_CLASS_NAME,
+               arguments: [todo1.id, todo2.id, source_table, partitioned_table, source_column])
+      end
+
+      let!(:pending_job2) do
+        create(:background_migration_job,
+               class_name: described_class::MIGRATION_CLASS_NAME,
+               arguments: [todo3.id, todo3.id, source_table, partitioned_table, source_column])
+      end
+
+      let!(:succeeded_job) do
+        create(:background_migration_job, :succeeded,
+               class_name: described_class::MIGRATION_CLASS_NAME,
+               arguments: [todo4.id, todo4.id, source_table, partitioned_table, source_column])
+      end
+
+      before do
+        partitioned_model.primary_key = :id
+        partitioned_model.table_name = partitioned_table
+
+        allow(migration).to receive(:queue_background_migration_jobs_by_range_at_intervals)
+
+        migration.partition_table_by_date source_table, partition_column, min_date: min_date, max_date: max_date
+
+        allow(Gitlab::BackgroundMigration).to receive(:steal)
+        allow(migration).to receive(:execute).with(/VACUUM/)
+      end
+
+      it 'idempotently cleans up after failed background migrations' do
+        expect(partitioned_model.count).to eq(0)
+
+        partitioned_model.insert!(todo2.attributes)
+
+        expect_next_instance_of(Gitlab::Database::PartitioningMigrationHelpers::BackfillPartitionedTable) do |backfill|
+          allow(backfill).to receive(:transaction_open?).and_return(false)
+
+          expect(backfill).to receive(:perform)
+            .with(todo1.id, todo2.id, source_table, partitioned_table, source_column)
+            .and_call_original
+
+          expect(backfill).to receive(:perform)
+            .with(todo3.id, todo3.id, source_table, partitioned_table, source_column)
+            .and_call_original
+        end
+
+        migration.finalize_backfilling_partitioned_table source_table
+
+        expect(partitioned_model.count).to eq(3)
+
+        [todo1, todo2, todo3].each do |original|
+          copy = partitioned_model.find(original.id)
+          expect(copy.attributes).to eq(original.attributes)
+        end
+
+        expect(partitioned_model.find_by_id(todo4.id)).to be_nil
+
+        [pending_job1, pending_job2].each do |job|
+          expect(job.reload).to be_succeeded
+        end
+      end
+
+      it 'raises an error if no job tracking records are marked as succeeded' do
+        expect_next_instance_of(Gitlab::Database::PartitioningMigrationHelpers::BackfillPartitionedTable) do |backfill|
+          allow(backfill).to receive(:transaction_open?).and_return(false)
+
+          expect(backfill).to receive(:perform).and_return(0)
+        end
+
+        expect do
+          migration.finalize_backfilling_partitioned_table source_table
+        end.to raise_error(/failed to update tracking record/)
+      end
+
+      it 'vacuums the table after loading is complete' do
+        expect_next_instance_of(Gitlab::Database::PartitioningMigrationHelpers::BackfillPartitionedTable) do |backfill|
+          allow(backfill).to receive(:perform).and_return(1)
+        end
+
+        expect(migration).to receive(:disable_statement_timeout).and_call_original
+        expect(migration).to receive(:execute).with("VACUUM FREEZE ANALYZE #{partitioned_table}")
+
+        migration.finalize_backfilling_partitioned_table source_table
+      end
+    end
+  end
+
   def filter_columns_by_name(columns, names)
     columns.reject { |c| names.include?(c.name) }
   end