Fix zeroFields setting when merging slices into maps (#83)

* Fix zeroFields setting when merging slices into maps Signed-off-by: Haytham Abuelfutuh <[email protected]> * PR Feedback Signed-off-by: Haytham Abuelfutuh <[email protected]>
flyteorg · May 4, 2021 · 99f4b20 · 99f4b20
1 parent 9681e7a
commit 99f4b20
Show file tree

Hide file tree

Showing 5 changed files with 113 additions and 2 deletions.
diff --git a/config/tests/accessor_test.go b/config/tests/accessor_test.go
@@ -331,6 +331,57 @@ func TestAccessor_UpdateConfig(t *testing.T) {
 			assert.Equal(t, "default_3", r.OtherItem.ID)
 		})
 
+		t.Run(fmt.Sprintf("[%v] Override default map config", provider(config.Options{}).ID()), func(t *testing.T) {
+			t.Run("Simple", func(t *testing.T) {
+				root := config.NewRootSection()
+				_, err := root.RegisterSection(MyComponentSectionKey, &ItemMap{
+					Items: map[string]Item{
+						"1": {
+							ID:   "default_1",
+							Name: "default_Name",
+						},
+						"2": {
+							ID:   "default_2",
+							Name: "default_2_Name",
+						},
+					},
+				})
+				assert.NoError(t, err)
+
+				v := provider(config.Options{
+					SearchPaths: []string{filepath.Join("testdata", "map_config.yaml")},
+					RootSection: root,
+				})
+
+				assert.NoError(t, v.UpdateConfig(context.TODO()))
+				r := root.GetSection(MyComponentSectionKey).GetConfig().(*ItemMap)
+				assert.Len(t, r.Items, 2)
+				assert.Equal(t, "abc", r.Items["1"].ID)
+			})
+
+			t.Run("NestedMaps", func(t *testing.T) {
+				root := config.NewRootSection()
+				_, err := root.RegisterSection(MyComponentSectionKey, &ItemMap{
+					ItemsMap: map[string]map[string]Item{},
+				})
+				assert.NoError(t, err)
+
+				v := provider(config.Options{
+					SearchPaths: []string{filepath.Join("testdata", "map_config_nested.yaml")},
+					RootSection: root,
+				})
+
+				assert.NoError(t, v.UpdateConfig(context.TODO()))
+				r := root.GetSection(MyComponentSectionKey).GetConfig().(*ItemMap)
+				assert.Len(t, r.ItemsMap, 2)
+				assert.Equal(t, "abc1", r.ItemsMap["itemA"]["itemAa"].ID)
+				assert.Equal(t, "hello world", r.ItemsMap["itemA"]["itemAa"].RandomValue)
+				assert.Equal(t, "abc2", r.ItemsMap["itemB"]["itemBa"].ID)
+				assert.Equal(t, "xyz1", r.ItemsMap["itemA"]["itemAb"].ID)
+				assert.Equal(t, "xyz2", r.ItemsMap["itemB"]["itemBb"].ID)
+			})
+		})
+
 		t.Run(fmt.Sprintf("[%v] Override in Env Var", provider(config.Options{}).ID()), func(t *testing.T) {
 			reg := config.NewRootSection()
 			_, err := reg.RegisterSection(MyComponentSectionKey, &MyComponentConfig{})

diff --git a/config/tests/testdata/map_config.yaml b/config/tests/testdata/map_config.yaml
@@ -0,0 +1,8 @@
+my-component:
+  items:
+    1:
+      id: abc
+      name: "A b c"
+    2:
+      id: xyz
+      name: "x y z"
diff --git a/config/tests/testdata/map_config_nested.yaml b/config/tests/testdata/map_config_nested.yaml
@@ -0,0 +1,17 @@
+my-component:
+  itemsMap:
+    - itemA:
+        - itemAa:
+            id: abc1
+            name: "A b c"
+            randomValue: "hello world"
+        - itemAb:
+            id: xyz1
+            name: "x y z"
+    - itemB:
+        - itemBa:
+            id: abc2
+            name: "A b c"
+        - itemBb:
+            id: xyz2
+            name: "x y z"
diff --git a/config/tests/types_test.go b/config/tests/types_test.go
@@ -28,15 +28,21 @@ type OtherComponentConfig struct {
 }
 
 type Item struct {
-	ID   string `json:"id"`
-	Name string `json:"name"`
+	ID          string `json:"id"`
+	Name        string `json:"name"`
+	RandomValue string `json:"randomValue"`
 }
 
 type ItemArray struct {
 	Items     []Item `json:"items"`
 	OtherItem Item   `json:"otherItem"`
 }
 
+type ItemMap struct {
+	Items    map[string]Item            `json:"items"`
+	ItemsMap map[string]map[string]Item `json:"itemsMap"`
+}
+
 func (MyComponentConfig) GetPFlagSet(prefix string) *pflag.FlagSet {
 	cmdFlags := pflag.NewFlagSet("MyComponentConfig", pflag.ExitOnError)
 	cmdFlags.String(fmt.Sprintf("%v%v", prefix, "str"), "hello world", "life is short")

diff --git a/config/viper/viper.go b/config/viper/viper.go
@@ -165,6 +165,34 @@ func canGetElement(t reflect.Kind) bool {
 	return exists
 }
 
+// sliceToMapHook allows the conversion from slices to maps. This is used as a hack due to the lack of support of case
+// sensitive keys in viper (see: https://github.com/spf13/viper#does-viper-support-case-sensitive-keys). The way we work
+// around that is by filling in fields that should be maps as slices in yaml config files. This hook then takes care of
+// reverting that process.
+func sliceToMapHook(f reflect.Kind, t reflect.Kind, data interface{}) (interface{}, error) {
+	// Only handle slice -> map conversion
+	if f == reflect.Slice && t == reflect.Map {
+		// this will be the target result
+		res := map[interface{}]interface{}{}
+		// It's safe to convert data into a slice since we did the type assertion above.
+		asSlice := data.([]interface{})
+		for _, item := range asSlice {
+			asMap, casted := item.(map[interface{}]interface{})
+			if !casted {
+				return data, nil
+			}
+
+			for key, value := range asMap {
+				res[key] = value
+			}
+		}
+
+		return res, nil
+	}
+
+	return data, nil
+}
+
 // This decoder hook tests types for json unmarshaling capability. If implemented, it uses json unmarshal to build the
 // object. Otherwise, it'll just pass on the original data.
 func jsonUnmarshallerHook(_, to reflect.Type, data interface{}) (interface{}, error) {
@@ -269,6 +297,7 @@ func defaultDecoderConfig(output interface{}, opts ...viperLib.DecoderConfigOpti
 			jsonUnmarshallerHook,
 			mapstructure.StringToTimeDurationHookFunc(),
 			mapstructure.StringToSliceHookFunc(","),
+			sliceToMapHook,
 		),
 		// Empty/zero fields before applying provided values. This avoids potentially undesired/unexpected merging logic.
 		ZeroFields: true,