Extract ProjectRepositoryStorageMove logic into concern

In this commit we extract the logic of the
ProjectRepositoryStorageMove class to a concern.

app/models/concerns/repository_storage_movable.rb

+# frozen_string_literal: true
+
+module RepositoryStorageMovable
+  extend ActiveSupport::Concern
+  include AfterCommitQueue
+
+  included do
+    scope :order_created_at_desc, -> { order(created_at: :desc) }
+
+    validates :container, presence: true
+    validates :state, presence: true
+    validates :source_storage_name,
+      on: :create,
+      presence: true,
+      inclusion: { in: ->(_) { Gitlab.config.repositories.storages.keys } }
+    validates :destination_storage_name,
+      on: :create,
+      presence: true,
+      inclusion: { in: ->(_) { Gitlab.config.repositories.storages.keys } }
+    validate :container_repository_writable, on: :create
+
+    default_value_for(:destination_storage_name, allows_nil: false) do
+      pick_repository_storage
+    end
+
+    state_machine initial: :initial do
+      event :schedule do
+        transition initial: :scheduled
+      end
+
+      event :start do
+        transition scheduled: :started
+      end
+
+      event :finish_replication do
+        transition started: :replicated
+      end
+
+      event :finish_cleanup do
+        transition replicated: :finished
+      end
+
+      event :do_fail do
+        transition [:initial, :scheduled, :started] => :failed
+        transition replicated: :cleanup_failed
+      end
+
+      around_transition initial: :scheduled do |storage_move, block|
+        block.call
+
+        begin
+          storage_move.container.set_repository_read_only!(skip_git_transfer_check: true)
+        rescue => err
+          storage_move.add_error(err.message)
+          next false
+        end
+
+        storage_move.run_after_commit do
+          storage_move.schedule_repository_storage_update_worker
+        end
+
+        true
+      end
+
+      before_transition started: :replicated do |storage_move|
+        storage_move.container.set_repository_writable!
+
+        storage_move.update_repository_storage(storage_move.destination_storage_name)
+      end
+
+      before_transition started: :failed do |storage_move|
+        storage_move.container.set_repository_writable!
+      end
+
+      state :initial, value: 1
+      state :scheduled, value: 2
+      state :started, value: 3
+      state :finished, value: 4
+      state :failed, value: 5
+      state :replicated, value: 6
+      state :cleanup_failed, value: 7
+    end
+  end
+
+  class_methods do
+    private
+
+    def pick_repository_storage
+      container_klass = reflect_on_association(:container).class_name.constantize
+
+      container_klass.pick_repository_storage
+    end
+  end
+
+  # Projects, snippets, and group wikis has different db structure. In projects,
+  # we need to update some columns in this step, but we don't with the other resources.
+  #
+  # Therefore, we create this No-op method for snippets and wikis and let project
+  # overwrite it in their implementation.
+  def update_repository_storage(new_storage)
+    # No-op
+  end
+
+  def schedule_repository_storage_update_worker
+    raise NotImplementedError
+  end
+
+  def add_error(message)
+    errors.add(error_key, message)
+  end
+
+  private
+
+  def container_repository_writable
+    add_error(_('is read only')) if container&.repository_read_only?
+  end
+
+  def error_key
+    raise NotImplementedError
+  end
+end

app/models/project.rb

@@ -342,7 +342,7 @@ class Project < ApplicationRecord
 
   has_many :daily_build_group_report_results, class_name: 'Ci::DailyBuildGroupReportResult'
 
-  has_many :repository_storage_moves, class_name: 'ProjectRepositoryStorageMove'
+  has_many :repository_storage_moves, class_name: 'ProjectRepositoryStorageMove', inverse_of: :container
 
   has_many :webide_pipelines, -> { webide_source }, class_name: 'Ci::Pipeline', inverse_of: :project
   has_many :reviews, inverse_of: :project

app/models/project_repository_storage_move.rb

@@ -4,100 +4,31 @@
 # project. For example, moving a project to another gitaly node to help
 # balance storage capacity.
 class ProjectRepositoryStorageMove < ApplicationRecord
-  include AfterCommitQueue
+  extend ::Gitlab::Utils::Override
+  include RepositoryStorageMovable
 
-  belongs_to :project, inverse_of: :repository_storage_moves
+  belongs_to :container, class_name: 'Project', inverse_of: :repository_storage_moves, foreign_key: :project_id
+  alias_attribute :project, :container
+  scope :with_projects, -> { includes(container: :route) }
 
-  validates :project, presence: true
-  validates :state, presence: true
-  validates :source_storage_name,
-    on: :create,
-    presence: true,
-    inclusion: { in: ->(_) { Gitlab.config.repositories.storages.keys } }
-  validates :destination_storage_name,
-    on: :create,
-    presence: true,
-    inclusion: { in: ->(_) { Gitlab.config.repositories.storages.keys } }
-  validate :project_repository_writable, on: :create
-
-  default_value_for(:destination_storage_name, allows_nil: false) do
-    pick_repository_storage
-  end
-
-  state_machine initial: :initial do
-    event :schedule do
-      transition initial: :scheduled
-    end
-
-    event :start do
-      transition scheduled: :started
-    end
-
-    event :finish_replication do
-      transition started: :replicated
-    end
-
-    event :finish_cleanup do
-      transition replicated: :finished
-    end
-
-    event :do_fail do
-      transition [:initial, :scheduled, :started] => :failed
-      transition replicated: :cleanup_failed
-    end
-
-    around_transition initial: :scheduled do |storage_move, block|
-      block.call
-
-      begin
-        storage_move.project.set_repository_read_only!(skip_git_transfer_check: true)
-      rescue => err
-        storage_move.errors.add(:project, err.message)
-        next false
+  override :update_repository_storage
+  def update_repository_storage(new_storage)
+    container.update_column(:repository_storage, new_storage)
   end
 
-      storage_move.run_after_commit do
+  override :schedule_repository_storage_update_worker
+  def schedule_repository_storage_update_worker
     ProjectUpdateRepositoryStorageWorker.perform_async(
-          storage_move.project_id,
-          storage_move.destination_storage_name,
-          storage_move.id
+      project_id,
+      destination_storage_name,
+      id
     )
   end
 
-      true
-    end
-
-    before_transition started: :replicated do |storage_move|
-      storage_move.project.set_repository_writable!
-
-      storage_move.project.update_column(:repository_storage, storage_move.destination_storage_name)
-    end
-
-    before_transition started: :failed do |storage_move|
-      storage_move.project.set_repository_writable!
-    end
-
-    state :initial, value: 1
-    state :scheduled, value: 2
-    state :started, value: 3
-    state :finished, value: 4
-    state :failed, value: 5
-    state :replicated, value: 6
-    state :cleanup_failed, value: 7
-  end
-
-  scope :order_created_at_desc, -> { order(created_at: :desc) }
-  scope :with_projects, -> { includes(project: :route) }
-
-  class << self
-    def pick_repository_storage
-      Project.pick_repository_storage
-    end
-  end
-
   private
 
-  def project_repository_writable
-    errors.add(:project, _('is read only')) if project&.repository_read_only?
+  override :error_key
+  def error_key
+    :project
   end
 end

spec/factories/project_repository_storage_moves.rb

@@ -2,7 +2,7 @@
 
 FactoryBot.define do
   factory :project_repository_storage_move, class: 'ProjectRepositoryStorageMove' do
-    project
+    container { association(:project) }
 
     source_storage_name { 'default' }
 

spec/models/project_repository_storage_move_spec.rb

@@ -3,116 +3,11 @@
 require 'spec_helper'
 
 RSpec.describe ProjectRepositoryStorageMove, type: :model do
-  describe 'associations' do
-    it { is_expected.to belong_to(:project) }
-  end
-
-  describe 'validations' do
-    it { is_expected.to validate_presence_of(:project) }
-    it { is_expected.to validate_presence_of(:state) }
-    it { is_expected.to validate_presence_of(:source_storage_name) }
-    it { is_expected.to validate_presence_of(:destination_storage_name) }
-
-    context 'source_storage_name inclusion' do
-      subject { build(:project_repository_storage_move, source_storage_name: 'missing') }
-
-      it "does not allow repository storages that don't match a label in the configuration" do
-        expect(subject).not_to be_valid
-        expect(subject.errors[:source_storage_name].first).to match(/is not included in the list/)
-      end
-    end
-
-    context 'destination_storage_name inclusion' do
-      subject { build(:project_repository_storage_move, destination_storage_name: 'missing') }
-
-      it "does not allow repository storages that don't match a label in the configuration" do
-        expect(subject).not_to be_valid
-        expect(subject.errors[:destination_storage_name].first).to match(/is not included in the list/)
-      end
-    end
-
-    context 'project repository read-only' do
-      subject { build(:project_repository_storage_move, project: project) }
-
-      let(:project) { build(:project, repository_read_only: true) }
-
-      it "does not allow the project to be read-only on create" do
-        expect(subject).not_to be_valid
-        expect(subject.errors[:project].first).to match(/is read only/)
-      end
-    end
-  end
-
-  describe 'defaults' do
-    context 'destination_storage_name' do
-      subject { build(:project_repository_storage_move) }
-
-      it 'picks storage from ApplicationSetting' do
-        expect(Gitlab::CurrentSettings).to receive(:pick_repository_storage).and_return('picked').at_least(:once)
-
-        expect(subject.destination_storage_name).to eq('picked')
-      end
-    end
-  end
-
-  describe 'state transitions' do
-    let(:project) { create(:project) }
-
-    before do
-      stub_storage_settings('test_second_storage' => { 'path' => 'tmp/tests/extra_storage' })
-    end
+  it_behaves_like 'handles repository moves' do
+    let_it_be_with_refind(:container) { create(:project) }
 
-    context 'when in the default state' do
-      subject(:storage_move) { create(:project_repository_storage_move, project: project, destination_storage_name: 'test_second_storage') }
-
-      context 'and transits to scheduled' do
-        it 'triggers ProjectUpdateRepositoryStorageWorker' do
-          expect(ProjectUpdateRepositoryStorageWorker).to receive(:perform_async).with(project.id, 'test_second_storage', storage_move.id)
-
-          storage_move.schedule!
-
-          expect(project).to be_repository_read_only
-        end
-
-        context 'when the transition fails' do
-          it 'does not trigger ProjectUpdateRepositoryStorageWorker and adds an error' do
-            allow(storage_move.project).to receive(:set_repository_read_only!).and_raise(StandardError, 'foobar')
-            expect(ProjectUpdateRepositoryStorageWorker).not_to receive(:perform_async)
-
-            storage_move.schedule!
-
-            expect(storage_move.errors[:project]).to include('foobar')
-          end
-        end
-      end
-
-      context 'and transits to started' do
-        it 'does not allow the transition' do
-          expect { storage_move.start! }
-            .to raise_error(StateMachines::InvalidTransition)
-        end
-      end
-    end
-
-    context 'when started' do
-      subject(:storage_move) { create(:project_repository_storage_move, :started, project: project, destination_storage_name: 'test_second_storage') }
-
-      context 'and transits to replicated' do
-        it 'sets the repository storage and marks the project as writable' do
-          storage_move.finish_replication!
-
-          expect(project.repository_storage).to eq('test_second_storage')
-          expect(project).not_to be_repository_read_only
-        end
-      end
-
-      context 'and transits to failed' do
-        it 'marks the project as writable' do
-          storage_move.do_fail!
-
-          expect(project).not_to be_repository_read_only
-        end
-      end
-    end
+    let(:repository_storage_factory_key) { :project_repository_storage_move }
+    let(:error_key) { :project }
+    let(:repository_storage_worker) { ProjectUpdateRepositoryStorageWorker }
   end
 end

spec/requests/api/project_repository_storage_moves_spec.rb

@@ -7,7 +7,7 @@ RSpec.describe API::ProjectRepositoryStorageMoves do
 
   let_it_be(:user) { create(:admin) }
   let_it_be(:project) { create(:project, :repository).tap { |project| project.track_project_repository } }
-  let_it_be(:storage_move) { create(:project_repository_storage_move, :scheduled, project: project) }
+  let_it_be(:storage_move) { create(:project_repository_storage_move, :scheduled, container: project) }
 
   shared_examples 'get single project repository storage move' do
     let(:project_repository_storage_move_id) { storage_move.id }
@@ -63,14 +63,14 @@ RSpec.describe API::ProjectRepositoryStorageMoves do
 
       control = ActiveRecord::QueryRecorder.new { get_project_repository_storage_moves }
 
-      create(:project_repository_storage_move, :scheduled, project: project)
+      create(:project_repository_storage_move, :scheduled, container: project)
 
       expect { get_project_repository_storage_moves }.not_to exceed_query_limit(control)
     end
 
     it 'returns the most recently created first' do
-      storage_move_oldest = create(:project_repository_storage_move, :scheduled, project: project, created_at: 2.days.ago)
-      storage_move_middle = create(:project_repository_storage_move, :scheduled, project: project, created_at: 1.day.ago)
+      storage_move_oldest = create(:project_repository_storage_move, :scheduled, container: project, created_at: 2.days.ago)
+      storage_move_middle = create(:project_repository_storage_move, :scheduled, container: project, created_at: 1.day.ago)
 
       get_project_repository_storage_moves
 

spec/services/projects/fork_service_spec.rb

@@ -325,7 +325,7 @@ RSpec.describe Projects::ForkService do
       storage_move = create(
         :project_repository_storage_move,
         :scheduled,
-        project: project,
+        container: project,
         destination_storage_name: 'test_second_storage'
       )
       Projects::UpdateRepositoryStorageService.new(storage_move).execute

spec/services/projects/update_repository_storage_service_spec.rb

@@ -18,7 +18,7 @@ RSpec.describe Projects::UpdateRepositoryStorageService do
     context 'without wiki and design repository' do
       let(:project) { create(:project, :repository, wiki_enabled: false) }
       let(:destination) { 'test_second_storage' }
-      let(:repository_storage_move) { create(:project_repository_storage_move, :scheduled, project: project, destination_storage_name: destination) }
+      let(:repository_storage_move) { create(:project_repository_storage_move, :scheduled, container: project, destination_storage_name: destination) }
       let!(:checksum) { project.repository.checksum }
       let(:project_repository_double) { double(:repository) }
       let(:original_project_repository_double) { double(:repository) }
@@ -144,7 +144,7 @@ RSpec.describe Projects::UpdateRepositoryStorageService do
       end
 
       context 'when the repository move is finished' do
-        let(:repository_storage_move) { create(:project_repository_storage_move, :finished, project: project, destination_storage_name: destination) }
+        let(:repository_storage_move) { create(:project_repository_storage_move, :finished, container: project, destination_storage_name: destination) }
 
         it 'is idempotent' do
           expect do
@@ -156,7 +156,7 @@ RSpec.describe Projects::UpdateRepositoryStorageService do
       end
 
       context 'when the repository move is failed' do
-        let(:repository_storage_move) { create(:project_repository_storage_move, :failed, project: project, destination_storage_name: destination) }
+        let(:repository_storage_move) { create(:project_repository_storage_move, :failed, container: project, destination_storage_name: destination) }
 
         it 'is idempotent' do
           expect do
@@ -170,7 +170,7 @@ RSpec.describe Projects::UpdateRepositoryStorageService do
 
     context 'project with no repositories' do
       let(:project) { create(:project) }
-      let(:repository_storage_move) { create(:project_repository_storage_move, :scheduled, project: project, destination_storage_name: 'test_second_storage') }
+      let(:repository_storage_move) { create(:project_repository_storage_move, :scheduled, container: project, destination_storage_name: 'test_second_storage') }
 
       it 'updates the database' do
         allow(Gitlab::GitalyClient).to receive(:filesystem_id).with('default').and_call_original
@@ -191,7 +191,7 @@ RSpec.describe Projects::UpdateRepositoryStorageService do
         let(:project) { create(:project, :repository, wiki_enabled: true) }
         let(:repository) { project.wiki.repository }
         let(:destination) { 'test_second_storage' }
-        let(:repository_storage_move) { create(:project_repository_storage_move, :scheduled, project: project, destination_storage_name: destination) }
+        let(:repository_storage_move) { create(:project_repository_storage_move, :scheduled, container: project, destination_storage_name: destination) }
 
         before do
           project.create_wiki
@@ -204,7 +204,7 @@ RSpec.describe Projects::UpdateRepositoryStorageService do
         let(:project) { create(:project, :repository) }
         let(:repository) { project.design_repository }
         let(:destination) { 'test_second_storage' }
-        let(:repository_storage_move) { create(:project_repository_storage_move, :scheduled, project: project, destination_storage_name: destination) }
+        let(:repository_storage_move) { create(:project_repository_storage_move, :scheduled, container: project, destination_storage_name: destination) }
 
         before do
           project.design_repository.create_if_not_exists

spec/support/shared_examples/models/concerns/repository_storage_movable_shared_examples.rb

+# frozen_string_literal: true
+
+RSpec.shared_examples 'handles repository moves' do
+  describe 'associations' do
+    it { is_expected.to belong_to(:container) }
+  end
+
+  describe 'validations' do
+    it { is_expected.to validate_presence_of(:container) }
+    it { is_expected.to validate_presence_of(:state) }
+    it { is_expected.to validate_presence_of(:source_storage_name) }
+    it { is_expected.to validate_presence_of(:destination_storage_name) }
+
+    context 'source_storage_name inclusion' do
+      subject { build(repository_storage_factory_key, source_storage_name: 'missing') }
+
+      it "does not allow repository storages that don't match a label in the configuration" do
+        expect(subject).not_to be_valid
+        expect(subject.errors[:source_storage_name].first).to match(/is not included in the list/)
+      end
+    end
+
+    context 'destination_storage_name inclusion' do
+      subject { build(repository_storage_factory_key, destination_storage_name: 'missing') }
+
+      it "does not allow repository storages that don't match a label in the configuration" do
+        expect(subject).not_to be_valid
+        expect(subject.errors[:destination_storage_name].first).to match(/is not included in the list/)
+      end
+    end
+
+    context 'container repository read-only' do
+      subject { build(repository_storage_factory_key, container: container) }
+
+      it "does not allow the container to be read-only on create" do
+        container.update!(repository_read_only: true)
+
+        expect(subject).not_to be_valid
+        expect(subject.errors[error_key].first).to match(/is read only/)
+      end
+    end
+  end
+
+  describe 'defaults' do
+    context 'destination_storage_name' do
+      subject { build(repository_storage_factory_key) }
+
+      it 'picks storage from ApplicationSetting' do
+        expect(Gitlab::CurrentSettings).to receive(:pick_repository_storage).and_return('picked').at_least(:once)
+
+        expect(subject.destination_storage_name).to eq('picked')
+      end
+    end
+  end
+
+  describe 'state transitions' do
+    before do
+      stub_storage_settings('test_second_storage' => { 'path' => 'tmp/tests/extra_storage' })
+    end
+
+    context 'when in the default state' do
+      subject(:storage_move) { create(repository_storage_factory_key, container: container, destination_storage_name: 'test_second_storage') }
+
+      context 'and transits to scheduled' do
+        it 'triggers the corresponding repository storage worker' do
+          expect(repository_storage_worker).to receive(:perform_async).with(container.id, 'test_second_storage', storage_move.id)
+
+          storage_move.schedule!
+
+          expect(container).to be_repository_read_only
+        end
+
+        context 'when the transition fails' do
+          it 'does not trigger ProjectUpdateRepositoryStorageWorker and adds an error' do
+            allow(storage_move.container).to receive(:set_repository_read_only!).and_raise(StandardError, 'foobar')
+            expect(repository_storage_worker).not_to receive(:perform_async)
+
+            storage_move.schedule!
+
+            expect(storage_move.errors[error_key]).to include('foobar')
+          end
+        end
+      end
+
+      context 'and transits to started' do
+        it 'does not allow the transition' do
+          expect { storage_move.start! }
+            .to raise_error(StateMachines::InvalidTransition)
+        end
+      end
+    end
+
+    context 'when started' do
+      subject(:storage_move) { create(repository_storage_factory_key, :started, container: container, destination_storage_name: 'test_second_storage') }
+
+      context 'and transits to replicated' do
+        it 'sets the repository storage and marks the container as writable' do
+          storage_move.finish_replication!
+
+          expect(container.repository_storage).to eq('test_second_storage')
+          expect(container).not_to be_repository_read_only
+        end
+      end
+
+      context 'and transits to failed' do
+        it 'marks the container as writable' do
+          storage_move.do_fail!
+
+          expect(container).not_to be_repository_read_only
+        end
+      end
+    end
+  end
+end