new dbaction: createfkconstraint

Author: kvch

pkg/migrations/dbactions.go

@@ -394,3 +394,49 @@ func (a *validateConstraintAction) Execute(ctx context.Context) error {
 		pq.QuoteIdentifier(a.constraint)))
 	return err
 }
+
+// createFKConstraintAction is a DBAction that creates a new foreign key constraint
+type createFKConstraintAction struct {
+	conn              db.DB
+	table             string
+	constraint        string
+	columns           []string
+	initiallyDeferred bool
+	deferrable        bool
+	reference         *TableForeignKeyReference
+	skipValidation    bool
+}
+
+func NewCreateFKConstraintAction(conn db.DB, table, constraint string, columns []string, reference *TableForeignKeyReference, initiallyDeferred, deferrable, skipValidation bool) *createFKConstraintAction {
+	return &createFKConstraintAction{
+		conn:              conn,
+		table:             table,
+		constraint:        constraint,
+		columns:           columns,
+		reference:         reference,
+		initiallyDeferred: initiallyDeferred,
+		deferrable:        deferrable,
+		skipValidation:    skipValidation,
+	}
+}
+
+func (a *createFKConstraintAction) Execute(ctx context.Context) error {
+	sql := fmt.Sprintf("ALTER TABLE %s ADD ", pq.QuoteIdentifier(a.table))
+	writer := &ConstraintSQLWriter{
+		Name:              a.constraint,
+		Columns:           a.columns,
+		InitiallyDeferred: a.initiallyDeferred,
+		Deferrable:        a.deferrable,
+		SkipValidation:    a.skipValidation,
+	}
+	sql += writer.WriteForeignKey(
+		a.reference.Table,
+		a.reference.Columns,
+		a.reference.OnDelete,
+		a.reference.OnUpdate,
+		a.reference.OnDeleteSetColumns,
+		a.reference.MatchType)
+
+	_, err := a.conn.ExecContext(ctx, sql)
+	return err
+}

pkg/migrations/op_create_constraint.go

@@ -95,7 +95,7 @@ func (o *OpCreateConstraint) Start(ctx context.Context, l Logger, conn db.DB, la
 	case OpCreateConstraintTypeCheck:
 		return table, o.addCheckConstraint(ctx, conn, table.Name)
 	case OpCreateConstraintTypeForeignKey:
-		return table, o.addForeignKeyConstraint(ctx, conn, table)
+		return table, NewCreateFKConstraintAction(conn, o.Name, table.Name, temporaryNames(o.Columns), o.References, false, false, true).Execute(ctx)
 	}
 
 	return table, nil
@@ -296,28 +296,6 @@ func (o *OpCreateConstraint) addCheckConstraint(ctx context.Context, conn db.DB,
 	return err
 }
 
-func (o *OpCreateConstraint) addForeignKeyConstraint(ctx context.Context, conn db.DB, table *schema.Table) error {
-	sql := fmt.Sprintf("ALTER TABLE %s ADD ", pq.QuoteIdentifier(table.Name))
-
-	writer := &ConstraintSQLWriter{
-		Name:           o.Name,
-		Columns:        temporaryNames(o.Columns),
-		SkipValidation: true,
-	}
-	sql += writer.WriteForeignKey(
-		o.References.Table,
-		o.References.Columns,
-		o.References.OnDelete,
-		o.References.OnUpdate,
-		o.References.OnDeleteSetColumns,
-		o.References.MatchType,
-	)
-
-	_, err := conn.ExecContext(ctx, sql)
-
-	return err
-}
-
 func temporaryNames(columns []string) []string {
 	names := make([]string, len(columns))
 	for i, col := range columns {

pkg/migrations/op_set_fk.go

@@ -6,8 +6,6 @@ import (
 	"context"
 	"fmt"
 
-	"github.com/lib/pq"
-
 	"github.com/xataio/pgroll/pkg/db"
 	"github.com/xataio/pgroll/pkg/schema"
 )
@@ -26,9 +24,39 @@ func (o *OpSetForeignKey) Start(ctx context.Context, l Logger, conn db.DB, lates
 	l.LogOperationStart(o)
 
 	table := s.GetTable(o.Table)
+	if table == nil {
+		return nil, TableDoesNotExistError{Name: o.Table}
+	}
+	column := table.GetColumn(o.Column)
+	if column == nil {
+		return nil, ColumnDoesNotExistError{Table: o.Table, Name: o.Column}
+	}
+	referencedTable := s.GetTable(o.References.Table)
+	if referencedTable == nil {
+		return nil, TableDoesNotExistError{Name: o.References.Table}
+	}
+
+	referencedColumn := referencedTable.GetColumn(o.References.Column)
+	if referencedColumn == nil {
+		return nil, ColumnDoesNotExistError{Table: o.References.Table, Name: o.References.Column}
+	}
 
 	// Create a NOT VALID foreign key constraint on the new column.
-	if err := o.addForeignKeyConstraint(ctx, conn, s); err != nil {
+	if err := NewCreateFKConstraintAction(conn,
+		o.Table,
+		o.References.Name,
+		[]string{o.Column},
+		&TableForeignKeyReference{
+			Table:     o.References.Table,
+			Columns:   []string{o.References.Column},
+			MatchType: o.References.MatchType,
+			OnDelete:  o.References.OnDelete,
+			OnUpdate:  o.References.OnUpdate,
+		},
+		o.References.InitiallyDeferred,
+		o.References.Deferrable,
+		true,
+	).Execute(ctx); err != nil {
 		return nil, fmt.Errorf("failed to add foreign key constraint: %w", err)
 	}
 
@@ -84,40 +112,3 @@ func (o *OpSetForeignKey) Validate(ctx context.Context, s *schema.Schema) error
 
 	return nil
 }
-
-func (o *OpSetForeignKey) addForeignKeyConstraint(ctx context.Context, conn db.DB, s *schema.Schema) error {
-	table := s.GetTable(o.Table)
-	if table == nil {
-		return TableDoesNotExistError{Name: o.Table}
-	}
-	column := table.GetColumn(o.Column)
-	if column == nil {
-		return ColumnDoesNotExistError{Table: o.Table, Name: o.Column}
-	}
-	referencedTable := s.GetTable(o.References.Table)
-	if referencedTable == nil {
-		return TableDoesNotExistError{Name: o.References.Table}
-	}
-
-	referencedColumn := referencedTable.GetColumn(o.References.Column)
-	if referencedColumn == nil {
-		return ColumnDoesNotExistError{Table: o.References.Table, Name: o.References.Column}
-	}
-
-	sql := fmt.Sprintf("ALTER TABLE %s ADD ", pq.QuoteIdentifier(table.Name))
-	writer := &ConstraintSQLWriter{
-		Name:           o.References.Name,
-		Columns:        []string{column.Name},
-		SkipValidation: true,
-	}
-	sql += writer.WriteForeignKey(
-		referencedTable.Name,
-		[]string{referencedColumn.Name},
-		o.References.OnDelete,
-		o.References.OnUpdate,
-		nil,
-		o.References.MatchType)
-
-	_, err := conn.ExecContext(ctx, sql)
-	return err
-}