From e53673c8d79ce89dee6cb842e8da747860f0effd Mon Sep 17 00:00:00 2001
From: Andrew Mason <amason@slack-corp.com>
Date: Wed, 15 Sep 2021 19:38:27 -0400
Subject: [PATCH] Add cli entrypoints

Signed-off-by: Andrew Mason <amason@slack-corp.com>
---
 .../vtctldclient/internal/command/tablets.go  | 54 +++++++++++++++++++
 1 file changed, 54 insertions(+)

diff --git a/go/cmd/vtctldclient/internal/command/tablets.go b/go/cmd/vtctldclient/internal/command/tablets.go
index bfd9ce9c33e..bedcad8b428 100644
--- a/go/cmd/vtctldclient/internal/command/tablets.go
+++ b/go/cmd/vtctldclient/internal/command/tablets.go
@@ -20,10 +20,12 @@ import (
 	"fmt"
 	"strconv"
 	"strings"
+	"time"
 
 	"github.com/spf13/cobra"
 
 	"vitess.io/vitess/go/cmd/vtctldclient/cli"
+	"vitess.io/vitess/go/protoutil"
 	"vitess.io/vitess/go/vt/topo/topoproto"
 
 	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
@@ -81,6 +83,14 @@ Valid output formats are "awk" and "json".`,
 		Args:                  cobra.NoArgs,
 		RunE:                  commandGetTablets,
 	}
+	// PingTablet makes a PingTablet gRPC call to a vtctld.
+	PingTablet = &cobra.Command{
+		Use:                   "PingTablet <alias>",
+		Short:                 "Checks that the specified tablet is awake and responding to RPCs. This command can be blocked by other in-flight operations.",
+		DisableFlagsInUseLine: true,
+		Args:                  cobra.ExactArgs(1),
+		RunE:                  commandPingTablet,
+	}
 	// RefreshState makes a RefreshState gRPC call to a vtctld.
 	RefreshState = &cobra.Command{
 		Use:                   "RefreshState <alias>",
@@ -105,6 +115,14 @@ Valid output formats are "awk" and "json".`,
 		Args:                  cobra.ExactArgs(2),
 		RunE:                  commandSetWritable,
 	}
+	// SleepTablet makes a SleepTablet gRPC call to a vtctld.
+	SleepTablet = &cobra.Command{
+		Use:                   "SleepTablet <alias> <duration>",
+		Short:                 "Blocks the action queue on the specified tablet for the specified amount of time. This is typically used for testing.",
+		DisableFlagsInUseLine: true,
+		Args:                  cobra.ExactArgs(2),
+		RunE:                  commandSleepTablet,
+	}
 	// StartReplication makes a StartReplication gRPC call to a vtctld.
 	StartReplication = &cobra.Command{
 		Use:                   "StartReplication <alias>",
@@ -285,6 +303,20 @@ func commandGetTablets(cmd *cobra.Command, args []string) error {
 	return nil
 }
 
+func commandPingTablet(cmd *cobra.Command, args []string) error {
+	alias, err := topoproto.ParseTabletAlias(cmd.Flags().Arg(0))
+	if err != nil {
+		return err
+	}
+
+	cli.FinishedParsing(cmd)
+
+	_, err = client.PingTablet(commandCtx, &vtctldatapb.PingTabletRequest{
+		TabletAlias: alias,
+	})
+	return err
+}
+
 func commandRefreshState(cmd *cobra.Command, args []string) error {
 	alias, err := topoproto.ParseTabletAlias(cmd.Flags().Arg(0))
 	if err != nil {
@@ -359,6 +391,26 @@ func commandSetWritable(cmd *cobra.Command, args []string) error {
 	return err
 }
 
+func commandSleepTablet(cmd *cobra.Command, args []string) error {
+	alias, err := topoproto.ParseTabletAlias(cmd.Flags().Arg(0))
+	if err != nil {
+		return err
+	}
+
+	duration, err := time.ParseDuration(cmd.Flags().Arg(1))
+	if err != nil {
+		return err
+	}
+
+	cli.FinishedParsing(cmd)
+
+	_, err = client.SleepTablet(commandCtx, &vtctldatapb.SleepTabletRequest{
+		TabletAlias: alias,
+		Duration:    protoutil.DurationToProto(duration),
+	})
+	return err
+}
+
 func commandStartReplication(cmd *cobra.Command, args []string) error {
 	alias, err := topoproto.ParseTabletAlias(cmd.Flags().Arg(0))
 	if err != nil {
@@ -404,12 +456,14 @@ func init() {
 	GetTablets.Flags().BoolVar(&getTabletsOptions.Strict, "strict", false, "Require all cells to return successful tablet data. Without --strict, tablet listings may be partial.")
 	Root.AddCommand(GetTablets)
 
+	Root.AddCommand(PingTablet)
 	Root.AddCommand(RefreshState)
 
 	RefreshStateByShard.Flags().StringSliceVarP(&refreshStateByShardOptions.Cells, "cells", "c", nil, "If specified, only call RefreshState on tablets in the specified cells. If empty, all cells are considered.")
 	Root.AddCommand(RefreshStateByShard)
 
 	Root.AddCommand(SetWritable)
+	Root.AddCommand(SleepTablet)
 	Root.AddCommand(StartReplication)
 	Root.AddCommand(StopReplication)
 }