From 505b1ace38fd15f8b111248201cc1478de53cb1b Mon Sep 17 00:00:00 2001
From: Hidde Beydals <hiddeco@users.noreply.github.com>
Date: Wed, 24 Apr 2024 09:18:33 +0200
Subject: [PATCH] feat(cli): add `--wait` for `kargo promote` (#1890)

Signed-off-by: Hidde Beydals <hidde@hhh.computer>
---
 internal/cli/cmd/promote/promote.go | 66 +++++++++++++++++++++++++++++
 1 file changed, 66 insertions(+)

diff --git a/internal/cli/cmd/promote/promote.go b/internal/cli/cmd/promote/promote.go
index f5c6567b09..6946629fba 100644
--- a/internal/cli/cmd/promote/promote.go
+++ b/internal/cli/cmd/promote/promote.go
@@ -7,9 +7,11 @@ import (
 
 	"connectrpc.com/connect"
 	"github.com/spf13/cobra"
+	"golang.org/x/sync/errgroup"
 	"k8s.io/cli-runtime/pkg/genericclioptions"
 	"k8s.io/cli-runtime/pkg/genericiooptions"
 
+	kargoapi "github.com/akuity/kargo/api/v1alpha1"
 	"github.com/akuity/kargo/internal/cli/client"
 	"github.com/akuity/kargo/internal/cli/config"
 	"github.com/akuity/kargo/internal/cli/io"
@@ -17,6 +19,7 @@ import (
 	"github.com/akuity/kargo/internal/cli/option"
 	"github.com/akuity/kargo/internal/cli/templates"
 	v1alpha1 "github.com/akuity/kargo/pkg/api/service/v1alpha1"
+	"github.com/akuity/kargo/pkg/api/service/v1alpha1/svcv1alpha1connect"
 )
 
 type promotionOptions struct {
@@ -31,6 +34,7 @@ type promotionOptions struct {
 	FreightAlias  string
 	Stage         string
 	SubscribersOf string
+	Wait          bool
 }
 
 func NewCommand(cfg config.CLIConfig, streams genericiooptions.IOStreams) *cobra.Command {
@@ -117,6 +121,7 @@ func (o *promotionOptions) addFlags(cmd *cobra.Command) {
 			option.StageFlag,
 		),
 	)
+	option.Wait(cmd.Flags(), &o.Wait, false, "Wait for the promotion(s) to complete.")
 
 	cmd.MarkFlagsOneRequired(option.FreightFlag, option.FreightAliasFlag)
 	cmd.MarkFlagsMutuallyExclusive(option.FreightFlag, option.FreightAliasFlag)
@@ -177,6 +182,11 @@ func (o *promotionOptions) run(ctx context.Context) error {
 		if err != nil {
 			return fmt.Errorf("promote stage: %w", err)
 		}
+		if o.Wait {
+			if err = waitForPromotion(ctx, kargoSvcCli, res.Msg.GetPromotion()); err != nil {
+				return fmt.Errorf("wait for promotion: %w", err)
+			}
+		}
 		_ = printer.PrintObj(res.Msg.GetPromotion(), o.IOStreams.Out)
 		return nil
 	case o.SubscribersOf != "":
@@ -194,6 +204,11 @@ func (o *promotionOptions) run(ctx context.Context) error {
 		if err != nil {
 			return fmt.Errorf("promote stage subscribers: %w", err)
 		}
+		if o.Wait {
+			if err = waitForPromotions(ctx, kargoSvcCli, res.Msg.GetPromotions()...); err != nil {
+				return fmt.Errorf("wait for promotions: %w", err)
+			}
+		}
 		for _, p := range res.Msg.GetPromotions() {
 			_ = printer.PrintObj(p, o.IOStreams.Out)
 		}
@@ -201,3 +216,54 @@ func (o *promotionOptions) run(ctx context.Context) error {
 	}
 	return nil
 }
+
+func waitForPromotions(
+	ctx context.Context,
+	kargoSvcCli svcv1alpha1connect.KargoServiceClient,
+	p ...*kargoapi.Promotion,
+) error {
+	g, ctx := errgroup.WithContext(ctx)
+	for _, promo := range p {
+		promo := promo
+		g.Go(func() error {
+			return waitForPromotion(ctx, kargoSvcCli, promo)
+		})
+	}
+	return g.Wait()
+}
+
+func waitForPromotion(
+	ctx context.Context,
+	kargoSvcCli svcv1alpha1connect.KargoServiceClient,
+	p *kargoapi.Promotion,
+) error {
+	if p == nil || p.Status.Phase.IsTerminal() {
+		// No need to wait for a promotion that is already terminal.
+		return nil
+	}
+
+	res, err := kargoSvcCli.WatchPromotion(ctx, connect.NewRequest(&v1alpha1.WatchPromotionRequest{
+		Project: p.Namespace,
+		Name:    p.Name,
+	}))
+	if err != nil {
+		return fmt.Errorf("watch promotion: %w", err)
+	}
+	defer func() {
+		if conn, connErr := res.Conn(); connErr == nil {
+			_ = conn.CloseRequest()
+		}
+	}()
+	for {
+		if !res.Receive() {
+			if err = res.Err(); err != nil {
+				return fmt.Errorf("watch promotion: %w", err)
+			}
+			return errors.New("unexpected end of watch stream")
+		}
+		msg := res.Msg()
+		if msg.GetPromotion().Status.Phase.IsTerminal() {
+			return nil
+		}
+	}
+}