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