VReplication Workflow command: Allow stop/start on specific shards (#…

…17581)

Signed-off-by: Rohit Nayak <[email protected]>

go/cmd/vtctldclient/command/vreplication/workflow/state.go

@@ -60,14 +60,17 @@ func commandUpdateState(cmd *cobra.Command, args []string) error {
 	cli.FinishedParsing(cmd)
 
 	var state binlogdatapb.VReplicationWorkflowState
+	var shards []string
 	switch strings.ToLower(cmd.Name()) {
 	case "start":
 		if err := common.CanRestartWorkflow(baseOptions.Keyspace, baseOptions.Workflow); err != nil {
 			return err
 		}
 		state = binlogdatapb.VReplicationWorkflowState_Running
+		shards = baseOptions.Shards
 	case "stop":
 		state = binlogdatapb.VReplicationWorkflowState_Stopped
+		shards = baseOptions.Shards
 	default:
 		return fmt.Errorf("invalid workflow state: %s", args[0])
 	}
@@ -80,6 +83,7 @@ func commandUpdateState(cmd *cobra.Command, args []string) error {
 			Cells:       textutil.SimulatedNullStringSlice,
 			TabletTypes: textutil.SimulatedNullTabletTypeSlice,
 			State:       &state,
+			Shards:      shards,
 		},
 	}
 

go/test/endtoend/vreplication/vreplication_vtctldclient_cli_test.go

@@ -25,6 +25,7 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/require"
+	"github.com/tidwall/gjson"
 	"golang.org/x/exp/maps"
 	"google.golang.org/protobuf/encoding/protojson"
 
@@ -501,6 +502,28 @@ func splitShard(t *testing.T, keyspace, workflowName, sourceShards, targetShards
 	waitForWorkflowState(t, vc, ksWorkflow, binlogdatapb.VReplicationWorkflowState_Stopped.String())
 	rs.Start()
 	waitForWorkflowState(t, vc, fmt.Sprintf("%s.%s", keyspace, workflowName), binlogdatapb.VReplicationWorkflowState_Running.String())
+
+	t.Run("Test --shards in workflow start/stop", func(t *testing.T) {
+		// This subtest expects workflow to be running at the start and restarts it at the end.
+		type tCase struct {
+			shards   string
+			action   string
+			expected int
+		}
+		testCases := []tCase{
+			{"-40", "stop", 1},
+			{"40-80", "stop", 1},
+			{"-40,40-80", "start", 2},
+		}
+		for _, tc := range testCases {
+			output, err := vc.VtctldClient.ExecuteCommandWithOutput("workflow", "--keyspace", keyspace, tc.action, "--workflow", workflowName, "--shards", tc.shards)
+			require.NoError(t, err, "failed to %s workflow: %v", tc.action, err)
+			cnt := gjson.Get(output, "details.#").Int()
+			require.EqualValuesf(t, tc.expected, cnt, "expected %d shards, got %d for action %s, shards %s", tc.expected, cnt, tc.action, tc.shards)
+		}
+	})
+	waitForWorkflowState(t, vc, fmt.Sprintf("%s.%s", keyspace, workflowName), binlogdatapb.VReplicationWorkflowState_Running.String())
+
 	for _, targetTab := range targetTabs {
 		catchup(t, targetTab, workflowName, "Reshard")
 	}

go/vt/vtctl/workflow/server.go

@@ -1719,8 +1719,10 @@ func (s *Server) WorkflowUpdate(ctx context.Context, req *vtctldatapb.WorkflowUp
 	span.Annotate("on_ddl", req.TabletRequest.OnDdl)
 	span.Annotate("state", req.TabletRequest.State)
 	span.Annotate("config_overrides", req.TabletRequest.ConfigOverrides)
+	span.Annotate("shards", req.TabletRequest.Shards)
 
 	vx := vexec.NewVExec(req.Keyspace, req.TabletRequest.Workflow, s.ts, s.tmc, s.env.Parser())
+	vx.SetShardSubset(req.TabletRequest.Shards)
 	callback := func(ctx context.Context, tablet *topo.TabletInfo) (*querypb.QueryResult, error) {
 		res, err := s.tmc.UpdateVReplicationWorkflow(ctx, tablet.Tablet, req.TabletRequest)
 		if err != nil {

proto/tabletmanagerdata.proto

@@ -781,7 +781,7 @@ message UpdateVReplicationWorkflowRequest {
   optional TabletSelectionPreference tablet_selection_preference = 4;
   optional binlogdata.OnDDLAction on_ddl = 5;
   optional binlogdata.VReplicationWorkflowState state = 6;
-  reserved 7; // unused, was: repeated string shards
+  repeated string shards = 7;
   map<string, string> config_overrides = 8;
   optional string message = 9;
 }