Repository: ariga/atlas
Pull request #2393 (2023-12-21 07:37:30)

sql/sqlspec: add triggers to document


Patch:
diff --git a/sql/internal/specutil/convert.go b/sql/internal/specutil/convert.go
index 08b95cf9..394925de 100644
--- a/sql/internal/specutil/convert.go
+++ b/sql/internal/specutil/convert.go
@@ -48,6 +48,7 @@ type (
 		Materialized []*sqlspec.View
 		Funcs        []*sqlspec.Func
 		Procs        []*sqlspec.Func
+		Triggers     []*sqlspec.Trigger
 	}
 
 	// ScanFuncs represents a set of scan functions
@@ -57,11 +58,13 @@ type (
 		View  ConvertViewFunc
 		Func  func(*sqlspec.Func) (*schema.Func, error)
 		Proc  func(*sqlspec.Func) (*schema.Proc, error)
+		// The trigger add itself to the relevant table/view.
+		Trigger func(*schema.Realm, *sqlspec.Trigger) error
 	}
 
-	// Funcs represents a set of spec functions
-	// used to convert the Realm to an HCL document.
-	Funcs struct {
+	// SchemaFuncs represents a set of spec functions
+	// used to convert the Schema object to an HCL document.
+	SchemaFuncs struct {
 		Table TableSpecFunc
 		View  ViewSpecFunc
 		Func  func(*schema.Func) (*sqlspec.Func, error)
@@ -233,6 +236,13 @@ func Scan(r *schema.Realm, doc *ScanDoc, funcs *ScanFuncs) error {
 			}
 		}
 	}
+	if funcs.Trigger != nil {
+		for _, st := range doc.Triggers {
+			if err := funcs.Trigger(r, st); err != nil {
+				return fmt.Errorf("cannot convert trigger %q: %w", st.Name, err)
+			}
+		}
+	}
 	for o, refs := range deps {
 		var err error
 		switch o := o.(type) {
@@ -381,7 +391,7 @@ func Default(d cty.Value) (schema.Expr, error) {
 		}
 		x = &schema.RawExpr{X: raw.X}
 	default:
-		return nil, fmt.Errorf("unsupported value type for default: %T", d)
+		return nil, fmt.Errorf("unsupported type for default value: %T", d)
 	}
 	return x, nil
 }
@@ -513,7 +523,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 *SchemaFuncs) (*SchemaSpec, error) {
 	spec := &SchemaSpec{
 		Schema: &sqlspec.Schema{
 			Name: s.Name,
@@ -531,6 +541,7 @@ func FromSchema(s *schema.Schema, funcs *Funcs) (*SchemaSpec, error) {
 			table.Schema = SchemaRef(s.Name)
 		}
 		spec.Tables = append(spec.Tables, table)
+		spec.Triggers = append(spec.Triggers, t.Triggers...)
 	}
 	for _, v := range s.Views {
 		view, err := funcs.View(v)
@@ -545,6 +556,7 @@ func FromSchema(s *schema.Schema, funcs *Funcs) (*SchemaSpec, error) {
 		} else {
 			spec.Views = append(spec.Views, view)
 		}
+		spec.Triggers = append(spec.Triggers, v.Triggers...)
 	}
 	if funcs.Func != nil {
 		for _, f := range s.Funcs {
@@ -976,7 +988,7 @@ func SchemaName(ref *schemahcl.Ref) (string, error) {
 }
 
 // ColumnByRef returns a column from the table by its reference.
-func ColumnByRef(t *schema.Table, ref *schemahcl.Ref) (*schema.Column, error) {
+func ColumnByRef[T *schema.View | *schema.Table](tv T, ref *schemahcl.Ref) (*schema.Column, error) {
 	vs, err := ref.ByType(typeColumn)
 	if err != nil {
 		return nil, err
@@ -984,11 +996,22 @@ func ColumnByRef(t *schema.Table, ref *schemahcl.Ref) (*schema.Column, error) {
 	if len(vs) != 1 {
 		return nil, fmt.Errorf("expected 1 column ref, got %d", len(vs))
 	}
-	c, ok := t.Column(vs[0])
-	if !ok {
-		return nil, fmt.Errorf("unknown column %q in table %q", vs[0], t.Name)
+	switch tv := any(tv).(type) {
+	case *schema.Table:
+		c, ok := tv.Column(vs[0])
+		if !ok {
+			return nil, fmt.Errorf("column %q was not found in table %s", vs[0], tv.Name)
+		}
+		return c, nil
+	case *schema.View:
+		c, ok := tv.Column(vs[0])
+		if !ok {
+			return nil, fmt.Errorf("column %q was not found in view %s", vs[0], tv.Name)
+		}
+		return c, nil
+	default:
+		return nil, fmt.Errorf("unreachable %T", tv)
 	}
-	return c, nil
 }
 
 func externalRef(ref *schemahcl.Ref, sch *schema.Schema) (*schema.Table, *schema.Column, error) {
diff --git a/sql/internal/specutil/spec.go b/sql/internal/specutil/spec.go
index e71d03df..4a6b7a88 100644
--- a/sql/internal/specutil/spec.go
+++ b/sql/internal/specutil/spec.go
@@ -30,24 +30,37 @@ type (
 		Funcs        []*sqlspec.Func
 		Procs        []*sqlspec.Func
 		Materialized []*sqlspec.View
+		// Collected triggers to convert into spec.
+		Triggers []*schema.Trigger
 	}
-	doc struct {
-		Tables       []*sqlspec.Table  `spec:"table"`
-		Views        []*sqlspec.View   `spec:"view"`
-		Materialized []*sqlspec.View   `spec:"materialized"`
-		Funcs        []*sqlspec.Func   `spec:"function"`
-		Procs        []*sqlspec.Func   `spec:"procedure"`
-		Schemas      []*sqlspec.Schema `spec:"schema"`
+	// RealmFuncs represents the functions that used
+	// to convert the schema.Realm into HCL spec document.
+	RealmFuncs struct {
+		Schema   func(*schema.Schema) (*SchemaSpec, error)
+		Triggers func([]*schema.Trigger, *Doc) ([]*sqlspec.Trigger, error)
+	}
+	// Doc represents the common HCL spec document.
+	Doc struct {
+		Tables       []*sqlspec.Table   `spec:"table"`
+		Views        []*sqlspec.View    `spec:"view"`
+		Materialized []*sqlspec.View    `spec:"materialized"`
+		Funcs        []*sqlspec.Func    `spec:"function"`
+		Procs        []*sqlspec.Func    `spec:"procedure"`
+		Triggers     []*sqlspec.Trigger `spec:"trigger"`
+		Schemas      []*sqlspec.Schema  `spec:"schema"`
 	}
 )
 
 // Marshal marshals v into an Atlas DDL document using a schemahcl.Marshaler. Marshal uses the given
 // schemaSpec function to convert a *schema.Schema into *sqlspec.Schema, []*sqlspec.Table and []*sqlspec.View.
-func Marshal(v any, marshaler schemahcl.Marshaler, convertFunc func(*schema.Schema) (*SchemaSpec, error)) ([]byte, error) {
-	d := &doc{}
+func Marshal(v any, marshaler schemahcl.Marshaler, funcs RealmFuncs) ([]byte, error) {
+	var (
+		d  = &Doc{}
+		ts []*schema.Trigger
+	)
 	switch s := v.(type) {
 	case *schema.Schema:
-		spec, err := convertFunc(s)
+		spec, err := funcs.Schema(s)
 		if err != nil {
 			return nil, fmt.Errorf("specutil: failed converting schema to spec: %w", err)
 		}
@@ -57,9 +70,10 @@ func Marshal(v any, marshaler schemahcl.Marshaler, convertFunc func(*schema.Sche
 		d.Schemas = []*sqlspec.Schema{spec.Schema}
 		d.Funcs = spec.Funcs
 		d.Procs = spec.Procs
+		ts = spec.Triggers
 	case *schema.Realm:
 		for _, s := range s.Schemas {
-			spec, err := convertFunc(s)
+			spec, err := funcs.Schema(s)
 			if err != nil {
 				return nil, fmt.Errorf("specutil: failed converting schema to spec: %w", err)
 			}
@@ -69,6 +83,7 @@ func Marshal(v any, marshaler schemahcl.Marshaler, convertFunc func(*schema.Sche
 			d.Schemas = append(d.Schemas, spec.Schema)
 			d.Funcs = append(d.Funcs, spec.Funcs...)
 			d.Procs = append(d.Procs, spec.Procs...)
+			ts = append(ts, spec.Triggers...)
 		}
 		if err := QualifyObjects(d.Tables); err != nil {
 			return nil, err
@@ -91,6 +106,13 @@ func Marshal(v any, marshaler schemahcl.Marshaler, convertFunc func(*schema.Sche
 	default:
 		return nil, fmt.Errorf("specutil: failed marshaling spec. %T is not supported", v)
 	}
+	if funcs.Triggers != nil {
+		specs, err := funcs.Triggers(ts, d)
+		if err != nil {
+			return nil, err
+		}
+		d.Triggers = specs
+	}
 	return marshaler.MarshalSpec(d)
 }
 
diff --git a/sql/mysql/driver_oss.go b/sql/mysql/driver_oss.go
index dbfba40f..283f64ed 100644
--- a/sql/mysql/driver_oss.go
+++ b/sql/mysql/driver_oss.go
@@ -16,7 +16,7 @@ import (
 
 var (
 	specOptions []schemahcl.Option
-	specFuncs   = &specutil.Funcs{
+	specFuncs   = &specutil.SchemaFuncs{
 		Table: tableSpec,
 		View:  viewSpec,
 	}
diff --git a/sql/mysql/sqlspec.go b/sql/mysql/sqlspec.go
index 68a0b201..52b19ae4 100644
--- a/sql/mysql/sqlspec.go
+++ b/sql/mysql/sqlspec.go
@@ -80,7 +80,9 @@ func evalSpec(p *hclparse.Parser, v any, input map[string]cty.Value) error {
 
 // MarshalSpec marshals v into an Atlas DDL document using a schemahcl.Marshaler.
 func MarshalSpec(v any, marshaler schemahcl.Marshaler) ([]byte, error) {
-	return specutil.Marshal(v, marshaler, schemaSpec)
+	return specutil.Marshal(v, marshaler, specutil.RealmFuncs{
+		Schema: schemaSpec,
+	})
 }
 
 var (
diff --git a/sql/postgres/driver_oss.go b/sql/postgres/driver_oss.go
index d6c16735..9332240d 100644
--- a/sql/postgres/driver_oss.go
+++ b/sql/postgres/driver_oss.go
@@ -20,7 +20,7 @@ import (
 
 var (
 	specOptions []schemahcl.Option
-	specFuncs   = &specutil.Funcs{
+	specFuncs   = &specutil.SchemaFuncs{
 		Table: tableSpec,
 		View:  viewSpec,
 	}
@@ -30,6 +30,10 @@ var (
 	}
 )
 
+func triggersSpec([]*schema.Trigger, *doc) error {
+	return nil // unimplemented.
+}
+
 func (*inspect) inspectViews(context.Context, *schema.Realm, *schema.InspectOptions) error {
 	return nil // unimplemented.
 }
diff --git a/sql/postgres/sqlspec.go b/sql/postgres/sqlspec.go
index c8965703..6166ed5f 100644
--- a/sql/postgres/sqlspec.go
+++ b/sql/postgres/sqlspec.go
@@ -23,15 +23,16 @@ import (
 
 type (
 	doc struct {
-		Tables       []*sqlspec.Table  `spec:"table"`
-		Views        []*sqlspec.View   `spec:"view"`
-		Materialized []*sqlspec.View   `spec:"materialized"`
-		Enums        []*enum           `spec:"enum"`
-		Domains      []*domain         `spec:"domain"`
-		Sequences    []*sequence       `spec:"sequence"`
-		Funcs        []*sqlspec.Func   `spec:"function"`
-		Procs        []*sqlspec.Func   `spec:"procedure"`
-		Schemas      []*sqlspec.Schema `spec:"schema"`
+		Tables       []*sqlspec.Table   `spec:"table"`
+		Views        []*sqlspec.View    `spec:"view"`
+		Materialized []*sqlspec.View    `spec:"materialized"`
+		Enums        []*enum            `spec:"enum"`
+		Domains      []*domain          `spec:"domain"`
+		Sequences    []*sequence        `spec:"sequence"`
+		Funcs        []*sqlspec.Func    `spec:"function"`
+		Procs        []*sqlspec.Func    `spec:"procedure"`
+		Triggers     []*sqlspec.Trigger `spec:"trigger"`
+		Schemas      []*sqlspec.Schema  `spec:"schema"`
 	}
 
 	// Enum holds a specification for an enum type.
@@ -76,9 +77,22 @@ func (d *doc) merge(d1 *doc) {
 	d.Domains = append(d.Domains, d1.Domains...)
 	d.Schemas = append(d.Schemas, d1.Schemas...)
 	d.Sequences = append(d.Sequences, d1.Sequences...)
+	d.Triggers = append(d.Triggers, d1.Triggers...)
 	d.Materialized = append(d.Materialized, d1.Materialized...)
 }
 
+func (d *doc) ScanDoc() *specutil.ScanDoc {
+	return &specutil.ScanDoc{
+		Schemas:      d.Schemas,
+		Tables:       d.Tables,
+		Views:        d.Views,
+		Funcs:        d.Funcs,
+		Procs:        d.Procs,
+		Triggers:     d.Triggers,
+		Materialized: d.Materialized,
+	}
+}
+
 // Label returns the defaults label used for the enum resource.
 func (e *enum) Label() string { return e.Name }
 
@@ -129,10 +143,7 @@ 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,
-			&specutil.ScanDoc{Schemas: d.Schemas, Tables: d.Tables, Views: d.Views, Materialized: d.Materialized, Funcs: d.Funcs, Procs: d.Procs},
-			scanFuncs,
-		); err != nil {
+		if err := specutil.Scan(v, d.ScanDoc(), scanFuncs); err != nil {
 			return fmt.Errorf("specutil: failed converting to *schema.Realm: %w", err)
 		}
 		if err := convertEnums(d.Tables, d.Enums, v); err != nil {
@@ -153,10 +164,7 @@ 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,
-			&specutil.ScanDoc{Schemas: d.Schemas, Tables: d.Tables, Views: d.Views, Materialized: d.Materialized, Funcs: d.Funcs, Procs: d.Procs},
-			scanFuncs,
-		); err != nil {
+		if err := specutil.Scan(r, d.ScanDoc(), scanFuncs); err != nil {
 			return err
 		}
 		if err := convertEnums(d.Tables, d.Enums, r); err != nil {
@@ -179,21 +187,26 @@ func evalSpec(p *hclparse.Parser, v any, input map[string]cty.Value) error {
 
 // MarshalSpec marshals v into an Atlas DDL document using a schemahcl.Marshaler.
 func MarshalSpec(v any, marshaler schemahcl.Marshaler) ([]byte, error) {
-	var d doc
+	var (
+		d  doc
+		ts []*schema.Trigger
+	)
 	switch s := v.(type) {
 	case *schema.Schema:
-		d1, err := schemaSpec(s)
+		d1, trs, err := schemaSpec(s)
 		if err != nil {
 			return nil, fmt.Errorf("specutil: failed converting schema to spec: %w", err)
 		}
+		ts = trs
 		d.merge(d1)
 	case *schema.Realm:
 		for _, s := range s.Schemas {
-			d1, err := schemaSpec(s)
+			d1, trs, err := schemaSpec(s)
 			if err != nil {
 				return nil, fmt.Errorf("specutil: failed converting schema to spec: %w", err)
 			}
 			d.merge(d1)
+			ts = append(ts, trs...)
 		}
 		if err := specutil.QualifyObjects(d.Tables); err != nil {
 			return nil, err
@@ -228,6 +241,9 @@ func MarshalSpec(v any, marshaler schemahcl.Marshaler) ([]byte, error) {
 	default:
 		return nil, fmt.Errorf("specutil: failed marshaling spec. %T is not supported", v)
 	}
+	if err := triggersSpec(ts, &d); err != nil {
+		return nil, err
+	}
 	return marshaler.MarshalSpec(&d)
 }
 
@@ -658,10 +674,10 @@ func enumRef(n string) *schemahcl.Ref {
 }
 
 // schemaSpec converts from a concrete Postgres schema to Atlas specification.
-func schemaSpec(s *schema.Schema) (*doc, error) {
+func schemaSpec(s *schema.Schema) (*doc, []*schema.Trigger, error) {
 	spec, err := specutil.FromSchema(s, specFuncs)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 	d := &doc{
 		Tables:       spec.Tables,
@@ -674,9 +690,9 @@ func schemaSpec(s *schema.Schema) (*doc, error) {
 		Domains:      make([]*domain, 0, len(s.Objects)),
 	}
 	if err := objectSpec(d, spec, s); err != nil {
-		return nil, err
+		return nil, nil, err
 	}
-	return d, nil
+	return d, spec.Triggers, nil
 }
 
 // tableSpec converts from a concrete Postgres sqlspec.Table to a schema.Table.
diff --git a/sql/sqlite/sqlspec.go b/sql/sqlite/sqlspec.go
index 7cc339d5..c101f838 100644
--- a/sql/sqlite/sqlspec.go
+++ b/sql/sqlite/sqlspec.go
@@ -65,7 +65,9 @@ func evalSpec(p *hclparse.Parser, v any, input map[string]cty.Value) error {
 
 // MarshalSpec marshals v into an Atlas DDL document using a schemahcl.Marshaler.
 func MarshalSpec(v any, marshaler schemahcl.Marshaler) ([]byte, error) {
-	return specutil.Marshal(v, marshaler, schemaSpec)
+	return specutil.Marshal(v, marshaler, specutil.RealmFuncs{
+		Schema: schemaSpec,
+	})
 }
 
 // convertTable converts a sqlspec.Table to a schema.Table. Table conversion is done without converting
@@ -158,7 +160,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.SchemaFuncs{
 		Table: tableSpec,
 		View:  viewSpec,
 	})
diff --git a/sql/sqlspec/sqlspec.go b/sql/sqlspec/sqlspec.go
index 6c2e400b..e72e337e 100644
--- a/sql/sqlspec/sqlspec.go
+++ b/sql/sqlspec/sqlspec.go
@@ -114,6 +114,14 @@ type (
 		// as their definition can be either a string or an enum (ref).
 		schemahcl.DefaultExtension
 	}
+
+	// Trigger holds the specification for a trigger.
+	Trigger struct {
+		Name string         `spec:",name"`
+		On   *schemahcl.Ref `spec:"on"` // A table or a view.
+		// Attributes and blocks are different for each driver.
+		schemahcl.DefaultExtension
+	}
 )
 
 // Label returns the defaults label used for the table resource.
@@ -158,5 +166,6 @@ func init() {
 	schemahcl.Register("table", &Table{})
 	schemahcl.Register("function", &Func{})
 	schemahcl.Register("procedure", &Func{})
+	schemahcl.Register("trigger", &Trigger{})
 	schemahcl.Register("schema", &Schema{})
 }


Test patch:
diff --git a/sql/internal/specutil/convert_test.go b/sql/internal/specutil/convert_test.go
index 9bb609d6..e16695c3 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, &SchemaFuncs{
 		Table: func(*schema.Table) (*sqlspec.Table, error) {
 			return &sqlspec.Table{}, nil
 		},