praefect: Add list-storages subcommandjc-praefect-subcmd-list-virtual-storages

It would be handy for Praefct admins to easily print out the virtual
storages with its associated storages without having to go into the
config. This change adds a new subcommand list-storages that either
prints out all virtual storages with their associated storages, or takes
a virtual storage as a flag and prints out storages for that virtual
storage.

Changelog: added

3 files changed, 280 insertions, 2 deletions

diff --git a/cmd/praefect/subcmd.go b/cmd/praefect/subcmd.go
index e91090d8d..292da9a30 100644
--- a/cmd/praefect/subcmd.go
+++ b/cmd/praefect/subcmd.go
@@ -51,8 +51,9 @@ var subcommands = map[string]subcmd{
 		praefect.NewUnavailableReposCheck,
 		praefect.NewClockSyncCheck(helper.CheckClockSync),
 	),
-	metadataCmdName: newMetadataSubcommand(os.Stdout),
-	verifyCmdName:   newVerifySubcommand(os.Stdout),
+	metadataCmdName:     newMetadataSubcommand(os.Stdout),
+	verifyCmdName:       newVerifySubcommand(os.Stdout),
+	listStoragesCmdName: newListStorages(os.Stdout),
 }
 
 // subCommand returns an exit code, to be fed into os.Exit.
diff --git a/cmd/praefect/subcmd_list_storages.go b/cmd/praefect/subcmd_list_storages.go
new file mode 100644
index 000000000..105454336
--- /dev/null
+++ b/cmd/praefect/subcmd_list_storages.go
@@ -0,0 +1,91 @@
+package main
+
+import (
+	"flag"
+	"fmt"
+	"io"
+
+	"github.com/olekukonko/tablewriter"
+	"gitlab.com/gitlab-org/gitaly/v15/internal/praefect/config"
+)
+
+const (
+	listStoragesCmdName = "list-storages"
+)
+
+type listStorages struct {
+	virtualStorage string
+	w              io.Writer
+}
+
+func newListStorages(w io.Writer) *listStorages {
+	return &listStorages{w: w}
+}
+
+func (cmd *listStorages) FlagSet() *flag.FlagSet {
+	fs := flag.NewFlagSet(listStoragesCmdName, flag.ExitOnError)
+	fs.StringVar(
+		&cmd.virtualStorage,
+		paramVirtualStorage,
+		"",
+		"name of the virtual storage to list storages for",
+	)
+	fs.Usage = func() {
+		printfErr("Description:\n" +
+			"	This command lists virtual storages and their associated storages.\n" +
+			"	Passing a virtual-storage argument will print out the storage associated with\n" +
+			"       that particular virtual storage.\n")
+		fs.PrintDefaults()
+	}
+
+	return fs
+}
+
+func (cmd listStorages) Exec(flags *flag.FlagSet, cfg config.Config) error {
+	if flags.NArg() > 0 {
+		return unexpectedPositionalArgsError{Command: flags.Name()}
+	}
+
+	table := tablewriter.NewWriter(cmd.w)
+	table.SetHeader([]string{"VIRTUAL_STORAGE", "NODE", "ADDRESS"})
+	table.SetHeaderAlignment(tablewriter.ALIGN_LEFT)
+	table.SetAutoFormatHeaders(false)
+	table.SetAlignment(tablewriter.ALIGN_LEFT)
+	table.SetCenterSeparator("")
+	table.SetColumnSeparator("")
+	table.SetRowSeparator("")
+	table.SetHeaderLine(false)
+	table.SetBorder(false)
+	table.SetTablePadding("\t") // pad with tabs
+	table.SetNoWhiteSpace(true)
+
+	if cmd.virtualStorage != "" {
+		for _, virtualStorage := range cfg.VirtualStorages {
+			if virtualStorage.Name != cmd.virtualStorage {
+				continue
+			}
+
+			for _, node := range virtualStorage.Nodes {
+				table.Append([]string{virtualStorage.Name, node.Storage, node.Address})
+			}
+
+			table.Render()
+
+			return nil
+		}
+
+		fmt.Fprintf(cmd.w, "No virtual storages named %s.\n", cmd.virtualStorage)
+
+		return nil
+	}
+
+	for _, virtualStorage := range cfg.VirtualStorages {
+		for _, node := range virtualStorage.Nodes {
+			table.Append([]string{virtualStorage.Name, node.Storage, node.Address})
+		}
+	}
+
+	table.Render()
+
+	return nil
+}
diff --git a/cmd/praefect/subcmd_list_storages_test.go b/cmd/praefect/subcmd_list_storages_test.go
new file mode 100644
index 000000000..fda7446c8
--- /dev/null
+++ b/cmd/praefect/subcmd_list_storages_test.go
@@ -0,0 +1,186 @@
+package main
+
+import (
+	"bytes"
+	"testing"
+
+	"github.com/olekukonko/tablewriter"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"gitlab.com/gitlab-org/gitaly/v15/internal/praefect/config"
+)
+
+func TestListStorages_FlagSet(t *testing.T) {
+	t.Parallel()
+	cmd := &listStorages{}
+	fs := cmd.FlagSet()
+	require.NoError(t, fs.Parse([]string{"--virtual-storage", "vs"}))
+	require.Equal(t, "vs", cmd.virtualStorage)
+}
+
+func TestListStorages_Exec(t *testing.T) {
+	t.Parallel()
+
+	testCases := []struct {
+		desc            string
+		virtualStorages []*config.VirtualStorage
+		args            []string
+		expectedOutput  func(*tablewriter.Table)
+	}{
+		{
+			desc: "one virtual storage",
+			virtualStorages: []*config.VirtualStorage{
+				{
+					Name: "vs-1",
+					Nodes: []*config.Node{
+						{
+							Storage: "storage-1",
+							Address: "tcp://1.2.3.4",
+						},
+						{
+							Storage: "storage-2",
+							Address: "tcp://4.3.2.1",
+						},
+					},
+				},
+			},
+			args: []string{},
+			expectedOutput: func(t *tablewriter.Table) {
+				t.Append([]string{"vs-1", "storage-1", "tcp://1.2.3.4"})
+				t.Append([]string{"vs-1", "storage-2", "tcp://4.3.2.1"})
+			},
+		},
+		{
+			desc: "multiple virtual storages but only show one",
+			virtualStorages: []*config.VirtualStorage{
+				{
+					Name: "vs-1",
+					Nodes: []*config.Node{
+						{
+							Storage: "storage-1",
+							Address: "tcp://1.2.3.4",
+						},
+						{
+							Storage: "storage-2",
+							Address: "tcp://4.3.2.1",
+						},
+					},
+				},
+				{
+					Name: "vs-2",
+					Nodes: []*config.Node{
+						{
+							Storage: "storage-3",
+							Address: "tcp://1.1.3.4",
+						},
+						{
+							Storage: "storage-4",
+							Address: "tcp://1.3.2.1",
+						},
+					},
+				},
+				{
+					Name: "vs-3",
+					Nodes: []*config.Node{
+						{
+							Storage: "storage-5",
+							Address: "tcp://2.1.3.4",
+						},
+						{
+							Storage: "storage-6",
+							Address: "tcp://2.3.2.1",
+						},
+					},
+				},
+			},
+			args: []string{"-virtual-storage", "vs-2"},
+			expectedOutput: func(t *tablewriter.Table) {
+				t.Append([]string{"vs-2", "storage-3", "tcp://1.1.3.4"})
+				t.Append([]string{"vs-2", "storage-4", "tcp://1.3.2.1"})
+			},
+		},
+		{
+			desc: "one virtual storage with virtual storage arg",
+			virtualStorages: []*config.VirtualStorage{
+				{
+					Name: "vs-1",
+					Nodes: []*config.Node{
+						{
+							Storage: "storage-1",
+							Address: "tcp://1.2.3.4",
+						},
+						{
+							Storage: "storage-2",
+							Address: "tcp://4.3.2.1",
+						},
+					},
+				},
+			},
+			args: []string{"-virtual-storage", "vs-1"},
+			expectedOutput: func(t *tablewriter.Table) {
+				t.Append([]string{"vs-1", "storage-1", "tcp://1.2.3.4"})
+				t.Append([]string{"vs-1", "storage-2", "tcp://4.3.2.1"})
+			},
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.desc, func(t *testing.T) {
+			var expectedOutput bytes.Buffer
+			table := tablewriter.NewWriter(&expectedOutput)
+			table.SetHeader([]string{"VIRTUAL_STORAGE", "NODE", "ADDRESS"})
+			table.SetHeaderAlignment(tablewriter.ALIGN_LEFT)
+			table.SetAutoFormatHeaders(false)
+			table.SetAlignment(tablewriter.ALIGN_LEFT)
+			table.SetCenterSeparator("")
+			table.SetColumnSeparator("")
+			table.SetRowSeparator("")
+			table.SetHeaderLine(false)
+			table.SetBorder(false)
+			table.SetTablePadding("\t") // pad with tabs
+			table.SetNoWhiteSpace(true)
+			tc.expectedOutput(table)
+			table.Render()
+
+			var out bytes.Buffer
+			cmd := &listStorages{
+				w: &out,
+			}
+
+			fs := cmd.FlagSet()
+			require.NoError(t, fs.Parse(tc.args))
+			require.NoError(t, cmd.Exec(fs, config.Config{
+				VirtualStorages: tc.virtualStorages,
+			}))
+			assert.Equal(t, expectedOutput.String(), out.String())
+		})
+	}
+
+	t.Run("virtual storage arg matches no virtual storages", func(t *testing.T) {
+		var out bytes.Buffer
+		cmd := &listStorages{
+			w: &out,
+		}
+
+		fs := cmd.FlagSet()
+		require.NoError(t, fs.Parse([]string{"-virtual-storage", "vs-2"}))
+		require.NoError(t, cmd.Exec(fs, config.Config{
+			VirtualStorages: []*config.VirtualStorage{
+				{
+					Name: "vs-1",
+					Nodes: []*config.Node{
+						{
+							Storage: "storage-1",
+							Address: "tcp://1.2.3.4",
+						},
+						{
+							Storage: "storage-2",
+							Address: "tcp://4.3.2.1",
+						},
+					},
+				},
+			},
+		}))
+		assert.Equal(t, "No virtual storages named vs-2.\n", out.String())
+	})
+}