1package main
2
3import (
4	"bytes"
5	"testing"
6
7	"github.com/stretchr/testify/require"
8	"gitlab.com/gitlab-org/gitaly/v14/internal/praefect/config"
9	"gitlab.com/gitlab-org/gitaly/v14/internal/praefect/datastore"
10	"gitlab.com/gitlab-org/gitaly/v14/internal/praefect/datastore/glsql"
11	"gitlab.com/gitlab-org/gitaly/v14/internal/praefect/service/info"
12	"gitlab.com/gitlab-org/gitaly/v14/internal/testhelper"
13	"gitlab.com/gitlab-org/gitaly/v14/proto/go/gitalypb"
14	"google.golang.org/grpc"
15)
16
17func registerPraefectInfoServer(impl gitalypb.PraefectInfoServiceServer) svcRegistrar {
18	return func(srv *grpc.Server) {
19		gitalypb.RegisterPraefectInfoServiceServer(srv, impl)
20	}
21}
22
23func TestDatalossSubcommand(t *testing.T) {
24	t.Parallel()
25	cfg := config.Config{
26		VirtualStorages: []*config.VirtualStorage{
27			{
28				Name: "virtual-storage-1",
29				Nodes: []*config.Node{
30					{Storage: "gitaly-1"},
31					{Storage: "gitaly-2"},
32					{Storage: "gitaly-3"},
33				},
34			},
35			{
36				Name: "virtual-storage-2",
37				Nodes: []*config.Node{
38					{Storage: "gitaly-4"},
39				},
40			},
41		},
42	}
43
44	tx := glsql.NewDB(t).Begin(t)
45	defer tx.Rollback(t)
46
47	ctx, cancel := testhelper.Context()
48	defer cancel()
49
50	testhelper.SetHealthyNodes(t, ctx, tx, map[string]map[string][]string{"praefect-0": {
51		"virtual-storage-1": {"gitaly-1", "gitaly-3"},
52	}})
53	gs := datastore.NewPostgresRepositoryStore(tx, cfg.StorageNames())
54
55	for _, q := range []string{
56		`
57				INSERT INTO repositories (virtual_storage, relative_path, "primary")
58				VALUES
59					('virtual-storage-1', 'repository-1', 'gitaly-1'),
60					('virtual-storage-1', 'repository-2', 'gitaly-3')
61				`,
62		`
63				INSERT INTO repository_assignments (virtual_storage, relative_path, storage)
64				VALUES
65					('virtual-storage-1', 'repository-1', 'gitaly-1'),
66					('virtual-storage-1', 'repository-1', 'gitaly-2'),
67					('virtual-storage-1', 'repository-2', 'gitaly-1'),
68					('virtual-storage-1', 'repository-2', 'gitaly-3')
69				`,
70	} {
71		_, err := tx.ExecContext(ctx, q)
72		require.NoError(t, err)
73	}
74
75	require.NoError(t, gs.SetGeneration(ctx, "virtual-storage-1", "repository-1", "gitaly-1", 1))
76	require.NoError(t, gs.SetGeneration(ctx, "virtual-storage-1", "repository-1", "gitaly-2", 0))
77	require.NoError(t, gs.SetGeneration(ctx, "virtual-storage-1", "repository-1", "gitaly-3", 0))
78
79	require.NoError(t, gs.SetGeneration(ctx, "virtual-storage-1", "repository-2", "gitaly-2", 1))
80	require.NoError(t, gs.SetGeneration(ctx, "virtual-storage-1", "repository-2", "gitaly-3", 0))
81
82	ln, clean := listenAndServe(t, []svcRegistrar{
83		registerPraefectInfoServer(info.NewServer(cfg, gs, nil, nil, nil)),
84	})
85	defer clean()
86	for _, tc := range []struct {
87		desc            string
88		args            []string
89		virtualStorages []*config.VirtualStorage
90		output          string
91		error           error
92	}{
93		{
94			desc:  "positional arguments",
95			args:  []string{"-virtual-storage=virtual-storage-1", "positional-arg"},
96			error: unexpectedPositionalArgsError{Command: "dataloss"},
97		},
98		{
99			desc: "data loss with unavailable repositories",
100			args: []string{"-virtual-storage=virtual-storage-1"},
101			output: `Virtual storage: virtual-storage-1
102  Repositories:
103    repository-2 (unavailable):
104      Primary: gitaly-3
105      In-Sync Storages:
106        gitaly-2, unhealthy
107      Outdated Storages:
108        gitaly-1 is behind by 2 changes or less, assigned host
109        gitaly-3 is behind by 1 change or less, assigned host
110`,
111		},
112		{
113			desc: "data loss with partially unavailable repositories",
114			args: []string{"-virtual-storage=virtual-storage-1", "-partially-unavailable"},
115			output: `Virtual storage: virtual-storage-1
116  Repositories:
117    repository-1:
118      Primary: gitaly-1
119      In-Sync Storages:
120        gitaly-1, assigned host
121      Outdated Storages:
122        gitaly-2 is behind by 1 change or less, assigned host, unhealthy
123        gitaly-3 is behind by 1 change or less
124    repository-2 (unavailable):
125      Primary: gitaly-3
126      In-Sync Storages:
127        gitaly-2, unhealthy
128      Outdated Storages:
129        gitaly-1 is behind by 2 changes or less, assigned host
130        gitaly-3 is behind by 1 change or less, assigned host
131`,
132		},
133		{
134			desc:            "multiple virtual storages with unavailable repositories",
135			virtualStorages: []*config.VirtualStorage{{Name: "virtual-storage-2"}, {Name: "virtual-storage-1"}},
136			output: `Virtual storage: virtual-storage-1
137  Repositories:
138    repository-2 (unavailable):
139      Primary: gitaly-3
140      In-Sync Storages:
141        gitaly-2, unhealthy
142      Outdated Storages:
143        gitaly-1 is behind by 2 changes or less, assigned host
144        gitaly-3 is behind by 1 change or less, assigned host
145Virtual storage: virtual-storage-2
146  All repositories are available!
147`,
148		},
149		{
150			desc:            "multiple virtual storages with partially unavailable repositories",
151			args:            []string{"-partially-unavailable"},
152			virtualStorages: []*config.VirtualStorage{{Name: "virtual-storage-2"}, {Name: "virtual-storage-1"}},
153			output: `Virtual storage: virtual-storage-1
154  Repositories:
155    repository-1:
156      Primary: gitaly-1
157      In-Sync Storages:
158        gitaly-1, assigned host
159      Outdated Storages:
160        gitaly-2 is behind by 1 change or less, assigned host, unhealthy
161        gitaly-3 is behind by 1 change or less
162    repository-2 (unavailable):
163      Primary: gitaly-3
164      In-Sync Storages:
165        gitaly-2, unhealthy
166      Outdated Storages:
167        gitaly-1 is behind by 2 changes or less, assigned host
168        gitaly-3 is behind by 1 change or less, assigned host
169Virtual storage: virtual-storage-2
170  All repositories are fully available on all assigned storages!
171`,
172		},
173	} {
174		t.Run(tc.desc, func(t *testing.T) {
175			cmd := newDatalossSubcommand()
176			output := &bytes.Buffer{}
177			cmd.output = output
178
179			fs := cmd.FlagSet()
180			require.NoError(t, fs.Parse(tc.args))
181			err := cmd.Exec(fs, config.Config{
182				VirtualStorages: tc.virtualStorages,
183				SocketPath:      ln.Addr().String(),
184			})
185			require.Equal(t, tc.error, err, err)
186			require.Equal(t, tc.output, output.String())
187		})
188	}
189}
190