Add latest changes from gitlab-org/gitlab@master

2 files changed, 69 insertions, 2 deletions

diff --git a/lib/gitlab/database/background_migration_job.rb b/lib/gitlab/database/background_migration_job.rb
index 445735b232a..1b9d7cbc9a1 100644
--- a/lib/gitlab/database/background_migration_job.rb
+++ b/lib/gitlab/database/background_migration_job.rb
@@ -3,6 +3,8 @@
 module Gitlab
   module Database
     class BackgroundMigrationJob < ActiveRecord::Base # rubocop:disable Rails/ApplicationRecord
+      include EachBatch
+
       self.table_name = :background_migration_jobs
 
       scope :for_migration_class, -> (class_name) { where(class_name: normalize_class_name(class_name)) }
diff --git a/lib/gitlab/database/partitioning_migration_helpers/table_management_helpers.rb b/lib/gitlab/database/partitioning_migration_helpers/table_management_helpers.rb
index e6d8ec55319..04c63d27f9e 100644
--- a/lib/gitlab/database/partitioning_migration_helpers/table_management_helpers.rb
+++ b/lib/gitlab/database/partitioning_migration_helpers/table_management_helpers.rb
@@ -6,6 +6,7 @@ module Gitlab
       module TableManagementHelpers
         include ::Gitlab::Database::SchemaHelpers
         include ::Gitlab::Database::DynamicModelHelpers
+        include ::Gitlab::Database::MigrationHelpers
         include ::Gitlab::Database::Migrations::BackgroundMigrationHelpers
 
         ALLOWED_TABLES = %w[audit_events].freeze
@@ -15,6 +16,12 @@ module Gitlab
         BATCH_INTERVAL = 2.minutes.freeze
         BATCH_SIZE = 50_000
 
+        JobArguments = Struct.new(:start_id, :stop_id, :source_table_name, :partitioned_table_name, :source_column) do
+          def self.from_array(arguments)
+            self.new(*arguments)
+          end
+        end
+
         # Creates a partitioned copy of an existing table, using a RANGE partitioning strategy on a timestamp column.
         # One partition is created per month between the given `min_date` and `max_date`. Also installs a trigger on
         # the original table to copy writes into the partitioned table. To copy over historic data from before creation
@@ -132,6 +139,42 @@ module Gitlab
           end
         end
 
+        # Executes cleanup tasks from a previous BackgroundMigration to backfill a partitioned table by finishing
+        # pending jobs and performing a final data synchronization.
+        # This performs two steps:
+        #   1. Wait to finish any pending BackgroundMigration jobs that have not succeeded
+        #   2. Inline copy any missed rows from the original table to the partitioned table
+        #
+        # **NOTE** Migrations using this method cannot be scheduled in the same release as the migration that
+        # schedules the background migration using the `enqueue_background_migration` helper, or else the
+        # background migration jobs will be force-executed.
+        #
+        # Example:
+        #
+        #   finalize_backfilling_partitioned_table :audit_events
+        #
+        def finalize_backfilling_partitioned_table(table_name)
+          assert_table_is_allowed(table_name)
+          assert_not_in_transaction_block(scope: ERROR_SCOPE)
+
+          partitioned_table_name = make_partitioned_table_name(table_name)
+          unless table_exists?(partitioned_table_name)
+            raise "could not find partitioned table for #{table_name}, " \
+              "this could indicate the previous partitioning migration has been rolled back."
+          end
+
+          Gitlab::BackgroundMigration.steal(MIGRATION_CLASS_NAME) do |raw_arguments|
+            JobArguments.from_array(raw_arguments).source_table_name == table_name.to_s
+          end
+
+          primary_key = connection.primary_key(table_name)
+          copy_missed_records(table_name, partitioned_table_name, primary_key)
+
+          disable_statement_timeout do
+            execute("VACUUM FREEZE ANALYZE #{partitioned_table_name}")
+          end
+        end
+
         private
 
         def assert_table_is_allowed(table_name)
@@ -284,7 +327,7 @@ module Gitlab
           create_trigger(table_name, trigger_name, function_name, fires: 'AFTER INSERT OR UPDATE OR DELETE')
         end
 
-        def enqueue_background_migration(source_table_name, partitioned_table_name, source_key)
+        def enqueue_background_migration(source_table_name, partitioned_table_name, source_column)
           source_model = define_batchable_model(source_table_name)
 
           queue_background_migration_jobs_by_range_at_intervals(
@@ -292,13 +335,35 @@ module Gitlab
             MIGRATION_CLASS_NAME,
             BATCH_INTERVAL,
             batch_size: BATCH_SIZE,
-            other_job_arguments: [source_table_name.to_s, partitioned_table_name, source_key],
+            other_job_arguments: [source_table_name.to_s, partitioned_table_name, source_column],
             track_jobs: true)
         end
 
         def cleanup_migration_jobs(table_name)
           ::Gitlab::Database::BackgroundMigrationJob.for_partitioning_migration(MIGRATION_CLASS_NAME, table_name).delete_all
         end
+
+        def copy_missed_records(source_table_name, partitioned_table_name, source_column)
+          backfill_table = BackfillPartitionedTable.new
+          relation = ::Gitlab::Database::BackgroundMigrationJob.pending
+            .for_partitioning_migration(MIGRATION_CLASS_NAME, source_table_name)
+
+          relation.each_batch do |batch|
+            batch.each do |pending_migration_job|
+              job_arguments = JobArguments.from_array(pending_migration_job.arguments)
+              start_id = job_arguments.start_id
+              stop_id = job_arguments.stop_id
+
+              say("Backfilling data into partitioned table for ids from #{start_id} to #{stop_id}")
+              job_updated_count = backfill_table.perform(start_id, stop_id, source_table_name,
+                partitioned_table_name, source_column)
+
+              unless job_updated_count > 0
+                raise "failed to update tracking record for ids from #{start_id} to #{stop_id}"
+              end
+            end
+          end
+        end
       end
     end
   end