Fix bug tracking snippet shard name

app/models/snippet.rb

@@ -296,12 +296,12 @@ class Snippet < ApplicationRecord
     return if repository_exists? && snippet_repository
 
     repository.create_if_not_exists
-    track_snippet_repository
+    track_snippet_repository(repository.storage)
   end
 
-  def track_snippet_repository
-    repository = snippet_repository || build_snippet_repository
-    repository.update!(shard_name: repository_storage, disk_path: disk_path)
+  def track_snippet_repository(shard)
+    snippet_repo = snippet_repository || build_snippet_repository
+    snippet_repo.update!(shard_name: shard, disk_path: disk_path)
   end
 
   def can_cache_field?(field)

changelogs/unreleased/fj-212399-fix-bug-tracking-snippet-shard-name.yml

+---
+title: Fix bug tracking snippet shard name
+merge_request: 27979
+author:
+type: fixed

lib/gitlab/import_export/snippet_repo_restorer.rb

@@ -30,7 +30,7 @@ module Gitlab
 
       def create_repository_from_bundle
         repository.create_from_bundle(path_to_bundle)
-        snippet.track_snippet_repository
+        snippet.track_snippet_repository(repository.storage)
       end
 
       def create_repository_from_db

spec/factories/snippets.rb

@@ -28,7 +28,7 @@ FactoryBot.define do
           bare_repo: TestEnv.factory_repo_path_bare,
           refs: TestEnv::BRANCH_SHA)
 
-        snippet.track_snippet_repository
+        snippet.track_snippet_repository(snippet.repository.storage)
       end
     end
 

spec/lib/gitlab/import_export/snippet_repo_restorer_spec.rb

@@ -55,9 +55,11 @@ describe Gitlab::ImportExport::SnippetRepoRestorer do
     let(:snippet_bundle_path) { File.join(bundle_path, "#{snippet_with_repo.hexdigest}.bundle") }
     let(:result) { exporter.save }
 
-    it 'creates the repository from the bundle' do
+    before do
       expect(exporter.save).to be_truthy
+    end
 
+    it 'creates the repository from the bundle' do
       expect(snippet.repository_exists?).to be_falsey
       expect(snippet.snippet_repository).to be_nil
       expect(snippet.repository).to receive(:create_from_bundle).and_call_original
@@ -66,5 +68,14 @@ describe Gitlab::ImportExport::SnippetRepoRestorer do
       expect(snippet.repository_exists?).to be_truthy
       expect(snippet.snippet_repository).not_to be_nil
     end
+
+    it 'sets same shard in snippet repository as in the repository storage' do
+      expect(snippet).to receive(:repository_storage).and_return('picked')
+      expect(snippet.repository).to receive(:create_from_bundle)
+
+      restorer.restore
+
+      expect(snippet.snippet_repository.shard_name).to eq 'picked'
+    end
   end
 end

spec/models/snippet_repository_spec.rb

@@ -16,7 +16,7 @@ describe SnippetRepository do
   describe '.find_snippet' do
     it 'finds snippet by disk path' do
       snippet = create(:snippet, author: user)
-      snippet.track_snippet_repository
+      snippet.track_snippet_repository(snippet.repository.storage)
 
       expect(described_class.find_snippet(snippet.disk_path)).to eq(snippet)
     end

spec/models/snippet_spec.rb

@@ -567,18 +567,21 @@ describe Snippet do
 
   describe '#track_snippet_repository' do
     let(:snippet) { create(:snippet) }
+    let(:shard_name) { 'foo' }
+
+    subject { snippet.track_snippet_repository(shard_name) }
 
     context 'when a snippet repository entry does not exist' do
       it 'creates a new entry' do
-        expect { snippet.track_snippet_repository }.to change(snippet, :snippet_repository)
+        expect { subject }.to change(snippet, :snippet_repository)
       end
 
       it 'tracks the snippet storage location' do
-        snippet.track_snippet_repository
+        subject
 
         expect(snippet.snippet_repository).to have_attributes(
           disk_path: snippet.disk_path,
-          shard_name: snippet.repository_storage
+          shard_name: shard_name
         )
       end
     end
@@ -586,21 +589,20 @@ describe Snippet do
     context 'when a tracking entry exists' do
       let!(:snippet) { create(:snippet, :repository) }
       let(:snippet_repository) { snippet.snippet_repository }
-      let!(:shard) { create(:shard, name: 'foo') }
+      let(:shard_name) { 'bar' }
 
       it 'does not create a new entry in the database' do
-        expect { snippet.track_snippet_repository }.not_to change(snippet, :snippet_repository)
+        expect { subject }.not_to change(snippet, :snippet_repository)
       end
 
       it 'updates the snippet storage location' do
         allow(snippet).to receive(:disk_path).and_return('fancy/new/path')
-        allow(snippet).to receive(:repository_storage).and_return('foo')
 
-        snippet.track_snippet_repository
+        subject
 
         expect(snippet.snippet_repository).to have_attributes(
           disk_path: 'fancy/new/path',
-          shard_name: 'foo'
+          shard_name: shard_name
         )
       end
     end
@@ -609,19 +611,31 @@ describe Snippet do
   describe '#create_repository' do
     let(:snippet) { create(:snippet) }
 
+    subject { snippet.create_repository }
+
     it 'creates the repository' do
       expect(snippet.repository).to receive(:after_create).and_call_original
 
-      expect(snippet.create_repository).to be_truthy
+      expect(subject).to be_truthy
       expect(snippet.repository.exists?).to be_truthy
     end
 
     it 'tracks snippet repository' do
       expect do
-        snippet.create_repository
+        subject
       end.to change(SnippetRepository, :count).by(1)
     end
 
+    it 'sets same shard in snippet repository as in the repository storage' do
+      expect(snippet).to receive(:repository_storage).and_return('picked')
+      expect(snippet).to receive(:repository_exists?).and_return(false)
+      expect(snippet.repository).to receive(:create_if_not_exists)
+
+      subject
+
+      expect(snippet.snippet_repository.shard_name).to eq 'picked'
+    end
+
     context 'when repository exists' do
       let!(:snippet) { create(:snippet, :repository) }