Admin list clients

cli/mongodb-backup-admin/main.go

@@ -2,61 +2,121 @@ package main
 
 import (
 	"context"
-	"flag"
 	"fmt"
+	"io"
 	"log"
+	"os"
+	"strings"
+	"time"
 
+	"github.com/alecthomas/kingpin"
 	pbapi "github.com/percona/mongodb-backup/proto/api"
+	"github.com/pkg/errors"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials"
 	"google.golang.org/grpc/testdata"
 )
 
-var (
-	clientID           = flag.String("id", "", "Client ID")
-	tls                = flag.Bool("tls", false, "Connection uses TLS if true, else plain TCP")
-	caFile             = flag.String("ca_file", "", "The file containning the CA root cert file")
-	serverAddr         = flag.String("server_addr", "127.0.0.1:10001", "The server address in the format of host:port")
-	serverHostOverride = flag.String("server_host_override", "x.test.youtube.com", "The server name use to verify the hostname returned by TLS handshake")
-)
+type cliOptions struct {
+	app        *kingpin.Application
+	clientID   *string
+	tls        *bool
+	caFile     *string
+	serverAddr *string
+
+	listClients *kingpin.CmdClause
+}
 
 func main() {
-	flag.Parse()
-	var opts []grpc.DialOption
-	if *tls {
-		if *caFile == "" {
-			*caFile = testdata.Path("ca.pem")
+	cmd, opts, err := processCliArgs(os.Args[1:])
+	if err != nil {
+		if opts != nil {
+			kingpin.Usage()
+		}
+		log.Fatal(err)
+	}
+
+	var grpcOpts []grpc.DialOption
+	if *opts.tls {
+		if *opts.caFile == "" {
+			*opts.caFile = testdata.Path("ca.pem")
 		}
-		creds, err := credentials.NewClientTLSFromFile(*caFile, *serverHostOverride)
+		creds, err := credentials.NewClientTLSFromFile(*opts.caFile, "")
 		if err != nil {
 			log.Fatalf("Failed to create TLS credentials %v", err)
 		}
-		opts = append(opts, grpc.WithTransportCredentials(creds))
+		grpcOpts = append(grpcOpts, grpc.WithTransportCredentials(creds))
 	} else {
-		opts = append(opts, grpc.WithInsecure())
+		grpcOpts = append(grpcOpts, grpc.WithInsecure())
 	}
 
-	conn, err := grpc.Dial(*serverAddr, opts...)
+	conn, err := grpc.Dial(*opts.serverAddr, grpcOpts...)
 	if err != nil {
 		log.Fatalf("fail to dial: %v", err)
 	}
 	defer conn.Close()
 
-	processMessages(conn)
+	switch cmd {
+	case "list-agents":
+		clients, err := getConnectedAgents(conn)
+		if err != nil {
+			log.Fatal(err)
+		}
+		printConnectedAgents(clients)
+	}
 
 }
 
-func processMessages(conn *grpc.ClientConn) {
+func getConnectedAgents(conn *grpc.ClientConn) ([]*pbapi.Client, error) {
 	apiClient := pbapi.NewApiClient(conn)
 	stream, err := apiClient.GetClients(context.Background(), &pbapi.Empty{})
 	if err != nil {
-		log.Fatal(err)
+		return nil, err
 	}
+	clients := []*pbapi.Client{}
 	for {
 		msg, err := stream.Recv()
 		if err != nil {
-			log.Fatal(err)
+			if err == io.EOF {
+				break
+			}
+			return nil, errors.Wrap(err, "Cannot get the connected agents list")
 		}
-		fmt.Printf("%+v\n", msg)
+		clients = append(clients, msg)
 	}
+	return clients, nil
+}
+
+func printConnectedAgents(clients []*pbapi.Client) {
+	if len(clients) == 0 {
+		fmt.Println("There are no agents connected to the backup master")
+		return
+	}
+	fmt.Println(strings.Repeat("-", 100))
+	for _, client := range clients {
+		fmt.Printf("%s  -  %s  -  %v\n", client.ClientID, client.Status, time.Unix(client.LastSeen, 0))
+	}
+}
+
+func processCliArgs(args []string) (string, *cliOptions, error) {
+
+	app := kingpin.New("mongodb-backup-admin", "MongoDB backup admin")
+	opts := &cliOptions{
+		clientID:    kingpin.Flag("client-id", "Client ID").Required().String(),
+		tls:         kingpin.Flag("tls", "Connection uses TLS if true, else plain TCP").Default("false").Bool(),
+		caFile:      kingpin.Flag("ca-file", "The file containning the CA root cert file").String(),
+		serverAddr:  kingpin.Flag("server-addr", "The server address in the format of host:port").Default("127.0.0.1:10001").String(),
+		listClients: kingpin.Command("list-agents", "List all agents connected to the server"),
+	}
+
+	cmd, err := app.Parse(args)
+	if err != nil {
+		return "", nil, err
+	}
+
+	if cmd == "" {
+		return "", opts, fmt.Errorf("Invalid command")
+	}
+
+	return cmd, opts, nil
 }

tests/general_test.go

@@ -78,7 +78,7 @@ func TestOne(t *testing.T) {
 	t1 := time.Now().Truncate(time.Second)
 	clientID := "ABC123"
 
-	time.Sleep(2 * time.Second)
+	time.Sleep(200 * time.Millisecond)
 	agentClient := pb.NewMessagesClient(agentConn)
 	agentStream, err := agentClient.MessagesChat(context.Background())
 	registerMsg := &pb.ClientMessage{