sql/internal/sqlx: allow extending diff with fk attributes (#3148)

sql/internal/specutil/convert.go

@@ -64,6 +64,8 @@ type (
 		Triggers func(*schema.Realm, []*sqlspec.Trigger) error
 		// Objects add themselves to the realm.
 		Objects func(*schema.Realm) error
+		// Optional function to extend the foreign keys.
+		ForeignKey func(*sqlspec.ForeignKey, *schema.ForeignKey) error
 	}
 
 	// SchemaFuncs represents a set of spec functions
@@ -159,7 +161,7 @@ func Scan(r *schema.Realm, doc *ScanDoc, funcs *ScanFuncs) error {
 	}
 	// Link the foreign keys.
 	for t, fks := range fks {
-		if err := linkForeignKeys(t, fks); err != nil {
+		if err := linkForeignKeys(funcs, t, fks); err != nil {
 			return err
 		}
 	}
@@ -511,7 +513,7 @@ 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, fks []*sqlspec.ForeignKey) error {
+func linkForeignKeys(funcs *ScanFuncs, tbl *schema.Table, fks []*sqlspec.ForeignKey) error {
 	for _, spec := range fks {
 		fk := &schema.ForeignKey{Symbol: spec.Symbol, Table: tbl}
 		if spec.OnUpdate != nil {
@@ -546,6 +548,11 @@ func linkForeignKeys(tbl *schema.Table, fks []*sqlspec.ForeignKey) error {
 			fk.RefColumns = append(fk.RefColumns, c)
 		}
 		tbl.ForeignKeys = append(tbl.ForeignKeys, fk)
+		if funcs.ForeignKey != nil {
+			if err := funcs.ForeignKey(spec, fk); err != nil {
+				return err
+			}
+		}
 	}
 	return nil
 }

sql/internal/sqlx/diff.go

@@ -67,6 +67,9 @@ type (
 		// ReferenceChanged reports if the foreign key referential action was
 		// changed. For example, action was changed from RESTRICT to CASCADE.
 		ReferenceChanged(from, to schema.ReferenceOption) bool
+
+		// ForeignKeyAttrChanged reports if any of the foreign-key attributes were changed.
+		ForeignKeyAttrChanged(from, to []schema.Attr) bool
 	}
 
 	// DropSchemaChanger is an optional interface allows DiffDriver to drop
@@ -614,6 +617,9 @@ func (d *Diff) fkChange(from, to *schema.ForeignKey) schema.ChangeKind {
 	if d.ReferenceChanged(from.OnDelete, to.OnDelete) {
 		change |= schema.ChangeDeleteAction
 	}
+	if d.ForeignKeyAttrChanged(from.Attrs, to.Attrs) {
+		change |= schema.ChangeAttr
+	}
 	return change
 }
 

sql/internal/sqlx/sqlx.go

@@ -112,15 +112,33 @@ func SchemaFKs(s *schema.Schema, rows *sql.Rows) error {
 	return TypedSchemaFKs[*nullString](s, rows)
 }
 
+// FKAttrScanner allows extending foreign-keys
+// scanning with extra attributes.
+type FKAttrScanner struct {
+	Columns  func() []any
+	ScanFunc func(*schema.ForeignKey, []any) error
+}
+
 // TypedSchemaFKs is a version of SchemaFKs that allows to specify the type of
 // used to scan update and delete actions from the database.
-func TypedSchemaFKs[T ScanStringer](s *schema.Schema, rows *sql.Rows) error {
+func TypedSchemaFKs[T ScanStringer](s *schema.Schema, rows *sql.Rows, attr ...*FKAttrScanner) error {
 	for rows.Next() {
 		var (
 			updateAction, deleteAction                                   = V(new(T)), V(new(T))
 			name, table, column, tSchema, refTable, refColumn, refSchema string
+			columns                                                      = []any{
+				&name, &table, &column, &tSchema, &refTable, &refColumn, &refSchema, &updateAction, &deleteAction,
+			}
+			origin = len(columns)
 		)
-		if err := rows.Scan(&name, &table, &column, &tSchema, &refTable, &refColumn, &refSchema, &updateAction, &deleteAction); err != nil {
+		// Allows to scan extra attributes.
+		switch {
+		case len(attr) > 1:
+			return fmt.Errorf("unexpected number of attributes scanners: %d", len(attr))
+		case len(attr) == 1:
+			columns = append(columns, attr[0].Columns()...)
+		}
+		if err := rows.Scan(columns...); err != nil {
 			return err
 		}
 		t, ok := s.Table(table)
@@ -172,6 +190,11 @@ func TypedSchemaFKs[T ScanStringer](s *schema.Schema, rows *sql.Rows) error {
 		if _, ok := fk.RefColumn(rc.Name); !ok {
 			fk.RefColumns = append(fk.RefColumns, rc)
 		}
+		if len(attr) == 1 {
+			if err := attr[0].ScanFunc(fk, columns[origin:]); err != nil {
+				return err
+			}
+		}
 	}
 	return nil
 }

sql/mysql/diff.go

@@ -212,6 +212,11 @@ func (*diff) ReferenceChanged(from, to schema.ReferenceOption) bool {
 	return from != to
 }
 
+// ForeignKeyAttrChanged reports if any of the foreign-key attributes were changed.
+func (*diff) ForeignKeyAttrChanged(_, _ []schema.Attr) bool {
+	return false
+}
+
 // Normalize implements the sqlx.Normalizer interface.
 func (d *diff) Normalize(from, to *schema.Table, opts *schema.DiffOptions) error {
 	if opts.Mode.Is(schema.DiffModeNormalized) {

sql/schema/dsl.go

@@ -926,6 +926,12 @@ func (f *ForeignKey) SetOnDelete(o ReferenceOption) *ForeignKey {
 	return f
 }
 
+// AddAttrs adds additional attributes to the schema.
+func (f *ForeignKey) AddAttrs(attrs ...Attr) *ForeignKey {
+	f.Attrs = append(f.Attrs, attrs...)
+	return f
+}
+
 // AddDeps adds the given objects as dependencies to the function.
 func (f *Func) AddDeps(objs ...Object) *Func {
 	f.Deps = append(f.Deps, objs...)

sql/schema/schema.go

@@ -112,6 +112,7 @@ type (
 		RefColumns []*Column
 		OnUpdate   ReferenceOption
 		OnDelete   ReferenceOption
+		Attrs      []Attr
 	}
 
 	// A Trigger represents a trigger definition.

sql/sqlite/diff.go

@@ -166,6 +166,11 @@ func (*diff) ReferenceChanged(from, to schema.ReferenceOption) bool {
 	return from != to
 }
 
+// ForeignKeyAttrChanged reports if any of the foreign-key attributes were changed.
+func (*diff) ForeignKeyAttrChanged(_, _ []schema.Attr) bool {
+	return false
+}
+
 // Normalize implements the sqlx.Normalizer interface.
 func (d *diff) Normalize(from, to *schema.Table, _ *schema.DiffOptions) error {
 	used := make([]bool, len(to.ForeignKeys))