diff --git a/flytestdlib/config/tests/accessor_test.go b/flytestdlib/config/tests/accessor_test.go index 899938be644..842dac9afe7 100644 --- a/flytestdlib/config/tests/accessor_test.go +++ b/flytestdlib/config/tests/accessor_test.go @@ -89,6 +89,16 @@ type ComplexType struct { type ComplexTypeArray []ComplexType +type Plugin struct { + Name string `json:"name"` + Annotations map[string]string `json:"annotations"` +} + +type DottedKeysConfig struct { + Annotations map[string]string `json:"annotations"` + Plugins []Plugin `json:"plugins"` +} + type ConfigWithLists struct { ListOfStuff []ComplexType `json:"list"` StringValue string `json:"string-val"` @@ -403,6 +413,27 @@ func TestAccessor_UpdateConfig(t *testing.T) { assert.Equal(t, "xyz1", r.ItemsMap["itemA"]["itemAb"].ID) assert.Equal(t, "xyz2", r.ItemsMap["itemB"]["itemBb"].ID) }) + + t.Run("DottedKeysEndToEnd", func(t *testing.T) { + root := config.NewRootSection() + _, err := root.RegisterSection("dotted-keys", &DottedKeysConfig{}) + assert.NoError(t, err) + + v := provider(config.Options{ + SearchPaths: []string{filepath.Join("testdata", "dotted_keys_config.yaml")}, + RootSection: root, + }) + + assert.NoError(t, v.UpdateConfig(context.TODO())) + r := root.GetSection("dotted-keys").GetConfig().(*DottedKeysConfig) + + assert.Equal(t, "false", r.Annotations["cluster-autoscaler.kubernetes.io/safe-to-evict"]) + assert.Equal(t, "true", r.Annotations["sidecar.istio.io/inject"]) + + assert.Len(t, r.Plugins, 1) + assert.Equal(t, "KeepMyCase", r.Plugins[0].Name) + assert.Equal(t, "/etc/plugin", r.Plugins[0].Annotations["config.path"]) + }) }) t.Run(fmt.Sprintf("[%v] Override in Env Var", provider(config.Options{}).ID()), func(t *testing.T) { diff --git a/flytestdlib/config/tests/testdata/dotted_keys_config.yaml b/flytestdlib/config/tests/testdata/dotted_keys_config.yaml new file mode 100644 index 00000000000..d2ffeee831f --- /dev/null +++ b/flytestdlib/config/tests/testdata/dotted_keys_config.yaml @@ -0,0 +1,8 @@ +dotted-keys: + annotations: + cluster-autoscaler.kubernetes.io/safe-to-evict: "false" + sidecar.istio.io/inject: "true" + plugins: + - name: KeepMyCase + annotations: + config.path: /etc/plugin diff --git a/flytestdlib/config/viper/viper.go b/flytestdlib/config/viper/viper.go index ad3c461860b..8440776b8ae 100644 --- a/flytestdlib/config/viper/viper.go +++ b/flytestdlib/config/viper/viper.go @@ -289,6 +289,7 @@ func (v viperAccessor) parseViperConfig(root config.Section) error { } restoreCaseSensitiveArrayKeys(settings, rawSettings) + restoreDottedMapKeys(settings, rawSettings) } return v.parseViperConfigRecursive(root, settings) @@ -324,6 +325,84 @@ func restoreCaseSensitiveArrayKeys(viperData, rawData map[string]interface{}) { } } +// restoreDottedMapKeys undoes viper's splitting of dotted YAML keys. +// +// - Viper hard-codes "." as its key delimiter and splits every key on it, +// so a YAML leaf key like "test.annotation" arrives in viperData as the +// nested map {"test": {"annotation": ...}}. +// - rawData is the authoritative source for what the user wrote. +// - We walk rawData, detect dotted keys, drop the nested skeleton from +// viperData, and reinsert the value under its original dotted key. +func restoreDottedMapKeys(viperData, rawData map[string]interface{}) { + for rawKey, rawVal := range rawData { + if strings.Contains(rawKey, keyDelim) { + // Drop the nested skeleton viper built from this dotted key, then + // reinsert the raw value under the original key. Lowercase the + // path because viper lowercases all keys. + pruneSplitPath(viperData, strings.Split(strings.ToLower(rawKey), keyDelim)) + viperData[rawKey] = rawVal + continue + } + + // rawVal may itself be a map containing dotted keys deeper down + // (e.g. annotations: {test.annotation: ...}); recurse to fix them. + viperKey, ok := findViperKey(viperData, rawKey) + if !ok { + continue + } + if rawMap, ok := rawVal.(map[string]interface{}); ok { + if viperMap, ok := viperData[viperKey].(map[string]interface{}); ok { + restoreDottedMapKeys(viperMap, rawMap) + } + } + } +} + +// findViperKey returns the key in viperData that case-insensitively matches +// rawKey. Viper lowercases all keys internally, but raw YAML preserves the +// original casing, so we must scan. +func findViperKey(viperData map[string]interface{}, rawKey string) (string, bool) { + for k := range viperData { + if strings.EqualFold(k, rawKey) { + return k, true + } + } + return "", false +} + +// pruneSplitPath removes the nested-map skeleton viper created when it split a +// dotted key on ".". It walks the path top-down, deleting only nodes that +// have no remaining children, so sibling keys that share a prefix are left +// intact (e.g. removing "a.b.c" must not drop "a.b.d"). +func pruneSplitPath(m map[string]interface{}, parts []string) { + if len(parts) == 0 { + return + } + head := parts[0] + child, ok := m[head] + if !ok { + // head not in the map m, directly return. + return + } + if len(parts) == 1 { + // Leaf of the split path — this is the value viper produced from the + // dotted key. Drop it; the caller will reinsert the raw value. + delete(m, head) + return + } + childMap, ok := child.(map[string]interface{}) + if !ok { + // Path diverges from what viper would have produced; leave alone. + return + } + pruneSplitPath(childMap, parts[1:]) + // Drop the parent if recursion emptied it — keeps the cleanup tight so we + // don't leave behind empty intermediate maps from the split. + if len(childMap) == 0 { + delete(m, head) + } +} + func (v viperAccessor) parseViperConfigRecursive(root config.Section, settings interface{}) error { errs := stdLibErrs.ErrorCollection{} var mine interface{} diff --git a/flytestdlib/config/viper/viper_test.go b/flytestdlib/config/viper/viper_test.go index 881d4436f2c..ecbbbacd1f7 100644 --- a/flytestdlib/config/viper/viper_test.go +++ b/flytestdlib/config/viper/viper_test.go @@ -52,3 +52,171 @@ func Test_stringToByteArray(t *testing.T) { assert.NotEqual(t, []byte("hello"), res) }) } + +func Test_restoreDottedMapKeys(t *testing.T) { + t.Run("siblings sharing dotted prefix", func(t *testing.T) { + viper := map[string]interface{}{ + "annotations": map[string]interface{}{ + "test": map[string]interface{}{ + "annotation": "true", + "other": "false", + }, + }, + } + raw := map[string]interface{}{ + "annotations": map[string]interface{}{ + "test.annotation": "true", + "test.other": "false", + }, + } + + restoreDottedMapKeys(viper, raw) + + assert.Equal(t, map[string]interface{}{ + "annotations": map[string]interface{}{ + "test.annotation": "true", + "test.other": "false", + }, + }, viper) + }) + + t.Run("case-insensitive recursion into nested map", func(t *testing.T) { + viper := map[string]interface{}{ + "annotations": map[string]interface{}{ + "test": map[string]interface{}{"annotation": "true"}, + }, + } + raw := map[string]interface{}{ + "Annotations": map[string]interface{}{ + "test.annotation": "true", + }, + } + + restoreDottedMapKeys(viper, raw) + + assert.Equal(t, map[string]interface{}{ + "annotations": map[string]interface{}{ + "test.annotation": "true", + }, + }, viper) + }) + + t.Run("dotted key with map value", func(t *testing.T) { + viper := map[string]interface{}{ + "a": map[string]interface{}{ + "b": map[string]interface{}{ + "c": map[string]interface{}{"d": "v"}, + }, + }, + } + raw := map[string]interface{}{ + "a.b": map[string]interface{}{ + "c.d": "v", + }, + } + + restoreDottedMapKeys(viper, raw) + + assert.Equal(t, map[string]interface{}{ + "a.b": map[string]interface{}{ + "c.d": "v", + }, + }, viper) + }) + + t.Run("dotted and non-dotted siblings at top level", func(t *testing.T) { + viper := map[string]interface{}{ + "a": map[string]interface{}{ + "b": "from-dotted", + "x": "from-non-dotted", + }, + } + raw := map[string]interface{}{ + "a.b": "from-dotted", + "a": map[string]interface{}{"x": "from-non-dotted"}, + } + + restoreDottedMapKeys(viper, raw) + + assert.Equal(t, map[string]interface{}{ + "a.b": "from-dotted", + "a": map[string]interface{}{"x": "from-non-dotted"}, + }, viper) + }) + + // This test case covers the usecases mentioned in issue #6166 + t.Run("k8s annotation key", func(t *testing.T) { + viper := map[string]interface{}{ + "default-annotations": map[string]interface{}{ + "cluster-autoscaler": map[string]interface{}{ + "kubernetes": map[string]interface{}{ + "io/safe-to-evict": "false", + }, + }, + }, + } + raw := map[string]interface{}{ + "default-annotations": map[string]interface{}{ + "cluster-autoscaler.kubernetes.io/safe-to-evict": "false", + }, + } + + restoreDottedMapKeys(viper, raw) + + assert.Equal(t, map[string]interface{}{ + "default-annotations": map[string]interface{}{ + "cluster-autoscaler.kubernetes.io/safe-to-evict": "false", + }, + }, viper) + }) +} + +func Test_pruneSplitPath(t *testing.T) { + t.Run("preserves sibling under shared prefix", func(t *testing.T) { + m := map[string]interface{}{ + "test": map[string]interface{}{ + "annotation": "v1", + "other": "v2", + }, + } + + pruneSplitPath(m, []string{"test", "annotation"}) + + assert.Equal(t, map[string]interface{}{ + "test": map[string]interface{}{"other": "v2"}, + }, m) + }) + + t.Run("collapses entire chain when emptied", func(t *testing.T) { + m := map[string]interface{}{ + "a": map[string]interface{}{ + "b": map[string]interface{}{"c": "v"}, + }, + } + + pruneSplitPath(m, []string{"a", "b", "c"}) + + assert.Equal(t, map[string]interface{}{}, m) + }) + + t.Run("path diverges from viper structure", func(t *testing.T) { + m := map[string]interface{}{"a": "scalar"} + + pruneSplitPath(m, []string{"a", "b"}) + + assert.Equal(t, map[string]interface{}{"a": "scalar"}, m) + }) +} + +func Test_findViperKey(t *testing.T) { + t.Run("case-insensitive match", func(t *testing.T) { + k, ok := findViperKey(map[string]interface{}{"foo": 1}, "Foo") + assert.True(t, ok) + assert.Equal(t, "foo", k) + }) + + t.Run("not found", func(t *testing.T) { + _, ok := findViperKey(map[string]interface{}{"foo": 1}, "bar") + assert.False(t, ok) + }) +}