1// +build postgres
2
3package datastore
4
5import (
6	"testing"
7
8	"github.com/stretchr/testify/require"
9	"gitlab.com/gitlab-org/gitaly/v14/internal/testhelper"
10)
11
12func TestAssignmentStore_GetHostAssignments(t *testing.T) {
13	type assignment struct {
14		virtualStorage string
15		relativePath   string
16		storage        string
17	}
18
19	configuredStorages := []string{"storage-1", "storage-2", "storage-3"}
20	for _, tc := range []struct {
21		desc                string
22		virtualStorage      string
23		existingAssignments []assignment
24		expectedAssignments []string
25		error               error
26	}{
27		{
28			desc:           "virtual storage not found",
29			virtualStorage: "invalid-virtual-storage",
30			error:          newVirtualStorageNotFoundError("invalid-virtual-storage"),
31		},
32		{
33			desc:                "configured storages fallback when no records",
34			virtualStorage:      "virtual-storage",
35			expectedAssignments: configuredStorages,
36		},
37		{
38			desc:           "configured storages fallback when a repo exists in different virtual storage",
39			virtualStorage: "virtual-storage",
40			existingAssignments: []assignment{
41				{virtualStorage: "other-virtual-storage", relativePath: "relative-path", storage: "storage-1"},
42			},
43			expectedAssignments: configuredStorages,
44		},
45		{
46			desc:           "configured storages fallback when a different repo exists in the virtual storage ",
47			virtualStorage: "virtual-storage",
48			existingAssignments: []assignment{
49				{virtualStorage: "virtual-storage", relativePath: "other-relative-path", storage: "storage-1"},
50			},
51			expectedAssignments: configuredStorages,
52		},
53		{
54			desc:           "unconfigured storages are ignored",
55			virtualStorage: "virtual-storage",
56			existingAssignments: []assignment{
57				{virtualStorage: "virtual-storage", relativePath: "relative-path", storage: "unconfigured-storage"},
58			},
59			expectedAssignments: configuredStorages,
60		},
61		{
62			desc:           "assignments found",
63			virtualStorage: "virtual-storage",
64			existingAssignments: []assignment{
65				{virtualStorage: "virtual-storage", relativePath: "relative-path", storage: "storage-1"},
66				{virtualStorage: "virtual-storage", relativePath: "relative-path", storage: "storage-2"},
67				{virtualStorage: "virtual-storage", relativePath: "relative-path", storage: "unconfigured"},
68			},
69			expectedAssignments: []string{"storage-1", "storage-2"},
70		},
71	} {
72		t.Run(tc.desc, func(t *testing.T) {
73			ctx, cancel := testhelper.Context()
74			defer cancel()
75
76			db := getDB(t)
77
78			for _, assignment := range tc.existingAssignments {
79				_, err := db.ExecContext(ctx, `
80					INSERT INTO repositories (virtual_storage, relative_path)
81					VALUES ($1, $2)
82					ON CONFLICT DO NOTHING
83				`, assignment.virtualStorage, assignment.relativePath)
84				require.NoError(t, err)
85
86				_, err = db.ExecContext(ctx, `
87					INSERT INTO repository_assignments VALUES ($1, $2, $3)
88				`, assignment.virtualStorage, assignment.relativePath, assignment.storage)
89				require.NoError(t, err)
90			}
91
92			actualAssignments, err := NewAssignmentStore(
93				db,
94				map[string][]string{"virtual-storage": configuredStorages},
95			).GetHostAssignments(ctx, tc.virtualStorage, "relative-path")
96			require.Equal(t, tc.error, err)
97			require.ElementsMatch(t, tc.expectedAssignments, actualAssignments)
98		})
99	}
100}
101
102func TestAssignmentStore_SetReplicationFactor(t *testing.T) {
103	type matcher func(testing.TB, []string)
104
105	equal := func(expected []string) matcher {
106		return func(t testing.TB, actual []string) {
107			t.Helper()
108			require.Equal(t, expected, actual)
109		}
110	}
111
112	contains := func(expecteds ...[]string) matcher {
113		return func(t testing.TB, actual []string) {
114			t.Helper()
115			require.Contains(t, expecteds, actual)
116		}
117	}
118
119	for _, tc := range []struct {
120		desc                  string
121		existingAssignments   []string
122		nonExistentRepository bool
123		replicationFactor     int
124		requireStorages       matcher
125		error                 error
126	}{
127		{
128			desc:                  "increase replication factor of non-existent repository",
129			nonExistentRepository: true,
130			replicationFactor:     1,
131			error:                 newRepositoryNotFoundError("virtual-storage", "relative-path"),
132		},
133		{
134			desc:              "primary prioritized when setting the first assignments",
135			replicationFactor: 1,
136			requireStorages:   equal([]string{"primary"}),
137		},
138		{
139			desc:                "increasing replication factor ignores unconfigured storages",
140			existingAssignments: []string{"unconfigured-storage"},
141			replicationFactor:   1,
142			requireStorages:     equal([]string{"primary"}),
143		},
144		{
145			desc:                "replication factor already achieved",
146			existingAssignments: []string{"primary", "secondary-1"},
147			replicationFactor:   2,
148			requireStorages:     equal([]string{"primary", "secondary-1"}),
149		},
150		{
151			desc:                "increase replication factor by a step",
152			existingAssignments: []string{"primary"},
153			replicationFactor:   2,
154			requireStorages:     contains([]string{"primary", "secondary-1"}, []string{"primary", "secondary-2"}),
155		},
156		{
157			desc:                "increase replication factor to maximum",
158			existingAssignments: []string{"primary"},
159			replicationFactor:   3,
160			requireStorages:     equal([]string{"primary", "secondary-1", "secondary-2"}),
161		},
162		{
163			desc:                "increased replication factor unattainable",
164			existingAssignments: []string{"primary"},
165			replicationFactor:   4,
166			error:               newUnattainableReplicationFactorError(4, 3),
167		},
168		{
169			desc:                "decreasing replication factor ignores unconfigured storages",
170			existingAssignments: []string{"secondary-1", "unconfigured-storage"},
171			replicationFactor:   1,
172			requireStorages:     equal([]string{"secondary-1"}),
173		},
174		{
175			desc:                "decrease replication factor by a step",
176			existingAssignments: []string{"primary", "secondary-1", "secondary-2"},
177			replicationFactor:   2,
178			requireStorages:     contains([]string{"primary", "secondary-1"}, []string{"primary", "secondary-2"}),
179		},
180		{
181			desc:                "decrease replication factor to minimum",
182			existingAssignments: []string{"primary", "secondary-1", "secondary-2"},
183			replicationFactor:   1,
184			requireStorages:     equal([]string{"primary"}),
185		},
186		{
187			desc:              "minimum replication factor is enforced",
188			replicationFactor: 0,
189			error:             newMinimumReplicationFactorError(0),
190		},
191	} {
192		t.Run(tc.desc, func(t *testing.T) {
193			ctx, cancel := testhelper.Context()
194			defer cancel()
195
196			db := getDB(t)
197
198			configuredStorages := map[string][]string{"virtual-storage": {"primary", "secondary-1", "secondary-2"}}
199
200			if !tc.nonExistentRepository {
201				_, err := db.ExecContext(ctx, `
202					INSERT INTO repositories (virtual_storage, relative_path, "primary")
203					VALUES ('virtual-storage', 'relative-path', 'primary')
204				`)
205				require.NoError(t, err)
206			}
207
208			for _, storage := range tc.existingAssignments {
209				_, err := db.ExecContext(ctx, `
210					INSERT INTO repository_assignments VALUES ('virtual-storage', 'relative-path', $1)
211				`, storage)
212				require.NoError(t, err)
213			}
214
215			store := NewAssignmentStore(db, configuredStorages)
216
217			setStorages, err := store.SetReplicationFactor(ctx, "virtual-storage", "relative-path", tc.replicationFactor)
218			require.Equal(t, tc.error, err)
219			if tc.error != nil {
220				return
221			}
222
223			tc.requireStorages(t, setStorages)
224
225			assignedStorages, err := store.GetHostAssignments(ctx, "virtual-storage", "relative-path")
226			require.NoError(t, err)
227			tc.requireStorages(t, assignedStorages)
228		})
229	}
230}
231