From 5bc720326a8b9236aa76d5dd74e9fd2e26e312d9 Mon Sep 17 00:00:00 2001
From: Kimmo Lehto <klehto@mirantis.com>
Date: Mon, 30 Sep 2024 15:10:24 +0300
Subject: [PATCH] Make Apply action phase list modifiable before Run()

Signed-off-by: Kimmo Lehto <klehto@mirantis.com>
---
 action/apply.go | 117 ++++++++++++++++++++++++++++--------------------
 cmd/apply.go    |   4 +-
 2 files changed, 72 insertions(+), 49 deletions(-)

diff --git a/action/apply.go b/action/apply.go
index a379ee60..bfe18a89 100644
--- a/action/apply.go
+++ b/action/apply.go
@@ -15,7 +15,7 @@ import (
 	log "github.com/sirupsen/logrus"
 )
 
-type Apply struct {
+type ApplyOptions struct {
 	// Manager is the phase manager
 	Manager *phase.Manager
 	// DisableDowngradeCheck skips the downgrade check
@@ -36,59 +36,80 @@ type Apply struct {
 	ConfigPath string
 }
 
-func (a Apply) Run() error {
-	start := time.Now()
-
-	phase.NoWait = a.NoWait
-	phase.Force = a.Force
+type Apply struct {
+	ApplyOptions
+	Phases phase.Phases
+}
 
+// NewApply creates a new Apply action. The list of phases can be modified via the Phases field, for example:
+//
+//	apply := NewApply(opts)
+//	gatherK0sFacts := &phase.GatherK0sFacts{} // advisable to get the title from the phase itself instead of hardcoding the title.
+//	apply.Phases.InsertBefore(gatherK0sFacts.Title(), &myCustomPhase{}) // insert a custom phase before the GatherK0sFacts phase
+func NewApply(opts ApplyOptions) *Apply {
 	lockPhase := &phase.Lock{}
-
-	a.Manager.AddPhase(
-		&phase.DefaultK0sVersion{},
-		&phase.Connect{},
-		&phase.DetectOS{},
-		lockPhase,
-		&phase.PrepareHosts{},
-		&phase.GatherFacts{},
-		&phase.ValidateHosts{},
-		&phase.GatherK0sFacts{},
-		&phase.ValidateFacts{SkipDowngradeCheck: a.DisableDowngradeCheck},
-
-		// if UploadBinaries: true
-		&phase.DownloadBinaries{}, // downloads k0s binaries to local cache
-		&phase.UploadK0s{},        // uploads k0s binaries to hosts from cache
-
-		// if UploadBinaries: false
-		&phase.DownloadK0s{}, // downloads k0s binaries directly from hosts
-
-		&phase.UploadFiles{},
-		&phase.InstallBinaries{},
-		&phase.PrepareArm{},
-		&phase.ConfigureK0s{},
-		&phase.Restore{
-			RestoreFrom: a.RestoreFrom,
+	unlockPhase := lockPhase.UnlockPhase()
+	apply := &Apply{
+		ApplyOptions: opts,
+		Phases: phase.Phases{
+			&phase.DefaultK0sVersion{},
+			&phase.Connect{},
+			&phase.DetectOS{},
+			lockPhase,
+			&phase.PrepareHosts{},
+			&phase.GatherFacts{},
+			&phase.ValidateHosts{},
+			&phase.GatherK0sFacts{},
+			&phase.ValidateFacts{SkipDowngradeCheck: opts.DisableDowngradeCheck},
+
+			// if UploadBinaries: true
+			&phase.DownloadBinaries{}, // downloads k0s binaries to local cache
+			&phase.UploadK0s{},        // uploads k0s binaries to hosts from cache
+
+			// if UploadBinaries: false
+			&phase.DownloadK0s{}, // downloads k0s binaries directly from hosts
+
+			&phase.UploadFiles{},
+			&phase.InstallBinaries{},
+			&phase.PrepareArm{},
+			&phase.ConfigureK0s{},
+			&phase.Restore{
+				RestoreFrom: opts.RestoreFrom,
+			},
+			&phase.RunHooks{Stage: "before", Action: "apply"},
+			&phase.InitializeK0s{},
+			&phase.InstallControllers{},
+			&phase.InstallWorkers{},
+			&phase.UpgradeControllers{},
+			&phase.UpgradeWorkers{NoDrain: opts.NoDrain},
+			&phase.Reinstall{},
+			&phase.ResetWorkers{NoDrain: opts.NoDrain},
+			&phase.ResetControllers{NoDrain: opts.NoDrain},
+			&phase.RunHooks{Stage: "after", Action: "apply"},
+			unlockPhase,
+			&phase.Disconnect{},
 		},
-		&phase.RunHooks{Stage: "before", Action: "apply"},
-		&phase.InitializeK0s{},
-		&phase.InstallControllers{},
-		&phase.InstallWorkers{},
-		&phase.UpgradeControllers{},
-		&phase.UpgradeWorkers{NoDrain: a.NoDrain},
-		&phase.Reinstall{},
-		&phase.ResetWorkers{NoDrain: a.NoDrain},
-		&phase.ResetControllers{NoDrain: a.NoDrain},
-		&phase.RunHooks{Stage: "after", Action: "apply"},
-	)
+	}
+	if opts.KubeconfigOut != nil {
+		apply.Phases.InsertBefore(unlockPhase.Title(), &phase.GetKubeconfig{APIAddress: opts.KubeconfigAPIAddress})
+	}
 
-	if a.KubeconfigOut != nil {
-		a.Manager.AddPhase(&phase.GetKubeconfig{APIAddress: a.KubeconfigAPIAddress})
+	return apply
+}
+
+// Run the Apply action
+func (a Apply) Run() error {
+	if len(a.Phases) == 0 {
+		// for backwards compatibility with the old Apply struct without NewApply(..)
+		tmpApply := NewApply(a.ApplyOptions)
+		a.Phases = tmpApply.Phases
 	}
+	start := time.Now()
+
+	phase.NoWait = a.NoWait
+	phase.Force = a.Force
 
-	a.Manager.AddPhase(
-		&phase.Unlock{Cancel: lockPhase.Cancel},
-		&phase.Disconnect{},
-	)
+	a.Manager.SetPhases(a.Phases)
 
 	analytics.Client.Publish("apply-start", map[string]interface{}{})
 
diff --git a/cmd/apply.go b/cmd/apply.go
index 0630bc8c..2a302341 100644
--- a/cmd/apply.go
+++ b/cmd/apply.go
@@ -72,7 +72,7 @@ var applyCommand = &cli.Command{
 			kubeconfigOut = out
 		}
 
-		applyAction := action.Apply{
+		applyOpts := action.ApplyOptions{
 			Force:                 ctx.Bool("force"),
 			Manager:               ctx.Context.Value(ctxManagerKey{}).(*phase.Manager),
 			KubeconfigOut:         kubeconfigOut,
@@ -84,6 +84,8 @@ var applyCommand = &cli.Command{
 			ConfigPath:            ctx.String("config"),
 		}
 
+		applyAction := action.NewApply(applyOpts)
+
 		if err := applyAction.Run(); err != nil {
 			return fmt.Errorf("apply failed - log file saved to %s: %w", ctx.Context.Value(ctxLogFileKey{}).(string), err)
 		}