Add method to create repositories

3 files changed, 97 insertions, 0 deletions

diff --git a/internal/praefect/datastore/repository_store.go b/internal/praefect/datastore/repository_store.go
index 288a4655d..254492f23 100644
--- a/internal/praefect/datastore/repository_store.go
+++ b/internal/praefect/datastore/repository_store.go
@@ -64,6 +64,26 @@ func (err RepositoryNotExistsError) Error() string {
 	)
 }
 
+// RepositoryExistsError is returned when trying to create a repository that already exists.
+type RepositoryExistsError struct {
+	virtualStorage string
+	relativePath   string
+	storage        string
+}
+
+// Is checks whether the other errors is of the same type.
+func (err RepositoryExistsError) Is(other error) bool {
+	_, ok := other.(RepositoryExistsError)
+	return ok
+}
+
+// Error returns the errors message.
+func (err RepositoryExistsError) Error() string {
+	return fmt.Sprintf("repository %q -> %q -> %q already exists",
+		err.virtualStorage, err.relativePath, err.storage,
+	)
+}
+
 // RepositoryStore provides access to repository state.
 type RepositoryStore interface {
 	// GetGeneration gets the repository's generation on a given storage.
@@ -76,6 +96,9 @@ 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
 	// 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.
@@ -275,6 +298,38 @@ AND storage = ANY($3)
 	return sourceGeneration, nil
 }
 
+func (rs *PostgresRepositoryStore) CreateRepository(ctx context.Context, virtualStorage, relativePath, storage string) error {
+	const q = `
+WITH repo AS (
+	INSERT INTO repositories (
+		virtual_storage,
+		relative_path,
+		generation
+	) VALUES ($1, $2, 0)
+)
+INSERT INTO storage_repositories (
+	virtual_storage,
+	relative_path,
+	storage,
+	generation
+)
+VALUES ($1, $2, $3, 0)
+`
+
+	_, err := rs.db.ExecContext(ctx, q, virtualStorage, relativePath, storage)
+
+	var pqerr *pq.Error
+	if errors.As(err, &pqerr) && pqerr.Code.Name() == "unique_violation" {
+		return RepositoryExistsError{
+			virtualStorage: virtualStorage,
+			relativePath:   relativePath,
+			storage:        storage,
+		}
+	}
+
+	return err
+}
+
 func (rs *PostgresRepositoryStore) DeleteRepository(ctx context.Context, virtualStorage, relativePath, storage string) error {
 	const q = `
 WITH repo AS (
diff --git a/internal/praefect/datastore/repository_store_mock.go b/internal/praefect/datastore/repository_store_mock.go
index d598a1692..faeaa3bfb 100644
--- a/internal/praefect/datastore/repository_store_mock.go
+++ b/internal/praefect/datastore/repository_store_mock.go
@@ -10,6 +10,7 @@ type MockRepositoryStore struct {
 	IsLatestGenerationFunc                 func(ctx context.Context, virtualStorage, relativePath, storage string) (bool, 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
 	DeleteRepositoryFunc                   func(ctx context.Context, virtualStorage, relativePath, storage string) error
 	RenameRepositoryFunc                   func(ctx context.Context, virtualStorage, relativePath, storage, newRelativePath string) error
 	GetConsistentSecondariesFunc           func(ctx context.Context, virtualStorage, relativePath, primary string) (map[string]struct{}, error)
@@ -58,6 +59,14 @@ func (m MockRepositoryStore) SetGeneration(ctx context.Context, virtualStorage,
 	return m.SetGenerationFunc(ctx, virtualStorage, relativePath, storage, generation)
 }
 
+func (m MockRepositoryStore) CreateRepository(ctx context.Context, virtualStorage, relativePath, storage string) error {
+	if m.CreateRepositoryFunc == nil {
+		return nil
+	}
+
+	return m.CreateRepositoryFunc(ctx, virtualStorage, relativePath, storage)
+}
+
 func (m MockRepositoryStore) DeleteRepository(ctx context.Context, virtualStorage, relativePath, storage string) error {
 	if m.DeleteRepositoryFunc == nil {
 		return nil
diff --git a/internal/praefect/datastore/repository_store_test.go b/internal/praefect/datastore/repository_store_test.go
index bd72c4729..a1543263e 100644
--- a/internal/praefect/datastore/repository_store_test.go
+++ b/internal/praefect/datastore/repository_store_test.go
@@ -285,6 +285,39 @@ 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: struct{}{},
+					},
+				},
+				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.Equal(t,
+				RepositoryExistsError{vs, repo, stor},
+				rs.CreateRepository(ctx, vs, repo, stor),
+			)
+		})
+	})
+
 	t.Run("DeleteRepository", func(t *testing.T) {
 		t.Run("delete non-existing", func(t *testing.T) {
 			rs, _ := newStore(t, nil)