Repository: ariga/atlas
Pull request #1816 (2023-07-03 18:32:31)

sql/postgres: unmarshal views with their dependencies 


Patch:
diff --git a/sql/internal/specutil/convert.go b/sql/internal/specutil/convert.go
index 10c47f21..e77aae7f 100644
--- a/sql/internal/specutil/convert.go
+++ b/sql/internal/specutil/convert.go
@@ -19,79 +19,136 @@ import (
 
 // List of convert function types.
 type (
-	ConvertTableFunc      func(*sqlspec.Table, *schema.Schema) (*schema.Table, error)
-	ConvertColumnFunc     func(*sqlspec.Column, *schema.Table) (*schema.Column, error)
-	ConvertTypeFunc       func(*sqlspec.Column) (schema.Type, error)
-	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)
-	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)
-	CheckSpecFunc         func(*schema.Check) *sqlspec.Check
+	ConvertTableFunc       func(*sqlspec.Table, *schema.Schema) (*schema.Table, error)
+	ConvertTableColumnFunc func(*sqlspec.Column, *schema.Table) (*schema.Column, error)
+	ConvertViewFunc        func(*sqlspec.View, *schema.Schema) (*schema.View, error)
+	ConvertViewColumnFunc  func(*sqlspec.Column, *schema.View) (*schema.Column, error)
+	ConvertTypeFunc        func(*sqlspec.Column) (schema.Type, error)
+	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)
+	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)
+	CheckSpecFunc          func(*schema.Check) *sqlspec.Check
+)
+
+type (
+	// ScanDoc represents a scanned HCL document.
+	ScanDoc struct {
+		Schemas []*sqlspec.Schema
+		Tables  []*sqlspec.Table
+		Views   []*sqlspec.View
+	}
+	// ScanFuncs represents a set of scan functions
+	// used to convert the HCL document to the Realm.
+	ScanFuncs struct {
+		Table ConvertTableFunc
+		View  ConvertViewFunc
+	}
 )
 
 // Scan populates the Realm from the schemas and table specs.
-func Scan(r *schema.Realm, schemas []*sqlspec.Schema, tables []*sqlspec.Table, convertTable ConvertTableFunc) error {
+func Scan(r *schema.Realm, doc *ScanDoc, funcs *ScanFuncs) error {
 	byName := make(map[string]*schema.Schema)
-	for _, spec := range schemas {
-		s := &schema.Schema{Name: spec.Name, Realm: r}
-		r.AddSchemas(s)
-		byName[spec.Name] = s
-	}
-	for _, spec := range tables {
-		name, err := SchemaName(spec.Schema)
+	for _, s := range doc.Schemas {
+		s1 := schema.New(s.Name)
+		r.AddSchemas(s1)
+		byName[s.Name] = s1
+	}
+	tableFKs := make(map[*schema.Table][]*sqlspec.ForeignKey)
+	for _, st := range doc.Tables {
+		name, err := SchemaName(st.Schema)
 		if err != nil {
-			return fmt.Errorf("specutil: cannot extract schema name for table %q: %w", spec.Name, err)
+			return fmt.Errorf("specutil: cannot extract schema name for table %q: %w", st.Name, err)
 		}
 		s, ok := byName[name]
 		if !ok {
-			return fmt.Errorf("specutil: schema %q not found for table %q", name, spec.Name)
+			return fmt.Errorf("specutil: schema %q not found for table %q", name, st.Name)
 		}
-		t, err := convertTable(spec, s)
+		t, err := funcs.Table(st, s)
 		if err != nil {
-			return fmt.Errorf("specutil: cannot convert table %q: %w", spec.Name, err)
+			return fmt.Errorf("specutil: cannot convert table %q: %w", st.Name, err)
 		}
+		tableFKs[t] = st.ForeignKeys
 		s.AddTables(t)
 	}
 	// Link the foreign keys.
-	for _, s := range r.Schemas {
-		for _, t := range s.Tables {
-			spec, err := findTableSpec(tables, s.Name, t.Name)
-			if err != nil {
-				return err
-			}
-			if err := linkForeignKeys(t, s, spec); err != nil {
-				return err
-			}
+	for t, fks := range tableFKs {
+		if err := linkForeignKeys(t, fks); err != nil {
+			return err
 		}
 	}
-	return nil
-}
-
-// findTableSpec searches tableSpecs for a spec of a table named tableName in a schema named schemaName.
-func findTableSpec(tableSpecs []*sqlspec.Table, schemaName, tableName string) (*sqlspec.Table, error) {
-	for _, tbl := range tableSpecs {
-		n, err := SchemaName(tbl.Schema)
+	viewDeps := make(map[*schema.View][]*schemahcl.Ref, len(doc.Views))
+	for _, sv := range doc.Views {
+		name, err := SchemaName(sv.Schema)
 		if err != nil {
-			return nil, err
+			return fmt.Errorf("specutil: cannot extract schema name for view %q: %w", sv.Name, err)
+		}
+		s, ok := byName[name]
+		if !ok {
+			return fmt.Errorf("specutil: schema %q not found for view %q", name, sv.Name)
 		}
-		if n == schemaName && tbl.Name == tableName {
-			return tbl, nil
+		v, err := funcs.View(sv, s)
+		if err != nil {
+			return fmt.Errorf("specutil: cannot convert view %q: %w", sv.Name, err)
+		}
+		s.AddViews(v)
+		if deps, ok := sv.Attr("depends_on"); ok {
+			refs, err := deps.Refs()
+			if err != nil {
+				return fmt.Errorf("specutil: expect list of references for attribute view.%s.depends_on: %w", sv.Name, err)
+			}
+			viewDeps[v] = refs
+		}
+	}
+	// Link views' dependencies.
+	for v, refs := range viewDeps {
+		for i, r := range refs {
+			switch p, err := r.Path(); {
+			case err != nil:
+				return fmt.Errorf("specutil: extract reference for view.%s: %w", v.Name, err)
+			case len(p) == 0:
+				return fmt.Errorf("specutil: empty reference for view.%s", v.Name)
+			case p[0].T == "view":
+				q, n, err := viewName(r)
+				if err != nil {
+					return fmt.Errorf("specutil: extract view name from view.%s.depends_on[%d]: %w", v.Name, i, err)
+				}
+				v1, err := findT(v.Schema, q, n, func(s *schema.Schema, name string) (*schema.View, bool) {
+					return s.View(name)
+				})
+				if err != nil {
+					return fmt.Errorf("specutil: find view refrence for view.%s.depends_on[%d]: %w", v.Name, i, err)
+				}
+				v.AddDeps(v1)
+			case p[0].T == "table":
+				q, n, err := tableName(r)
+				if err != nil {
+					return fmt.Errorf("specutil: extract view name from view.%s.depends_on[%d]: %w", v.Name, i, err)
+				}
+				t1, err := findT(v.Schema, q, n, func(s *schema.Schema, name string) (*schema.Table, bool) {
+					return s.Table(name)
+				})
+				if err != nil {
+					return fmt.Errorf("specutil: find table refrence for view.%s.depends_on[%d]: %w", v.Name, i, err)
+				}
+				v.AddDeps(t1)
+			}
 		}
 	}
-	return nil, fmt.Errorf("table %s.%s not found", schemaName, tableName)
+	return nil
 }
 
 // Table converts a sqlspec.Table to a schema.Table. Table conversion is done without converting
 // ForeignKeySpecs into ForeignKeys, as the target tables do not necessarily exist in the schema
 // at this point. Instead, the linking is done by the Schema function.
-func Table(spec *sqlspec.Table, parent *schema.Schema, convertColumn ConvertColumnFunc,
+func Table(spec *sqlspec.Table, parent *schema.Schema, convertColumn ConvertTableColumnFunc,
 	convertPK ConvertPrimaryKeyFunc, convertIndex ConvertIndexFunc, convertCheck ConvertCheckFunc) (*schema.Table, error) {
 	tbl := &schema.Table{
 		Name:   spec.Name,
@@ -131,6 +188,30 @@ func Table(spec *sqlspec.Table, parent *schema.Schema, convertColumn ConvertColu
 	return tbl, nil
 }
 
+// View converts a sqlspec.View to a schema.View.
+func View(spec *sqlspec.View, parent *schema.Schema, convertColumn ConvertViewColumnFunc) (*schema.View, error) {
+	as, ok := spec.Extra.Attr("as")
+	if !ok {
+		return nil, fmt.Errorf("specutil: missing 'as' definition for view %q", spec.Name)
+	}
+	def, err := as.String()
+	if err != nil {
+		return nil, fmt.Errorf("specutil: expect string definition for attribute view.%s.as: %w", spec.Name, err)
+	}
+	v := schema.NewView(spec.Name, def).SetSchema(parent)
+	for _, c := range spec.Columns {
+		c, err := convertColumn(c, v)
+		if err != nil {
+			return nil, err
+		}
+		v.AddColumns(c)
+	}
+	if err := convertCommentFromSpec(spec, &v.Attrs); err != nil {
+		return nil, err
+	}
+	return v, nil
+}
+
 // Column converts a sqlspec.Column into a schema.Column.
 func Column(spec *sqlspec.Column, conv ConvertTypeFunc) (*schema.Column, error) {
 	out := &schema.Column{
@@ -255,8 +336,8 @@ func PrimaryKey(spec *sqlspec.PrimaryKey, parent *schema.Table) (*schema.Index,
 // linkForeignKeys creates the foreign keys defined in the Table's spec by creating references
 // to column in the provided Schema. It is assumed that all tables referenced FK definitions in the spec
 // are reachable from the provided schema or its connected realm.
-func linkForeignKeys(tbl *schema.Table, sch *schema.Schema, table *sqlspec.Table) error {
-	for _, spec := range table.ForeignKeys {
+func linkForeignKeys(tbl *schema.Table, fks []*sqlspec.ForeignKey) error {
+	for _, spec := range fks {
 		fk := &schema.ForeignKey{Symbol: spec.Symbol, Table: tbl}
 		if spec.OnUpdate != nil {
 			fk.OnUpdate = schema.ReferenceOption(FromVar(spec.OnUpdate.V))
@@ -275,7 +356,7 @@ func linkForeignKeys(tbl *schema.Table, sch *schema.Schema, table *sqlspec.Table
 			fk.Columns = append(fk.Columns, c)
 		}
 		for i, ref := range spec.RefColumns {
-			t, c, err := externalRef(ref, sch)
+			t, c, err := externalRef(ref, tbl.Schema)
 			if isLocalRef(ref) {
 				t = fk.Table
 				c, err = ColumnByRef(fk.Table, ref)
@@ -646,27 +727,29 @@ func ColumnByRef(t *schema.Table, ref *schemahcl.Ref) (*schema.Column, error) {
 }
 
 func externalRef(ref *schemahcl.Ref, sch *schema.Schema) (*schema.Table, *schema.Column, error) {
-	tbl, err := findTable(ref, sch)
+	qualifier, name, err := tableName(ref)
 	if err != nil {
 		return nil, nil, err
 	}
-	c, err := ColumnByRef(tbl, ref)
+	t, err := findT(sch, qualifier, name, func(s *schema.Schema, name string) (*schema.Table, bool) {
+		return s.Table(name)
+	})
 	if err != nil {
 		return nil, nil, err
 	}
-	return tbl, c, nil
-}
-
-// findTable finds the table referenced by ref in the provided schema. If the table
-// is not in the provided schema.Schema other schemas in the connected schema.Realm
-// are searched as well.
-func findTable(ref *schemahcl.Ref, sch *schema.Schema) (*schema.Table, error) {
-	qualifier, name, err := tableName(ref)
+	c, err := ColumnByRef(t, ref)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
+	return t, c, nil
+}
+
+// findT finds the table/view referenced by ref in the provided schema. If the table/view
+// is not in the provided schema.Schema other schemas in the connected schema.Realm are
+// searched as well.
+func findT[T schema.View | schema.Table](sch *schema.Schema, qualifier, name string, findT func(*schema.Schema, string) (*T, bool)) (*T, error) {
 	var (
-		matches []*schema.Table  // Found references.
+		matches []*T             // Found references.
 		schemas []*schema.Schema // Schemas to search.
 	)
 	switch {
@@ -681,7 +764,7 @@ func findTable(ref *schemahcl.Ref, sch *schema.Schema) (*schema.Table, error) {
 		}
 	}
 	for _, s := range schemas {
-		t, ok := s.Table(name)
+		t, ok := findT(s, name)
 		if ok {
 			matches = append(matches, t)
 		}
@@ -690,9 +773,9 @@ func findTable(ref *schemahcl.Ref, sch *schema.Schema) (*schema.Table, error) {
 	case 1:
 		return matches[0], nil
 	case 0:
-		return nil, fmt.Errorf("sqlspec: table %q not found", name)
+		return nil, fmt.Errorf("specutil: refrenced table/view %q not found", name)
 	default:
-		return nil, fmt.Errorf("specutil: multiple tables found for %q", name)
+		return nil, fmt.Errorf("specutil: multiple refrences tables/views found for %q", name)
 	}
 }
 
@@ -712,6 +795,22 @@ func tableName(ref *schemahcl.Ref) (qualifier, name string, err error) {
 	return
 }
 
+func viewName(ref *schemahcl.Ref) (qualifier, name string, err error) {
+	vs, err := ref.ByType("view")
+	if err != nil {
+		return "", "", err
+	}
+	switch len(vs) {
+	case 1:
+		name = vs[0]
+	case 2:
+		qualifier, name = vs[0], vs[1]
+	default:
+		return "", "", fmt.Errorf("sqlspec: unexpected number of references in %q", vs)
+	}
+	return
+}
+
 func isLocalRef(r *schemahcl.Ref) bool {
 	return strings.HasPrefix(r.V, "$column")
 }
diff --git a/sql/mysql/sqlspec.go b/sql/mysql/sqlspec.go
index ee85390f..b4843862 100644
--- a/sql/mysql/sqlspec.go
+++ b/sql/mysql/sqlspec.go
@@ -33,8 +33,10 @@ func evalSpec(p *hclparse.Parser, v any, input map[string]cty.Value) error {
 		if err := hclState.Eval(p, &d, input); err != nil {
 			return err
 		}
-		err := specutil.Scan(v, d.Schemas, d.Tables, convertTable)
-		if err != nil {
+		if err := specutil.Scan(v,
+			&specutil.ScanDoc{Schemas: d.Schemas, Tables: d.Tables},
+			&specutil.ScanFuncs{Table: convertTable},
+		); err != nil {
 			return fmt.Errorf("mysql: failed converting to *schema.Realm: %w", err)
 		}
 		for _, spec := range d.Schemas {
@@ -54,8 +56,11 @@ func evalSpec(p *hclparse.Parser, v any, input map[string]cty.Value) error {
 		if len(d.Schemas) != 1 {
 			return fmt.Errorf("mysql: expecting document to contain a single schema, got %d", len(d.Schemas))
 		}
-		var r schema.Realm
-		if err := specutil.Scan(&r, d.Schemas, d.Tables, convertTable); err != nil {
+		r := &schema.Realm{}
+		if err := specutil.Scan(r,
+			&specutil.ScanDoc{Schemas: d.Schemas, Tables: d.Tables},
+			&specutil.ScanFuncs{Table: convertTable},
+		); err != nil {
 			return err
 		}
 		if err := convertCharset(d.Schemas[0], &r.Schemas[0].Attrs); err != nil {
diff --git a/sql/postgres/sqlspec.go b/sql/postgres/sqlspec.go
index 7987a7dc..c2d3af8c 100644
--- a/sql/postgres/sqlspec.go
+++ b/sql/postgres/sqlspec.go
@@ -49,7 +49,10 @@ func evalSpec(p *hclparse.Parser, v any, input map[string]cty.Value) error {
 		if err := hclState.Eval(p, &d, input); err != nil {
 			return err
 		}
-		if err := specutil.Scan(v, d.Schemas, d.Tables, convertTable); err != nil {
+		if err := specutil.Scan(v,
+			&specutil.ScanDoc{Schemas: d.Schemas, Tables: d.Tables, Views: d.Views},
+			&specutil.ScanFuncs{Table: convertTable, View: convertView},
+		); err != nil {
 			return fmt.Errorf("specutil: failed converting to *schema.Realm: %w", err)
 		}
 		if len(d.Enums) > 0 {
@@ -66,7 +69,10 @@ func evalSpec(p *hclparse.Parser, v any, input map[string]cty.Value) error {
 			return fmt.Errorf("specutil: expecting document to contain a single schema, got %d", len(d.Schemas))
 		}
 		r := &schema.Realm{}
-		if err := specutil.Scan(r, d.Schemas, d.Tables, convertTable); err != nil {
+		if err := specutil.Scan(r,
+			&specutil.ScanDoc{Schemas: d.Schemas, Tables: d.Tables, Views: d.Views},
+			&specutil.ScanFuncs{Table: convertTable, View: convertView},
+		); err != nil {
 			return err
 		}
 		if err := convertEnums(d.Tables, d.Enums, r); err != nil {
@@ -120,6 +126,7 @@ func MarshalSpec(v any, marshaler schemahcl.Marshaler) ([]byte, error) {
 
 var (
 	hclState = schemahcl.New(
+		schemahcl.WithTypes("view.column.type", TypeRegistry.Specs()),
 		schemahcl.WithTypes("table.column.type", TypeRegistry.Specs()),
 		schemahcl.WithScopedEnums("table.index.type", IndexTypeBTree, IndexTypeBRIN, IndexTypeHash, IndexTypeGIN, IndexTypeGiST, "GiST", IndexTypeSPGiST, "SPGiST"),
 		schemahcl.WithScopedEnums("table.partition.type", PartitionTypeRange, PartitionTypeList, PartitionTypeHash),
@@ -160,6 +167,17 @@ func convertTable(spec *sqlspec.Table, parent *schema.Schema) (*schema.Table, er
 	return t, nil
 }
 
+// convertView converts a sqlspec.View to a schema.View.
+func convertView(spec *sqlspec.View, parent *schema.Schema) (*schema.View, error) {
+	v, err := specutil.View(spec, parent, func(c *sqlspec.Column, _ *schema.View) (*schema.Column, error) {
+		return specutil.Column(c, convertColumnType)
+	})
+	if err != nil {
+		return nil, err
+	}
+	return v, nil
+}
+
 // convertPartition converts and appends the partition block into the table attributes if exists.
 func convertPartition(spec schemahcl.Resource, table *schema.Table) error {
 	r, ok := spec.Resource("partition")
diff --git a/sql/sqlite/sqlspec.go b/sql/sqlite/sqlspec.go
index 3bbcb493..d039ddb8 100644
--- a/sql/sqlite/sqlspec.go
+++ b/sql/sqlite/sqlspec.go
@@ -27,8 +27,10 @@ func evalSpec(p *hclparse.Parser, v any, input map[string]cty.Value) error {
 		if err := hclState.Eval(p, &d, input); err != nil {
 			return err
 		}
-		err := specutil.Scan(v, d.Schemas, d.Tables, convertTable)
-		if err != nil {
+		if err := specutil.Scan(v,
+			&specutil.ScanDoc{Schemas: d.Schemas, Tables: d.Tables},
+			&specutil.ScanFuncs{Table: convertTable},
+		); err != nil {
 			return fmt.Errorf("specutil: failed converting to *schema.Realm: %w", err)
 		}
 	case *schema.Schema:
@@ -39,8 +41,11 @@ func evalSpec(p *hclparse.Parser, v any, input map[string]cty.Value) error {
 		if len(d.Schemas) != 1 {
 			return fmt.Errorf("specutil: expecting document to contain a single schema, got %d", len(d.Schemas))
 		}
-		var r schema.Realm
-		if err := specutil.Scan(&r, d.Schemas, d.Tables, convertTable); err != nil {
+		r := &schema.Realm{}
+		if err := specutil.Scan(r,
+			&specutil.ScanDoc{Schemas: d.Schemas, Tables: d.Tables},
+			&specutil.ScanFuncs{Table: convertTable},
+		); err != nil {
 			return err
 		}
 		r.Schemas[0].Realm = nil


Test patch:
diff --git a/sql/postgres/sqlspec_test.go b/sql/postgres/sqlspec_test.go
index e254d713..3348053c 100644
--- a/sql/postgres/sqlspec_test.go
+++ b/sql/postgres/sqlspec_test.go
@@ -330,6 +330,63 @@ schema "public" {
 	require.Equal(t, f, string(buf))
 }
 
+func TestUnmarshalViews(t *testing.T) {
+	f := `table "t1" {
+  schema = schema.public
+  column "id" {
+    null = false
+    type = int
+  }
+}
+view "v1" {
+  schema = schema.public
+  as     = "SELECT * FROM t2 WHERE id IS NOT NULL"
+}
+view "v2" {
+ schema = schema.public
+ column "id" {
+   null = false
+   type = int
+ }
+ as      = "SELECT * FROM t3 WHERE id IS NOT NULL ORDER BY id"
+ comment = "view comment"
+}
+view "v3" {
+ schema     = schema.public
+ as         = "SELECT * FROM v2 JOIN t1 USING (id)"
+ depends_on = [view.v1, table.t1]
+}
+schema "public" {
+}
+`
+	var (
+		r   schema.Realm
+		s   schema.Schema
+		exp = schema.New("public").
+			AddTables(
+				schema.NewTable("t1").
+					AddColumns(
+						schema.NewIntColumn("id", "int"),
+					),
+			).
+			AddViews(
+				schema.NewView("v1", "SELECT * FROM t2 WHERE id IS NOT NULL"),
+				schema.NewView("v2", "SELECT * FROM t3 WHERE id IS NOT NULL ORDER BY id").
+					AddColumns(
+						schema.NewIntColumn("id", "int"),
+					).
+					SetComment("view comment"),
+			)
+	)
+	exp.AddViews(
+		schema.NewView("v3", "SELECT * FROM v2 JOIN t1 USING (id)").
+			AddDeps(exp.Views[0], exp.Tables[0]),
+	)
+	r.AddSchemas(exp)
+	require.NoError(t, EvalHCLBytes([]byte(f), &s, nil))
+	require.EqualValues(t, exp, &s)
+}
+
 func TestUnmarshalSpec_IndexType(t *testing.T) {
 	f := `
 schema "s" {}