From 0acc02624136c1e8a546592ef2e90be74e38f8cd Mon Sep 17 00:00:00 2001
From: Kobi Samoray <kobi.samoray@broadcom.com>
Date: Thu, 10 Oct 2024 16:45:08 +0200
Subject: [PATCH] Stage host upgrade in vLCM when flagged to do so

Signed-off-by: Kobi Samoray <kobi.samoray@broadcom.com>
---
 nsxt/resource_nsxt_upgrade_run.go | 23 +++++++++++++++++------
 1 file changed, 17 insertions(+), 6 deletions(-)

diff --git a/nsxt/resource_nsxt_upgrade_run.go b/nsxt/resource_nsxt_upgrade_run.go
index fdfc27e36..a291f895e 100644
--- a/nsxt/resource_nsxt_upgrade_run.go
+++ b/nsxt/resource_nsxt_upgrade_run.go
@@ -415,7 +415,8 @@ func upgradeRunCreateOrUpdate(d *schema.ResourceData, m interface{}) error {
 	}
 
 	log.Printf("[INFO] Updating UpgradeUnitGroup and UpgradePlanSetting.")
-	err = prepareUpgrade(upgradeClientSet, d, targetVersion)
+	var hasVLCM bool
+	err = prepareUpgrade(upgradeClientSet, d, targetVersion, &hasVLCM)
 	if err != nil {
 		return handleCreateError("NsxtUpgradeRun", id, err)
 	}
@@ -430,7 +431,7 @@ func upgradeRunCreateOrUpdate(d *schema.ResourceData, m interface{}) error {
 		}
 	}
 
-	err = runUpgrade(upgradeClientSet, getPartialUpgradeMap(d, targetVersion), targetVersion, finalizeUpgrade)
+	err = runUpgrade(upgradeClientSet, getPartialUpgradeMap(d, targetVersion), targetVersion, hasVLCM, finalizeUpgrade)
 	if err != nil {
 		return handleCreateError("NsxtUpgradeRun", id, err)
 	}
@@ -441,7 +442,7 @@ func upgradeRunCreateOrUpdate(d *schema.ResourceData, m interface{}) error {
 	return resourceNsxtUpgradeRunRead(d, m)
 }
 
-func prepareUpgrade(upgradeClientSet *upgradeClientSet, d *schema.ResourceData, targetVersion string) error {
+func prepareUpgrade(upgradeClientSet *upgradeClientSet, d *schema.ResourceData, targetVersion string, hasVLCM *bool) error {
 	for _, component := range getUpgradeComponentList(targetVersion) {
 		// Customize MP upgrade is not allowed
 		if component == mpUpgradeGroup || component == finalizeUpgradeGroup {
@@ -487,7 +488,7 @@ func prepareUpgrade(upgradeClientSet *upgradeClientSet, d *schema.ResourceData,
 			return err
 		}
 
-		err = updateUpgradeUnitGroups(upgradeClientSet, d, component, preResetGroupList)
+		err = updateUpgradeUnitGroups(upgradeClientSet, d, component, preResetGroupList, hasVLCM)
 		if err != nil {
 			return err
 		}
@@ -591,8 +592,9 @@ func waitUpgradeForStatus(upgradeClientSet *upgradeClientSet, component *string,
 	return nil
 }
 
-func updateUpgradeUnitGroups(upgradeClientSet *upgradeClientSet, d *schema.ResourceData, component string, preResetGroupList model.UpgradeUnitGroupListResult) error {
+func updateUpgradeUnitGroups(upgradeClientSet *upgradeClientSet, d *schema.ResourceData, component string, preResetGroupList model.UpgradeUnitGroupListResult, hasVLCM *bool) error {
 	isBefore := false
+	*hasVLCM = false
 	getReorderAfterReq := func(id string) model.ReorderRequest {
 		return model.ReorderRequest{
 			Id:       &id,
@@ -686,6 +688,9 @@ func updateUpgradeUnitGroups(upgradeClientSet *upgradeClientSet, d *schema.Resou
 			}
 
 			if upgradeMode != "" {
+				if upgradeMode == "stage_in_vlcm" {
+					*hasVLCM = true
+				}
 				upgradeModeKey := "upgrade_mode"
 				extendConfig = util.KeyValuePairsReplaceOrAppend(extendConfig, upgradeModeKey, upgradeMode)
 			}
@@ -778,7 +783,7 @@ func updateComponentUpgradePlanSetting(settingClient plan.SettingsClient, d *sch
 	return err
 }
 
-func runUpgrade(upgradeClientSet *upgradeClientSet, partialUpgradeMap map[string]bool, targetVersion string, finalizeUpgrade bool) error {
+func runUpgrade(upgradeClientSet *upgradeClientSet, partialUpgradeMap map[string]bool, targetVersion string, hasVLCM, finalizeUpgrade bool) error {
 	partialUpgradeExist := false
 	prevComponent := ""
 	for _, c := range getUpgradeComponentList(targetVersion) {
@@ -828,6 +833,12 @@ func runUpgrade(upgradeClientSet *upgradeClientSet, partialUpgradeMap map[string
 			completeLog = fmt.Sprintf("[INFO] %s upgrade is partially completed.", component)
 		}
 		//#nosec G601 Ignore implicit memory aliasing in for loop temporarily
+		if hasVLCM && component == hostUpgradeGroup {
+			err = upgradeClientSet.PlanClient.Stageupgrade(&component)
+			if err != nil {
+				return err
+			}
+		}
 		err = upgradeClientSet.PlanClient.Upgrade(&component)
 		if err != nil {
 			return err