Link created assignments via the repository ID to the repository

The assignments are already being linked with the repository ID when
they are being created through CreateRepository as part of repository
creation. Assignment records can also be created when SetReplicationFactor
is called. This commit updates SetReplicationFactor to link the created
assignments records via the repository ID so we can later do a one swoop
migration to link all existing records to their repositories via the id.

2 files changed, 14 insertions, 5 deletions

diff --git a/internal/praefect/datastore/assignment.go b/internal/praefect/datastore/assignment.go
index b097272c4..7fc48eedc 100644
--- a/internal/praefect/datastore/assignment.go
+++ b/internal/praefect/datastore/assignment.go
@@ -123,7 +123,7 @@ func (s AssignmentStore) SetReplicationFactor(ctx context.Context, virtualStorag
 	//    current assignments.
 	rows, err := s.db.QueryContext(ctx, `
 WITH repository AS (
-	SELECT virtual_storage, relative_path, "primary"
+	SELECT repository_id, virtual_storage, relative_path, "primary"
 	FROM repositories
 	WHERE virtual_storage = $1
 	AND   relative_path   = $2
@@ -139,7 +139,7 @@ existing_assignments AS (
 
 created_assignments AS (
 	INSERT INTO repository_assignments
-	SELECT virtual_storage, relative_path, storage
+	SELECT virtual_storage, relative_path, storage, repository_id
 	FROM repository
 	CROSS JOIN ( SELECT unnest($4::text[]) AS storage ) AS configured_storages
 	WHERE storage NOT IN ( SELECT storage FROM existing_assignments )
diff --git a/internal/praefect/datastore/assignment_test.go b/internal/praefect/datastore/assignment_test.go
index c8ac5ad84..40dcb1c59 100644
--- a/internal/praefect/datastore/assignment_test.go
+++ b/internal/praefect/datastore/assignment_test.go
@@ -3,6 +3,7 @@ package datastore
 import (
 	"testing"
 
+	"github.com/lib/pq"
 	"github.com/stretchr/testify/require"
 	"gitlab.com/gitlab-org/gitaly/v14/internal/praefect/datastore/glsql"
 	"gitlab.com/gitlab-org/gitaly/v14/internal/testhelper"
@@ -204,15 +205,15 @@ func TestAssignmentStore_SetReplicationFactor(t *testing.T) {
 
 			if !tc.nonExistentRepository {
 				_, err := db.ExecContext(ctx, `
-					INSERT INTO repositories (virtual_storage, relative_path, "primary")
-					VALUES ('virtual-storage', 'relative-path', 'primary')
+					INSERT INTO repositories (virtual_storage, relative_path, "primary", repository_id)
+					VALUES ('virtual-storage', 'relative-path', 'primary', 1)
 				`)
 				require.NoError(t, err)
 			}
 
 			for _, storage := range tc.existingAssignments {
 				_, err := db.ExecContext(ctx, `
-					INSERT INTO repository_assignments VALUES ('virtual-storage', 'relative-path', $1)
+					INSERT INTO repository_assignments VALUES ('virtual-storage', 'relative-path', $1, 1)
 				`, storage)
 				require.NoError(t, err)
 			}
@@ -230,6 +231,14 @@ func TestAssignmentStore_SetReplicationFactor(t *testing.T) {
 			assignedStorages, err := store.GetHostAssignments(ctx, "virtual-storage", "relative-path")
 			require.NoError(t, err)
 			tc.requireStorages(t, assignedStorages)
+
+			var storagesWithIncorrectRepositoryID pq.StringArray
+			require.NoError(t, db.QueryRowContext(ctx, `
+				SELECT array_agg(storage)
+				FROM repository_assignments
+				WHERE COALESCE(repository_id != 1, true)
+			`).Scan(&storagesWithIncorrectRepositoryID))
+			require.Empty(t, storagesWithIncorrectRepositoryID)
 		})
 	}
 }