Repository: vitessio/vitess
Pull request #6533 (2020-08-05 01:39:39)

Workflow: Add List-All Command
This PR adds a `list-all` subcommand to the workflow command. Choosing `list-all` will return a list of all active workflows for the given keyspace.

The justification for this is that consumers of Vitess (for instance `vitess-operator`) will need the ability to discover active workflows, and not be required to hold onto workflow state themselves. Being able to discover workflows means that an operator can easily uncover whether resharding is currently ongoing, and act accordingly.


Patch:
diff --git a/go/vt/discovery/healthcheck.go b/go/vt/discovery/healthcheck.go
index d0ca06a706..d1f4fefa69 100644
--- a/go/vt/discovery/healthcheck.go
+++ b/go/vt/discovery/healthcheck.go
@@ -313,16 +313,11 @@ func NewHealthCheck(ctx context.Context, retryDelay, healthCheckTimeout time.Dur
 // It does not block on making connection.
 // name is an optional tag for the tablet, e.g. an alternative address.
 func (hc *HealthCheckImpl) AddTablet(tablet *topodata.Tablet) {
+	log.Infof("Calling AddTablet for tablet: %v", tablet)
 	// check whether we should really add this tablet
 	if !hc.isIncluded(tablet) {
 		return
 	}
-	// check whether grpc port is present on tablet, if not return
-	if tablet.PortMap["grpc"] == 0 {
-		return
-	}
-
-	log.Infof("Adding tablet to healthcheck: %v", tablet)
 	hc.mu.Lock()
 	defer hc.mu.Unlock()
 	if hc.healthByAlias == nil {
@@ -376,12 +371,11 @@ func (hc *HealthCheckImpl) RemoveTablet(tablet *topodata.Tablet) {
 
 // ReplaceTablet removes the old tablet and adds the new tablet.
 func (hc *HealthCheckImpl) ReplaceTablet(old, new *topodata.Tablet) {
-	hc.RemoveTablet(old)
+	hc.deleteTablet(old)
 	hc.AddTablet(new)
 }
 
 func (hc *HealthCheckImpl) deleteTablet(tablet *topodata.Tablet) {
-	log.Infof("Removing tablet from healthcheck: %v", tablet)
 	hc.mu.Lock()
 	defer hc.mu.Unlock()
 
diff --git a/go/vt/discovery/tablet_health_check.go b/go/vt/discovery/tablet_health_check.go
index 8a3beecbfd..79a8acd40b 100644
--- a/go/vt/discovery/tablet_health_check.go
+++ b/go/vt/discovery/tablet_health_check.go
@@ -107,9 +107,8 @@ func (thc *tabletHealthCheck) setServingState(serving bool, reason string) {
 	if !thc.loggedServingState || (serving != thc.Serving) {
 		// Emit the log from a separate goroutine to avoid holding
 		// the th lock while logging is happening
-		log.Infof("HealthCheckUpdate(Serving State): tablet: %v serving %v => %v for %v/%v (%v) reason: %s",
+		go log.Infof("HealthCheckUpdate(Serving State): tablet: %v serving => %v for %v/%v (%v) reason: %s",
 			topotools.TabletIdent(thc.Tablet),
-			thc.Serving,
 			serving,
 			thc.Tablet.GetKeyspace(),
 			thc.Tablet.GetShard(),
diff --git a/go/vt/vtctl/vtctl.go b/go/vt/vtctl/vtctl.go
index 0be416ff66..0b97041207 100644
--- a/go/vt/vtctl/vtctl.go
+++ b/go/vt/vtctl/vtctl.go
@@ -478,7 +478,7 @@ var commands = []commandGroup{
 		"Workflow", []command{
 			{"Workflow", commandWorkflow,
 				"<ks.workflow> <action> --dry-run",
-				"Start/Stop/Delete Workflow on all target tablets in workflow. Example: Workflow merchant.morders Start",
+				"Start/Stop/Delete/List/ListAll Workflow on all target tablets in workflow. Example: Workflow merchant.morders Start",
 			},
 		},
 	},
@@ -2946,23 +2946,28 @@ func commandWorkflow(ctx context.Context, wr *wrangler.Wrangler, subFlags *flag.
 		return err
 	}
 	if subFlags.NArg() != 2 {
-		return fmt.Errorf("usage: Workflow --dry-run keyspace.workflow start/stop/delete")
+		return fmt.Errorf("usage: Workflow --dry-run keyspace.workflow start/stop/delete/list/list-all")
 	}
-	keyspace, workflow, err := splitKeyspaceWorkflow(subFlags.Arg(0))
-	if err != nil {
-		return err
+	keyspace := subFlags.Arg(0)
+	action := strings.ToLower(subFlags.Arg(1))
+	var workflow string
+	var err error
+	if action != "listall" {
+		keyspace, workflow, err = splitKeyspaceWorkflow(subFlags.Arg(0))
+		if err != nil {
+			return err
+		}
 	}
 	_, err = wr.TopoServer().GetKeyspace(ctx, keyspace)
 	if err != nil {
 		wr.Logger().Errorf("Keyspace %s not found", keyspace)
 	}
-	action := subFlags.Arg(1)
 
 	results, err := wr.WorkflowAction(ctx, workflow, keyspace, action, *dryRun)
 	if err != nil {
 		return err
 	}
-	if action == "list" {
+	if action == "list" || action == "listall" {
 		return nil
 	}
 	if len(results) == 0 {
diff --git a/go/vt/vtgate/vschema_manager.go b/go/vt/vtgate/vschema_manager.go
index 3d810af55f..87d0e0eb6f 100644
--- a/go/vt/vtgate/vschema_manager.go
+++ b/go/vt/vtgate/vschema_manager.go
@@ -71,7 +71,6 @@ func (vm *VSchemaManager) watchSrvVSchema(ctx context.Context, cell string) {
 		// we don't know the real value. In this case, we want
 		// to use the previous value if it was set, or an
 		// empty vschema if it wasn't.
-		log.Infof("Received vschema update")
 		switch {
 		case err == nil:
 			// Good case, we can try to save that value.
@@ -80,7 +79,6 @@ func (vm *VSchemaManager) watchSrvVSchema(ctx context.Context, cell string) {
 			// Otherwise, keep what we already had before.
 			v = nil
 		default:
-			log.Errorf("SrvVschema watch error: %v", err)
 			// Watch error, increment our counters.
 			if vschemaCounters != nil {
 				vschemaCounters.Add("WatchError", 1)
diff --git a/go/vt/wrangler/vexec.go b/go/vt/wrangler/vexec.go
index e978811e63..a2bbbdba91 100644
--- a/go/vt/wrangler/vexec.go
+++ b/go/vt/wrangler/vexec.go
@@ -24,6 +24,8 @@ import (
 	"sync"
 	"time"
 
+	"k8s.io/apimachinery/pkg/util/sets"
+
 	"vitess.io/vitess/go/vt/log"
 
 	"github.com/gogo/protobuf/proto"
@@ -188,10 +190,13 @@ func (vx *vexec) getMasterForShard(shard string) (*topo.TabletInfo, error) {
 	return master, nil
 }
 
-// WorkflowAction can start/stop/delete or list strams in _vt.vreplication on all masters in the target keyspace of the workflow
+// WorkflowAction can start/stop/delete or list streams in _vt.vreplication on all masters in the target keyspace of the workflow.
 func (wr *Wrangler) WorkflowAction(ctx context.Context, workflow, keyspace, action string, dryRun bool) (map[*topo.TabletInfo]*sqltypes.Result, error) {
 	if action == "list" {
 		return nil, wr.listStreams(ctx, workflow, keyspace)
+	} else if action == "listall" {
+		_, err := wr.ListAllWorkflows(ctx, keyspace)
+		return nil, err
 	}
 	results, err := wr.execWorkflowAction(ctx, workflow, keyspace, action, dryRun)
 	retResults := make(map[*topo.TabletInfo]*sqltypes.Result)
@@ -344,6 +349,31 @@ func (wr *Wrangler) getStreams(ctx context.Context, workflow, keyspace string) (
 	return &rsr, nil
 }
 
+// ListAllWorkflows will return a list of all active workflows for the given keyspace.
+func (wr *Wrangler) ListAllWorkflows(ctx context.Context, keyspace string) ([]string, error) {
+	query := "select distinct workflow from _vt.vreplication where state <> 'Stopped'"
+	results, err := wr.runVexec(ctx, "", keyspace, query, false)
+	if err != nil {
+		return nil, err
+	}
+	workflowsSet := sets.NewString()
+	for _, result := range results {
+		if len(result.Rows) == 0 {
+			continue
+		}
+		qr := sqltypes.Proto3ToResult(result)
+		for _, row := range qr.Rows {
+			for _, value := range row {
+				// Even though we query for distinct, we must de-dup because we query per master tablet.
+				workflowsSet.Insert(value.ToString())
+			}
+		}
+	}
+	workflows := workflowsSet.List()
+	wr.printWorkflowList(workflows)
+	return workflows, nil
+}
+
 func (wr *Wrangler) listStreams(ctx context.Context, workflow, keyspace string) error {
 	replStatus, err := wr.getStreams(ctx, workflow, keyspace)
 	if err != nil {
@@ -376,6 +406,11 @@ func dumpStreamListAsJSON(replStatus *replicationStatusResult, wr *Wrangler) err
 	return nil
 }
 
+func (wr *Wrangler) printWorkflowList(workflows []string) {
+	list := strings.Join(workflows, ", ")
+	wr.Logger().Printf("Workflows: %v", list)
+}
+
 func (wr *Wrangler) getCopyState(ctx context.Context, tablet *topo.TabletInfo, id int64) ([]copyState, error) {
 	var cs []copyState
 	query := fmt.Sprintf(`select table_name, lastpk from _vt.copy_state where vrepl_id = %d`, id)
diff --git a/go/vt/wrangler/vexec_plan.go b/go/vt/wrangler/vexec_plan.go
index a61a986d67..6499d71ebd 100644
--- a/go/vt/wrangler/vexec_plan.go
+++ b/go/vt/wrangler/vexec_plan.go
@@ -118,7 +118,7 @@ func (vx *vexec) addDefaultWheres(where *sqlparser.Where) *sqlparser.Where {
 			}
 		}
 	}
-	if !hasWorkflow {
+	if !hasWorkflow && vx.workflow != "" {
 		expr := &sqlparser.ComparisonExpr{
 			Left:     &sqlparser.ColName{Name: sqlparser.NewColIdent("workflow")},
 			Operator: sqlparser.EqualStr,


Test patch:
diff --git a/go/test/endtoend/cellalias/cell_alias_test.go b/go/test/endtoend/cellalias/cell_alias_test.go
index 8479254767..d85f9cdd4d 100644
--- a/go/test/endtoend/cellalias/cell_alias_test.go
+++ b/go/test/endtoend/cellalias/cell_alias_test.go
@@ -118,11 +118,11 @@ func TestMain(m *testing.M) {
 			return 1, err
 		}
 
-		shard1Master = localCluster.NewVttabletInstance("master", 0, cell1)
+		shard1Master = localCluster.NewVttabletInstance("master", 0, "")
 		shard1Replica = localCluster.NewVttabletInstance("replica", 0, cell2)
 		shard1Rdonly = localCluster.NewVttabletInstance("rdonly", 0, cell2)
 
-		shard2Master = localCluster.NewVttabletInstance("master", 0, cell1)
+		shard2Master = localCluster.NewVttabletInstance("master", 0, "")
 		shard2Replica = localCluster.NewVttabletInstance("replica", 0, cell2)
 		shard2Rdonly = localCluster.NewVttabletInstance("rdonly", 0, cell2)
 
@@ -256,8 +256,6 @@ func TestAlias(t *testing.T) {
 	vtgateInstance := localCluster.NewVtgateInstance()
 	vtgateInstance.CellsToWatch = allCells
 	vtgateInstance.TabletTypesToWait = "MASTER,REPLICA"
-	// Use legacy gateway. There's a separate test for tabletgateway in go/test/endtoend/tabletgateway/cellalias/cell_alias_test.go
-	vtgateInstance.GatewayImplementation = "discoverygateway"
 	err = vtgateInstance.Setup()
 	require.NoError(t, err)
 
@@ -311,12 +309,14 @@ func TestAddAliasWhileVtgateUp(t *testing.T) {
 
 	vtgateInstance := localCluster.NewVtgateInstance()
 	vtgateInstance.CellsToWatch = allCells
-	// only MASTER is in vtgate's "cell", other tablet types are not visible because they are in the other cell
-	vtgateInstance.TabletTypesToWait = "MASTER"
+	vtgateInstance.TabletTypesToWait = "MASTER,REPLICA,RDONLY"
+	vtgateInstance.GatewayImplementation = "discoverygateway"
 	err = vtgateInstance.Setup()
 	require.NoError(t, err)
 	defer vtgateInstance.TearDown()
 
+	waitTillAllTabletsAreHealthyInVtgate(t, *vtgateInstance, shard1.Name, shard2.Name)
+
 	// since replica and rdonly tablets of all shards in cell2, the last 2 assertion is expected to fail
 	testQueriesOnTabletType(t, "master", vtgateInstance.GrpcPort, false)
 	testQueriesOnTabletType(t, "replica", vtgateInstance.GrpcPort, true)
diff --git a/go/test/endtoend/vtgate/buffer/buffer_test.go b/go/test/endtoend/vtgate/buffer/buffer_test.go
index 89e3435cd9..83a4b556af 100644
--- a/go/test/endtoend/vtgate/buffer/buffer_test.go
+++ b/go/test/endtoend/vtgate/buffer/buffer_test.go
@@ -205,9 +205,7 @@ func createCluster() (*cluster.LocalProcessCluster, int) {
 		// Long timeout in case failover is slow.
 		"-buffer_window", "10m",
 		"-buffer_max_failover_duration", "10m",
-		"-buffer_min_time_between_failovers", "20m",
-		// Use legacy gateway. tabletgateway test is at go/test/endtoend/tabletgateway/buffer/buffer_test.go
-		"-gateway_implementation", "discoverygateway"}
+		"-buffer_min_time_between_failovers", "20m"}
 
 	// Start vtgate
 	if err := clusterInstance.StartVtgate(); err != nil {
diff --git a/go/vt/discovery/healthcheck_test.go b/go/vt/discovery/healthcheck_test.go
index e1bc0325cb..88db7b9fec 100644
--- a/go/vt/discovery/healthcheck_test.go
+++ b/go/vt/discovery/healthcheck_test.go
@@ -62,7 +62,10 @@ func TestHealthCheck(t *testing.T) {
 	hc := createTestHc(ts)
 	// close healthcheck
 	defer hc.Close()
-	tablet := createTestTablet(0, "cell", "a")
+	tablet := topo.NewTablet(0, "cell", "a")
+	tablet.Keyspace = "k"
+	tablet.Shard = "s"
+	tablet.PortMap["vt"] = 1
 	tablet.Type = topodatapb.TabletType_REPLICA
 	input := make(chan *querypb.StreamHealthResponse)
 	conn := createFakeConn(tablet, input)
@@ -200,7 +203,8 @@ func TestHealthCheckStreamError(t *testing.T) {
 	hc := createTestHc(ts)
 	defer hc.Close()
 
-	tablet := createTestTablet(0, "cell", "a")
+	tablet := topo.NewTablet(0, "cell", "a")
+	tablet.PortMap["vt"] = 1
 	input := make(chan *querypb.StreamHealthResponse)
 	resultChan := hc.Subscribe()
 	fc := createFakeConn(tablet, input)
@@ -210,7 +214,7 @@ func TestHealthCheckStreamError(t *testing.T) {
 	// Immediately after AddTablet() there will be the first notification.
 	want := &TabletHealth{
 		Tablet:              tablet,
-		Target:              &querypb.Target{Keyspace: "k", Shard: "s"},
+		Target:              &querypb.Target{},
 		Serving:             false,
 		MasterTermStartTime: 0,
 	}
@@ -256,7 +260,8 @@ func TestHealthCheckVerifiesTabletAlias(t *testing.T) {
 	hc := createTestHc(ts)
 	defer hc.Close()
 
-	tablet := createTestTablet(0, "cell", "a")
+	tablet := topo.NewTablet(1, "cell", "a")
+	tablet.PortMap["vt"] = 1
 	input := make(chan *querypb.StreamHealthResponse, 1)
 	fc := createFakeConn(tablet, input)
 	resultChan := hc.Subscribe()
@@ -266,7 +271,7 @@ func TestHealthCheckVerifiesTabletAlias(t *testing.T) {
 	// Immediately after AddTablet() there will be the first notification.
 	want := &TabletHealth{
 		Tablet:              tablet,
-		Target:              &querypb.Target{Keyspace: "k", Shard: "s"},
+		Target:              &querypb.Target{},
 		Serving:             false,
 		MasterTermStartTime: 0,
 	}
@@ -297,7 +302,8 @@ func TestHealthCheckVerifiesTabletAlias(t *testing.T) {
 func TestHealthCheckCloseWaitsForGoRoutines(t *testing.T) {
 	ts := memorytopo.NewServer("cell")
 	hc := createTestHc(ts)
-	tablet := createTestTablet(0, "cell", "a")
+	tablet := topo.NewTablet(0, "cell", "a")
+	tablet.PortMap["vt"] = 1
 	input := make(chan *querypb.StreamHealthResponse, 1)
 	createFakeConn(tablet, input)
 	resultChan := hc.Subscribe()
@@ -307,7 +313,7 @@ func TestHealthCheckCloseWaitsForGoRoutines(t *testing.T) {
 	// Immediately after AddTablet() there will be the first notification.
 	want := &TabletHealth{
 		Tablet:              tablet,
-		Target:              &querypb.Target{Keyspace: "k", Shard: "s"},
+		Target:              &querypb.Target{},
 		Serving:             false,
 		MasterTermStartTime: 0,
 	}
@@ -358,7 +364,9 @@ func TestHealthCheckTimeout(t *testing.T) {
 	hc := createTestHc(ts)
 	hc.healthCheckTimeout = 500 * time.Millisecond
 	defer hc.Close()
-	tablet := createTestTablet(0, "cell", "a")
+
+	tablet := topo.NewTablet(0, "cell", "a")
+	tablet.PortMap["vt"] = 1
 	input := make(chan *querypb.StreamHealthResponse)
 	fc := createFakeConn(tablet, input)
 	resultChan := hc.Subscribe()
@@ -366,7 +374,7 @@ func TestHealthCheckTimeout(t *testing.T) {
 	// Immediately after AddTablet() there will be the first notification.
 	want := &TabletHealth{
 		Tablet:              tablet,
-		Target:              &querypb.Target{Keyspace: "k", Shard: "s"},
+		Target:              &querypb.Target{},
 		Serving:             false,
 		MasterTermStartTime: 0,
 	}
@@ -424,7 +432,10 @@ func TestGetHealthyTablets(t *testing.T) {
 	ts := memorytopo.NewServer("cell")
 	hc := createTestHc(ts)
 	defer hc.Close()
-	tablet := createTestTablet(0, "cell", "a")
+	tablet := topo.NewTablet(0, "cell", "a")
+	tablet.Keyspace = "k"
+	tablet.Shard = "s"
+	tablet.PortMap["vt"] = 1
 	tablet.Type = topodatapb.TabletType_REPLICA
 	input := make(chan *querypb.StreamHealthResponse)
 	createFakeConn(tablet, input)
@@ -496,11 +507,16 @@ func TestGetHealthyTablets(t *testing.T) {
 	mustMatch(t, want, a, "unexpected result")
 
 	// add a second tablet
-	tablet2 := createTestTablet(11, "cell", "host2")
+	tablet2 := topo.NewTablet(11, "cell", "host2")
+	tablet2.Keyspace = "k"
+	tablet2.Shard = "s"
+	tablet2.PortMap["vt"] = 2
 	tablet2.Type = topodatapb.TabletType_REPLICA
 	input2 := make(chan *querypb.StreamHealthResponse)
 	createFakeConn(tablet2, input2)
+	t.Logf(`createFakeConn({Host: "a", PortMap: {"vt": 2}}, c)`)
 	hc.AddTablet(tablet2)
+	t.Logf(`hc = HealthCheck(); hc.AddTablet({Host: "a", PortMap: {"vt": 1}}, "")`)
 	// there will be a first result, get and discard it
 	<-resultChan
 
@@ -621,7 +637,10 @@ func TestAliases(t *testing.T) {
 	defer ts.DeleteCellsAlias(context.Background(), "region2")
 
 	// add a tablet as replica in diff cell, same region
-	tablet := createTestTablet(1, "cell1", "host3")
+	tablet := topo.NewTablet(1, "cell1", "host3")
+	tablet.Keyspace = "k"
+	tablet.Shard = "s"
+	tablet.PortMap["vt"] = 1
 	tablet.Type = topodatapb.TabletType_REPLICA
 	input := make(chan *querypb.StreamHealthResponse)
 	fc := createFakeConn(tablet, input)
@@ -669,7 +688,10 @@ func TestAliases(t *testing.T) {
 	mustMatch(t, want, a, "Wrong TabletHealth data")
 
 	// add another tablet in a diff cell, diff region
-	tablet2 := createTestTablet(2, "cell2", "host4")
+	tablet2 := topo.NewTablet(2, "cell2", "host4")
+	tablet2.Keyspace = "k"
+	tablet2.Shard = "s"
+	tablet2.PortMap["vt"] = 2
 	tablet2.Type = topodatapb.TabletType_REPLICA
 	input2 := make(chan *querypb.StreamHealthResponse)
 	fc = createFakeConn(tablet2, input2)
@@ -689,26 +711,6 @@ func TestAliases(t *testing.T) {
 	mustMatch(t, want, a, "Wrong TabletHealth data")
 }
 
-func TestHealthCheckChecksGrpcPort(t *testing.T) {
-	ts := memorytopo.NewServer("cell")
-	hc := createTestHc(ts)
-	defer hc.Close()
-
-	tablet := createTestTablet(0, "cell", "a")
-	tablet.PortMap["grpc"] = 0
-	resultChan := hc.Subscribe()
-
-	// AddTablet should not add the tablet because port is 0
-	hc.AddTablet(tablet)
-
-	select {
-	case result := <-resultChan:
-		assert.Nil(t, result, "healthCheck received result: %v", result)
-	case <-time.After(2 * time.Millisecond):
-		// No response after timeout. Success.
-	}
-}
-
 func TestTemplate(t *testing.T) {
 	tablet := topo.NewTablet(0, "cell", "a")
 	ts := []*TabletHealth{
@@ -734,6 +736,7 @@ func TestTemplate(t *testing.T) {
 }
 
 func TestDebugURLFormatting(t *testing.T) {
+
 	//log error
 	if err2 := flag.Set("tablet_url_template", "https://{{.GetHostNameLevel 0}}.bastion.{{.Tablet.Alias.Cell}}.corp"); err2 != nil {
 		log.Errorf("flag.Set(\"tablet_url_template\", \"https://{{.GetHostNameLevel 0}}.bastion.{{.Tablet.Alias.Cell}}.corp\") failed : %v", err2)
@@ -870,15 +873,6 @@ func createFixedHealthConn(tablet *topodatapb.Tablet, fixedResult *querypb.Strea
 	return conn
 }
 
-func createTestTablet(uid uint32, cell, host string) *topodatapb.Tablet {
-	tablet := topo.NewTablet(uid, cell, host)
-	tablet.PortMap["vt"] = 1
-	tablet.PortMap["grpc"] = 2
-	tablet.Keyspace = "k"
-	tablet.Shard = "s"
-	return tablet
-}
-
 var mustMatch = utils.MustMatchFn(
 	[]interface{}{ // types with unexported fields
 		TabletHealth{},
diff --git a/go/vt/wrangler/vexec_test.go b/go/vt/wrangler/vexec_test.go
index f0a15e0841..754bb2f041 100644
--- a/go/vt/wrangler/vexec_test.go
+++ b/go/vt/wrangler/vexec_test.go
@@ -292,6 +292,20 @@ will be run on the following streams in keyspace target for workflow wrWorkflow:
 	require.Equal(t, dryRunResult, logger.String())
 }
 
+func TestWorkflowListAll(t *testing.T) {
+	ctx := context.Background()
+	keyspace := "target"
+	workflow := "wrWorkflow"
+	env := newWranglerTestEnv([]string{"0"}, []string{"-80", "80-"}, "", nil)
+	defer env.close()
+	logger := logutil.NewMemoryLogger()
+	wr := New(logger, env.topoServ, env.tmc)
+
+	workflows, err := wr.ListAllWorkflows(ctx, keyspace)
+	require.Nil(t, err)
+	require.Equal(t, []string{workflow}, workflows)
+}
+
 func TestVExecValidations(t *testing.T) {
 	ctx := context.Background()
 	workflow := "wf"
diff --git a/go/vt/wrangler/wrangler_env_test.go b/go/vt/wrangler/wrangler_env_test.go
index fd2c9af4c3..92cd904e88 100644
--- a/go/vt/wrangler/wrangler_env_test.go
+++ b/go/vt/wrangler/wrangler_env_test.go
@@ -149,7 +149,6 @@ func newWranglerTestEnv(sourceShards, targetShards []string, query string, posit
 			fmt.Sprintf("1|%v|pos||0|Running|vt_target|1234|", bls),
 		)
 		env.tmc.setVRResults(master.tablet, "select id, source, pos, stop_pos, max_replication_lag, state, db_name, time_updated, message from _vt.vreplication where db_name = 'vt_target' and workflow = 'wrWorkflow'", result)
-		//"
 		env.tmc.setVRResults(
 			master.tablet,
 			"select source, pos from _vt.vreplication where db_name='vt_target' and workflow='wrWorkflow'",
@@ -159,6 +158,12 @@ func newWranglerTestEnv(sourceShards, targetShards []string, query string, posit
 				posRows...,
 			),
 		)
+		result = sqltypes.MakeTestResult(sqltypes.MakeTestFields(
+			"workflow",
+			"varchar"),
+			"wrWorkflow",
+		)
+		env.tmc.setVRResults(master.tablet, "select distinct workflow from _vt.vreplication where state != 'Stopped' and db_name = 'vt_target'", result)
 
 		result = sqltypes.MakeTestResult(sqltypes.MakeTestFields(
 			"table|lastpk",