Use repository ID in GetGeneration

To prepare for Praefect generating disk paths for replicas, this commit
changes GetGeneration to identify the repository by its ID rather
than the virtual storage and relative path.

5 files changed, 13 insertions, 14 deletions

diff --git a/cmd/praefect/subcmd_accept_dataloss_test.go b/cmd/praefect/subcmd_accept_dataloss_test.go
index 2945b43f1..a3b9940df 100644
--- a/cmd/praefect/subcmd_accept_dataloss_test.go
+++ b/cmd/praefect/subcmd_accept_dataloss_test.go
@@ -145,7 +145,7 @@ func TestAcceptDatalossSubcommand(t *testing.T) {
 			require.NoError(t, fs.Parse(tc.args))
 			tc.matchError(t, cmd.Exec(fs, conf))
 			for storage, expected := range tc.expectedGenerations {
-				actual, err := rs.GetGeneration(ctx, vs, repo, storage)
+				actual, err := rs.GetGeneration(ctx, 1, storage)
 				require.NoError(t, err)
 				require.Equal(t, expected, actual, storage)
 			}
diff --git a/internal/praefect/coordinator_pg_test.go b/internal/praefect/coordinator_pg_test.go
index 14c281229..1d81b1989 100644
--- a/internal/praefect/coordinator_pg_test.go
+++ b/internal/praefect/coordinator_pg_test.go
@@ -312,7 +312,7 @@ func TestStreamDirectorMutator_Transaction(t *testing.T) {
 			// Nodes that did not successfully commit or did not participate should remain on their
 			// existing generation.
 			for i, n := range tc.nodes {
-				gen, err := rs.GetGeneration(ctx, repo.StorageName, repo.RelativePath, storageNodes[i].Storage)
+				gen, err := rs.GetGeneration(ctx, 1, storageNodes[i].Storage)
 				require.NoError(t, err)
 				require.Equal(t, n.expectedGeneration, gen, "node %d has wrong generation", i)
 			}
diff --git a/internal/praefect/datastore/repository_store.go b/internal/praefect/datastore/repository_store.go
index 765e0d6fe..834be79f3 100644
--- a/internal/praefect/datastore/repository_store.go
+++ b/internal/praefect/datastore/repository_store.go
@@ -83,7 +83,7 @@ var ErrNoRowsAffected = errors.New("no rows were affected by the query")
 // RepositoryStore provides access to repository state.
 type RepositoryStore interface {
 	// GetGeneration gets the repository's generation on a given storage.
-	GetGeneration(ctx context.Context, virtualStorage, relativePath, storage string) (int, error)
+	GetGeneration(ctx context.Context, repositoryID int64, storage string) (int, error)
 	// IncrementGeneration increments the generations of up to date nodes.
 	IncrementGeneration(ctx context.Context, repositoryID int64, primary string, secondaries []string) error
 	// SetGeneration sets the repository's generation on the given storage. If the generation is higher
@@ -145,17 +145,16 @@ func NewPostgresRepositoryStore(db glsql.Querier, configuredStorages map[string]
 	return &PostgresRepositoryStore{db: db, storages: storages(configuredStorages)}
 }
 
-func (rs *PostgresRepositoryStore) GetGeneration(ctx context.Context, virtualStorage, relativePath, storage string) (int, error) {
+func (rs *PostgresRepositoryStore) GetGeneration(ctx context.Context, repositoryID int64, storage string) (int, error) {
 	const q = `
 SELECT generation
 FROM storage_repositories
-WHERE virtual_storage = $1
-AND relative_path = $2
-AND storage = $3
+WHERE repository_id = $1
+AND storage = $2
 `
 
 	var gen int
-	if err := rs.db.QueryRowContext(ctx, q, virtualStorage, relativePath, storage).Scan(&gen); err != nil {
+	if err := rs.db.QueryRowContext(ctx, q, repositoryID, storage).Scan(&gen); err != nil {
 		if errors.Is(err, sql.ErrNoRows) {
 			return GenerationUnknown, nil
 		}
diff --git a/internal/praefect/datastore/repository_store_mock.go b/internal/praefect/datastore/repository_store_mock.go
index b9c9f674d..6a6975677 100644
--- a/internal/praefect/datastore/repository_store_mock.go
+++ b/internal/praefect/datastore/repository_store_mock.go
@@ -5,7 +5,7 @@ import "context"
 // MockRepositoryStore allows for mocking a RepositoryStore by parametrizing its behavior. All methods
 // default to what could be considered success if not set.
 type MockRepositoryStore struct {
-	GetGenerationFunc                     func(ctx context.Context, virtualStorage, relativePath, storage string) (int, error)
+	GetGenerationFunc                     func(ctx context.Context, repositoryID int64, storage string) (int, error)
 	IncrementGenerationFunc               func(ctx context.Context, repositoryID int64, primary string, secondaries []string) error
 	GetReplicatedGenerationFunc           func(ctx context.Context, repositoryID int64, source, target string) (int, error)
 	SetGenerationFunc                     func(ctx context.Context, virtualStorage, relativePath, storage string, generation int) error
@@ -22,12 +22,12 @@ type MockRepositoryStore struct {
 	GetRepositoryIDFunc                   func(ctx context.Context, virtualStorage, relativePath string) (int64, error)
 }
 
-func (m MockRepositoryStore) GetGeneration(ctx context.Context, virtualStorage, relativePath, storage string) (int, error) {
+func (m MockRepositoryStore) GetGeneration(ctx context.Context, repositoryID int64, storage string) (int, error) {
 	if m.GetGenerationFunc == nil {
 		return GenerationUnknown, nil
 	}
 
-	return m.GetGenerationFunc(ctx, virtualStorage, relativePath, storage)
+	return m.GetGenerationFunc(ctx, repositoryID, storage)
 }
 
 func (m MockRepositoryStore) IncrementGeneration(ctx context.Context, repositoryID int64, primary string, secondaries []string) error {
diff --git a/internal/praefect/datastore/repository_store_test.go b/internal/praefect/datastore/repository_store_test.go
index d635737ea..2368a5cea 100644
--- a/internal/praefect/datastore/repository_store_test.go
+++ b/internal/praefect/datastore/repository_store_test.go
@@ -526,13 +526,13 @@ func testRepositoryStore(t *testing.T, newStore repositoryStoreFactory) {
 	t.Run("GetGeneration", func(t *testing.T) {
 		rs, _ := newStore(t, nil)
 
-		generation, err := rs.GetGeneration(ctx, vs, repo, stor)
+		generation, err := rs.GetGeneration(ctx, 1, stor)
 		require.NoError(t, err)
 		require.Equal(t, GenerationUnknown, generation)
 
-		require.NoError(t, rs.SetGeneration(ctx, vs, repo, stor, 0))
+		require.NoError(t, rs.CreateRepository(ctx, 1, vs, repo, stor, nil, nil, false, false))
 
-		generation, err = rs.GetGeneration(ctx, vs, repo, stor)
+		generation, err = rs.GetGeneration(ctx, 1, stor)
 		require.NoError(t, err)
 		require.Equal(t, 0, generation)
 	})