Merge branch 'smh-database-assignments' into 'master'

Fetch host node assignments from the database

See merge request gitlab-org/gitaly!2789

3 files changed, 205 insertions, 1 deletions

diff --git a/cmd/praefect/main.go b/cmd/praefect/main.go
index 78f597443..df4505bb9 100644
--- a/cmd/praefect/main.go
+++ b/cmd/praefect/main.go
@@ -300,7 +300,7 @@ func run(cfgs []starter.Config, conf config.Config) error {
 			hm,
 			praefect.NewLockedRandom(rand.New(rand.NewSource(time.Now().UnixNano()))),
 			rs,
-			praefect.StaticStorageAssignments(conf.StorageNames()),
+			datastore.NewAssignmentStore(db, conf.StorageNames()),
 		)
 	} else {
 		healthChecker = praefect.HealthChecker(nodeManager)
diff --git a/internal/praefect/datastore/assignment.go b/internal/praefect/datastore/assignment.go
new file mode 100644
index 000000000..012fc8838
--- /dev/null
+++ b/internal/praefect/datastore/assignment.go
@@ -0,0 +1,69 @@
+package datastore
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/lib/pq"
+	"gitlab.com/gitlab-org/gitaly/internal/praefect/datastore/glsql"
+)
+
+func newVirtualStorageNotFoundError(virtualStorage string) error {
+	return fmt.Errorf("virtual storage %q not found", virtualStorage)
+}
+
+func newAssignmentsNotFoundError(virtualStorage, relativePath string) error {
+	return fmt.Errorf("host assignments for repository %q/%q not found", virtualStorage, relativePath)
+}
+
+// AssignmentStore manages host assignments in Postgres.
+type AssignmentStore struct {
+	db                 glsql.Querier
+	configuredStorages map[string][]string
+}
+
+// NewAssignmentsStore returns a new AssignmentStore using the passed in database.
+func NewAssignmentStore(db glsql.Querier, configuredStorages map[string][]string) AssignmentStore {
+	return AssignmentStore{db: db, configuredStorages: configuredStorages}
+}
+
+func (s AssignmentStore) GetHostAssignments(ctx context.Context, virtualStorage, relativePath string) ([]string, error) {
+	configuredStorages, ok := s.configuredStorages[virtualStorage]
+	if !ok {
+		return nil, newVirtualStorageNotFoundError(virtualStorage)
+	}
+
+	rows, err := s.db.QueryContext(ctx, `
+SELECT storage
+FROM repositories
+JOIN storage_repositories USING (virtual_storage, relative_path)
+WHERE virtual_storage = $1
+AND   relative_path = $2
+AND   assigned
+AND   storage = ANY($3::text[])
+	`, virtualStorage, relativePath, pq.StringArray(configuredStorages))
+	if err != nil {
+		return nil, fmt.Errorf("query: %w", err)
+	}
+	defer rows.Close()
+
+	var assignedStorages []string
+	for rows.Next() {
+		var storage string
+		if err := rows.Scan(&storage); err != nil {
+			return nil, fmt.Errorf("scan: %w", err)
+		}
+
+		assignedStorages = append(assignedStorages, storage)
+	}
+
+	if err := rows.Err(); err != nil {
+		return nil, fmt.Errorf("iterating rows: %w", err)
+	}
+
+	if len(assignedStorages) == 0 {
+		return nil, newAssignmentsNotFoundError(virtualStorage, relativePath)
+	}
+
+	return assignedStorages, nil
+}
diff --git a/internal/praefect/datastore/assignment_test.go b/internal/praefect/datastore/assignment_test.go
new file mode 100644
index 000000000..8ac8592c7
--- /dev/null
+++ b/internal/praefect/datastore/assignment_test.go
@@ -0,0 +1,135 @@
+// +build postgres
+
+package datastore
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	"gitlab.com/gitlab-org/gitaly/internal/testhelper"
+)
+
+func TestAssignmentStore_GetHostAssignments(t *testing.T) {
+	type repository struct {
+		virtualStorage string
+		relativePath   string
+	}
+
+	type storage struct {
+		virtualStorage string
+		relativePath   string
+		storage        string
+		assigned       bool
+	}
+
+	for _, tc := range []struct {
+		desc           string
+		virtualStorage string
+		repositories   []repository
+		storages       []storage
+		assignments    []string
+		error          error
+	}{
+		{
+			desc:           "virtual storage not found",
+			virtualStorage: "invalid-virtual-storage",
+			error:          newVirtualStorageNotFoundError("invalid-virtual-storage"),
+		},
+		{
+			desc:           "not found when no records",
+			virtualStorage: "virtual-storage",
+			error:          newAssignmentsNotFoundError("virtual-storage", "relative-path"),
+		},
+		{
+			desc:           "not found when only repository record",
+			virtualStorage: "virtual-storage",
+			repositories: []repository{
+				{virtualStorage: "virtual-storage", relativePath: "relative-path"},
+			},
+			error: newAssignmentsNotFoundError("virtual-storage", "relative-path"),
+		},
+		{
+			desc:           "not found with incorrect virtual storage",
+			virtualStorage: "virtual-storage",
+			repositories: []repository{
+				{virtualStorage: "other-virtual-storage", relativePath: "relative-path"},
+			},
+			storages: []storage{
+				{virtualStorage: "other-virtual-storage", relativePath: "relative-path", storage: "storage-1", assigned: true},
+			},
+			error: newAssignmentsNotFoundError("virtual-storage", "relative-path"),
+		},
+		{
+			desc:           "not found with incorrect relative path",
+			virtualStorage: "virtual-storage",
+			repositories: []repository{
+				{virtualStorage: "virtual-storage", relativePath: "other-relative-path"},
+			},
+			storages: []storage{
+				{virtualStorage: "virtual-storage", relativePath: "other-relative-path", storage: "storage-1", assigned: true},
+			},
+			error: newAssignmentsNotFoundError("virtual-storage", "relative-path"),
+		},
+
+		{
+			desc:           "not found when only storage record",
+			virtualStorage: "virtual-storage",
+			storages: []storage{
+				{virtualStorage: "virtual-storage", relativePath: "relative-path", storage: "storage-1", assigned: true},
+			},
+			error: newAssignmentsNotFoundError("virtual-storage", "relative-path"),
+		},
+		{
+			desc:           "unconfigured storages are ignored",
+			virtualStorage: "virtual-storage",
+			storages: []storage{
+				{virtualStorage: "virtual-storage", relativePath: "relative-path", storage: "unconfigured-storage", assigned: true},
+			},
+			error: newAssignmentsNotFoundError("virtual-storage", "relative-path"),
+		},
+		{
+			desc:           "assignments found",
+			virtualStorage: "virtual-storage",
+			repositories: []repository{
+				{virtualStorage: "virtual-storage", relativePath: "relative-path"},
+			},
+			storages: []storage{
+				{virtualStorage: "virtual-storage", relativePath: "relative-path", storage: "storage-1", assigned: true},
+				{virtualStorage: "virtual-storage", relativePath: "relative-path", storage: "storage-2", assigned: true},
+				{virtualStorage: "virtual-storage", relativePath: "relative-path", storage: "storage-3", assigned: false},
+				{virtualStorage: "virtual-storage", relativePath: "relative-path", storage: "unconfigured", assigned: true},
+			},
+			assignments: []string{"storage-1", "storage-2"},
+		},
+	} {
+		t.Run(tc.desc, func(t *testing.T) {
+			ctx, cancel := testhelper.Context()
+			defer cancel()
+
+			db := getDB(t)
+
+			for _, repository := range tc.repositories {
+				_, err := db.ExecContext(ctx, `
+					INSERT INTO repositories (virtual_storage, relative_path)
+					VALUES ($1, $2)
+				`, repository.virtualStorage, repository.relativePath)
+				require.NoError(t, err)
+			}
+
+			for _, storage := range tc.storages {
+				_, err := db.ExecContext(ctx, `
+					INSERT INTO storage_repositories (virtual_storage, relative_path, storage, assigned, generation)
+					VALUES ($1, $2, $3, $4, 0)
+				`, storage.virtualStorage, storage.relativePath, storage.storage, storage.assigned)
+				require.NoError(t, err)
+			}
+
+			assignments, err := NewAssignmentStore(
+				db,
+				map[string][]string{"virtual-storage": {"storage-1", "storage-2", "storage-3"}},
+			).GetHostAssignments(ctx, tc.virtualStorage, "relative-path")
+			require.Equal(t, tc.error, err)
+			require.ElementsMatch(t, tc.assignments, assignments)
+		})
+	}
+}