Add migration_plan to container_repositories

Add column to store migration_plan in container_repositories
to be used when scheduling the phase 2 registry migration.

Changelog: other

db/migrate/20220316202200_add_migration_plan_to_container_repositories.rb

+# frozen_string_literal: true
+
+class AddMigrationPlanToContainerRepositories < Gitlab::Database::Migration[1.0]
+  # rubocop:disable Migration/AddLimitToTextColumns
+  # limit is added in 20220316202402_add_text_limit_to_container_repositories_migration_plan
+  def change
+    add_column(:container_repositories, :migration_plan, :text)
+  end
+  # rubocop:enable Migration/AddLimitToTextColumns
+end

db/migrate/20220316202402_add_text_limit_to_container_repositories_migration_plan.rb

+# frozen_string_literal: true
+
+class AddTextLimitToContainerRepositoriesMigrationPlan < Gitlab::Database::Migration[1.0]
+  disable_ddl_transaction!
+
+  def up
+    add_text_limit :container_repositories, :migration_plan, 255
+  end
+
+  def down
+    remove_text_limit :container_repositories, :migration_plan
+  end
+end

db/post_migrate/20220316202640_populate_container_repositories_migration_plan.rb

+# frozen_string_literal: true
+
+class PopulateContainerRepositoriesMigrationPlan < Gitlab::Database::Migration[1.0]
+  MIGRATION = 'PopulateContainerRepositoryMigrationPlan'
+  DELAY_INTERVAL = 2.minutes.to_i
+  BATCH_SIZE = 1500
+
+  disable_ddl_transaction!
+
+  def up
+    queue_background_migration_jobs_by_range_at_intervals(
+      define_batchable_model('container_repositories'),
+      MIGRATION,
+      DELAY_INTERVAL,
+      batch_size: BATCH_SIZE,
+      track_jobs: true
+    )
+  end
+
+  def down
+    # no-op
+  end
+end

db/schema_migrations/20220316202200

+28722182d4ead079c8c0c36a5c075a29c5da93b4315af277311f561fc945d65f
\ No newline at end of file

db/schema_migrations/20220316202402

+006ef4f0c63e39e3874d39528ca6cfd233c17d6256d07bd2c0e7b2262d0a4825
\ No newline at end of file

db/schema_migrations/20220316202640

+7ea2288c45d497e2fde56b1d7c1e82360ed40442a2501a24e9795380adf5b911
\ No newline at end of file

db/structure.sql

@@ -13724,6 +13724,8 @@ CREATE TABLE container_repositories (
     migration_skipped_reason smallint,
     migration_state text DEFAULT 'default'::text NOT NULL,
     migration_aborted_in_state text,
+    migration_plan text,
+    CONSTRAINT check_05e9012f36 CHECK ((char_length(migration_plan) <= 255)),
     CONSTRAINT check_13c58fe73a CHECK ((char_length(migration_state) <= 255)),
     CONSTRAINT check_97f0249439 CHECK ((char_length(migration_aborted_in_state) <= 255))
 );

lib/gitlab/background_migration/populate_container_repository_migration_plan.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module BackgroundMigration
+    # The class to populates the migration_plan column of container_repositories
+    # with the current plan of the namespaces that owns the container_repository
+    #
+    # The plan can be NULL, in which case no UPDATE
+    # will be executed.
+    class PopulateContainerRepositoryMigrationPlan
+      def perform(start_id, end_id)
+        (start_id..end_id).each do |id|
+          execute(<<~SQL)
+            WITH selected_plan AS (
+              SELECT "plans"."name"
+              FROM "container_repositories"
+              INNER JOIN "projects" ON "projects"."id" = "container_repositories"."project_id"
+              INNER JOIN "namespaces" ON "namespaces"."id" = "projects"."namespace_id"
+              INNER JOIN "gitlab_subscriptions" ON "gitlab_subscriptions"."namespace_id" = "namespaces"."traversal_ids"[1]
+              INNER JOIN "plans" ON "plans"."id" = "gitlab_subscriptions"."hosted_plan_id"
+              WHERE "container_repositories"."id" = #{id}
+            )
+            UPDATE container_repositories
+            SET migration_plan = selected_plan.name
+            FROM selected_plan
+            WHERE container_repositories.id = #{id};
+          SQL
+        end
+
+        mark_job_as_succeeded(start_id, end_id)
+      end
+
+      private
+
+      def connection
+        @connection ||= ::ActiveRecord::Base.connection
+      end
+
+      def execute(sql)
+        connection.execute(sql)
+      end
+
+      def mark_job_as_succeeded(*arguments)
+        Gitlab::Database::BackgroundMigrationJob.mark_all_as_succeeded(
+          self.class.name.demodulize,
+          arguments
+        )
+      end
+    end
+  end
+end

spec/lib/gitlab/background_migration/populate_container_repository_migration_plan_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::BackgroundMigration::PopulateContainerRepositoryMigrationPlan, schema: 20220316202640 do
+  let_it_be(:container_repositories) { table(:container_repositories) }
+  let_it_be(:projects) { table(:projects) }
+  let_it_be(:namespaces) { table(:namespaces) }
+  let_it_be(:gitlab_subscriptions) { table(:gitlab_subscriptions) }
+  let_it_be(:plans) { table(:plans) }
+  let_it_be(:namespace_statistics) { table(:namespace_statistics) }
+
+  let!(:namepace1) { namespaces.create!(id: 1, type: 'Group', name: 'group1', path: 'group1', traversal_ids: [1]) }
+  let!(:namepace2) { namespaces.create!(id: 2, type: 'Group', name: 'group2', path: 'group2', traversal_ids: [2]) }
+  let!(:namepace3) { namespaces.create!(id: 3, type: 'Group', name: 'group3', path: 'group3', traversal_ids: [3]) }
+  let!(:sub_namespace) { namespaces.create!(id: 4, type: 'Group', name: 'group3', path: 'group3', parent_id: 1, traversal_ids: [1, 4]) }
+  let!(:plan1) { plans.create!(id: 1, name: 'plan1') }
+  let!(:plan2) { plans.create!(id: 2, name: 'plan2') }
+  let!(:gitlab_subscription1) { gitlab_subscriptions.create!(id: 1, namespace_id: 1, hosted_plan_id: 1) }
+  let!(:gitlab_subscription2) { gitlab_subscriptions.create!(id: 2, namespace_id: 2, hosted_plan_id: 2) }
+  let!(:project1) { projects.create!(id: 1, name: 'project1', path: 'project1', namespace_id: 4) }
+  let!(:project2) { projects.create!(id: 2, name: 'project2', path: 'project2', namespace_id: 2) }
+  let!(:project3) { projects.create!(id: 3, name: 'project3', path: 'project3', namespace_id: 3) }
+  let!(:container_repository1) { container_repositories.create!(id: 1, name: 'cr1', project_id: 1) }
+  let!(:container_repository2) { container_repositories.create!(id: 2, name: 'cr2', project_id: 2) }
+  let!(:container_repository3) { container_repositories.create!(id: 3, name: 'cr3', project_id: 3) }
+
+  let(:migration) { described_class.new }
+
+  subject do
+    migration.perform(1, 4)
+  end
+
+  it 'updates the migration_plan to match the actual plan', :aggregate_failures do
+    expect(Gitlab::Database::BackgroundMigrationJob).to receive(:mark_all_as_succeeded)
+      .with('PopulateContainerRepositoryMigrationPlan', [1, 4]).and_return(true)
+
+    subject
+
+    expect(container_repository1.reload.migration_plan).to eq('plan1')
+    expect(container_repository2.reload.migration_plan).to eq('plan2')
+    expect(container_repository3.reload.migration_plan).to eq(nil)
+  end
+end

spec/migrations/20220316202640_populate_container_repositories_migration_plan_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+require_migration!
+
+RSpec.describe PopulateContainerRepositoriesMigrationPlan, :aggregate_failures do
+  let_it_be(:namespaces) { table(:namespaces) }
+  let_it_be(:projects) { table(:projects) }
+  let_it_be(:container_repositories) { table(:container_repositories) }
+
+  let!(:namespace) { namespaces.create!(id: 1, name: 'namespace', path: 'namespace') }
+  let!(:project) { projects.create!(id: 1, name: 'project', path: 'project', namespace_id: 1) }
+  let!(:container_repository1) { container_repositories.create!(name: 'container_repository1', project_id: 1) }
+  let!(:container_repository2) { container_repositories.create!(name: 'container_repository2', project_id: 1) }
+  let!(:container_repository3) { container_repositories.create!(name: 'container_repository3', project_id: 1) }
+
+  before do
+    stub_const("#{described_class.name}::BATCH_SIZE", 2)
+  end
+
+  it 'schedules jobs for container_repositories to populate migration_state' do
+    Sidekiq::Testing.fake! do
+      freeze_time do
+        migrate!
+
+        expect(described_class::MIGRATION).to be_scheduled_delayed_migration(
+          2.minutes, container_repository1.id, container_repository2.id)
+        expect(described_class::MIGRATION).to be_scheduled_delayed_migration(
+          4.minutes, container_repository3.id, container_repository3.id)
+        expect(BackgroundMigrationWorker.jobs.size).to eq(2)
+      end
+    end
+  end
+end