update setProgress and setResults

replicatedhq · Jun 14, 2024 · bdc4458 · bdc4458
1 parent 0912496
commit bdc4458
Show file tree

Hide file tree

Showing 3 changed files with 52 additions and 48 deletions.
diff --git a/pkg/preflight/execute.go b/pkg/preflight/execute.go
@@ -1,7 +1,6 @@
 package preflight
 
 import (
-	"encoding/json"
 	"strings"
 	"sync"
 	"time"
@@ -19,7 +18,7 @@ import (
 
 // Execute will Execute the preflights using spec in preflightSpec.
 // This spec should be rendered, no template functions remaining
-func Execute(preflightSpec *troubleshootv1beta2.Preflight, ignorePermissionErrors bool, setProgress, setResults func(b []byte) error) (*types.PreflightResults, error) {
+func Execute(preflightSpec *troubleshootv1beta2.Preflight, ignorePermissionErrors bool, setProgress func(progress map[string]interface{}) error, setResults func(results *types.PreflightResults) error) (*types.PreflightResults, error) {
 	progressChan := make(chan interface{}, 0) // non-zero buffer will result in missed messages
 	defer close(progressChan)
 
@@ -44,26 +43,23 @@ func Execute(preflightSpec *troubleshootv1beta2.Preflight, ignorePermissionError
 				}
 			}
 
-			progress, ok := msg.(preflight.CollectProgress)
+			collectProgress, ok := msg.(preflight.CollectProgress)
 			if !ok {
 				continue
 			}
 
 			// TODO: We need a nice title to display
-			progressBytes, err := json.Marshal(map[string]interface{}{
-				"completedCount": progress.CompletedCount,
-				"totalCount":     progress.TotalCount,
-				"currentName":    progress.CurrentName,
-				"currentStatus":  progress.CurrentStatus,
+			progress := map[string]interface{}{
+				"completedCount": collectProgress.CompletedCount,
+				"totalCount":     collectProgress.TotalCount,
+				"currentName":    collectProgress.CurrentName,
+				"currentStatus":  collectProgress.CurrentStatus,
 				"updatedAt":      time.Now().Format(time.RFC3339),
-			})
-			if err != nil {
-				continue
 			}
 
 			completeMx.Lock()
 			if !isComplete {
-				if err := setProgress(progressBytes); err != nil {
+				if err := setProgress(progress); err != nil {
 					logger.Error(errors.Wrap(err, "failed to set preflight progress"))
 				}
 			}
@@ -88,13 +84,7 @@ func Execute(preflightSpec *troubleshootv1beta2.Preflight, ignorePermissionError
 			})
 		}
 
-		b, err := json.Marshal(uploadPreflightResults)
-		if err != nil {
-			logger.Error(errors.Wrap(err, "failed to marshal preflight results"))
-			return
-		}
-
-		if err := setResults(b); err != nil {
+		if err := setResults(uploadPreflightResults); err != nil {
 			logger.Error(errors.Wrap(err, "failed to set preflight results"))
 			return
 		}

diff --git a/pkg/preflight/preflight.go b/pkg/preflight/preflight.go
@@ -190,10 +190,18 @@ func Run(appID string, appSlug string, sequence int64, isAirgap bool, archiveDir
 				zap.String("appID", appID),
 				zap.Int64("sequence", sequence))
 
-			setProgress := func(b []byte) error {
+			setProgress := func(progress map[string]interface{}) error {
+				b, err := json.Marshal(progress)
+				if err != nil {
+					return errors.Wrap(err, "failed to marshal preflight progress")
+				}
 				return store.GetStore().SetPreflightProgress(appID, sequence, string(b))
 			}
-			setResults := func(b []byte) error {
+			setResults := func(results *types.PreflightResults) error {
+				b, err := json.Marshal(results)
+				if err != nil {
+					return errors.Wrap(err, "failed to marshal preflight results")
+				}
 				return store.GetStore().SetPreflightResults(appID, sequence, b)
 			}
 			uploadPreflightResults, err := Execute(preflight, ignoreRBAC, setProgress, setResults)

diff --git a/pkg/upgradeservice/preflight/preflight.go b/pkg/upgradeservice/preflight/preflight.go
@@ -23,9 +23,10 @@ import (
 	"go.uber.org/zap"
 )
 
+// TODO: change this so that it's not byte arrays, but the real types
 type PreflightData struct {
-	Progress []byte `json:"progress"`
-	Results  []byte `json:"results"`
+	Progress map[string]interface{}  `json:"progress"`
+	Results  *types.PreflightResults `json:"results"`
 }
 
 var PreflightDataFilepath string
@@ -109,13 +110,7 @@ func Run(app *apptypes.App, archiveDir string, sequence int64, registrySettings
 				},
 			}
 
-			b, err := json.Marshal(preflightResults)
-			if err != nil {
-				logger.Error(errors.Wrap(err, "failed to marshal preflight results"))
-				return
-			}
-
-			if err := setPreflightResult(b); err != nil {
+			if err := setPreflightResults(preflightResults); err != nil {
 				logger.Error(errors.Wrap(err, "failed to set preflight results"))
 				return
 			}
@@ -134,7 +129,7 @@ func Run(app *apptypes.App, archiveDir string, sequence int64, registrySettings
 			zap.String("appID", app.ID),
 			zap.Int64("sequence", sequence))
 
-		_, err := preflightpkg.Execute(preflight, ignoreRBAC, setPreflightProgress, setPreflightResult)
+		_, err := preflightpkg.Execute(preflight, ignoreRBAC, setPreflightProgress, setPreflightResults)
 		if err != nil {
 			logger.Error(errors.Wrap(err, "failed to run preflight checks"))
 			return
@@ -150,6 +145,30 @@ func Run(app *apptypes.App, archiveDir string, sequence int64, registrySettings
 	return nil
 }
 
+func setPreflightResults(results *types.PreflightResults) error {
+	preflightData, err := getPreflightData()
+	if err != nil {
+		return errors.Wrap(err, "failed to get preflight data")
+	}
+	preflightData.Results = results
+	if err := setPreflightData(preflightData); err != nil {
+		return errors.Wrap(err, "failed to set preflight results")
+	}
+	return nil
+}
+
+func setPreflightProgress(progress map[string]interface{}) error {
+	preflightData, err := getPreflightData()
+	if err != nil {
+		return errors.Wrap(err, "failed to get preflight data")
+	}
+	preflightData.Progress = progress
+	if err := setPreflightData(preflightData); err != nil {
+		return errors.Wrap(err, "failed to set preflight progress")
+	}
+	return nil
+}
+
 func getPreflightData() (*PreflightData, error) {
 	var preflightData *PreflightData
 	if _, err := os.Stat(PreflightDataFilepath); err != nil {
@@ -169,26 +188,13 @@ func getPreflightData() (*PreflightData, error) {
 	return preflightData, nil
 }
 
-func setPreflightResult(b []byte) error {
-	preflightData, err := getPreflightData()
+func setPreflightData(preflightData *PreflightData) error {
+	b, err := json.Marshal(preflightData)
 	if err != nil {
-		return errors.Wrap(err, "failed to get preflight data")
-	}
-	preflightData.Results = b
-	if err := os.WriteFile(PreflightDataFilepath, b, 0644); err != nil {
-		return errors.Wrap(err, "failed to write preflight results")
-	}
-	return nil
-}
-
-func setPreflightProgress(b []byte) error {
-	preflightData, err := getPreflightData()
-	if err != nil {
-		return errors.Wrap(err, "failed to get preflight data")
+		return errors.Wrap(err, "failed to marshal preflight data")
 	}
-	preflightData.Progress = b
 	if err := os.WriteFile(PreflightDataFilepath, b, 0644); err != nil {
-		return errors.Wrap(err, "failed to write preflight progress")
+		return errors.Wrap(err, "failed to write preflight data")
 	}
 	return nil
 }