1 files changed, 49 insertions, 36 deletions
diff --git a/lib/gitlab/background_migration/batched_migration_job.rb b/lib/gitlab/background_migration/batched_migration_job.rb
index 64401bc0674..973ab20f547 100644
--- a/lib/gitlab/background_migration/batched_migration_job.rb
+++ b/lib/gitlab/background_migration/batched_migration_job.rb
@@ -9,55 +9,68 @@ module Gitlab
     # see https://docs.gitlab.com/ee/development/database/batched_background_migrations.html#job-arguments.
     class BatchedMigrationJob
       include Gitlab::Database::DynamicModelHelpers
+      include Gitlab::ClassAttributes
 
-      def initialize(
-        start_id:, end_id:, batch_table:, batch_column:, sub_batch_size:, pause_ms:, job_arguments: [], connection:
-      )
+      DEFAULT_FEATURE_CATEGORY = :database
 
-        @start_id = start_id
-        @end_id = end_id
-        @batch_table = batch_table
-        @batch_column = batch_column
-        @sub_batch_size = sub_batch_size
-        @pause_ms = pause_ms
-        @job_arguments = job_arguments
-        @connection = connection
-      end
+      class << self
+        def generic_instance(batch_table:, batch_column:, job_arguments: [], connection:)
+          new(
+            batch_table: batch_table, batch_column: batch_column,
+            job_arguments: job_arguments, connection: connection,
+            start_id: 0, end_id: 0, sub_batch_size: 0, pause_ms: 0
+          )
+        end
 
-      def self.generic_instance(batch_table:, batch_column:, job_arguments: [], connection:)
-        new(
-          batch_table: batch_table, batch_column: batch_column,
-          job_arguments: job_arguments, connection: connection,
-          start_id: 0, end_id: 0, sub_batch_size: 0, pause_ms: 0
-        )
-      end
+        def job_arguments_count
+          0
+        end
 
-      def self.job_arguments_count
-        0
-      end
+        def operation_name(operation)
+          define_method('operation_name') do
+            operation
+          end
+        end
 
-      def self.operation_name(operation)
-        define_method('operation_name') do
-          operation
+        def job_arguments(*args)
+          args.each.with_index do |arg, index|
+            define_method(arg) do
+              @job_arguments[index]
+            end
+          end
+
+          define_singleton_method(:job_arguments_count) do
+            args.count
+          end
         end
-      end
 
-      def self.job_arguments(*args)
-        args.each.with_index do |arg, index|
-          define_method(arg) do
-            @job_arguments[index]
+        def scope_to(scope)
+          define_method(:filter_batch) do |relation|
+            instance_exec(relation, &scope)
           end
         end
 
-        define_singleton_method(:job_arguments_count) do
-          args.count
+        def feature_category(feature_category_name = nil)
+          if feature_category_name.present?
+            set_class_attribute(:feature_category, feature_category_name)
+          else
+            get_class_attribute(:feature_category) || DEFAULT_FEATURE_CATEGORY
+          end
         end
       end
 
-      def self.scope_to(scope)
-        define_method(:filter_batch) do |relation|
-          instance_exec(relation, &scope)
-        end
+      def initialize(
+        start_id:, end_id:, batch_table:, batch_column:, sub_batch_size:, pause_ms:, job_arguments: [], connection:
+      )
+
+        @start_id = start_id
+        @end_id = end_id
+        @batch_table = batch_table
+        @batch_column = batch_column
+        @sub_batch_size = sub_batch_size
+        @pause_ms = pause_ms
+        @job_arguments = job_arguments
+        @connection = connection
       end
 
       def filter_batch(relation)