Get host assignment by repository ID

Praefect's routing should use repository ID consistently in all of the
queries in order to avoid races where the external key of the repository
changes between the various queries. This commit updates GetHostAssignments
to get assignments by the repository's ID rather than its virtual storage
and relative path.

6 files changed, 19 insertions, 12 deletions

diff --git a/internal/praefect/assignment.go b/internal/praefect/assignment.go
index f0f38e5af..1aa6447ad 100644
--- a/internal/praefect/assignment.go
+++ b/internal/praefect/assignment.go
@@ -23,7 +23,7 @@ func NewDisabledAssignmentStore(storages map[string][]string) AssignmentStore {
 }
 
 // GetHostAssigments simply returns all of the storages configured for the virtual storage.
-func (storages disabledAssignments) GetHostAssignments(ctx context.Context, virtualStorage, relativePath string) ([]string, error) {
+func (storages disabledAssignments) GetHostAssignments(ctx context.Context, virtualStorage string, repositoryID int64) ([]string, error) {
 	return storages[virtualStorage], nil
 }
 
diff --git a/internal/praefect/datastore/assignment.go b/internal/praefect/datastore/assignment.go
index 7fc48eedc..b4524b568 100644
--- a/internal/praefect/datastore/assignment.go
+++ b/internal/praefect/datastore/assignment.go
@@ -38,7 +38,7 @@ func NewAssignmentStore(db glsql.Querier, configuredStorages map[string][]string
 	return AssignmentStore{db: db, configuredStorages: configuredStorages}
 }
 
-func (s AssignmentStore) GetHostAssignments(ctx context.Context, virtualStorage, relativePath string) ([]string, error) {
+func (s AssignmentStore) GetHostAssignments(ctx context.Context, virtualStorage string, repositoryID int64) ([]string, error) {
 	configuredStorages, ok := s.configuredStorages[virtualStorage]
 	if !ok {
 		return nil, newVirtualStorageNotFoundError(virtualStorage)
@@ -47,10 +47,9 @@ func (s AssignmentStore) GetHostAssignments(ctx context.Context, virtualStorage,
 	rows, err := s.db.QueryContext(ctx, `
 SELECT storage
 FROM repository_assignments
-WHERE virtual_storage = $1
-AND   relative_path = $2
-AND   storage = ANY($3)
-`, virtualStorage, relativePath, pq.StringArray(configuredStorages))
+WHERE repository_id = $1
+AND   storage = ANY($2)
+`, repositoryID, pq.StringArray(configuredStorages))
 	if err != nil {
 		return nil, fmt.Errorf("query: %w", err)
 	}
diff --git a/internal/praefect/datastore/assignment_test.go b/internal/praefect/datastore/assignment_test.go
index 67f83e7fc..3062ab125 100644
--- a/internal/praefect/datastore/assignment_test.go
+++ b/internal/praefect/datastore/assignment_test.go
@@ -2,6 +2,7 @@ package datastore
 
 import (
 	"errors"
+	"sort"
 	"testing"
 
 	"github.com/lib/pq"
@@ -97,10 +98,15 @@ func TestAssignmentStore_GetHostAssignments(t *testing.T) {
 				require.NoError(t, err)
 			}
 
+			repositoryID, err := rs.GetRepositoryID(ctx, tc.virtualStorage, "relative-path")
+			if err != nil {
+				require.Equal(t, commonerr.NewRepositoryNotFoundError(tc.virtualStorage, "relative-path"), err)
+			}
+
 			actualAssignments, err := NewAssignmentStore(
 				db,
 				map[string][]string{"virtual-storage": configuredStorages},
-			).GetHostAssignments(ctx, tc.virtualStorage, "relative-path")
+			).GetHostAssignments(ctx, tc.virtualStorage, repositoryID)
 			require.Equal(t, tc.error, err)
 			require.ElementsMatch(t, tc.expectedAssignments, actualAssignments)
 		})
@@ -233,8 +239,10 @@ func TestAssignmentStore_SetReplicationFactor(t *testing.T) {
 
 			tc.requireStorages(t, setStorages)
 
-			assignedStorages, err := store.GetHostAssignments(ctx, "virtual-storage", "relative-path")
+			assignedStorages, err := store.GetHostAssignments(ctx, "virtual-storage", 1)
 			require.NoError(t, err)
+
+			sort.Strings(assignedStorages)
 			tc.requireStorages(t, assignedStorages)
 
 			var storagesWithIncorrectRepositoryID pq.StringArray
diff --git a/internal/praefect/router_per_repository.go b/internal/praefect/router_per_repository.go
index 431406acb..e3455c3a9 100644
--- a/internal/praefect/router_per_repository.go
+++ b/internal/praefect/router_per_repository.go
@@ -25,7 +25,7 @@ var errPrimaryUnassigned = errors.New("primary node is not assigned")
 type AssignmentGetter interface {
 	// GetHostAssignments returns the names of the storages assigned to host the repository.
 	// The primary node must always be assigned.
-	GetHostAssignments(ctx context.Context, virtualStorage, relativePath string) ([]string, error)
+	GetHostAssignments(ctx context.Context, virtualStorage string, repositoryID int64) ([]string, error)
 }
 
 // ErrNoSuitableNode is returned when there is not suitable node to serve a request.
@@ -207,7 +207,7 @@ func (r *PerRepositoryRouter) RouteRepositoryMutator(ctx context.Context, virtua
 		return RepositoryMutatorRoute{}, ErrRepositoryReadOnly
 	}
 
-	assignedStorages, err := r.ag.GetHostAssignments(ctx, virtualStorage, relativePath)
+	assignedStorages, err := r.ag.GetHostAssignments(ctx, virtualStorage, repositoryID)
 	if err != nil {
 		return RepositoryMutatorRoute{}, fmt.Errorf("get host assignments: %w", err)
 	}
diff --git a/internal/praefect/service/info/repositories.go b/internal/praefect/service/info/repositories.go
index e701329d7..bfb5f0fe2 100644
--- a/internal/praefect/service/info/repositories.go
+++ b/internal/praefect/service/info/repositories.go
@@ -25,7 +25,7 @@ func (s *Server) RepositoryReplicas(ctx context.Context, in *gitalypb.Repository
 		return nil, fmt.Errorf("get primary: %w", err)
 	}
 
-	assignments, err := s.assignmentStore.GetHostAssignments(ctx, virtualStorage, relativePath)
+	assignments, err := s.assignmentStore.GetHostAssignments(ctx, virtualStorage, repositoryID)
 	if err != nil {
 		return nil, fmt.Errorf("get host assignments: %w", err)
 	}
diff --git a/internal/praefect/service/info/server.go b/internal/praefect/service/info/server.go
index 26703a411..2bc7852dd 100644
--- a/internal/praefect/service/info/server.go
+++ b/internal/praefect/service/info/server.go
@@ -19,7 +19,7 @@ import (
 type AssignmentStore interface {
 	// GetHostAssignments returns the names of the storages assigned to host the repository.
 	// The primary node must always be assigned.
-	GetHostAssignments(ctx context.Context, virtualStorage, relativePath string) ([]string, error)
+	GetHostAssignments(ctx context.Context, virtualStorage string, repositoryID int64) ([]string, error)
 	// SetReplicationFactor sets a repository's replication factor and returns the current assignments.
 	SetReplicationFactor(ctx context.Context, virtualStorage, relativePath string, replicationFactor int) ([]string, error)
 }