Repository: ariga/atlas
Pull request #1725 (2023-06-08 08:23:23)

sql/internal/specutil: marshal view spec


Patch:
diff --git a/schemahcl/hcl.go b/schemahcl/hcl.go
index 209c311f..3fc2b8b6 100644
--- a/schemahcl/hcl.go
+++ b/schemahcl/hcl.go
@@ -498,6 +498,22 @@ func (s *State) writeAttr(attr *Attr, body *hclwrite.Body) error {
 			}
 		}
 		body.SetAttributeRaw(attr.K, hclList(tokens))
+	// Heredoc is a special case that currently is not handled by hclwrite:
+	// https://github.com/hashicorp/hcl/blob/main/hclwrite/generate.go#L218-L219.
+	case attr.V.Type() == cty.String && strings.Count(attr.V.AsString(), "\n") > 1 && strings.HasPrefix(attr.V.AsString(), "<<"):
+		v := strings.TrimLeft(attr.V.AsString(), "<-")
+		// Heredoc begins with << (or <<-), followed by a token that
+		// specifies the terminator and ends with the \n + terminator.
+		if lines := strings.Split(v, "\n"); len(lines) > 2 && strings.TrimSpace(lines[0]) == strings.TrimSpace(lines[len(lines)-1]) {
+			body.SetAttributeRaw(attr.K, hclwrite.Tokens{
+				&hclwrite.Token{
+					Type:  hclsyntax.TokenOHeredoc,
+					Bytes: []byte(attr.V.AsString()),
+				},
+			})
+		} else {
+			body.SetAttributeValue(attr.K, attr.V)
+		}
 	default:
 		body.SetAttributeValue(attr.K, attr.V)
 	}
diff --git a/sql/internal/specutil/convert.go b/sql/internal/specutil/convert.go
index e11cce95..5d94f5fa 100644
--- a/sql/internal/specutil/convert.go
+++ b/sql/internal/specutil/convert.go
@@ -25,9 +25,11 @@ type (
 	ConvertPrimaryKeyFunc func(*sqlspec.PrimaryKey, *schema.Table) (*schema.Index, error)
 	ConvertIndexFunc      func(*sqlspec.Index, *schema.Table) (*schema.Index, error)
 	ConvertCheckFunc      func(*sqlspec.Check) (*schema.Check, error)
-	ColumnSpecFunc        func(*schema.Column, *schema.Table) (*sqlspec.Column, error)
 	ColumnTypeSpecFunc    func(schema.Type) (*sqlspec.Column, error)
 	TableSpecFunc         func(*schema.Table) (*sqlspec.Table, error)
+	TableColumnSpecFunc   func(*schema.Column, *schema.Table) (*sqlspec.Column, error)
+	ViewSpecFunc          func(*schema.View) (*sqlspec.View, error)
+	ViewColumnSpecFunc    func(*schema.Column, *schema.View) (*sqlspec.Column, error)
 	PrimaryKeySpecFunc    func(*schema.Index) (*sqlspec.PrimaryKey, error)
 	IndexSpecFunc         func(*schema.Index) (*sqlspec.Index, error)
 	ForeignKeySpecFunc    func(*schema.ForeignKey) (*sqlspec.ForeignKey, error)
@@ -293,13 +295,16 @@ func linkForeignKeys(tbl *schema.Table, sch *schema.Schema, table *sqlspec.Table
 }
 
 // FromSchema converts a schema.Schema into sqlspec.Schema and []sqlspec.Table.
-func FromSchema(s *schema.Schema, fn TableSpecFunc) (*SchemaSpec, error) {
-	spec := &sqlspec.Schema{
-		Name: s.Name,
-	}
-	tables := make([]*sqlspec.Table, 0, len(s.Tables))
+func FromSchema(s *schema.Schema, specT TableSpecFunc, specV ViewSpecFunc) (*SchemaSpec, error) {
+	var (
+		spec = &sqlspec.Schema{
+			Name: s.Name,
+		}
+		views  = make([]*sqlspec.View, 0, len(s.Views))
+		tables = make([]*sqlspec.Table, 0, len(s.Tables))
+	)
 	for _, t := range s.Tables {
-		table, err := fn(t)
+		table, err := specT(t)
 		if err != nil {
 			return nil, err
 		}
@@ -308,14 +313,25 @@ func FromSchema(s *schema.Schema, fn TableSpecFunc) (*SchemaSpec, error) {
 		}
 		tables = append(tables, table)
 	}
+	for _, v := range s.Views {
+		view, err := specV(v)
+		if err != nil {
+			return nil, err
+		}
+		if s.Name != "" {
+			view.Schema = SchemaRef(s.Name)
+		}
+		views = append(views, view)
+	}
 	return &SchemaSpec{
 		Schema: spec,
 		Tables: tables,
+		Views:  views,
 	}, nil
 }
 
 // FromTable converts a schema.Table to a sqlspec.Table.
-func FromTable(t *schema.Table, colFn ColumnSpecFunc, pkFn PrimaryKeySpecFunc, idxFn IndexSpecFunc,
+func FromTable(t *schema.Table, colFn TableColumnSpecFunc, pkFn PrimaryKeySpecFunc, idxFn IndexSpecFunc,
 	fkFn ForeignKeySpecFunc, ckFn CheckSpecFunc) (*sqlspec.Table, error) {
 	spec := &sqlspec.Table{
 		Name: t.Name,
@@ -357,6 +373,28 @@ func FromTable(t *schema.Table, colFn ColumnSpecFunc, pkFn PrimaryKeySpecFunc, i
 	return spec, nil
 }
 
+// FromView converts a schema.View to a sqlspec.View.
+func FromView(v *schema.View, colFn ViewColumnSpecFunc) (*sqlspec.View, error) {
+	spec := &sqlspec.View{
+		Name: v.Name,
+		As:   v.Def,
+	}
+	// In case the view definition is multi-line,
+	// format it as indented heredoc with two spaces.
+	if lines := strings.Split(v.Def, "\n"); len(lines) > 1 {
+		spec.As = fmt.Sprintf("<<-SQL\n  %s\n  SQL", strings.Join(lines, "\n  "))
+	}
+	for _, c := range v.Columns {
+		col, err := colFn(c, v)
+		if err != nil {
+			return nil, err
+		}
+		spec.Columns = append(spec.Columns, col)
+	}
+	convertCommentFromSchema(v.Attrs, &spec.Extra.Attrs)
+	return spec, nil
+}
+
 // FromPrimaryKey converts schema.Index to a sqlspec.PrimaryKey.
 func FromPrimaryKey(s *schema.Index) (*sqlspec.PrimaryKey, error) {
 	c := make([]*schemahcl.Ref, 0, len(s.Parts))
@@ -702,10 +740,10 @@ func convertCommentFromSpec(spec Attrer, attrs *[]schema.Attr) error {
 }
 
 // convertCommentFromSchema converts a schema element comment attribute to a spec comment attribute.
-func convertCommentFromSchema(src []schema.Attr, trgt *[]*schemahcl.Attr) {
+func convertCommentFromSchema(src []schema.Attr, target *[]*schemahcl.Attr) {
 	var c schema.Comment
 	if sqlx.Has(src, &c) {
-		*trgt = append(*trgt, schemahcl.StringAttr("comment", c.Text))
+		*target = append(*target, schemahcl.StringAttr("comment", c.Text))
 	}
 }
 
diff --git a/sql/internal/specutil/spec.go b/sql/internal/specutil/spec.go
index 6fbd8b61..2b8bcd5a 100644
--- a/sql/internal/specutil/spec.go
+++ b/sql/internal/specutil/spec.go
@@ -26,6 +26,7 @@ type (
 	SchemaSpec struct {
 		Schema *sqlspec.Schema
 		Tables []*sqlspec.Table
+		Views  []*sqlspec.View
 	}
 	doc struct {
 		Tables  []*sqlspec.Table  `spec:"table"`
diff --git a/sql/mysql/sqlspec.go b/sql/mysql/sqlspec.go
index ae61db86..40e7110f 100644
--- a/sql/mysql/sqlspec.go
+++ b/sql/mysql/sqlspec.go
@@ -231,7 +231,7 @@ func convertColumnType(spec *sqlspec.Column) (schema.Type, error) {
 
 // schemaSpec converts from a concrete MySQL schema to Atlas specification.
 func schemaSpec(s *schema.Schema) (*specutil.SchemaSpec, error) {
-	spec, err := specutil.FromSchema(s, tableSpec)
+	spec, err := specutil.FromSchema(s, tableSpec, viewSpec)
 	if err != nil {
 		return nil, err
 	}
@@ -277,6 +277,17 @@ func tableSpec(t *schema.Table) (*sqlspec.Table, error) {
 	return ts, nil
 }
 
+// viewSpec converts from a concrete MySQL schema.View to a sqlspec.View.
+func viewSpec(view *schema.View) (*sqlspec.View, error) {
+	spec, err := specutil.FromView(view, func(c *schema.Column, _ *schema.View) (*sqlspec.Column, error) {
+		return specutil.FromColumn(c, columnTypeSpec)
+	})
+	if err != nil {
+		return nil, err
+	}
+	return spec, nil
+}
+
 func pkSpec(idx *schema.Index) (*sqlspec.PrimaryKey, error) {
 	spec, err := specutil.FromPrimaryKey(idx)
 	if err != nil {
diff --git a/sql/postgres/sqlspec.go b/sql/postgres/sqlspec.go
index ded410aa..bfd1eb3e 100644
--- a/sql/postgres/sqlspec.go
+++ b/sql/postgres/sqlspec.go
@@ -24,6 +24,7 @@ import (
 type (
 	doc struct {
 		Tables  []*sqlspec.Table  `spec:"table"`
+		Views   []*sqlspec.View   `spec:"view"`
 		Enums   []*Enum           `spec:"enum"`
 		Schemas []*sqlspec.Schema `spec:"schema"`
 	}
@@ -91,6 +92,7 @@ func MarshalSpec(v any, marshaler schemahcl.Marshaler) ([]byte, error) {
 			return nil, fmt.Errorf("specutil: failed converting schema to spec: %w", err)
 		}
 		d.Tables = doc.Tables
+		d.Views = doc.Views
 		d.Schemas = doc.Schemas
 		d.Enums = doc.Enums
 	case *schema.Realm:
@@ -100,6 +102,7 @@ func MarshalSpec(v any, marshaler schemahcl.Marshaler) ([]byte, error) {
 				return nil, fmt.Errorf("specutil: failed converting schema to spec: %w", err)
 			}
 			d.Tables = append(d.Tables, doc.Tables...)
+			d.Views = append(d.Views, doc.Views...)
 			d.Schemas = append(d.Schemas, doc.Schemas...)
 			d.Enums = append(d.Enums, doc.Enums...)
 		}
@@ -512,25 +515,34 @@ func enumRef(n string) *schemahcl.Ref {
 
 // schemaSpec converts from a concrete Postgres schema to Atlas specification.
 func schemaSpec(schem *schema.Schema) (*doc, error) {
-	spec, err := specutil.FromSchema(schem, tableSpec)
+	spec, err := specutil.FromSchema(schem, tableSpec, viewSpec)
 	if err != nil {
 		return nil, err
 	}
 	d := &doc{
 		Tables:  spec.Tables,
+		Views:   spec.Views,
 		Schemas: []*sqlspec.Schema{spec.Schema},
 	}
 	enums := make(map[string]bool)
+	mayAdd := func(c *schema.Column) {
+		if e, ok := hasEnumType(c); ok && !enums[e.T] {
+			d.Enums = append(d.Enums, &Enum{
+				Name:   e.T,
+				Schema: specutil.SchemaRef(spec.Schema.Name),
+				Values: e.Values,
+			})
+			enums[e.T] = true
+		}
+	}
 	for _, t := range schem.Tables {
 		for _, c := range t.Columns {
-			if e, ok := hasEnumType(c); ok && !enums[e.T] {
-				d.Enums = append(d.Enums, &Enum{
-					Name:   e.T,
-					Schema: specutil.SchemaRef(spec.Schema.Name),
-					Values: e.Values,
-				})
-				enums[e.T] = true
-			}
+			mayAdd(c)
+		}
+	}
+	for _, t := range schem.Views {
+		for _, c := range t.Columns {
+			mayAdd(c)
 		}
 	}
 	return d, nil
@@ -540,7 +552,7 @@ func schemaSpec(schem *schema.Schema) (*doc, error) {
 func tableSpec(table *schema.Table) (*sqlspec.Table, error) {
 	spec, err := specutil.FromTable(
 		table,
-		columnSpec,
+		tableColumnSpec,
 		pkSpec,
 		indexSpec,
 		specutil.FromForeignKey,
@@ -555,6 +567,17 @@ func tableSpec(table *schema.Table) (*sqlspec.Table, error) {
 	return spec, nil
 }
 
+// viewSpec converts from a concrete PostgreSQL schema.View to a sqlspec.View.
+func viewSpec(view *schema.View) (*sqlspec.View, error) {
+	spec, err := specutil.FromView(view, func(c *schema.Column, _ *schema.View) (*sqlspec.Column, error) {
+		return specutil.FromColumn(c, columnTypeSpec)
+	})
+	if err != nil {
+		return nil, err
+	}
+	return spec, nil
+}
+
 func pkSpec(idx *schema.Index) (*sqlspec.PrimaryKey, error) {
 	spec, err := specutil.FromPrimaryKey(idx)
 	if err != nil {
@@ -611,8 +634,8 @@ func partAttr(idx *schema.Index, part *schema.IndexPart, spec *sqlspec.IndexPart
 	return nil
 }
 
-// columnSpec converts from a concrete Postgres schema.Column into a sqlspec.Column.
-func columnSpec(c *schema.Column, _ *schema.Table) (*sqlspec.Column, error) {
+// tableColumnSpec converts from a concrete Postgres schema.Column into a sqlspec.Column.
+func tableColumnSpec(c *schema.Column, _ *schema.Table) (*sqlspec.Column, error) {
 	s, err := specutil.FromColumn(c, columnTypeSpec)
 	if err != nil {
 		return nil, err
diff --git a/sql/sqlite/sqlspec.go b/sql/sqlite/sqlspec.go
index a9556ad4..f3d965a9 100644
--- a/sql/sqlite/sqlspec.go
+++ b/sql/sqlite/sqlspec.go
@@ -110,7 +110,7 @@ func convertColumnType(spec *sqlspec.Column) (schema.Type, error) {
 
 // schemaSpec converts from a concrete SQLite schema to Atlas specification.
 func schemaSpec(s *schema.Schema) (*specutil.SchemaSpec, error) {
-	return specutil.FromSchema(s, tableSpec)
+	return specutil.FromSchema(s, tableSpec, viewSpec)
 }
 
 // tableSpec converts from a concrete SQLite sqlspec.Table to a schema.Table.
@@ -125,6 +125,17 @@ func tableSpec(tab *schema.Table) (*sqlspec.Table, error) {
 	)
 }
 
+// viewSpec converts from a concrete SQLite schema.View to a sqlspec.View.
+func viewSpec(view *schema.View) (*sqlspec.View, error) {
+	spec, err := specutil.FromView(view, func(c *schema.Column, _ *schema.View) (*sqlspec.Column, error) {
+		return specutil.FromColumn(c, columnTypeSpec)
+	})
+	if err != nil {
+		return nil, err
+	}
+	return spec, nil
+}
+
 func indexSpec(idx *schema.Index) (*sqlspec.Index, error) {
 	spec, err := specutil.FromIndex(idx)
 	if err != nil {
diff --git a/sql/sqlspec/sqlspec.go b/sql/sqlspec/sqlspec.go
index 7b98495d..06de446a 100644
--- a/sql/sqlspec/sqlspec.go
+++ b/sql/sqlspec/sqlspec.go
@@ -36,6 +36,7 @@ type (
 		Qualifier string         `spec:",qualifier"`
 		Schema    *schemahcl.Ref `spec:"schema"`
 		Columns   []*Column      `spec:"column"`
+		As        string         `spec:"as"`
 		schemahcl.DefaultExtension
 	}
 


Test patch:
diff --git a/schemahcl/hcl_test.go b/schemahcl/hcl_test.go
index f1b0564f..925db613 100644
--- a/schemahcl/hcl_test.go
+++ b/schemahcl/hcl_test.go
@@ -25,6 +25,10 @@ b  = true
 s  = "hello, world"
 sl = ["hello", "world"]
 bl = [true, false]
+hd = <<-EOT
+  hello
+  world
+EOT
 `
 	var test struct {
 		Int        int      `spec:"i"`
@@ -32,6 +36,7 @@ bl = [true, false]
 		Str        string   `spec:"s"`
 		StringList []string `spec:"sl"`
 		BoolList   []bool   `spec:"bl"`
+		Heredoc    string   `spec:"hd"`
 	}
 	err := New().EvalBytes([]byte(f), &test, nil)
 	require.NoError(t, err)
@@ -40,6 +45,9 @@ bl = [true, false]
 	require.EqualValues(t, "hello, world", test.Str)
 	require.EqualValues(t, []string{"hello", "world"}, test.StringList)
 	require.EqualValues(t, []bool{true, false}, test.BoolList)
+	require.EqualValues(t, "hello\nworld\n", test.Heredoc)
+	// Heredoc needs to be explicitly formatted this way.
+	test.Heredoc = "<<-EOT\n  hello\n  world\nEOT"
 	marshal, err := Marshal(&test)
 	require.NoError(t, err)
 	require.EqualValues(t, f, string(marshal))
diff --git a/sql/internal/specutil/convert_test.go b/sql/internal/specutil/convert_test.go
index b71b4775..6879a5a3 100644
--- a/sql/internal/specutil/convert_test.go
+++ b/sql/internal/specutil/convert_test.go
@@ -22,9 +22,15 @@ func TestFromSpec_SchemaName(t *testing.T) {
 		},
 	}
 	sc.Tables[0].Schema = sc
-	spec, err := FromSchema(sc, func(table *schema.Table) (*sqlspec.Table, error) {
-		return &sqlspec.Table{}, nil
-	})
+	spec, err := FromSchema(
+		sc,
+		func(*schema.Table) (*sqlspec.Table, error) {
+			return &sqlspec.Table{}, nil
+		},
+		func(*schema.View) (*sqlspec.View, error) {
+			return &sqlspec.View{}, nil
+		},
+	)
 	require.NoError(t, err)
 	require.Equal(t, sc.Name, spec.Schema.Name)
 	require.Equal(t, "$schema."+sc.Name, spec.Tables[0].Schema.V)
diff --git a/sql/postgres/sqlspec_test.go b/sql/postgres/sqlspec_test.go
index d1204cfc..629e5b74 100644
--- a/sql/postgres/sqlspec_test.go
+++ b/sql/postgres/sqlspec_test.go
@@ -263,6 +263,47 @@ enum "account_type" {
 	require.EqualValues(t, exp, &s)
 }
 
+func TestMarshalViews(t *testing.T) {
+	s := schema.New("public").
+		AddViews(
+			schema.NewView("v1", "SELECT 1"),
+			schema.NewView("v2", "SELECT * FROM t2\n\tWHERE id IS NOT NULL"),
+			schema.NewView("v3", "SELECT * FROM t3\n\tWHERE id IS NOT NULL\n\tORDER BY id").
+				AddColumns(
+					schema.NewIntColumn("id", "id"),
+				),
+		)
+	buf, err := MarshalHCL(s)
+	require.NoError(t, err)
+	f := `view "v1" {
+  schema = schema.public
+  as     = "SELECT 1"
+}
+view "v2" {
+  schema = schema.public
+  as     = <<-SQL
+  SELECT * FROM t2
+  	WHERE id IS NOT NULL
+  SQL
+}
+view "v3" {
+  schema = schema.public
+  as     = <<-SQL
+  SELECT * FROM t3
+  	WHERE id IS NOT NULL
+  	ORDER BY id
+  SQL
+  column "id" {
+    null = false
+    type = sql("id")
+  }
+}
+schema "public" {
+}
+`
+	require.Equal(t, f, string(buf))
+}
+
 func TestUnmarshalSpec_IndexType(t *testing.T) {
 	f := `
 schema "s" {}