Repository: ariga/atlas
Pull request #3477 (2025-04-28 18:12:36)

schemahcl: `yamlmerge` to merge documents


Patch:
diff --git a/go.mod b/go.mod
index d0538aee..0e5f553b 100644
--- a/go.mod
+++ b/go.mod
@@ -11,6 +11,7 @@ require (
 	github.com/zclconf/go-cty v1.14.4
 	github.com/zclconf/go-cty-yaml v1.1.0
 	golang.org/x/mod v0.17.0
+	gopkg.in/yaml.v3 v3.0.1
 )
 
 require (
@@ -24,5 +25,4 @@ require (
 	github.com/mitchellh/go-wordwrap v0.0.0-20150314170334-ad45545899c7 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
 	golang.org/x/text v0.21.0 // indirect
-	gopkg.in/yaml.v3 v3.0.1 // indirect
 )
diff --git a/schemahcl/stdlib.go b/schemahcl/stdlib.go
index 8f84eb56..46a30adc 100644
--- a/schemahcl/stdlib.go
+++ b/schemahcl/stdlib.go
@@ -23,6 +23,7 @@ import (
 	"github.com/zclconf/go-cty/cty/function"
 	"github.com/zclconf/go-cty/cty/function/stdlib"
 	"github.com/zclconf/go-cty/cty/json"
+	"gopkg.in/yaml.v3"
 )
 
 func stdTypes(ctx *hcl.EvalContext) *hcl.EvalContext {
@@ -135,6 +136,7 @@ func stdFuncs() map[string]function.Function {
 		"jsonencode":      stdlib.JSONEncodeFunc,
 		"yamldecode":      ctyyaml.YAMLDecodeFunc,
 		"yamlencode":      ctyyaml.YAMLEncodeFunc,
+		"yamlmerge":       yamlMergeFunc,
 		"keys":            stdlib.KeysFunc,
 		"length":          stdlib.LengthFunc,
 		"log":             stdlib.LogFunc,
@@ -522,6 +524,53 @@ var (
 			return cty.True, nil
 		},
 	})
+
+	yamlMergeFunc = function.New(&function.Spec{
+		Description: "Merges multiple YAML documents into one.",
+		VarParam: &function.Parameter{
+			Name:      "values",
+			Type:      cty.String,
+			AllowNull: true,
+		},
+		Type:         function.StaticReturnType(cty.String),
+		RefineResult: refineNonNull,
+		Impl: func(args []cty.Value, _ cty.Type) (ret cty.Value, err error) {
+			switch len(args) {
+			case 0:
+				return cty.NullVal(cty.String), nil
+			case 1:
+				if args[0].IsNull() {
+					return cty.NullVal(cty.String), nil
+				}
+				if args[0].Type() != cty.String {
+					return cty.NilVal, function.NewArgErrorf(0, "yamlmerge: expected string value, got %s", args[0].Type())
+				}
+				return args[0], nil
+			default:
+				dst := make(map[string]any)
+				for i, v := range args {
+					if v.IsNull() {
+						continue
+					}
+					if v.Type() != cty.String {
+						return cty.NilVal, function.NewArgErrorf(i, "yamlmerge: expected string value, got %s", v.Type())
+					}
+					var obj map[string]any
+					if err = yaml.Unmarshal([]byte(v.AsString()), &obj); err != nil {
+						return cty.NilVal, function.NewArgErrorf(i, "yamlmerge: failed to unmarshal yaml: %s", err)
+					}
+					if err = deepMerge(dst, obj); err != nil {
+						return cty.NilVal, function.NewArgErrorf(i, "yamlmerge: failed to merge yaml: %s", err)
+					}
+				}
+				out, err := yaml.Marshal(dst)
+				if err != nil {
+					return cty.NilVal, fmt.Errorf("yamlmerge: failed to marshal yaml: %s", err)
+				}
+				return cty.StringVal(string(out)), nil
+			}
+		},
+	})
 )
 
 // MakeFileFunc returns a function that reads a file
@@ -626,3 +675,33 @@ func MakeFileSetFunc(base string) function.Function {
 		},
 	})
 }
+
+func deepMerge(dst, src map[string]any) error {
+	for k, v := range src {
+		switch val := v.(type) {
+		case map[string]any:
+			if _, ok := dst[k]; !ok {
+				dst[k] = make(map[string]any)
+			}
+			d, ok := dst[k].(map[string]any)
+			if !ok {
+				return fmt.Errorf("key %s is defined in both src and dst, but has a different type", k)
+			}
+			if err := deepMerge(d, val); err != nil {
+				return err
+			}
+		case []any:
+			if _, ok := dst[k]; !ok {
+				dst[k] = []any{}
+			}
+			d, ok := dst[k].([]any)
+			if !ok {
+				return fmt.Errorf("key %s is defined in both src and dst, but has a different type", k)
+			}
+			dst[k] = append(d, val...)
+		default:
+			dst[k] = val
+		}
+	}
+	return nil
+}


Test patch:
diff --git a/schemahcl/stdlib_test.go b/schemahcl/stdlib_test.go
index c6f101a5..8926e73e 100644
--- a/schemahcl/stdlib_test.go
+++ b/schemahcl/stdlib_test.go
@@ -353,3 +353,45 @@ func Example_printFunc() {
 	// cty.ListVal([]cty.Value{cty.StringVal("hello"), cty.StringVal("world")})
 	//
 }
+
+func TestYAMLMerge(t *testing.T) {
+	got, err := yamlMergeFunc.Call([]cty.Value{
+		cty.StringVal("map:\n  zzz: true\n  foo:\n    - bar"),
+		cty.StringVal("map:\n  bar:\n    - baz\n  foo:\n    - bar"),
+		cty.StringVal("map:\n  map:\n    baz: true\n    foo: false"),
+	})
+	require.NoError(t, err)
+	require.Equal(t, cty.StringVal(`map:
+    bar:
+        - baz
+    foo:
+        - bar
+        - bar
+    map:
+        baz: true
+        foo: false
+    zzz: true
+`), got)
+
+	got, err = yamlMergeFunc.Call([]cty.Value{
+		cty.StringVal("map:\n  zzz: true\n  foo:\n    - bar"),
+	})
+	require.NoError(t, err)
+	require.Equal(t, cty.StringVal(`map:
+  zzz: true
+  foo:
+    - bar`), got)
+
+	_, err = yamlMergeFunc.Call([]cty.Value{
+		cty.StringVal("map:\n  zzz: true\n  foo:\n    - bar"),
+		cty.NullVal(cty.String),
+		cty.StringVal("map:\n  zzz:\n    - baz"),
+	})
+	require.EqualError(t, err, "yamlmerge: failed to merge yaml: key zzz is defined in both src and dst, but has a different type")
+
+	_, err = yamlMergeFunc.Call([]cty.Value{
+		cty.StringVal("map:\n  zzz: true\n  foo:\n    - bar"),
+		cty.StringVal("map:\n  zzz:\n    foo: true"),
+	})
+	require.EqualError(t, err, "yamlmerge: failed to merge yaml: key zzz is defined in both src and dst, but has a different type")
+}