Repository: ariga/atlas
Pull request #447 (2022-01-20 07:43:36)

sql/postgres: marshal HCL enum
This is the `Marshal` part of the postgres support for enums.


Patch:
diff --git a/doc/md/getting-started/01-introduction.mdx b/doc/md/getting-started/01-introduction.mdx
index fa960942..37d818b0 100644
--- a/doc/md/getting-started/01-introduction.mdx
+++ b/doc/md/getting-started/01-introduction.mdx
@@ -62,11 +62,8 @@ sudo install -o root -g root -m 0755 ./atlas-linux-amd64-v0.3.0 /usr/local/bin/a
 </TabItem>
 <TabItem value="windows">
 
-Download latest release.
-```shell
-curl -LO https://release.ariga.io/atlas/atlas-windows-amd64-v0.3.0.exe
-```
-Move the atlas binary to a file location on your system PATH.
+Download the [latest release](https://release.ariga.io/atlas/atlas-windows-amd64-v0.3.0.exe) and
+move the atlas binary to a file location on your system PATH.
 
 
 </TabItem>
diff --git a/schema/schemaspec/schemahcl/hcl.go b/schema/schemaspec/schemahcl/hcl.go
index ffc0b428..21e6b6be 100644
--- a/schema/schemaspec/schemahcl/hcl.go
+++ b/schema/schemaspec/schemahcl/hcl.go
@@ -255,6 +255,11 @@ func (s *state) writeAttr(attr *schemaspec.Attr, body *hclwrite.Body) error {
 		expr := strings.ReplaceAll(v.V, "$", "")
 		body.SetAttributeRaw(attr.K, hclRawTokens(expr))
 	case *schemaspec.Type:
+		if v.IsRef {
+			expr := strings.ReplaceAll(v.T, "$", "")
+			body.SetAttributeRaw(attr.K, hclRawTokens(expr))
+			break
+		}
 		spec, ok := s.findTypeSpec(v.T)
 		if !ok {
 			v := fmt.Sprintf("sql(%q)", v.T)
diff --git a/sql/internal/specutil/convert.go b/sql/internal/specutil/convert.go
index 03f242c7..6e3727a5 100644
--- a/sql/internal/specutil/convert.go
+++ b/sql/internal/specutil/convert.go
@@ -302,7 +302,7 @@ func FromSchema(s *schema.Schema, fn TableSpecFunc) (*sqlspec.Schema, []*sqlspec
 			return nil, nil, err
 		}
 		if s.Name != "" {
-			table.Schema = schemaRef(s.Name)
+			table.Schema = SchemaRef(s.Name)
 		}
 		tables = append(tables, table)
 	}
@@ -501,7 +501,8 @@ func colRef(cName string, tName string) *schemaspec.Ref {
 	}
 }
 
-func schemaRef(n string) *schemaspec.Ref {
+// SchemaRef returns the schemaspec.Ref to the schema with the given name.
+func SchemaRef(n string) *schemaspec.Ref {
 	return &schemaspec.Ref{V: "$schema." + n}
 }
 
diff --git a/sql/postgres/sqlspec.go b/sql/postgres/sqlspec.go
index 8b062b0d..d29f458d 100644
--- a/sql/postgres/sqlspec.go
+++ b/sql/postgres/sqlspec.go
@@ -63,7 +63,31 @@ func UnmarshalSpec(data []byte, unmarshaler schemaspec.Unmarshaler, v interface{
 
 // MarshalSpec marshals v into an Atlas DDL document using a schemaspec.Marshaler.
 func MarshalSpec(v interface{}, marshaler schemaspec.Marshaler) ([]byte, error) {
-	return specutil.Marshal(v, marshaler, schemaSpec)
+	var d doc
+	switch s := v.(type) {
+	case *schema.Schema:
+		var err error
+		doc, err := schemaSpec(s)
+		if err != nil {
+			return nil, fmt.Errorf("specutil: failed converting schema to spec: %w", err)
+		}
+		d.Tables = doc.Tables
+		d.Schemas = doc.Schemas
+		d.Enums = doc.Enums
+	case *schema.Realm:
+		for _, s := range s.Schemas {
+			doc, err := schemaSpec(s)
+			if err != nil {
+				return nil, fmt.Errorf("specutil: failed converting schema to spec: %w", err)
+			}
+			d.Tables = append(d.Tables, doc.Tables...)
+			d.Schemas = append(d.Schemas, doc.Schemas...)
+			d.Enums = append(d.Enums, doc.Enums...)
+		}
+	default:
+		return nil, fmt.Errorf("specutil: failed marshaling spec. %T is not supported", v)
+	}
+	return marshaler.MarshalSpec(&d)
 }
 
 // Realm converts the schemas and tables of the doc into a schema.Realm.
@@ -216,9 +240,34 @@ func enumName(ref *schemaspec.Type) (string, error) {
 	return s[1], nil
 }
 
+// enumRef returns a reference string to the given enum name.
+func enumRef(n string) *schemaspec.Ref {
+	return &schemaspec.Ref{
+		V: "$enum." + n,
+	}
+}
+
 // schemaSpec converts from a concrete Postgres schema to Atlas specification.
-func schemaSpec(schem *schema.Schema) (*sqlspec.Schema, []*sqlspec.Table, error) {
-	return specutil.FromSchema(schem, tableSpec)
+func schemaSpec(schem *schema.Schema) (*doc, error) {
+	var d doc
+	s, tbls, err := specutil.FromSchema(schem, tableSpec)
+	if err != nil {
+		return nil, err
+	}
+	d.Schemas = []*sqlspec.Schema{s}
+	d.Tables = tbls
+	for _, t := range schem.Tables {
+		for _, c := range t.Columns {
+			if t, ok := c.Type.Type.(*schema.EnumType); ok {
+				d.Enums = append(d.Enums, &Enum{
+					Name:   t.T,
+					Schema: specutil.SchemaRef(s.Name),
+					Values: t.Values,
+				})
+			}
+		}
+	}
+	return &d, nil
 }
 
 // tableSpec converts from a concrete Postgres sqlspec.Table to a schema.Table.
@@ -240,6 +289,13 @@ func columnSpec(col *schema.Column, _ *schema.Table) (*sqlspec.Column, error) {
 
 // columnTypeSpec converts from a concrete Postgres schema.Type into sqlspec.Column Type.
 func columnTypeSpec(t schema.Type) (*sqlspec.Column, error) {
+	// Handle postgres enum types. They cannot be put into the TypeRegistry since their name is dynamic.
+	if e, ok := t.(*schema.EnumType); ok {
+		return &sqlspec.Column{Type: &schemaspec.Type{
+			T:     enumRef(e.T).V,
+			IsRef: true,
+		}}, nil
+	}
 	st, err := TypeRegistry.Convert(t)
 	if err != nil {
 		return nil, err


Test patch:
diff --git a/sql/postgres/sqlspec_test.go b/sql/postgres/sqlspec_test.go
index 80682f4e..694af026 100644
--- a/sql/postgres/sqlspec_test.go
+++ b/sql/postgres/sqlspec_test.go
@@ -201,6 +201,36 @@ enum "account_type" {
 	require.EqualValues(t, exp, &s)
 }
 
+func TestMarshalSpec_Enum(t *testing.T) {
+	s := schema.New("test").
+		AddTables(
+			schema.NewTable("account").
+				AddColumns(
+					schema.NewEnumColumn("account_type",
+						schema.EnumName("account_type"),
+						schema.EnumValues("private", "business"),
+					),
+				),
+		)
+	buf, err := MarshalSpec(s, hclState)
+	require.NoError(t, err)
+	const expected = `table "account" {
+  schema = schema.test
+  column "account_type" {
+    null = false
+    type = enum.account_type
+  }
+}
+schema "test" {
+}
+enum "account_type" {
+  schema = schema.test
+  values = ["private", "business", ]
+}
+`
+	require.EqualValues(t, expected, string(buf))
+}
+
 func TestTypes(t *testing.T) {
 	// TODO(rotemtam): enum, timestamptz, interval
 	for _, tt := range []struct {