Add latest changes from gitlab-org/gitlab@master

3 files changed, 78 insertions, 11 deletions

diff --git a/lib/generators/batched_background_migration/batched_background_migration_generator.rb b/lib/generators/batched_background_migration/batched_background_migration_generator.rb
index c68ed52c1a0..44aff6fe17a 100644
--- a/lib/generators/batched_background_migration/batched_background_migration_generator.rb
+++ b/lib/generators/batched_background_migration/batched_background_migration_generator.rb
@@ -6,9 +6,12 @@ module BatchedBackgroundMigration
   class BatchedBackgroundMigrationGenerator < ActiveRecord::Generators::Base
     source_root File.expand_path('templates', __dir__)
 
-    class_option :table_name
-    class_option :column_name
-    class_option :feature_category
+    class_option :table_name, type: :string, required: true, desc: "Table from which records we will be batching"
+    class_option :column_name, type: :string, required: true, desc: "Column to use for batching", default: :id
+    class_option :feature_category, type: :string, required: true,
+      desc: "Feature category to which this batched background migration belongs to"
+    class_option :ee_only, type: :boolean, desc: "Generate files for EE-only batched background migration",
+      default: false
 
     def validate!
       raise ArgumentError, "table_name is required" unless table_name.present?
@@ -29,15 +32,32 @@ module BatchedBackgroundMigration
     end
 
     def create_batched_background_migration_class_and_specs
-      template(
-        "batched_background_migration_job.template",
-        File.join("lib/gitlab/background_migration/#{file_name}.rb")
-      )
+      if ee_only?
+        template(
+          "ee_batched_background_migration_job.template",
+          File.join("ee/lib/ee/gitlab/background_migration/#{file_name}.rb")
+        )
 
-      template(
-        "batched_background_migration_job_spec.template",
-        File.join("spec/lib/gitlab/background_migration/#{file_name}_spec.rb")
-      )
+        template(
+          "foss_batched_background_migration_job.template",
+          File.join("lib/gitlab/background_migration/#{file_name}.rb")
+        )
+
+        template(
+          "batched_background_migration_job_spec.template",
+          File.join("ee/spec/lib/ee/gitlab/background_migration/#{file_name}_spec.rb")
+        )
+      else
+        template(
+          "batched_background_migration_job.template",
+          File.join("lib/gitlab/background_migration/#{file_name}.rb")
+        )
+
+        template(
+          "batched_background_migration_job_spec.template",
+          File.join("spec/lib/gitlab/background_migration/#{file_name}_spec.rb")
+        )
+      end
     end
 
     def create_dictionary_file
@@ -65,6 +85,10 @@ module BatchedBackgroundMigration
       options[:feature_category]
     end
 
+    def ee_only?
+      options[:ee_only]
+    end
+
     def current_milestone
       version = Gem::Version.new(File.read('VERSION'))
       version.release.segments.first(2).join('.')
diff --git a/lib/generators/batched_background_migration/templates/ee_batched_background_migration_job.template b/lib/generators/batched_background_migration/templates/ee_batched_background_migration_job.template
new file mode 100644
index 00000000000..b36fc216acd
--- /dev/null
+++ b/lib/generators/batched_background_migration/templates/ee_batched_background_migration_job.template
@@ -0,0 +1,29 @@
+# frozen_string_literal: true
+
+# See https://docs.gitlab.com/ee/development/database/batched_background_migrations.html
+# for more information on how to use batched background migrations
+
+# Update below commented lines with appropriate values.
+
+module EE
+  module Gitlab
+    module BackgroundMigration
+      module <%= class_name %>
+        extend ActiveSupport::Concern
+        extend ::Gitlab::Utils::Override
+
+        prepended do
+          # operation_name :my_operation
+          # scope_to ->(relation) { relation.where(column: "value") }
+        end
+
+        override :perform
+        def perform
+          each_sub_batch do |sub_batch|
+            # Your action on each sub_batch
+          end
+        end
+      end
+    end
+  end
+end
diff --git a/lib/generators/batched_background_migration/templates/foss_batched_background_migration_job.template b/lib/generators/batched_background_migration/templates/foss_batched_background_migration_job.template
new file mode 100644
index 00000000000..6a39c32955d
--- /dev/null
+++ b/lib/generators/batched_background_migration/templates/foss_batched_background_migration_job.template
@@ -0,0 +1,14 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module BackgroundMigration
+    # TODO Add a top-level documentation comment for the class
+    class <%= class_name %> < BatchedMigrationJob
+      feature_category :<%= feature_category %>
+
+      def perform; end
+    end
+  end
+end
+
+Gitlab::BackgroundMigration::<%= class_name %>.prepend_mod