Add latest changes from gitlab-org/gitlab@master

3 files changed, 56 insertions, 50 deletions

diff --git a/lib/gitlab/database/partitioning_migration_helpers/backfill_partitioned_table.rb b/lib/gitlab/database/partitioning_migration_helpers/backfill_partitioned_table.rb
index dcf457b9d63..e87707953ae 100644
--- a/lib/gitlab/database/partitioning_migration_helpers/backfill_partitioned_table.rb
+++ b/lib/gitlab/database/partitioning_migration_helpers/backfill_partitioned_table.rb
@@ -21,7 +21,7 @@ module Gitlab
             return
           end
 
-          bulk_copy = BulkCopy.new(source_table, partitioned_table, source_column, connection: connection)
+          bulk_copy = Gitlab::Database::PartitioningMigrationHelpers::BulkCopy.new(source_table, partitioned_table, source_column, connection: connection)
           parent_batch_relation = relation_scoped_to_range(source_table, source_column, start_id, stop_id)
 
           parent_batch_relation.each_batch(of: SUB_BATCH_SIZE) do |sub_batch|
@@ -56,41 +56,6 @@ module Gitlab
         def mark_jobs_as_succeeded(*arguments)
           BackgroundMigrationJob.mark_all_as_succeeded(self.class.name, arguments)
         end
-
-        # Helper class to copy data between two tables via upserts
-        class BulkCopy
-          DELIMITER = ', '
-
-          attr_reader :source_table, :destination_table, :source_column, :connection
-
-          def initialize(source_table, destination_table, source_column, connection:)
-            @source_table = source_table
-            @destination_table = destination_table
-            @source_column = source_column
-            @connection = connection
-          end
-
-          def copy_between(start_id, stop_id)
-            connection.execute(<<~SQL)
-              INSERT INTO #{destination_table} (#{column_listing})
-              SELECT #{column_listing}
-              FROM #{source_table}
-              WHERE #{source_column} BETWEEN #{start_id} AND #{stop_id}
-              FOR UPDATE
-              ON CONFLICT (#{conflict_targets}) DO NOTHING
-            SQL
-          end
-
-          private
-
-          def column_listing
-            @column_listing ||= connection.columns(source_table).map(&:name).join(DELIMITER)
-          end
-
-          def conflict_targets
-            connection.primary_key(destination_table).join(DELIMITER)
-          end
-        end
       end
     end
   end
diff --git a/lib/gitlab/database/partitioning_migration_helpers/bulk_copy.rb b/lib/gitlab/database/partitioning_migration_helpers/bulk_copy.rb
new file mode 100644
index 00000000000..b8f5a2e3ad4
--- /dev/null
+++ b/lib/gitlab/database/partitioning_migration_helpers/bulk_copy.rb
@@ -0,0 +1,42 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    module PartitioningMigrationHelpers
+      # Helper class to copy data between two tables via upserts
+      class BulkCopy
+        DELIMITER = ', '
+
+        attr_reader :source_table, :destination_table, :source_column, :connection
+
+        def initialize(source_table, destination_table, source_column, connection:)
+          @source_table = source_table
+          @destination_table = destination_table
+          @source_column = source_column
+          @connection = connection
+        end
+
+        def copy_between(start_id, stop_id)
+          connection.execute(<<~SQL)
+            INSERT INTO #{destination_table} (#{column_listing})
+            SELECT #{column_listing}
+            FROM #{source_table}
+            WHERE #{source_column} BETWEEN #{start_id} AND #{stop_id}
+            FOR UPDATE
+            ON CONFLICT (#{conflict_targets}) DO NOTHING
+          SQL
+        end
+
+        private
+
+        def column_listing
+          @column_listing ||= connection.columns(source_table).map(&:name).join(DELIMITER)
+        end
+
+        def conflict_targets
+          connection.primary_keys(destination_table).join(DELIMITER)
+        end
+      end
+    end
+  end
+end
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 5a942577006..e3cf1298df6 100644
--- a/lib/gitlab/database/partitioning_migration_helpers/table_management_helpers.rb
+++ b/lib/gitlab/database/partitioning_migration_helpers/table_management_helpers.rb
@@ -12,8 +12,10 @@ module Gitlab
         ERROR_SCOPE = 'table partitioning'
 
         MIGRATION_CLASS_NAME = "::#{module_parent_name}::BackfillPartitionedTable"
+        MIGRATION = "BackfillPartitionedTable"
         BATCH_INTERVAL = 2.minutes.freeze
         BATCH_SIZE = 50_000
+        SUB_BATCH_SIZE = 2_500
 
         JobArguments = Struct.new(:start_id, :stop_id, :source_table_name, :partitioned_table_name, :source_column) do
           def self.from_array(arguments)
@@ -108,7 +110,16 @@ module Gitlab
 
           partitioned_table_name = make_partitioned_table_name(table_name)
           primary_key = connection.primary_key(table_name)
-          enqueue_background_migration(table_name, partitioned_table_name, primary_key)
+
+          queue_batched_background_migration(
+            MIGRATION,
+            table_name,
+            primary_key,
+            partitioned_table_name,
+            batch_size: BATCH_SIZE,
+            sub_batch_size: SUB_BATCH_SIZE,
+            job_interval: BATCH_INTERVAL
+          )
         end
 
         # Cleanup a previously enqueued background migration to copy data into a partitioned table. This will not
@@ -150,7 +161,7 @@ module Gitlab
         #   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
+        # schedules the background migration using the `enqueue_partitioning_data_migration` helper, or else the
         # background migration jobs will be force-executed.
         #
         # Example:
@@ -445,18 +456,6 @@ 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_column)
-          source_model = define_batchable_model(source_table_name)
-
-          queue_background_migration_jobs_by_range_at_intervals(
-            source_model,
-            MIGRATION_CLASS_NAME,
-            BATCH_INTERVAL,
-            batch_size: BATCH_SIZE,
-            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