Use primary for internal registry migration API

Update the internal API endpoints for the registry
migration to read from primary to prevent any race
conditions with the related workers.

Changelog: fixed

lib/api/internal/container_registry/migration.rb

@@ -36,23 +36,25 @@ def find_repository!(path)
           requires :status, type: String, values: POSSIBLE_VALUES, desc: 'The migration step status'
         end
         put 'internal/registry/repositories/*repository_path/migration/status' do
-          repository = find_repository!(declared_params[:repository_path])
+          ::Gitlab::Database::LoadBalancing::Session.current.use_primary do
+            repository = find_repository!(declared_params[:repository_path])
 
-          unless repository.migration_in_active_state?
-            bad_request!("Wrong migration state (#{repository.migration_state})")
-          end
-
-          case declared_params[:status]
-          when STATUS_PRE_IMPORT_COMPLETE
-            unless repository.finish_pre_import_and_start_import
-              bad_request!("Couldn't transition from pre_importing to importing")
+            unless repository.migration_in_active_state?
+              bad_request!("Wrong migration state (#{repository.migration_state})")
             end
-          when STATUS_IMPORT_COMPLETE
-            unless repository.finish_import
-              bad_request!("Couldn't transition from importing to import_done")
+
+            case declared_params[:status]
+            when STATUS_PRE_IMPORT_COMPLETE
+              unless repository.finish_pre_import_and_start_import
+                bad_request!("Couldn't transition from pre_importing to importing")
+              end
+            when STATUS_IMPORT_COMPLETE
+              unless repository.finish_import
+                bad_request!("Couldn't transition from importing to import_done")
+              end
+            when STATUS_IMPORT_FAILED, STATUS_PRE_IMPORT_FAILED
+              repository.abort_import
             end
-          when STATUS_IMPORT_FAILED, STATUS_PRE_IMPORT_FAILED
-            repository.abort_import
           end
 
           status 200

spec/requests/api/internal/container_registry/migration_spec.rb

@@ -2,7 +2,7 @@
 
 require 'spec_helper'
 
-RSpec.describe API::Internal::ContainerRegistry::Migration do
+RSpec.describe API::Internal::ContainerRegistry::Migration, :aggregate_failures do
   let_it_be_with_reload(:repository) { create(:container_repository) }
 
   let(:secret_token) { 'secret_token' }
@@ -147,6 +147,17 @@
 
         it_behaves_like 'returning an error', returning_status: :not_found
       end
+
+      context 'query read location' do
+        it 'reads from the primary' do
+          expect(ContainerRepository).to receive(:find_by_path!).and_wrap_original do |m, *args|
+            expect(::Gitlab::Database::LoadBalancing::Session.current.use_primary?).to eq(true)
+            m.call(*args)
+          end
+
+          subject
+        end
+      end
     end
 
     context 'with an invalid sent token' do