3 files changed, 64 insertions, 11 deletions

diff --git a/lib/container_registry/base_client.rb b/lib/container_registry/base_client.rb
index 22d4510fe71..bb9422ae048 100644
--- a/lib/container_registry/base_client.rb
+++ b/lib/container_registry/base_client.rb
@@ -37,14 +37,24 @@ module ContainerRegistry
     class << self
       private
 
-      def with_dummy_client(return_value_if_disabled: nil)
+      def with_dummy_client(return_value_if_disabled: nil, token_config: { type: :full_access_token, path: nil })
         registry_config = Gitlab.config.registry
         unless registry_config.enabled && registry_config.api_url.present?
           return return_value_if_disabled
         end
 
-        token = Auth::ContainerRegistryAuthenticationService.access_token([], [])
-        yield new(registry_config.api_url, token: token)
+        yield new(registry_config.api_url, token: token_from(token_config))
+      end
+
+      def token_from(config)
+        case config[:type]
+        when :full_access_token
+          Auth::ContainerRegistryAuthenticationService.access_token([], [])
+        when :nested_repositories_token
+          return unless config[:path]
+
+          Auth::ContainerRegistryAuthenticationService.pull_nested_repositories_access_token(config[:path])
+        end
       end
     end
 
diff --git a/lib/container_registry/gitlab_api_client.rb b/lib/container_registry/gitlab_api_client.rb
index 3cd7003d1f8..0cd8f8509f6 100644
--- a/lib/container_registry/gitlab_api_client.rb
+++ b/lib/container_registry/gitlab_api_client.rb
@@ -5,10 +5,12 @@ module ContainerRegistry
     include Gitlab::Utils::StrongMemoize
 
     JSON_TYPE = 'application/json'
+    CANCEL_RESPONSE_STATUS_HEADER = 'status'
 
     IMPORT_RESPONSES = {
       200 => :already_imported,
       202 => :ok,
+      400 => :bad_request,
       401 => :unauthorized,
       404 => :not_found,
       409 => :already_being_imported,
@@ -25,6 +27,12 @@ module ContainerRegistry
       end
     end
 
+    def self.deduplicated_size(path)
+      with_dummy_client(token_config: { type: :nested_repositories_token, path: path }) do |client|
+        client.repository_details(path, sizing: :self_with_descendants)['size_bytes']
+      end
+    end
+
     # https://gitlab.com/gitlab-org/container-registry/-/blob/master/docs-gitlab/api.md#compliance-check
     def supports_gitlab_api?
       strong_memoize(:supports_gitlab_api) do
@@ -50,18 +58,38 @@ module ContainerRegistry
       IMPORT_RESPONSES.fetch(response.status, :error)
     end
 
+    # https://gitlab.com/gitlab-org/container-registry/-/blob/master/docs-gitlab/api.md#cancel-repository-import
+    def cancel_repository_import(path, force: false)
+      response = with_import_token_faraday do |faraday_client|
+        faraday_client.delete(import_url_for(path)) do |req|
+          req.params['force'] = true if force
+        end
+      end
+
+      status = IMPORT_RESPONSES.fetch(response.status, :error)
+      actual_state = response.body[CANCEL_RESPONSE_STATUS_HEADER]
+
+      { status: status, migration_state: actual_state }
+    end
+
     # https://gitlab.com/gitlab-org/container-registry/-/blob/master/docs-gitlab/api.md#get-repository-import-status
     def import_status(path)
       with_import_token_faraday do |faraday_client|
-        body_hash = response_body(faraday_client.get(import_url_for(path)))
-        body_hash['status'] || 'error'
+        response = faraday_client.get(import_url_for(path))
+
+        # Temporary solution for https://gitlab.com/gitlab-org/gitlab/-/issues/356085#solutions
+        # this will trigger a `retry_pre_import`
+        break 'pre_import_failed' unless response.success?
+
+        body_hash = response_body(response)
+        body_hash&.fetch('status') || 'error'
       end
     end
 
-    def repository_details(path, with_size: false)
+    def repository_details(path, sizing: nil)
       with_token_faraday do |faraday_client|
         req = faraday_client.get("/gitlab/v1/repositories/#{path}/") do |req|
-          req.params['size'] = 'self' if with_size
+          req.params['size'] = sizing if sizing
         end
 
         break {} unless req.success?
diff --git a/lib/container_registry/migration.rb b/lib/container_registry/migration.rb
index b03c94e5ebf..005ef880034 100644
--- a/lib/container_registry/migration.rb
+++ b/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
@@ -28,9 +39,9 @@ module ContainerRegistry
 
     def self.enqueue_waiting_time
       return 0 if Feature.enabled?(:container_registry_migration_phase2_enqueue_speed_fast)
-      return 6.hours if Feature.enabled?(:container_registry_migration_phase2_enqueue_speed_slow)
+      return 165.minutes if Feature.enabled?(:container_registry_migration_phase2_enqueue_speed_slow)
 
-      1.hour
+      45.minutes
     end
 
     def self.capacity
@@ -46,8 +57,12 @@ 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?
+      Feature.enabled?(:container_registry_migration_phase2_all_plans)
     end
   end
 end