add option for storing primary when creating a repository

Repository creation when using repository specific primaries does not
work as the creation goes through the usual repository mutator route
which fails due to the repository not having a primary. As a preliminary
step towards fixing that, this commit adds an option that can be used
to store a repository's primary node in the database when creating its
record.

6 files changed, 57 insertions, 34 deletions

diff --git a/internal/praefect/coordinator.go b/internal/praefect/coordinator.go
index 63a894a9e..5bd99ada7 100644
--- a/internal/praefect/coordinator.go
+++ b/internal/praefect/coordinator.go
@@ -754,7 +754,7 @@ func (c *Coordinator) newRequestFinalizer(
 				ctxlogrus.Extract(ctx).WithError(err).Info("deleted repository does not have a store entry")
 			}
 		case datastore.CreateRepo:
-			if err := c.rs.CreateRepository(ctx, virtualStorage, targetRepo.GetRelativePath(), primary); err != nil {
+			if err := c.rs.CreateRepository(ctx, virtualStorage, targetRepo.GetRelativePath(), primary, false); err != nil {
 				if !errors.Is(err, datastore.RepositoryExistsError{}) {
 					return fmt.Errorf("create repository: %w", err)
 				}
diff --git a/internal/praefect/coordinator_test.go b/internal/praefect/coordinator_test.go
index f4057e3c3..28cec30a1 100644
--- a/internal/praefect/coordinator_test.go
+++ b/internal/praefect/coordinator_test.go
@@ -638,11 +638,12 @@ func TestStreamDirector_repo_creation(t *testing.T) {
 
 	var createRepositoryCalled int64
 	rs := datastore.MockRepositoryStore{
-		CreateRepositoryFunc: func(ctx context.Context, virtualStorage, relativePath, storage string) error {
+		CreateRepositoryFunc: func(ctx context.Context, virtualStorage, relativePath, primary string, storePrimary bool) error {
 			atomic.AddInt64(&createRepositoryCalled, 1)
 			assert.Equal(t, targetRepo.StorageName, virtualStorage)
 			assert.Equal(t, targetRepo.RelativePath, relativePath)
-			assert.Equal(t, rewrittenStorage, storage)
+			assert.Equal(t, rewrittenStorage, primary)
+			assert.False(t, storePrimary)
 			return nil
 		},
 	}
diff --git a/internal/praefect/datastore/repository_store.go b/internal/praefect/datastore/repository_store.go
index d3cf1e514..2cf36be9a 100644
--- a/internal/praefect/datastore/repository_store.go
+++ b/internal/praefect/datastore/repository_store.go
@@ -97,9 +97,13 @@ type RepositoryStore interface {
 	// GetReplicatedGeneration returns the generation propagated by applying the replication. If the generation would
 	// downgrade, a DowngradeAttemptedError is returned.
 	GetReplicatedGeneration(ctx context.Context, virtualStorage, relativePath, source, target string) (int, error)
-	// CreateRepository creates the repository for the virtual storage and the storage. Returns
-	// RepositoryExistsError when trying to create a repository which already has a matching record.
-	CreateRepository(ctx context.Context, virtualStorage, relativePath, storage string) error
+	// CreateRepository creates a record for a repository in the specified virtual storage and relative path.
+	// Primary is the storage the repository was created on. Returns RepositoryExistsError when trying to create
+	// a repository which already exists in the store.
+	//
+	// storePrimary should be set when repository specific primaries are enabled. When set, the primary is stored as
+	// the repository's primary.
+	CreateRepository(ctx context.Context, virtualStorage, relativePath, primary string, storePrimary bool) error
 	// DeleteRepository deletes the repository from the virtual storage and the storage. Returns
 	// RepositoryNotExistsError when trying to delete a repository which has no record in the virtual storage
 	// or the storage.
@@ -297,14 +301,15 @@ AND storage = ANY($3)
 
 //nolint:stylecheck
 //nolint:golint
-func (rs *PostgresRepositoryStore) CreateRepository(ctx context.Context, virtualStorage, relativePath, storage string) error {
+func (rs *PostgresRepositoryStore) CreateRepository(ctx context.Context, virtualStorage, relativePath, primary string, storePrimary bool) error {
 	const q = `
 WITH repo AS (
 	INSERT INTO repositories (
 		virtual_storage,
 		relative_path,
-		generation
-	) VALUES ($1, $2, 0)
+		generation,
+		"primary"
+	) VALUES ($1, $2, 0, CASE WHEN $4 THEN $3 END)
 )
 INSERT INTO storage_repositories (
 	virtual_storage,
@@ -315,14 +320,14 @@ INSERT INTO storage_repositories (
 VALUES ($1, $2, $3, 0)
 `
 
-	_, err := rs.db.ExecContext(ctx, q, virtualStorage, relativePath, storage)
+	_, err := rs.db.ExecContext(ctx, q, virtualStorage, relativePath, primary, storePrimary)
 
 	var pqerr *pq.Error
 	if errors.As(err, &pqerr) && pqerr.Code.Name() == "unique_violation" {
 		return RepositoryExistsError{
 			virtualStorage: virtualStorage,
 			relativePath:   relativePath,
-			storage:        storage,
+			storage:        primary,
 		}
 	}
 
diff --git a/internal/praefect/datastore/repository_store_mock.go b/internal/praefect/datastore/repository_store_mock.go
index 5809be6ee..ea051d1d0 100644
--- a/internal/praefect/datastore/repository_store_mock.go
+++ b/internal/praefect/datastore/repository_store_mock.go
@@ -9,7 +9,7 @@ type MockRepositoryStore struct {
 	IncrementGenerationFunc                func(ctx context.Context, virtualStorage, relativePath, primary string, secondaries []string) error
 	GetReplicatedGenerationFunc            func(ctx context.Context, virtualStorage, relativePath, source, target string) (int, error)
 	SetGenerationFunc                      func(ctx context.Context, virtualStorage, relativePath, storage string, generation int) error
-	CreateRepositoryFunc                   func(ctx context.Context, virtualStorage, relativePath, storage string) error
+	CreateRepositoryFunc                   func(ctx context.Context, virtualStorage, relativePath, storage string, storePrimary bool) error
 	DeleteRepositoryFunc                   func(ctx context.Context, virtualStorage, relativePath, storage string) error
 	RenameRepositoryFunc                   func(ctx context.Context, virtualStorage, relativePath, storage, newRelativePath string) error
 	GetConsistentStoragesFunc              func(ctx context.Context, virtualStorage, relativePath string) (map[string]struct{}, error)
@@ -52,12 +52,12 @@ func (m MockRepositoryStore) SetGeneration(ctx context.Context, virtualStorage,
 
 //nolint:stylecheck
 //nolint:golint
-func (m MockRepositoryStore) CreateRepository(ctx context.Context, virtualStorage, relativePath, storage string) error {
+func (m MockRepositoryStore) CreateRepository(ctx context.Context, virtualStorage, relativePath, storage string, storePrimary bool) error {
 	if m.CreateRepositoryFunc == nil {
 		return nil
 	}
 
-	return m.CreateRepositoryFunc(ctx, virtualStorage, relativePath, storage)
+	return m.CreateRepositoryFunc(ctx, virtualStorage, relativePath, storage, storePrimary)
 }
 
 func (m MockRepositoryStore) DeleteRepository(ctx context.Context, virtualStorage, relativePath, storage string) error {
diff --git a/internal/praefect/datastore/repository_store_test.go b/internal/praefect/datastore/repository_store_test.go
index 68751fc81..efcfbe41f 100644
--- a/internal/praefect/datastore/repository_store_test.go
+++ b/internal/praefect/datastore/repository_store_test.go
@@ -297,34 +297,51 @@ func testRepositoryStore(t *testing.T, newStore repositoryStoreFactory) {
 	})
 
 	t.Run("CreateRepository", func(t *testing.T) {
-		t.Run("create", func(t *testing.T) {
-			rs, requireState := newStore(t, nil)
-
-			require.NoError(t, rs.CreateRepository(ctx, vs, repo, stor))
-
-			requireState(t, ctx,
-				virtualStorageState{
-					vs: {
-						repo: repositoryRecord{},
-					},
+		t.Run("successfully created", func(t *testing.T) {
+			for _, tc := range []struct {
+				desc         string
+				storePrimary bool
+				primary      string
+			}{
+				{
+					desc:         "primary not stored",
+					storePrimary: false,
 				},
-				storageState{
-					vs: {
-						repo: {
-							stor: 0,
-						},
-					},
+				{
+					desc:         "primary stored",
+					storePrimary: true,
+					primary:      stor,
 				},
-			)
+			} {
+				t.Run(tc.desc, func(t *testing.T) {
+					rs, requireState := newStore(t, nil)
+
+					require.NoError(t, rs.CreateRepository(ctx, vs, repo, stor, tc.storePrimary))
+					requireState(t, ctx,
+						virtualStorageState{
+							vs: {
+								repo: repositoryRecord{primary: tc.primary},
+							},
+						},
+						storageState{
+							vs: {
+								repo: {
+									stor: 0,
+								},
+							},
+						},
+					)
+				})
+			}
 		})
 
 		t.Run("conflict", func(t *testing.T) {
 			rs, _ := newStore(t, nil)
 
-			require.NoError(t, rs.CreateRepository(ctx, vs, repo, stor))
+			require.NoError(t, rs.CreateRepository(ctx, vs, repo, stor, false))
 			require.Equal(t,
 				RepositoryExistsError{vs, repo, stor},
-				rs.CreateRepository(ctx, vs, repo, stor),
+				rs.CreateRepository(ctx, vs, repo, stor, false),
 			)
 		})
 	})
diff --git a/internal/praefect/repository_exists_test.go b/internal/praefect/repository_exists_test.go
index f0bb53a76..1179354c0 100644
--- a/internal/praefect/repository_exists_test.go
+++ b/internal/praefect/repository_exists_test.go
@@ -72,7 +72,7 @@ func TestRepositoryExistsStreamInterceptor(t *testing.T) {
 			ctx, cancel := testhelper.Context()
 			defer cancel()
 
-			require.NoError(t, rs.CreateRepository(ctx, "virtual-storage", "relative-path", "storage"))
+			require.NoError(t, rs.CreateRepository(ctx, "virtual-storage", "relative-path", "storage", false))
 
 			electionStrategy := config.ElectionStrategyPerRepository
 			if tc.routeToGitaly {