Add workflow metadata to keyspace topo object. Save it currently only…

… for movetables and use the cached information, if available, to only contact primaries of participating shards for all workflow commands

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

go/vt/topo/keyspace.go

@@ -203,7 +203,7 @@ func (ts *Server) GetKeyspace(ctx context.Context, keyspace string) (*KeyspaceIn
 	if err = k.UnmarshalVT(data); err != nil {
 		return nil, vterrors.Wrap(err, "bad keyspace data")
 	}
-
+	log.Infof("GetKeyspace:%v:%v:%v:", keyspace, version, k)
 	return &KeyspaceInfo{
 		keyspace: keyspace,
 		version:  version,

go/vt/topo/keyspace_workflow.go

@@ -0,0 +1,98 @@
+/*
+Copyright 2023 The Vitess Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package topo
+
+import (
+	"context"
+	"fmt"
+
+	"vitess.io/vitess/go/vt/proto/topodata"
+	"vitess.io/vitess/go/vt/proto/vtrpc"
+	"vitess.io/vitess/go/vt/vterrors"
+)
+
+func (ts *Server) SaveWorkflowMetadata(ctx context.Context, targetKeyspace string, wm *topodata.WorkflowMetadata) (err error) {
+	ctx, unlock, lockErr := ts.LockKeyspace(ctx, targetKeyspace, "GetWorkflowMetadata")
+	if lockErr != nil {
+		err = lockErr
+		return err
+	}
+	defer unlock(&err)
+
+	ki, err := ts.GetKeyspace(ctx, targetKeyspace)
+	for _, w := range ki.Workflows {
+		if w.Name == wm.Name {
+			return vterrors.New(vtrpc.Code_ALREADY_EXISTS,
+				fmt.Sprintf("workflow %s already exists in %s", w.Name, ki.KeyspaceName()))
+		}
+	}
+	ki.Workflows = append(ki.Workflows, &topodata.WorkflowMetadata{
+		Name:           wm.Name,
+		Type:           wm.Type,
+		TargetShards:   wm.TargetShards,
+		SourceKeyspace: wm.SourceKeyspace,
+		SourceShards:   wm.SourceShards,
+	})
+	if err := ts.UpdateKeyspace(ctx, ki); err != nil {
+		return err
+	}
+	return nil
+}
+
+func (ts *Server) DeleteWorkflowMetadata(ctx context.Context, targetKeyspace, workflowName string) (err error) {
+	ctx, unlock, lockErr := ts.LockKeyspace(ctx, targetKeyspace, "GetWorkflowMetadata")
+	if lockErr != nil {
+		err = lockErr
+		return err
+	}
+	defer unlock(&err)
+
+	ki, err := ts.GetKeyspace(ctx, targetKeyspace)
+	if err != nil {
+		return err
+	}
+	for i, w := range ki.Workflows {
+		if w.Name == workflowName {
+			ki.Workflows = append(ki.Workflows[:i], ki.Workflows[i+1:]...)
+			if err := ts.UpdateKeyspace(ctx, ki); err != nil {
+				return err
+			}
+			return nil
+		}
+	}
+	return nil
+}
+
+func (ts *Server) GetWorkflowMetadata(ctx context.Context, targetKeyspace, workflowName string) (wm *topodata.WorkflowMetadata, err error) {
+	ctx, unlock, lockErr := ts.LockKeyspace(ctx, targetKeyspace, "GetWorkflowMetadata")
+	if lockErr != nil {
+		err = lockErr
+		return nil, err
+	}
+	defer unlock(&err)
+
+	ki, err := ts.GetKeyspace(ctx, targetKeyspace)
+	if err != nil {
+		return nil, err
+	}
+	for _, w := range ki.Workflows {
+		if w.Name == workflowName {
+			return w, nil
+		}
+	}
+	return nil, nil
+}

go/vt/vtctl/workflow/utils.go

@@ -647,6 +647,40 @@ func areTabletsAvailableToStreamFrom(ctx context.Context, req *vtctldatapb.Workf
 	return nil
 }
 
+func FilterShardsForWorkflow(ctx context.Context, ts *topo.Server, keyspace, workflow string, shards []string) ([]string, error) {
+	log.Infof("Filtering shards for workflow %s in keyspace %s", workflow, keyspace)
+	ks, err := ts.GetKeyspace(ctx, keyspace)
+	if ks.Keyspace == nil || err != nil {
+		return shards, nil // keyspace doesn't exist, so we can't filter shards, should only come here in tests.
+	}
+	wm, err := ts.GetWorkflowMetadata(ctx, keyspace, workflow)
+	if err != nil {
+		return nil, err
+	}
+	if wm == nil {
+		return shards, nil
+	}
+
+	var targetShards []string
+	for _, shard := range shards {
+		for _, s := range wm.TargetShards {
+			if shard == s {
+				targetShards = append(targetShards, shard)
+			}
+		}
+	}
+
+	if len(targetShards) != len(wm.TargetShards) {
+		return nil, fmt.Errorf("target shards %v for workflow %s does not match the shards in the metadata for keyspace %s",
+			wm.TargetShards, workflow, keyspace)
+	}
+
+	if len(shards) != len(targetShards) {
+		log.Infof("workflow has fewer shards: %d than the keyspace: %d", len(targetShards), len(shards))
+	}
+	return targetShards, nil
+}
+
 // LegacyBuildTargets collects MigrationTargets and other metadata (see TargetInfo)
 // from a workflow in the target keyspace. It uses VReplicationExec to get the workflow
 // details rather than the new TabletManager ReadVReplicationWorkflow RPC. This is
@@ -658,7 +692,12 @@ func areTabletsAvailableToStreamFrom(ctx context.Context, req *vtctldatapb.Workf
 //
 // It returns ErrNoStreams if there are no targets found for the workflow.
 func LegacyBuildTargets(ctx context.Context, ts *topo.Server, tmc tmclient.TabletManagerClient, targetKeyspace string, workflow string) (*TargetInfo, error) {
-	targetShards, err := ts.GetShardNames(ctx, targetKeyspace)
+	shards, err := ts.GetShardNames(ctx, targetKeyspace)
+	if err != nil {
+		return nil, err
+	}
+
+	targetShards, err := FilterShardsForWorkflow(ctx, ts, targetKeyspace, workflow, shards)
 	if err != nil {
 		return nil, err
 	}
@@ -687,6 +726,7 @@ func LegacyBuildTargets(ctx context.Context, ts *topo.Server, tmc tmclient.Table
 	// two target shards will have vreplication streams, and the other shards in
 	// the target keyspace will not.
 	for _, targetShard := range targetShards {
+		log.Infof("Checking shard %s for workflow %s", targetShard, workflow)
 		si, err := ts.GetShard(ctx, targetKeyspace, targetShard)
 		if err != nil {
 			return nil, err

go/vt/wrangler/materializer.go

@@ -27,6 +27,8 @@ import (
 	"text/template"
 	"time"
 
+	"vitess.io/vitess/go/vt/proto/topodata"
+
 	"google.golang.org/protobuf/encoding/prototext"
 	"google.golang.org/protobuf/proto"
 
@@ -136,7 +138,6 @@ func (wr *Wrangler) MoveTables(ctx context.Context, workflow, sourceKeyspace, ta
 	//FIXME validate tableSpecs, allTables, excludeTables
 	var tables []string
 	var externalTopo *topo.Server
-
 	if externalCluster != "" { // when the source is an external mysql cluster mounted using the Mount command
 		externalTopo, err = wr.ts.OpenExternalVitessClusterServer(ctx, externalCluster)
 		if err != nil {
@@ -359,6 +360,23 @@ func (wr *Wrangler) MoveTables(ctx context.Context, workflow, sourceKeyspace, ta
 			return fmt.Errorf(msg)
 		}
 	}
+	var participatingShards []string
+	for _, si := range mz.targetShards {
+		participatingShards = append(participatingShards, si.ShardName())
+	}
+
+	wm := &topodata.WorkflowMetadata{
+		Name:           workflow,
+		Type:           "MoveTables",
+		TargetShards:   participatingShards,
+		SourceKeyspace: sourceKeyspace,
+		SourceShards:   participatingShards,
+	}
+	log.Infof("Going to save workflow metadata for workflow %s: %+v", workflow, wm)
+	if err := wr.ts.SaveWorkflowMetadata(ctx, targetKeyspace, wm); err != nil {
+		return err
+	}
+
 	if autoStart {
 		return mz.startStreams(ctx)
 	}

go/vt/wrangler/vexec.go

@@ -220,6 +220,7 @@ func (vx *vexec) exec() (map[*topo.TabletInfo]*querypb.QueryResult, error) {
 		wg.Add(1)
 		go func(ctx context.Context, primary *topo.TabletInfo) {
 			defer wg.Done()
+			log.Infof("Executing query on tablet %s", primary.AliasString())
 			qr, err := vx.planner.exec(ctx, primary.Alias, vx.plannedQuery)
 			if err != nil {
 				allErrors.RecordError(err)
@@ -294,9 +295,15 @@ func (vx *vexec) getPrimaries() error {
 	if len(shards) == 0 {
 		return fmt.Errorf("no shards found in keyspace %s", vx.keyspace)
 	}
+
+	targetShards, err := workflow2.FilterShardsForWorkflow(vx.ctx, vx.wr.ts, vx.keyspace, vx.workflow, shards)
+	if err != nil {
+		return err
+	}
+
 	var allPrimaries []*topo.TabletInfo
 	var primary *topo.TabletInfo
-	for _, shard := range shards {
+	for _, shard := range targetShards {
 		if primary, err = vx.getPrimaryForShard(shard); err != nil {
 			return err
 		}

proto/topodata.proto

@@ -256,6 +256,14 @@ message Shard {
   reserved 5;
 }
 
+message WorkflowMetadata {
+  string name = 1;
+  string type = 2;
+  repeated string target_shards = 4;
+  string source_keyspace = 5;
+  repeated string source_shards = 6;
+}
+
 // A Keyspace contains data about a keyspace.
 message Keyspace {
   // OBSOLETE string sharding_column_name = 1;
@@ -312,6 +320,10 @@ message Keyspace {
   // used for various system metadata that is stored in each
   // tablet's mysqld instance.
   string sidecar_db_name = 10;
+
+  // WorkflowMetadata is the metadata for workflows in this keyspace.
+  repeated WorkflowMetadata workflows = 11;
+
 }
 
 // ShardReplication describes the MySQL replication relationships