Repository: ariga/atlas
Pull request #2928 (2024-07-07 18:54:06)

schemahcl: allow creating list with mixed strings/enums


Patch:
diff --git a/schemahcl/schemahcl.go b/schemahcl/schemahcl.go
index 167f6726..12d51125 100644
--- a/schemahcl/schemahcl.go
+++ b/schemahcl/schemahcl.go
@@ -849,18 +849,32 @@ func (s *State) writeAttr(attr *Attr, body *hclwrite.Body) error {
 		}
 		tokens := make([]hclwrite.Tokens, 0, attr.V.LengthInt())
 		for _, v := range attr.V.AsValueSlice() {
-			if v.Type().IsCapsuleType() {
-				ref, ok := v.EncapsulatedValue().(*Ref)
-				if !ok {
-					return fmt.Errorf("unsupported capsule type: %v", v.Type())
-				}
-				ts, err := hclRefTokens(ref.V)
+			if !v.Type().IsCapsuleType() {
+				tokens = append(tokens, hclwrite.TokensForValue(v))
+				continue
+			}
+			switch ev := v.EncapsulatedValue().(type) {
+			case *Ref:
+				ts, err := hclRefTokens(ev.V)
 				if err != nil {
 					return err
 				}
 				tokens = append(tokens, ts)
-			} else {
-				tokens = append(tokens, hclwrite.TokensForValue(v))
+			case *EnumString:
+				switch {
+				case ev.S != "" && ev.E != "":
+					return fmt.Errorf("enum string cannot have both a string and an expression, got: %#v", ev)
+				case ev.S != "":
+					tokens = append(tokens, hclwrite.TokensForValue(cty.StringVal(ev.S)))
+				case ev.E != "":
+					tokens = append(tokens, hclwrite.Tokens{
+						&hclwrite.Token{Type: hclsyntax.TokenIdent, Bytes: []byte(ev.E)},
+					})
+				default:
+					return fmt.Errorf("enum string must have either a string or an expression, got: %#v", ev)
+				}
+			default:
+				return fmt.Errorf("unsupported capsule type: %v", v.Type())
 			}
 		}
 		body.SetAttributeRaw(attr.K, hclList(tokens))
diff --git a/schemahcl/spec.go b/schemahcl/spec.go
index 505eda05..d3f2158c 100644
--- a/schemahcl/spec.go
+++ b/schemahcl/spec.go
@@ -578,3 +578,25 @@ func RawAttr(k string, x string) *Attr {
 func RawExprValue(x *RawExpr) cty.Value {
 	return cty.CapsuleVal(ctyRawExpr, x)
 }
+
+// ctyEnumString is a capsule type for EnumString.
+var ctyEnumString = cty.Capsule("enum_string", reflect.TypeOf(EnumString{}))
+
+// EnumString is a helper type that represents
+// either an enum or a string value.
+type EnumString struct {
+	E, S string // Enum or string value.
+}
+
+// StringEnumsAttr is a helper method for constructing *schemahcl.Attr instances
+// that contain list of elements that their values can be either enum or string.
+func StringEnumsAttr(k string, elems ...*EnumString) *Attr {
+	vv := make([]cty.Value, len(elems))
+	for i, e := range elems {
+		vv[i] = cty.CapsuleVal(ctyEnumString, e)
+	}
+	return &Attr{
+		K: k,
+		V: cty.ListVal(vv),
+	}
+}


Test patch:
diff --git a/schemahcl/schemahcl_test.go b/schemahcl/schemahcl_test.go
index f0c11629..f7d58bdd 100644
--- a/schemahcl/schemahcl_test.go
+++ b/schemahcl/schemahcl_test.go
@@ -1078,6 +1078,24 @@ bar "b1" {
 	}, doc.Bar[0])
 }
 
+func Test_MarshalAttr(t *testing.T) {
+	var doc struct {
+		DefaultExtension
+	}
+	doc.Extra.Attrs = append(
+		doc.Extra.Attrs,
+		StringEnumsAttr("mixed1", &EnumString{S: "string"}, &EnumString{E: "enum"}),
+		StringEnumsAttr("mixed2", &EnumString{E: "enum1"}, &EnumString{E: "enum2"}),
+		StringEnumsAttr("mixed3", &EnumString{S: "string1"}, &EnumString{S: "string1"}),
+	)
+	buf, err := Marshal(&doc)
+	require.NoError(t, err)
+	require.Equal(t, `mixed1 = ["string", enum]
+mixed2 = [enum1, enum2]
+mixed3 = ["string1", "string1"]
+`, string(buf))
+}
+
 func Test_WithPos(t *testing.T) {
 	var (
 		doc struct {