Registry import enqueuer uses migration_plan

Update the ready_for_import scope on ContainerRepository
to use the new migration_plan column to reduce the
number of JOINs in the query.

Adds an index on migration_plan.

EE: true
Changelog: performance

db/post_migrate/20220322035654_add_migration_plan_index_to_container_repositories.rb

+# frozen_string_literal: true
+
+class AddMigrationPlanIndexToContainerRepositories < Gitlab::Database::Migration[1.0]
+  INDEX_NAME = 'idx_container_repos_on_migration_state_migration_plan_created'
+
+  disable_ddl_transaction!
+
+  def up
+    add_concurrent_index :container_repositories,
+      [:migration_state, :migration_plan, :created_at],
+      name: INDEX_NAME
+  end
+
+  def down
+    remove_concurrent_index_by_name :container_repositories, INDEX_NAME
+  end
+end

db/schema_migrations/20220322035654

+b37315953ac3cdc6d75f5287a4c264ebdee932f19cfe975f5d2c57b353da35d9
\ No newline at end of file

db/structure.sql

@@ -26472,6 +26472,8 @@ CREATE INDEX idx_container_repos_on_exp_cleanup_status_project_id_start_date ON
 
 CREATE INDEX idx_container_repos_on_import_started_at_when_importing ON container_repositories USING btree (migration_import_started_at) WHERE (migration_state = 'importing'::text);
 
+CREATE INDEX idx_container_repos_on_migration_state_migration_plan_created ON container_repositories USING btree (migration_state, migration_plan, created_at);
+
 CREATE INDEX idx_container_repos_on_pre_import_done_at_when_pre_import_done ON container_repositories USING btree (migration_pre_import_done_at) WHERE (migration_state = 'pre_import_done'::text);
 
 CREATE INDEX idx_container_repos_on_pre_import_started_at_when_pre_importing ON container_repositories USING btree (migration_pre_import_started_at) WHERE (migration_state = 'pre_importing'::text);

ee/app/models/ee/container_repository.rb

@@ -31,14 +31,7 @@ module EE
         if ::ContainerRegistry::Migration.limit_gitlab_org?
           joins(project: [:namespace]).where(namespaces: { path: GITLAB_ORG_NAMESPACE })
         else
-          joins(
-            %{
-              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(plans: { id: ::ContainerRegistry::Migration.target_plan.id })
+          where(migration_plan: ::ContainerRegistry::Migration.target_plans)
         end
       end
     end

ee/spec/models/container_repository_spec.rb

@@ -4,25 +4,18 @@ require 'spec_helper'
 
 RSpec.describe ContainerRepository, :saas do
   describe '.with_target_import_tier' do
-    let_it_be(:root_group) { create(:group) }
-    let_it_be(:group) { create(:group, parent_id: root_group.id) }
-    let_it_be(:project) { create(:project, namespace: group) }
-    let_it_be(:valid_container_repository) { create(:container_repository, project: project) }
+    let_it_be(:valid_container_repository) { create(:container_repository, migration_plan: 'free') }
 
     let_it_be(:gitlab_namespace) { create(:namespace, path: 'gitlab-org') }
     let_it_be(:gitlab_project) { create(:project, namespace: gitlab_namespace) }
     let_it_be(:gitlab_container_repository) { create(:container_repository, project: gitlab_project) }
 
-    let_it_be(:ultimate_project) { create(:project) }
-    let_it_be(:ultimate_container_repository) { create(:container_repository, project: ultimate_project) }
-
-    let(:subscription) { create(:gitlab_subscription, :premium, namespace: root_group) }
-    let(:ultimate_subscription) { create(:gitlab_subscription, :ultimate, namespace: ultimate_project.namespace) }
+    let_it_be(:ultimate_container_repository) { create(:container_repository, migration_plan: 'ultimate') }
 
     subject { described_class.with_target_import_tier }
 
     before do
-      stub_application_setting(container_registry_import_target_plan: subscription.hosted_plan.name)
+      stub_application_setting(container_registry_import_target_plan: valid_container_repository.migration_plan)
     end
 
     context 'all_plans disabled' do
@@ -39,7 +32,7 @@ RSpec.describe ContainerRepository, :saas do
           stub_feature_flags(container_registry_migration_limit_gitlab_org: false)
         end
 
-        it { is_expected.to contain_exactly(valid_container_repository) }
+        it { is_expected.to contain_exactly(valid_container_repository, gitlab_container_repository) }
       end
     end
 
@@ -51,17 +44,12 @@ RSpec.describe ContainerRepository, :saas do
   describe '.ready_for_import' do
     include_context 'importable repositories'
 
-    let_it_be(:ultimate_project) { create(:project) }
-    let_it_be(:ultimate_container_repository) { create(:container_repository, project: ultimate_project, created_at: 2.days.ago) }
-
-    let_it_be(:subscription) { create(:gitlab_subscription, :premium, namespace: root_group) }
-    let_it_be(:denied_subscription) { create(:gitlab_subscription, :premium, namespace: denied_project.namespace) }
-    let_it_be(:ultimate_subscription) { create(:gitlab_subscription, :ultimate, namespace: ultimate_project.namespace) }
+    let_it_be(:ultimate_container_repository) { create(:container_repository, migration_plan: 'ultimate', created_at: 2.days.ago) }
 
     subject { described_class.ready_for_import }
 
     before do
-      stub_application_setting(container_registry_import_target_plan: subscription.hosted_plan.name)
+      stub_application_setting(container_registry_import_target_plan: valid_container_repository.migration_plan)
     end
 
     it { is_expected.to contain_exactly(valid_container_repository, valid_container_repository2) }

lib/container_registry/migration.rb

@@ -2,6 +2,17 @@
 
 module ContainerRegistry
   module Migration
+    # Some container repositories do not have a plan associated with them, they will be imported with
+    # the free tiers
+    FREE_TIERS = ['free', 'early_adopter', nil].freeze
+    PREMIUM_TIERS = %w[premium bronze silver premium_trial].freeze
+    ULTIMATE_TIERS = %w[ultimate gold ultimate_trial].freeze
+    PLAN_GROUPS = {
+      'free' => FREE_TIERS,
+      'premium' => PREMIUM_TIERS,
+      'ultimate' => ULTIMATE_TIERS
+    }.freeze
+
     class << self
       delegate :container_registry_import_max_tags_count, to: ::Gitlab::CurrentSettings
       delegate :container_registry_import_max_retries, to: ::Gitlab::CurrentSettings
@@ -46,8 +57,8 @@ module ContainerRegistry
       0
     end
 
-    def self.target_plan
-      Plan.find_by_name(target_plan_name)
+    def self.target_plans
+      PLAN_GROUPS[target_plan_name]
     end
 
     def self.all_plans?

spec/lib/container_registry/migration_spec.rb

@@ -154,15 +154,21 @@ RSpec.describe ContainerRegistry::Migration do
     end
   end
 
-  describe '.target_plan' do
-    let_it_be(:plan) { create(:plan) }
+  describe '.target_plans' do
+    subject { described_class.target_plans }
 
-    before do
-      stub_application_setting(container_registry_import_target_plan: plan.name)
+    where(:target_plan, :result) do
+      'free'     | described_class::FREE_TIERS
+      'premium'  | described_class::PREMIUM_TIERS
+      'ultimate' | described_class::ULTIMATE_TIERS
     end
 
-    it 'returns the matching application_setting' do
-      expect(described_class.target_plan).to eq(plan)
+    with_them do
+      before do
+        stub_application_setting(container_registry_import_target_plan: target_plan)
+      end
+
+      it { is_expected.to eq(result) }
     end
   end
 

spec/models/container_repository_spec.rb

@@ -1274,7 +1274,7 @@ RSpec.describe ContainerRepository, :aggregate_failures do
       subject { described_class.ready_for_import }
 
       before do
-        stub_application_setting(container_registry_import_target_plan: root_group.actual_plan_name)
+        stub_application_setting(container_registry_import_target_plan: valid_container_repository.migration_plan)
       end
 
       it 'works' do

spec/support/shared_contexts/container_repositories_shared_context.rb

 # frozen_string_literal: true
 
 RSpec.shared_context 'importable repositories' do
-  let_it_be(:root_group) { create(:group) }
-  let_it_be(:group) { create(:group, parent_id: root_group.id) }
-  let_it_be(:project) { create(:project, namespace: group) }
-  let_it_be(:valid_container_repository) { create(:container_repository, project: project, created_at: 2.days.ago) }
-  let_it_be(:valid_container_repository2) { create(:container_repository, project: project, created_at: 1.year.ago) }
-  let_it_be(:importing_container_repository) { create(:container_repository, :importing, project: project, created_at: 2.days.ago) }
-  let_it_be(:new_container_repository) { create(:container_repository, project: project) }
+  let_it_be(:valid_container_repository) { create(:container_repository, created_at: 2.days.ago, migration_plan: 'free') }
+  let_it_be(:valid_container_repository2) { create(:container_repository, created_at: 1.year.ago, migration_plan: 'free') }
+  let_it_be(:importing_container_repository) { create(:container_repository, :importing, created_at: 2.days.ago, migration_plan: 'free') }
+  let_it_be(:new_container_repository) { create(:container_repository, migration_plan: 'free') }
 
   let_it_be(:denied_root_group) { create(:group) }
   let_it_be(:denied_group) { create(:group, parent_id: denied_root_group.id) }