Update aborted registry import logic

Update Registry EnqueuerWorker to fetch the true
repository status from the registry when making
decisions about what to do with aborted imports.

app/models/container_repository.rb

@@ -15,6 +15,7 @@ class ContainerRepository < ApplicationRecord
   MIGRATION_STATES = (IDLE_MIGRATION_STATES + ACTIVE_MIGRATION_STATES).freeze
 
   TooManyImportsError = Class.new(StandardError)
+  NativeImportError = Class.new(StandardError)
 
   belongs_to :project
 
@@ -85,9 +86,7 @@ class ContainerRepository < ApplicationRecord
     end
 
     state :pre_import_done do
-      validates :migration_pre_import_started_at,
-                :migration_pre_import_done_at,
-                presence: true
+      validates :migration_pre_import_done_at, presence: true
     end
 
     state :importing do
@@ -113,7 +112,7 @@ class ContainerRepository < ApplicationRecord
     end
 
     event :finish_pre_import do
-      transition pre_importing: :pre_import_done
+      transition %i[pre_importing import_aborted] => :pre_import_done
     end
 
     event :start_import do
@@ -121,7 +120,7 @@ class ContainerRepository < ApplicationRecord
     end
 
     event :finish_import do
-      transition importing: :import_done
+      transition %i[importing import_aborted] => :import_done
     end
 
     event :already_migrated do
@@ -155,7 +154,7 @@ class ContainerRepository < ApplicationRecord
       end
     end
 
-    before_transition pre_importing: :pre_import_done do |container_repository|
+    before_transition %i[pre_importing import_aborted] => :pre_import_done do |container_repository|
       container_repository.migration_pre_import_done_at = Time.zone.now
     end
 
@@ -170,7 +169,7 @@ class ContainerRepository < ApplicationRecord
       end
     end
 
-    before_transition importing: :import_done do |container_repository|
+    before_transition %i[importing import_aborted] => :import_done do |container_repository|
       container_repository.migration_import_done_at = Time.zone.now
     end
 
@@ -272,13 +271,29 @@ class ContainerRepository < ApplicationRecord
     finish_pre_import && start_import
   end
 
-  def retry_migration
-    return if migration_import_done_at
+  def retry_aborted_migration
+    return unless migration_state == 'import_aborted'
+
+    import_status = gitlab_api_client.import_status(self.path)
 
-    if migration_pre_import_done_at
+    case import_status
+    when 'native'
+      raise NativeImportError
+    when 'import_in_progress'
+      nil
+    when 'import_complete'
+      finish_import
+    when 'import_failed'
       retry_import
-    else
+    when 'pre_import_in_progress'
+      nil
+    when 'pre_import_complete'
+      finish_pre_import_and_start_import
+    when 'pre_import_failed'
       retry_pre_import
+    else
+      # If the import_status request fails, use the timestamp to guess current state
+      migration_pre_import_done_at ? retry_import : retry_pre_import
     end
   end
 

app/workers/container_registry/migration/enqueuer_worker.rb

@@ -32,7 +32,7 @@ module ContainerRegistry
       private
 
       def handle_aborted_migration
-        return unless next_aborted_repository&.retry_migration
+        return unless next_aborted_repository&.retry_aborted_migration
 
         log_extra_metadata_on_done(:container_repository_id, next_aborted_repository.id)
         log_extra_metadata_on_done(:import_type, 'retry')

lib/container_registry/gitlab_api_client.rb

@@ -45,6 +45,11 @@ module ContainerRegistry
       IMPORT_RESPONSES.fetch(response.status, :error)
     end
 
+    def import_status(path)
+      body_hash = response_body(faraday.get(import_url_for(path)))
+      body_hash['status'] || 'error'
+    end
+
     private
 
     def start_import_for(path, pre:)

spec/lib/container_registry/gitlab_api_client_spec.rb

@@ -77,7 +77,7 @@ RSpec.describe ContainerRegistry::GitlabApiClient do
     end
   end
 
-  describe '#pre_import_repository' do
+  describe '#import_repository' do
     subject { client.import_repository(path) }
 
     where(:status_code, :expected_result) do
@@ -101,6 +101,26 @@ RSpec.describe ContainerRegistry::GitlabApiClient do
     end
   end
 
+  describe '#import_status' do
+    subject { client.import_status(path) }
+
+    before do
+      stub_import_status(path, status)
+    end
+
+    context 'with a status' do
+      let(:status) { 'this_is_a_test' }
+
+      it { is_expected.to eq(status) }
+    end
+
+    context 'with no status' do
+      let(:status) { nil }
+
+      it { is_expected.to eq('error') }
+    end
+  end
+
   describe '.supports_gitlab_api?' do
     subject { described_class.supports_gitlab_api? }
 
@@ -171,4 +191,14 @@ RSpec.describe ContainerRegistry::GitlabApiClient do
       .with(headers: { 'Accept' => described_class::JSON_TYPE })
       .to_return(status: status_code, body: '')
   end
+
+  def stub_import_status(path, status)
+    stub_request(:get, "#{registry_api_url}/gitlab/v1/import/#{path}/")
+      .with(headers: { 'Accept' => described_class::JSON_TYPE })
+      .to_return(
+        status: 200,
+        body: { status: status }.to_json,
+        headers: { content_type: 'application/json' }
+      )
+  end
 end

spec/models/container_repository_spec.rb

@@ -218,7 +218,7 @@ RSpec.describe ContainerRepository, :aggregate_failures do
 
       subject { repository.finish_pre_import }
 
-      it_behaves_like 'transitioning from allowed states', %w[pre_importing]
+      it_behaves_like 'transitioning from allowed states', %w[pre_importing import_aborted]
 
       it 'sets migration_pre_import_done_at' do
         expect { subject }.to change { repository.reload.migration_pre_import_done_at }
@@ -263,7 +263,7 @@ RSpec.describe ContainerRepository, :aggregate_failures do
 
       subject { repository.finish_import }
 
-      it_behaves_like 'transitioning from allowed states', %w[importing]
+      it_behaves_like 'transitioning from allowed states', %w[importing import_aborted]
       it_behaves_like 'queueing the next import'
 
       it 'sets migration_import_done_at and queues the next import' do
@@ -334,44 +334,118 @@ RSpec.describe ContainerRepository, :aggregate_failures do
         end
       end
 
-      it_behaves_like 'transitioning from allowed states', %w[pre_importing]
+      it_behaves_like 'transitioning from allowed states', %w[pre_importing import_aborted]
       it_behaves_like 'transitioning to importing'
     end
   end
 
-  describe '#retry_migration' do
-    subject { repository.retry_migration }
+  describe '#retry_aborted_migration' do
+    subject { repository.retry_aborted_migration }
 
-    it 'retries the pre_import' do
-      expect(repository).to receive(:retry_pre_import).and_return(true)
-      expect(repository).not_to receive(:retry_import)
+    shared_examples 'no action' do
+      it 'does nothing' do
+        expect { subject }.not_to change { repository.reload.migration_state }
 
-      expect(subject).to eq(true)
+        expect(subject).to eq(nil)
+      end
     end
 
-    context 'when migration is done pre-importing' do
-      before do
-        repository.update_columns(migration_pre_import_done_at: Time.zone.now)
+    shared_examples 'retrying the pre_import' do
+      it 'retries the pre_import' do
+        expect(repository).to receive(:migration_pre_import).and_return(:ok)
+
+        expect { subject }.to change { repository.reload.migration_state }.to('pre_importing')
       end
+    end
 
-      it 'returns' do
-        expect(repository).to receive(:retry_import).and_return(true)
-        expect(repository).not_to receive(:retry_pre_import)
+    shared_examples 'retrying the import' do
+      it 'retries the import' do
+        expect(repository).to receive(:migration_import).and_return(:ok)
 
-        expect(subject).to eq(true)
+        expect { subject }.to change { repository.reload.migration_state }.to('importing')
       end
     end
 
-    context 'when migration is already complete' do
+    context 'when migration_state is not aborted' do
+      it_behaves_like 'no action'
+    end
+
+    context 'when migration_state is aborted' do
       before do
-        repository.update_columns(migration_import_done_at: Time.zone.now)
+        repository.abort_import
+
+        allow(repository.gitlab_api_client)
+            .to receive(:import_status).with(repository.path).and_return(client_response)
       end
 
-      it 'returns' do
-        expect(repository).not_to receive(:retry_pre_import)
-        expect(repository).not_to receive(:retry_import)
+      context 'native response' do
+        let(:client_response) { 'native' }
 
-        expect(subject).to eq(nil)
+        it 'raises an error' do
+          expect { subject }.to raise_error(described_class::NativeImportError)
+        end
+      end
+
+      context 'import_in_progress response' do
+        let(:client_response) { 'import_in_progress' }
+
+        it_behaves_like 'no action'
+      end
+
+      context 'import_complete response' do
+        let(:client_response) { 'import_complete' }
+
+        it 'finishes the import' do
+          expect { subject }.to change { repository.reload.migration_state }.to('import_done')
+        end
+      end
+
+      context 'import_failed response' do
+        let(:client_response) { 'import_failed' }
+
+        it_behaves_like 'retrying the import'
+      end
+
+      context 'pre_import_in_progress response' do
+        let(:client_response) { 'pre_import_in_progress' }
+
+        it_behaves_like 'no action'
+      end
+
+      context 'pre_import_complete response' do
+        let(:client_response) { 'pre_import_complete' }
+
+        it 'finishes the pre_import and starts the import' do
+          expect(repository).to receive(:finish_pre_import).and_call_original
+          expect(repository).to receive(:migration_import).and_return(:ok)
+
+          expect { subject }.to change { repository.reload.migration_state }.to('importing')
+        end
+      end
+
+      context 'pre_import_failed response' do
+        let(:client_response) { 'pre_import_failed' }
+
+        it_behaves_like 'retrying the pre_import'
+      end
+
+      context 'error response' do
+        let(:client_response) { 'error' }
+
+        context 'migration_pre_import_done_at is NULL' do
+          it_behaves_like 'retrying the pre_import'
+        end
+
+        context 'migration_pre_import_done_at is not NULL' do
+          before do
+            repository.update_columns(
+              migration_pre_import_started_at: 5.minutes.ago,
+              migration_pre_import_done_at: Time.zone.now
+            )
+          end
+
+          it_behaves_like 'retrying the import'
+        end
       end
     end
   end

spec/workers/container_registry/migration/enqueuer_worker_spec.rb

@@ -113,6 +113,7 @@ RSpec.describe ContainerRegistry::Migration::EnqueuerWorker, :aggregate_failures
         allow(worker).to receive(:next_aborted_repository) do
           next_aborted_repository = method.call
           allow(next_aborted_repository).to receive(:migration_import).and_return(:ok)
+          allow(next_aborted_repository.gitlab_api_client).to receive(:import_status).and_return('import_failed')
           next_aborted_repository
         end