admin: add alter-partition-assignments, list-partition-reassignments

Author: twmb

commands/admin/admin.go

@@ -4,6 +4,7 @@ package admin
 import (
 	"context"
 	"fmt"
+	"sort"
 
 	"github.com/spf13/cobra"
 
@@ -25,6 +26,8 @@ func Command(cl *client.Client) *cobra.Command {
 	cmd.AddCommand(deleteRecordsCommand(cl))
 	cmd.AddCommand(logdirsDescribeCommand(cl))
 	cmd.AddCommand(logdirsAlterReplicasCommand(cl))
+	cmd.AddCommand(alterPartitionAssignments(cl))
+	cmd.AddCommand(listPartitionReassignments(cl))
 
 	return cmd
 }
@@ -118,3 +121,138 @@ To avoid accidental triggers, this command requires a --run flag to run.
 
 	return cmd
 }
+
+func alterPartitionAssignments(cl *client.Client) *cobra.Command {
+	var topicPartReplicas []string
+
+	cmd := &cobra.Command{
+		Use:   "alter-partition-assignments",
+		Short: "Alter partition assignments.",
+		Long: `Alter which brokers partitions are assigned to (Kafka 2.4.0+).
+
+The syntax for each topic is
+
+  topic: 1->2,3,4 ; 2->1,2,3
+
+where the first number is the partition, and -> points to the replicas you
+want to move the partition to. Note that since this contains a >, you likely
+need to quote your input to the flag.
+
+If a replica list is empty for a specific partition, this cancels any active
+reassignment for that partition.
+`,
+		Example: "alter-partition-assignments -t 'foo:1->1,2,3' -t 'bar:2->3,4,5;5->3,4,5'",
+		Run: func(_ *cobra.Command, args []string) {
+			tprs, err := flagutil.ParseTopicPartitionReplicas(topicPartReplicas)
+			out.MaybeDie(err, "unable to parse topic partitions replicas: %v", err)
+
+			req := &kmsg.AlterPartitionAssignmentsRequest{
+				TimeoutMillis: cl.TimeoutMillis(),
+			}
+			for topic, partitions := range tprs {
+				t := kmsg.AlterPartitionAssignmentsRequestTopic{
+					Topic: topic,
+				}
+				for partition, replicas := range partitions {
+					t.Partitions = append(t.Partitions, kmsg.AlterPartitionAssignmentsRequestTopicPartition{
+						Partition: partition,
+						Replicas:  replicas,
+					})
+				}
+				req.Topics = append(req.Topics, t)
+			}
+
+			kresp, err := cl.Client().Request(context.Background(), req)
+			out.MaybeDie(err, "unable to alter partition assignments: %v", err)
+			resp := kresp.(*kmsg.AlterPartitionAssignmentsResponse)
+			if cl.AsJSON() {
+				out.ExitJSON(resp)
+			}
+
+			if resp.ErrorCode != 0 {
+				out.ErrAndMsg(resp.ErrorCode, resp.ErrorMessage)
+				out.Exit()
+			}
+
+			tw := out.BeginTabWrite()
+			defer tw.Flush()
+			for _, topic := range resp.Topics {
+				for _, partition := range topic.Partitions {
+					msg := "OK"
+					if err := kerr.ErrorForCode(partition.ErrorCode); err != nil {
+						msg = err.Error()
+					}
+					detail := ""
+					if partition.ErrorMessage != nil {
+						detail = *partition.ErrorMessage
+					}
+					fmt.Fprintf(tw, "%s\t%d\t%s\t%s\n", topic.Topic, partition.Partition, msg, detail)
+				}
+			}
+		},
+	}
+
+	cmd.Flags().StringArrayVarP(&topicPartReplicas, "topic", "t", nil, "topic, partitions, and replica destinations; repeatable")
+	return cmd
+}
+
+func listPartitionReassignments(cl *client.Client) *cobra.Command {
+	var topicParts []string
+
+	cmd := &cobra.Command{
+		Use:   "list-partition-reassignments",
+		Short: "List partition reassignments.",
+		Long: `List which partitions are currently being reassigned (Kafka 2.4.0+).
+
+The syntax for each topic is
+
+  topic:1,2,3
+
+where the numbers correspond to partitions for a topic.
+
+If no topics are specified, this lists all active reassignments.
+`,
+		Run: func(_ *cobra.Command, args []string) {
+			tps, err := flagutil.ParseTopicPartitions(topicParts)
+			out.MaybeDie(err, "unable to parse topic partitions: %v", err)
+
+			req := &kmsg.ListPartitionReassignmentsRequest{
+				TimeoutMillis: cl.TimeoutMillis(),
+			}
+			for topic, partitions := range tps {
+				req.Topics = append(req.Topics, kmsg.ListPartitionReassignmentsRequestTopic{
+					Topic:      topic,
+					Partitions: partitions,
+				})
+			}
+
+			kresp, err := cl.Client().Request(context.Background(), req)
+			out.MaybeDie(err, "unable to list partition reassignments: %v", err)
+			resp := kresp.(*kmsg.ListPartitionReassignmentsResponse)
+			if cl.AsJSON() {
+				out.ExitJSON(resp)
+			}
+
+			if resp.ErrorCode != 0 {
+				out.ErrAndMsg(resp.ErrorCode, resp.ErrorMessage)
+				out.Exit()
+			}
+
+			tw := out.BeginTabWrite()
+			defer tw.Flush()
+			fmt.Fprint(tw, "TOPIC\tPARTITION\tCURRENT REPLICAS\tADDING\tREMOVING\n")
+			for _, topic := range resp.Topics {
+				for _, p := range topic.Partitions {
+					sort.Slice(p.Replicas, func(i, j int) bool { return p.Replicas[i] < p.Replicas[j] })
+					sort.Slice(p.AddingReplicas, func(i, j int) bool { return p.AddingReplicas[i] < p.AddingReplicas[j] })
+					sort.Slice(p.RemovingReplicas, func(i, j int) bool { return p.RemovingReplicas[i] < p.RemovingReplicas[j] })
+					fmt.Fprintf(tw, "%s\t%d\t%v\t%v\t%v\n", topic.Topic, p.Partition, p.Replicas, p.AddingReplicas, p.RemovingReplicas)
+				}
+			}
+		},
+	}
+
+	cmd.Flags().StringArrayVarP(&topicParts, "topic", "t", nil, "topic and partitions to list partition reassignments for; repeatable")
+
+	return cmd
+}

flagutil/flagutil.go

@@ -33,3 +33,45 @@ func ParseTopicPartitions(list []string) (map[string][]int32, error) {
 	}
 	return tps, nil
 }
+
+// ParseTopicPartitionReplicas parses a list of the following, spaces trimmed:
+//
+//     topic: 4->3,2,1 ; 5->3,2,1
+func ParseTopicPartitionReplicas(list []string) (map[string]map[int32][]int32, error) {
+	tprs := make(map[string]map[int32][]int32)
+	for _, item := range list {
+		tps := strings.SplitN(item, ":", 2)
+		if len(tps) != 2 {
+			return nil, fmt.Errorf("item %q invalid empty topic", item)
+		}
+
+		topic := strings.TrimSpace(tps[0])
+		prs := make(map[int32][]int32)
+		tprs[topic] = prs
+
+		for _, partitionReplicasRaw := range strings.Split(tps[1], ";") {
+			partitionReplicas := strings.SplitN(partitionReplicasRaw, "->", 2)
+			if len(partitionReplicas) != 2 {
+				return nil, fmt.Errorf("item %q invalid partition->replicas bit %q", item, partitionReplicasRaw)
+			}
+			partition, err := strconv.Atoi(strings.TrimSpace(partitionReplicas[0]))
+			if err != nil {
+				return nil, fmt.Errorf("item %q invalid partition in %q", item, partitionReplicasRaw)
+			}
+			p := int32(partition)
+			prs[p] = nil
+			for _, r := range strings.Split(partitionReplicas[1], ",") {
+				r = strings.TrimSpace(r)
+				if len(r) == 0 {
+					continue
+				}
+				replica, err := strconv.Atoi(r)
+				if err != nil {
+					return nil, fmt.Errorf("item %q invalid replica in %q", item, partitionReplicasRaw)
+				}
+				prs[p] = append(prs[p], int32(replica))
+			}
+		}
+	}
+	return tprs, nil
+}