Supporting wildcards for RebuildShardGraph and RebuildKeyspaceGraph.

LGTM Mike.

go/cmd/vtctl/vtctl.go

@@ -135,7 +135,7 @@ var commands = []commandGroup{
 	commandGroup{
 		"Shards", []command{
 			command{"RebuildShardGraph", commandRebuildShardGraph,
-				"<zk shard path> (/zk/global/vt/keyspaces/<keyspace>/shards/<shard>)",
+				"<zk shard path> ... (/zk/global/vt/keyspaces/<keyspace>/shards/<shard>)",
 				"Rebuild the replication graph and shard serving data in zk. This may trigger an update to all connected clients."},
 			command{"ReparentShard", commandReparentShard,
 				"[-force] [-leave-master-read-only] <zk shard path> <zk tablet path>",
@@ -160,7 +160,7 @@ var commands = []commandGroup{
 				"[-force] <zk keyspaces path>/<name>",
 				"e.g. CreateKeyspace /zk/global/vt/keyspaces/my_keyspace"},
 			command{"RebuildKeyspaceGraph", commandRebuildKeyspaceGraph,
-				"<zk keyspace path> (/zk/global/vt/keyspaces/<keyspace>)",
+				"<zk keyspace path> ... (/zk/global/vt/keyspaces/<keyspace>)",
 				"Rebuild the serving data for all shards in this keyspace. This may trigger an update to all connected clients."},
 			command{"ValidateKeyspace", commandValidateKeyspace,
 				"[-ping-tablets] <zk keyspace path> (/zk/global/vt/keyspaces/<keyspace>)",
@@ -176,7 +176,7 @@ var commands = []commandGroup{
 				"[-max-staleness=<duration>] <zk action path> ... (/zk/global/vt/keyspaces/<keyspace>/shards/<shard>/action)",
 				"List any queued actions that are considered stale."},
 			command{"PruneActionLogs", commandPruneActionLogs,
-				"[-keep-count=<count to keep>] <zk actionlog path>",
+				"[-keep-count=<count to keep>] <zk actionlog path> ...",
 				"e.g. PruneActionLogs -keep-count=10 /zk/global/vt/keyspaces/my_keyspace/shards/0/actionlog\n" +
 					"Removes older actionlog entries until at most <count to keep> are left."},
 			command{"WaitForAction", commandWaitForAction,
@@ -933,10 +933,24 @@ func commandExecuteHook(wrangler *wr.Wrangler, subFlags *flag.FlagSet, args []st
 
 func commandRebuildShardGraph(wrangler *wr.Wrangler, subFlags *flag.FlagSet, args []string) (string, error) {
 	subFlags.Parse(args)
-	if subFlags.NArg() != 1 {
-		relog.Fatal("action RebuildShardGraph requires <zk shard path>")
+	if subFlags.NArg() == 0 {
+		relog.Fatal("action RebuildShardGraph requires at least one <zk shard path>")
 	}
-	return wrangler.RebuildShardGraph(subFlags.Arg(0))
+
+	zkPaths, err := zk.ResolveWildcards(wrangler.ZkConn(), subFlags.Args())
+	if err != nil {
+		return "", err
+	}
+	if len(zkPaths) == 0 {
+		return "", nil
+	}
+
+	for _, zkPath := range zkPaths {
+		if err := wrangler.RebuildShardGraph(zkPath); err != nil {
+			return "", err
+		}
+	}
+	return "", nil
 }
 
 func commandReparentShard(wrangler *wr.Wrangler, subFlags *flag.FlagSet, args []string) (string, error) {
@@ -1014,10 +1028,24 @@ func commandCreateKeyspace(wrangler *wr.Wrangler, subFlags *flag.FlagSet, args [
 
 func commandRebuildKeyspaceGraph(wrangler *wr.Wrangler, subFlags *flag.FlagSet, args []string) (string, error) {
 	subFlags.Parse(args)
-	if subFlags.NArg() != 1 {
-		relog.Fatal("action RebuildKeyspaceGraph requires <zk keyspace path>")
+	if subFlags.NArg() == 0 {
+		relog.Fatal("action RebuildKeyspaceGraph requires at least one <zk keyspace path>")
 	}
-	return wrangler.RebuildKeyspaceGraph(subFlags.Arg(0))
+
+	zkPaths, err := zk.ResolveWildcards(wrangler.ZkConn(), subFlags.Args())
+	if err != nil {
+		return "", err
+	}
+	if len(zkPaths) == 0 {
+		return "", nil
+	}
+
+	for _, zkPath := range zkPaths {
+		if err := wrangler.RebuildKeyspaceGraph(zkPath); err != nil {
+			return "", err
+		}
+	}
+	return "", nil
 }
 
 func commandValidateKeyspace(wrangler *wr.Wrangler, subFlags *flag.FlagSet, args []string) (string, error) {

go/vt/wrangler/rebuild.go

@@ -20,16 +20,16 @@ import (
 
 // Rebuild the serving and replication rollup data data while locking
 // out other changes.
-func (wr *Wrangler) RebuildShardGraph(zkShardPath string) (actionPath string, err error) {
+func (wr *Wrangler) RebuildShardGraph(zkShardPath string) error {
 	tm.MustBeShardPath(zkShardPath)
-	actionPath, err = wr.ai.RebuildShard(zkShardPath)
+	actionPath, err := wr.ai.RebuildShard(zkShardPath)
 	if err != nil {
-		return
+		return err
 	}
 
 	// Make sure two of these don't get scheduled at the same time.
 	if err = wr.obtainActionLock(actionPath); err != nil {
-		return "", err
+		return err
 	}
 
 	rebuildErr := wr.rebuildShard(zkShardPath, false)
@@ -38,9 +38,9 @@ func (wr *Wrangler) RebuildShardGraph(zkShardPath string) (actionPath string, er
 		if err != nil {
 			relog.Warning("handleActionError failed: %v", err)
 		}
-		return actionPath, rebuildErr
+		return rebuildErr
 	}
-	return
+	return err
 }
 
 // Update shard file with new master, replicas, etc.
@@ -177,16 +177,16 @@ func (wr *Wrangler) rebuildShardSrvGraph(zkShardPath string, shardInfo *tm.Shard
 }
 
 // Rebuild the serving graph data while locking out other changes.
-func (wr *Wrangler) RebuildKeyspaceGraph(zkKeyspacePath string) (actionPath string, err error) {
+func (wr *Wrangler) RebuildKeyspaceGraph(zkKeyspacePath string) error {
 	tm.MustBeKeyspacePath(zkKeyspacePath)
-	actionPath, err = wr.ai.RebuildKeyspace(zkKeyspacePath)
+	actionPath, err := wr.ai.RebuildKeyspace(zkKeyspacePath)
 	if err != nil {
-		return
+		return err
 	}
 
 	// Make sure two of these don't get scheduled at the same time.
 	if err = wr.obtainActionLock(actionPath); err != nil {
-		return "", err
+		return err
 	}
 
 	rebuildErr := wr.rebuildKeyspace(zkKeyspacePath)
@@ -195,9 +195,9 @@ func (wr *Wrangler) RebuildKeyspaceGraph(zkKeyspacePath string) (actionPath stri
 		if err != nil {
 			relog.Warning("handleActionError failed: %v", err)
 		}
-		return actionPath, rebuildErr
+		return rebuildErr
 	}
-	return
+	return err
 }
 
 // This function should only be used with an action lock on the shard - otherwise the
@@ -222,7 +222,7 @@ func (wr *Wrangler) rebuildKeyspace(zkKeyspacePath string) error {
 		zkShardPath := tm.ShardPath(vtRoot, keyspace, shardName)
 		wg.Add(1)
 		go func() {
-			if _, err := wr.RebuildShardGraph(zkShardPath); err != nil {
+			if err := wr.RebuildShardGraph(zkShardPath); err != nil {
 				relog.Error("RebuildShardGraph failed: %v %v", zkShardPath, err)
 				mu.Lock()
 				rebuildErr = fmt.Errorf("RebuildShardGraph failed on some shards")
@@ -368,8 +368,7 @@ func (wr *Wrangler) RebuildReplicationGraph(zkVtPaths []string, keyspaces []stri
 		wg.Add(1)
 		go func(keyspacePath string) {
 			defer wg.Done()
-			_, err := wr.RebuildKeyspaceGraph(keyspacePath)
-			if err != nil {
+			if err := wr.RebuildKeyspaceGraph(keyspacePath); err != nil {
 				mu.Lock()
 				hasErr = true
 				mu.Unlock()

go/vt/wrangler/wrangler.go

@@ -100,7 +100,7 @@ func (wr *Wrangler) ChangeType(zkTabletPath string, dbType tm.TabletType, force
 	}
 
 	if rebuildRequired {
-		if _, err := wr.RebuildShardGraph(shardToRebuild); err != nil {
+		if err := wr.RebuildShardGraph(shardToRebuild); err != nil {
 			return err
 		}
 	}
@@ -263,5 +263,5 @@ func (wr *Wrangler) Scrap(zkTabletPath string, force, skipRebuild bool) (actionP
 	}
 
 	// and rebuild the original shard / keyspace
-	return wr.RebuildShardGraph(ti.ShardPath())
+	return "", wr.RebuildShardGraph(ti.ShardPath())
 }

test/sharded.py

@@ -97,10 +97,9 @@ def run_test_sharding():
   shard_1_master.init_tablet( 'master',  'test_keyspace', '8000000000000000-0000000000000000', key_start='8000000000000000')
   shard_1_replica.init_tablet('replica', 'test_keyspace', '8000000000000000-0000000000000000', key_start='8000000000000000')
 
-  utils.run_vtctl('RebuildShardGraph /zk/global/vt/keyspaces/test_keyspace/shards/0000000000000000-8000000000000000')
-  utils.run_vtctl('RebuildShardGraph /zk/global/vt/keyspaces/test_keyspace/shards/8000000000000000-0000000000000000')
+  utils.run_vtctl('RebuildShardGraph /zk/global/vt/keyspaces/test_keyspace/shards/*', auto_log=True)
 
-  utils.run_vtctl('RebuildKeyspaceGraph /zk/global/vt/keyspaces/test_keyspace')
+  utils.run_vtctl('RebuildKeyspaceGraph /zk/global/vt/keyspaces/*', auto_log=True)
 
   # run checks now before we start the tablets
   utils.zk_check()