Delete tracking records on partitioning rollback

When a partitioned migration is rolled back by calling the
drop_partitioned_table_for method, delete all records related to that
migration in the background_migration_jobs table.

changelogs/unreleased/pb-background-migration-tracking-cleanup-on-rollback.yml

+---
+title: Delete tracking records on partitioning migration rollback
+merge_request: 36743
+author:
+type: fixed

db/migrate/20200713152443_add_background_migration_job_index_for_partitioning_migrations.rb

+# frozen_string_literal: true
+
+class AddBackgroundMigrationJobIndexForPartitioningMigrations < ActiveRecord::Migration[6.0]
+  DOWNTIME = false
+
+  INDEX_NAME = 'index_background_migration_jobs_for_partitioning_migrations'
+
+  def up
+    # rubocop:disable Migration/AddIndex
+    add_index :background_migration_jobs, '((arguments ->> 2))', name: INDEX_NAME,
+      where: "class_name = 'Gitlab::Database::PartitioningMigrationHelpers::BackfillPartitionedTable'"
+    # rubocop:enable Migration/AddIndex
+  end
+
+  def down
+    remove_index :background_migration_jobs, name: INDEX_NAME # rubocop:disable Migration/RemoveIndex
+  end
+end

db/structure.sql

@@ -18692,6 +18692,8 @@ CREATE UNIQUE INDEX index_aws_roles_on_role_external_id ON public.aws_roles USIN
 
 CREATE UNIQUE INDEX index_aws_roles_on_user_id ON public.aws_roles USING btree (user_id);
 
+CREATE INDEX index_background_migration_jobs_for_partitioning_migrations ON public.background_migration_jobs USING btree (((arguments ->> 2))) WHERE (class_name = 'Gitlab::Database::PartitioningMigrationHelpers::BackfillPartitionedTable'::text);
+
 CREATE INDEX index_background_migration_jobs_on_class_name_and_arguments ON public.background_migration_jobs USING btree (class_name, arguments);
 
 CREATE INDEX index_background_migration_jobs_on_class_name_and_status_and_id ON public.background_migration_jobs USING btree (class_name, status, id);
@@ -23773,5 +23775,6 @@ COPY "schema_migrations" (version) FROM STDIN;
 20200710105332
 20200710130234
 20200713071042
+20200713152443
 \.
 

lib/gitlab/database/background_migration_job.rb

@@ -5,8 +5,13 @@ module Gitlab
     class BackgroundMigrationJob < ActiveRecord::Base # rubocop:disable Rails/ApplicationRecord
       self.table_name = :background_migration_jobs
 
+      scope :for_migration_class, -> (class_name) { where(class_name: normalize_class_name(class_name)) }
       scope :for_migration_execution, -> (class_name, arguments) do
-        where('class_name = ? AND arguments = ?', normalize_class_name(class_name), arguments.to_json)
+        for_migration_class(class_name).where('arguments = ?', arguments.to_json)
+      end
+
+      scope :for_partitioning_migration, -> (class_name, table_name) do
+        for_migration_class(class_name).where('arguments ->> 2 = ?', table_name)
       end
 
       enum status: {

lib/gitlab/database/partitioning_migration_helpers/table_management_helpers.rb

@@ -69,6 +69,8 @@ module Gitlab
           assert_table_is_allowed(table_name)
           assert_not_in_transaction_block(scope: ERROR_SCOPE)
 
+          cleanup_migration_jobs(table_name)
+
           with_lock_retries do
             trigger_name = make_sync_trigger_name(table_name)
             drop_trigger(table_name, trigger_name)
@@ -261,6 +263,10 @@ module Gitlab
             other_job_arguments: [source_table_name.to_s, partitioned_table_name, source_key],
             track_jobs: true)
         end
+
+        def cleanup_migration_jobs(table_name)
+          ::Gitlab::Database::BackgroundMigrationJob.for_partitioning_migration(MIGRATION_CLASS_NAME, table_name).delete_all
+        end
       end
     end
   end

spec/lib/gitlab/database/background_migration_job_spec.rb

@@ -25,6 +25,26 @@ RSpec.describe Gitlab::Database::BackgroundMigrationJob do
     end
   end
 
+  describe '.for_partitioning_migration' do
+    let!(:job1) { create(:background_migration_job, arguments: [1, 100, 'other_table']) }
+    let!(:job2) { create(:background_migration_job, arguments: [1, 100, 'audit_events']) }
+    let!(:job3) { create(:background_migration_job, class_name: 'OtherJob', arguments: [1, 100, 'audit_events']) }
+
+    it 'returns jobs matching class_name and the table_name job argument' do
+      relation = described_class.for_partitioning_migration('TestJob', 'audit_events')
+
+      expect(relation.count).to eq(1)
+      expect(relation.first).to have_attributes(class_name: 'TestJob', arguments: [1, 100, 'audit_events'])
+    end
+
+    it 'normalizes class names by removing leading ::' do
+      relation = described_class.for_partitioning_migration('::TestJob', 'audit_events')
+
+      expect(relation.count).to eq(1)
+      expect(relation.first).to have_attributes(class_name: 'TestJob', arguments: [1, 100, 'audit_events'])
+    end
+  end
+
   describe '.mark_all_as_succeeded' do
     let!(:job1) { create(:background_migration_job, arguments: [1, 100]) }
     let!(:job2) { create(:background_migration_job, arguments: [1, 100]) }

spec/lib/gitlab/database/partitioning_migration_helpers/table_management_helpers_spec.rb

@@ -351,6 +351,7 @@ RSpec.describe Gitlab::Database::PartitioningMigrationHelpers::TableManagementHe
     let(:expected_tables) do
       %w[000000 201912 202001 202002].map { |suffix| "#{Gitlab::Database::DYNAMIC_PARTITIONS_SCHEMA}.#{partitioned_table}_#{suffix}" }.unshift(partitioned_table)
     end
+    let(:migration_class) { 'Gitlab::Database::PartitioningMigrationHelpers::BackfillPartitionedTable' }
 
     context 'when the table is not allowed' do
       let(:source_table) { :this_table_is_not_allowed }
@@ -389,6 +390,22 @@ RSpec.describe Gitlab::Database::PartitioningMigrationHelpers::TableManagementHe
         expect(connection.table_exists?(table)).to be(false)
       end
     end
+
+    context 'cleaning up background migration tracking records' do
+      let!(:job1) { create(:background_migration_job, class_name: migration_class, arguments: [1, 10, source_table]) }
+      let!(:job2) { create(:background_migration_job, class_name: migration_class, arguments: [11, 20, source_table]) }
+      let!(:job3) { create(:background_migration_job, class_name: migration_class, arguments: [1, 10, 'other_table']) }
+
+      it 'deletes any tracking records from the background_migration_jobs table' do
+        migration.partition_table_by_date source_table, partition_column, min_date: min_date, max_date: max_date
+
+        expect { migration.drop_partitioned_table_for(source_table) }
+          .to change { ::Gitlab::Database::BackgroundMigrationJob.count }.from(3).to(1)
+
+        remaining_record = ::Gitlab::Database::BackgroundMigrationJob.first
+        expect(remaining_record).to have_attributes(class_name: migration_class, arguments: [1, 10, 'other_table'])
+      end
+    end
   end
 
   describe '#create_hash_partitions' do