Repository: ariga/atlas
Pull request #2146 (2023-10-01 20:03:56)

sql/internal/specutil: add procs/funcs scan functions


Patch:
diff --git a/sql/internal/specutil/convert.go b/sql/internal/specutil/convert.go
index 17c433eb..4c06540e 100644
--- a/sql/internal/specutil/convert.go
+++ b/sql/internal/specutil/convert.go
@@ -46,6 +46,8 @@ type (
 		Tables       []*sqlspec.Table
 		Views        []*sqlspec.View
 		Materialized []*sqlspec.View
+		Funcs        []*sqlspec.Func
+		Procs        []*sqlspec.Func
 	}
 
 	// ScanFuncs represents a set of scan functions
@@ -53,6 +55,8 @@ type (
 	ScanFuncs struct {
 		Table ConvertTableFunc
 		View  ConvertViewFunc
+		Func  func(*sqlspec.Func) (*schema.Func, error)
+		Proc  func(*sqlspec.Func) (*schema.Proc, error)
 	}
 
 	// Funcs represents a set of spec functions
@@ -203,6 +207,40 @@ func Scan(r *schema.Realm, doc *ScanDoc, funcs *ScanFuncs) error {
 			}
 		}
 	}
+	if funcs.Func != nil {
+		for _, sf := range doc.Funcs {
+			name, err := SchemaName(sf.Schema)
+			if err != nil {
+				return fmt.Errorf("specutil: cannot extract schema name for function %q: %w", sf.Name, err)
+			}
+			s, ok := byName[name]
+			if !ok {
+				return fmt.Errorf("specutil: schema %q not found for function %q", name, sf.Name)
+			}
+			f, err := funcs.Func(sf)
+			if err != nil {
+				return fmt.Errorf("specutil: cannot convert function %q: %w", sf.Name, err)
+			}
+			s.AddFuncs(f)
+		}
+	}
+	if funcs.Proc != nil {
+		for _, sf := range doc.Procs {
+			name, err := SchemaName(sf.Schema)
+			if err != nil {
+				return fmt.Errorf("specutil: cannot extract schema name for procedure %q: %w", sf.Name, err)
+			}
+			s, ok := byName[name]
+			if !ok {
+				return fmt.Errorf("specutil: schema %q not found for procedure %q", name, sf.Name)
+			}
+			f, err := funcs.Proc(sf)
+			if err != nil {
+				return fmt.Errorf("specutil: cannot convert procedure %q: %w", sf.Name, err)
+			}
+			s.AddProcs(f)
+		}
+	}
 	return nil
 }
 
@@ -295,24 +333,11 @@ func Column(spec *sqlspec.Column, conv ConvertTypeFunc) (*schema.Column, error)
 			Null: spec.Null,
 		},
 	}
-	if d := spec.Default; !d.IsNull() {
-		switch {
-		case d.Type() == cty.String:
-			out.Default = &schema.Literal{V: d.AsString()}
-		case d.Type() == cty.Number:
-			out.Default = &schema.Literal{V: d.AsBigFloat().String()}
-		case d.Type() == cty.Bool:
-			out.Default = &schema.Literal{V: strconv.FormatBool(d.True())}
-		case d.Type().IsCapsuleType():
-			x, ok := d.EncapsulatedValue().(*schemahcl.RawExpr)
-			if !ok {
-				return nil, fmt.Errorf("invalid default value %q", d.Type().FriendlyName())
-			}
-			out.Default = &schema.RawExpr{X: x.X}
-		default:
-			return nil, fmt.Errorf("unsupported value type for default: %T", d)
-		}
+	d, err := Default(spec.Default)
+	if err != nil {
+		return nil, err
 	}
+	out.Default = d
 	ct, err := conv(spec)
 	if err != nil {
 		return nil, err
@@ -324,6 +349,31 @@ func Column(spec *sqlspec.Column, conv ConvertTypeFunc) (*schema.Column, error)
 	return out, err
 }
 
+// Default converts a cty.Value (as defined in the spec) into a schema.Expr.
+func Default(d cty.Value) (schema.Expr, error) {
+	if d.IsNull() {
+		return nil, nil // no default.
+	}
+	var x schema.Expr
+	switch {
+	case d.Type() == cty.String:
+		x = &schema.Literal{V: d.AsString()}
+	case d.Type() == cty.Number:
+		x = &schema.Literal{V: d.AsBigFloat().String()}
+	case d.Type() == cty.Bool:
+		x = &schema.Literal{V: strconv.FormatBool(d.True())}
+	case d.Type().IsCapsuleType():
+		raw, ok := d.EncapsulatedValue().(*schemahcl.RawExpr)
+		if !ok {
+			return nil, fmt.Errorf("invalid default value %q", d.Type().FriendlyName())
+		}
+		x = &schema.RawExpr{X: raw.X}
+	default:
+		return nil, fmt.Errorf("unsupported value type for default: %T", d)
+	}
+	return x, nil
+}
+
 // Index converts a sqlspec.Index to a schema.Index. The optional arguments allow
 // passing functions for mutating the created index-part (e.g. add attributes).
 func Index(spec *sqlspec.Index, parent *schema.Table, partFns ...func(*sqlspec.IndexPart, *schema.IndexPart) error) (*schema.Index, error) {
@@ -451,7 +501,7 @@ func linkForeignKeys(tbl *schema.Table, fks []*sqlspec.ForeignKey) error {
 }
 
 // FromSchema converts a schema.Schema into sqlspec.Schema and []sqlspec.Table.
-func FromSchema(s *schema.Schema, funcs Funcs) (*SchemaSpec, error) {
+func FromSchema(s *schema.Schema, funcs *Funcs) (*SchemaSpec, error) {
 	spec := &SchemaSpec{
 		Schema: &sqlspec.Schema{
 			Name: s.Name,
diff --git a/sql/mysql/sqlspec.go b/sql/mysql/sqlspec.go
index b3d58f90..31ade3b4 100644
--- a/sql/mysql/sqlspec.go
+++ b/sql/mysql/sqlspec.go
@@ -276,7 +276,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, specutil.Funcs{
+	spec, err := specutil.FromSchema(s, &specutil.Funcs{
 		Table: tableSpec,
 		View:  viewSpec,
 	})
diff --git a/sql/postgres/driver_oss.go b/sql/postgres/driver_oss.go
index d1034a42..293a553e 100644
--- a/sql/postgres/driver_oss.go
+++ b/sql/postgres/driver_oss.go
@@ -16,10 +16,14 @@ import (
 
 var (
 	specOptions []schemahcl.Option
-	specFuncs   = specutil.Funcs{
+	specFuncs   = &specutil.Funcs{
 		Table: tableSpec,
 		View:  viewSpec,
 	}
+	scanFuncs = &specutil.ScanFuncs{
+		Table: convertTable,
+		View:  convertView,
+	}
 )
 
 func (*inspect) inspectViews(context.Context, *schema.Realm, *schema.InspectOptions) error {
diff --git a/sql/postgres/sqlspec.go b/sql/postgres/sqlspec.go
index 1be2197d..76810cbf 100644
--- a/sql/postgres/sqlspec.go
+++ b/sql/postgres/sqlspec.go
@@ -77,8 +77,8 @@ func evalSpec(p *hclparse.Parser, v any, input map[string]cty.Value) error {
 			return err
 		}
 		if err := specutil.Scan(v,
-			&specutil.ScanDoc{Schemas: d.Schemas, Tables: d.Tables, Views: d.Views, Materialized: d.Materialized},
-			&specutil.ScanFuncs{Table: convertTable, View: convertView},
+			&specutil.ScanDoc{Schemas: d.Schemas, Tables: d.Tables, Views: d.Views, Materialized: d.Materialized, Funcs: d.Funcs, Procs: d.Procs},
+			scanFuncs,
 		); err != nil {
 			return fmt.Errorf("specutil: failed converting to *schema.Realm: %w", err)
 		}
@@ -97,8 +97,8 @@ func evalSpec(p *hclparse.Parser, v any, input map[string]cty.Value) error {
 		}
 		r := &schema.Realm{}
 		if err := specutil.Scan(r,
-			&specutil.ScanDoc{Schemas: d.Schemas, Tables: d.Tables, Views: d.Views, Materialized: d.Materialized},
-			&specutil.ScanFuncs{Table: convertTable, View: convertView},
+			&specutil.ScanDoc{Schemas: d.Schemas, Tables: d.Tables, Views: d.Views, Materialized: d.Materialized, Funcs: d.Funcs, Procs: d.Procs},
+			scanFuncs,
 		); err != nil {
 			return err
 		}
diff --git a/sql/schema/dsl.go b/sql/schema/dsl.go
index b975866d..096119da 100644
--- a/sql/schema/dsl.go
+++ b/sql/schema/dsl.go
@@ -85,6 +85,24 @@ func (s *Schema) AddObjects(objs ...Object) *Schema {
 	return s
 }
 
+// AddFuncs appends the given functions to the schema.
+func (s *Schema) AddFuncs(funcs ...*Func) *Schema {
+	for _, f := range funcs {
+		f.Schema = s
+	}
+	s.Funcs = append(s.Funcs, funcs...)
+	return s
+}
+
+// AddProcs appends the given procedures to the schema.
+func (s *Schema) AddProcs(procs ...*Proc) *Schema {
+	for _, f := range procs {
+		f.Schema = s
+	}
+	s.Procs = append(s.Procs, procs...)
+	return s
+}
+
 // NewRealm creates a new Realm.
 func NewRealm(schemas ...*Schema) *Realm {
 	r := &Realm{Schemas: schemas}
diff --git a/sql/sqlite/sqlspec.go b/sql/sqlite/sqlspec.go
index 747f2fe3..7cc339d5 100644
--- a/sql/sqlite/sqlspec.go
+++ b/sql/sqlite/sqlspec.go
@@ -158,7 +158,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, specutil.Funcs{
+	return specutil.FromSchema(s, &specutil.Funcs{
 		Table: tableSpec,
 		View:  viewSpec,
 	})


Test patch:
diff --git a/sql/internal/specutil/convert_test.go b/sql/internal/specutil/convert_test.go
index 72528f14..9bb609d6 100644
--- a/sql/internal/specutil/convert_test.go
+++ b/sql/internal/specutil/convert_test.go
@@ -22,7 +22,7 @@ func TestFromSpec_SchemaName(t *testing.T) {
 		},
 	}
 	sc.Tables[0].Schema = sc
-	spec, err := FromSchema(sc, Funcs{
+	spec, err := FromSchema(sc, &Funcs{
 		Table: func(*schema.Table) (*sqlspec.Table, error) {
 			return &sqlspec.Table{}, nil
 		},