Link repository id to replicas created from replication jobs

SetGeneration is called to set a replica's generation following a replication
job. As it might create new replica records, it should ensure the record is
linked to the correct repository via the repository ID. This commit links the
replica's record by getting the correct repository ID from the repositories table.
As all replicas still have the same (virtual_storage, relative_path) as the
repository, this is still safe to do. It's also necessary as not every replication
job yet carries a repository ID. This commit also includes the repository ID in
replication jobs so the approach can later be updated to simply take the repository
ID from the replication job.

4 files changed, 13 insertions, 4 deletions

diff --git a/internal/praefect/coordinator.go b/internal/praefect/coordinator.go
index 0d837224d..218d5fbb3 100644
--- a/internal/praefect/coordinator.go
+++ b/internal/praefect/coordinator.go
@@ -1018,6 +1018,7 @@ func (c *Coordinator) newRequestFinalizer(
 		for _, secondary := range outdatedSecondaries {
 			event := datastore.ReplicationEvent{
 				Job: datastore.ReplicationJob{
+					RepositoryID:      repositoryID,
 					Change:            change,
 					RelativePath:      targetRepo.GetRelativePath(),
 					VirtualStorage:    virtualStorage,
diff --git a/internal/praefect/datastore/queue.go b/internal/praefect/datastore/queue.go
index dcff2567e..72d88f970 100644
--- a/internal/praefect/datastore/queue.go
+++ b/internal/praefect/datastore/queue.go
@@ -49,6 +49,10 @@ func allowToAck(state JobState) error {
 
 // ReplicationJob is a persistent representation of the replication job.
 type ReplicationJob struct {
+	// RepositoryID is the ID of the repository this job relates to. RepositoryID
+	// may be 0 if the job doesn't relate to any known repository. This can happen
+	// for example when the job is deleting an orphaned replica of a deleted repository.
+	RepositoryID      int64      `json:"repository_id"`
 	Change            ChangeType `json:"change"`
 	RelativePath      string     `json:"relative_path"`
 	TargetNodeStorage string     `json:"target_node_storage"`
diff --git a/internal/praefect/datastore/repository_store.go b/internal/praefect/datastore/repository_store.go
index 0b0650991..83571fa7e 100644
--- a/internal/praefect/datastore/repository_store.go
+++ b/internal/praefect/datastore/repository_store.go
@@ -235,13 +235,17 @@ WITH repository AS (
 )
 
 INSERT INTO storage_repositories (
+	repository_id,
 	virtual_storage,
 	relative_path,
 	storage,
 	generation
 )
-VALUES ($1, $2, $3, $4)
+SELECT
+	(SELECT repository_id FROM repositories WHERE virtual_storage = $1 AND relative_path = $2),
+	$1, $2, $3, $4
 ON CONFLICT (virtual_storage, relative_path, storage) DO UPDATE SET
+	repository_id = EXCLUDED.repository_id,
 	generation = EXCLUDED.generation
 `
 
diff --git a/internal/praefect/datastore/repository_store_test.go b/internal/praefect/datastore/repository_store_test.go
index 912241112..3940e26ff 100644
--- a/internal/praefect/datastore/repository_store_test.go
+++ b/internal/praefect/datastore/repository_store_test.go
@@ -917,7 +917,7 @@ func testRepositoryStore(t *testing.T, newStore repositoryStoreFactory) {
 					"repository-1": {
 						"primary":                {repositoryID: 1, generation: 1},
 						"consistent-secondary":   {repositoryID: 1, generation: 1},
-						"inconsistent-secondary": {generation: 0},
+						"inconsistent-secondary": {repositoryID: 1, generation: 0},
 					},
 				},
 			},
@@ -949,10 +949,10 @@ func testRepositoryStore(t *testing.T, newStore repositoryStoreFactory) {
 				storageState{
 					"virtual-storage-1": {
 						"repository-1": {
-							"unknown":                {generation: 2},
+							"unknown":                {repositoryID: 1, generation: 2},
 							"primary":                {repositoryID: 1, generation: 1},
 							"consistent-secondary":   {repositoryID: 1, generation: 1},
-							"inconsistent-secondary": {generation: 0},
+							"inconsistent-secondary": {repositoryID: 1, generation: 0},
 						},
 					},
 				},