From 81febdc179b8fc420dffae46c0a4f1e58bf6822f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Micha=C5=82=20Leszczy=C5=84ski?= <2000michal@wp.pl>
Date: Tue, 27 Aug 2024 19:56:26 +0200
Subject: [PATCH] feat(command/restore): allow for setting
 --stream-to-all-replicas flag

---
 pkg/command/restore/cmd.go | 34 +++++++++++++++++++++-------------
 1 file changed, 21 insertions(+), 13 deletions(-)

diff --git a/pkg/command/restore/cmd.go b/pkg/command/restore/cmd.go
index 5e8ab8bf4..0d1a400a9 100644
--- a/pkg/command/restore/cmd.go
+++ b/pkg/command/restore/cmd.go
@@ -23,18 +23,19 @@ type command struct {
 	flag.TaskBase
 	client *managerclient.Client
 
-	cluster       string
-	location      []string
-	keyspace      []string
-	snapshotTag   string
-	batchSize     int
-	parallel      int
-	tableParallel int
-	restoreSchema bool
-	restoreTables bool
-	unpinAgentCPU bool
-	dryRun        bool
-	showTables    bool
+	cluster             string
+	location            []string
+	keyspace            []string
+	snapshotTag         string
+	batchSize           int
+	parallel            int
+	tableParallel       int
+	restoreSchema       bool
+	restoreTables       bool
+	unpinAgentCPU       bool
+	streamToAllReplicas bool
+	dryRun              bool
+	showTables          bool
 }
 
 func NewCommand(client *managerclient.Client) *cobra.Command {
@@ -84,6 +85,7 @@ func (cmd *command) init() {
 	w.Unwrap().BoolVar(&cmd.restoreSchema, "restore-schema", false, "")
 	w.Unwrap().BoolVar(&cmd.restoreTables, "restore-tables", false, "")
 	w.Unwrap().BoolVar(&cmd.unpinAgentCPU, "unpin-agent-cpu", false, "")
+	w.Unwrap().BoolVar(&cmd.streamToAllReplicas, "stream-to-all-replicas", false, "")
 	w.Unwrap().BoolVar(&cmd.dryRun, "dry-run", false, "")
 	w.Unwrap().BoolVar(&cmd.showTables, "show-tables", false, "")
 }
@@ -170,7 +172,13 @@ func (cmd *command) run(args []string) error {
 		props["unpin_agent_cpu"] = cmd.unpinAgentCPU
 		ok = true
 	}
-
+	if cmd.Flag("stream-to-all-replicas").Changed {
+		if cmd.Update() {
+			return wrapper("stream-to-all-replicas")
+		}
+		props["stream_to_all_replicas"] = cmd.streamToAllReplicas
+		ok = true
+	}
 	if cmd.dryRun {
 		res, err := cmd.client.GetRestoreTarget(cmd.Context(), cmd.cluster, task)
 		if err != nil {