From 5ad91a98e21f71c52ee65f5e1349dca1d5f8c0f9 Mon Sep 17 00:00:00 2001
From: Zachary Blasczyk <zachary.blasczyk@wandb.com>
Date: Tue, 3 Sep 2024 15:42:56 -0500
Subject: [PATCH] fix: Merge specs deeply

---
 controllers/weightsandbiases_controller.go | 21 +++++++++----
 pkg/wandb/spec/spec.go                     | 35 +++++++++++++---------
 2 files changed, 37 insertions(+), 19 deletions(-)

diff --git a/controllers/weightsandbiases_controller.go b/controllers/weightsandbiases_controller.go
index 5a15a06..5617058 100644
--- a/controllers/weightsandbiases_controller.go
+++ b/controllers/weightsandbiases_controller.go
@@ -159,23 +159,34 @@ func (r *WeightsAndBiasesReconciler) Reconcile(ctx context.Context, req ctrl.Req
 		log.Info("Initial desired spec", "spec", desiredSpec.SensitiveValuesMasked())
 	}
 
-	// First takes precedence
-	desiredSpec.Merge(crdSpec)
+	if err := desiredSpec.Merge(crdSpec); err != nil {
+		log.Error(err, "Failed to merge CRD spec into desired spec")
+		return ctrlqueue.RequeueWithError(err)
+	}
 	if r.Debug {
 		log.Info("Desired spec after merging crdSpec", "spec", desiredSpec.SensitiveValuesMasked())
 	}
 
-	desiredSpec.Merge(userInputSpec)
+	if err := desiredSpec.Merge(userInputSpec); err != nil {
+		log.Error(err, "Failed to merge user input spec into desired spec")
+		return ctrlqueue.RequeueWithError(err)
+	}
 	if r.Debug {
 		log.Info("Desired spec after merging userInputSpec", "spec", desiredSpec.SensitiveValuesMasked())
 	}
 
-	desiredSpec.Merge(deployerSpec)
+	if err := desiredSpec.Merge(deployerSpec); err != nil {
+		log.Error(err, "Failed to merge deployer spec into desired spec")
+		return ctrlqueue.RequeueWithError(err)
+	}
 	if r.Debug {
 		log.Info("Desired spec after merging deployerSpec", "spec", desiredSpec.SensitiveValuesMasked())
 	}
 
-	desiredSpec.Merge(operator.Defaults(wandb, r.Scheme))
+	if err := desiredSpec.Merge(operator.Defaults(wandb, r.Scheme)); err != nil {
+		log.Error(err, "Failed to merge operator defaults into desired spec")
+		return ctrlqueue.RequeueWithError(err)
+	}
 	if r.Debug {
 		log.Info("Desired spec after merging operator defaults", "spec", desiredSpec.SensitiveValuesMasked())
 	}
diff --git a/pkg/wandb/spec/spec.go b/pkg/wandb/spec/spec.go
index 7117493..a9a8f01 100644
--- a/pkg/wandb/spec/spec.go
+++ b/pkg/wandb/spec/spec.go
@@ -88,26 +88,30 @@ func (s *Spec) IsEqual(spec *Spec) bool {
 	return isReleaseEqual && isValuesEqual && isMetadataEqual
 }
 
-func (s *Spec) mergeConfig(values Values) (err error) {
+func (s *Spec) mergeConfig(newValues map[string]interface{}) error {
 	if s.Values == nil {
-		s.Values = values
-		return nil
+		s.Values = make(map[string]interface{})
 	}
-	mergedValues, err := s.Values.Merge(values)
-	if err != nil {
-		return err
+
+	for key, newValue := range newValues {
+		if existingValue, exists := s.Values[key]; exists {
+			if _, ok := existingValue.(map[string]interface{}); ok {
+				if newMap, ok := newValue.(map[string]interface{}); ok {
+					if err := s.mergeConfig(newMap); err != nil {
+						return err
+					}
+					continue
+				}
+			}
+		}
+		s.Values[key] = newValue
 	}
-	s.Values = mergedValues
 	return nil
 }
 
-// Merge merges the given spec into the current spec. The currently spec will
-// take precedence. This means, if the current spec has a value for a key, the
-// new spec will not override it. Likewise, charts & metadata cannot be merged,
-// so if a chart/metadata value is set, the one passed in will be ignored.
-func (s *Spec) Merge(spec *Spec) {
+func (s *Spec) Merge(spec *Spec) error {
 	if spec == nil {
-		return
+		return nil
 	}
 
 	if s.Metadata == nil {
@@ -117,8 +121,11 @@ func (s *Spec) Merge(spec *Spec) {
 		s.Chart = spec.Chart
 	}
 	if spec.Values != nil {
-		s.mergeConfig(spec.Values)
+		if err := s.mergeConfig(spec.Values); err != nil {
+			return err
+		}
 	}
+	return nil
 }
 
 func (s *Spec) SensitiveValuesMasked() *Spec {