Duplicate unique and check constraints correctly (#466)

Previously, unique and check constraints with multiple columns were not
duplicated correctly.

We had two issues:
1. Pgroll tried to create the same duplicated index for each renamed
column. From now on pgroll creates the index once with the new column
names.
2. Pgroll tried to convert the unique index to a constraint multiple
times. From now on, the index is only converted once.

Required by #464

---------

Co-authored-by: Andrew Farries <[email protected]>

Author: kvch

pkg/migrations/duplicate.go

@@ -17,174 +17,267 @@ import (
 // Duplicator duplicates a column in a table, including all constraints and
 // comments.
 type Duplicator struct {
+	stmtBuilder       *duplicatorStmtBuilder
 	conn              db.DB
-	table             *schema.Table
-	column            *schema.Column
-	asName            string
-	withoutNotNull    bool
-	withType          string
-	withoutConstraint string
+	columns           map[string]*columnToDuplicate
+	withoutConstraint []string
+}
+
+type columnToDuplicate struct {
+	column         *schema.Column
+	asName         string
+	withoutNotNull bool
+	withType       string
+}
+
+// duplicatorStmtBuilder is a helper for building SQL statements to duplicate
+// columns and constraints in a table.
+type duplicatorStmtBuilder struct {
+	table *schema.Table
 }
 
 const (
 	dataTypeMismatchErrorCode  pq.ErrorCode = "42804"
 	undefinedFunctionErrorCode pq.ErrorCode = "42883"
+
+	cCreateUniqueIndexSQL            = `CREATE UNIQUE INDEX CONCURRENTLY %s ON %s (%s)`
+	cSetDefaultSQL                   = `ALTER TABLE %s ALTER COLUMN %s SET DEFAULT %s`
+	cAlterTableAddCheckConstraintSQL = `ALTER TABLE %s ADD CONSTRAINT %s %s NOT VALID`
 )
 
 // NewColumnDuplicator creates a new Duplicator for a column.
-func NewColumnDuplicator(conn db.DB, table *schema.Table, column *schema.Column) *Duplicator {
+func NewColumnDuplicator(conn db.DB, table *schema.Table, columns ...*schema.Column) *Duplicator {
+	cols := make(map[string]*columnToDuplicate, len(columns))
+	for _, column := range columns {
+		cols[column.Name] = &columnToDuplicate{
+			column:   column,
+			asName:   TemporaryName(column.Name),
+			withType: column.Type,
+		}
+	}
 	return &Duplicator{
-		conn:     conn,
-		table:    table,
-		column:   column,
-		asName:   TemporaryName(column.Name),
-		withType: column.Type,
+		stmtBuilder: &duplicatorStmtBuilder{
+			table: table,
+		},
+		conn:              conn,
+		columns:           cols,
+		withoutConstraint: make([]string, 0),
 	}
 }
 
 // WithType sets the type of the new column.
-func (d *Duplicator) WithType(t string) *Duplicator {
-	d.withType = t
+func (d *Duplicator) WithType(columnName, t string) *Duplicator {
+	d.columns[columnName].withType = t
 	return d
 }
 
 // WithoutConstraint excludes a constraint from being duplicated.
 func (d *Duplicator) WithoutConstraint(c string) *Duplicator {
-	d.withoutConstraint = c
+	d.withoutConstraint = append(d.withoutConstraint, c)
 	return d
 }
 
 // WithoutNotNull excludes the NOT NULL constraint from being duplicated.
-func (d *Duplicator) WithoutNotNull() *Duplicator {
-	d.withoutNotNull = true
+func (d *Duplicator) WithoutNotNull(columnName string) *Duplicator {
+	d.columns[columnName].withoutNotNull = true
 	return d
 }
 
 // Duplicate duplicates a column in the table, including all constraints and
 // comments.
 func (d *Duplicator) Duplicate(ctx context.Context) error {
+	colNames := make([]string, 0, len(d.columns))
+	for name, c := range d.columns {
+		colNames = append(colNames, name)
+
+		// Duplicate the column with the new type
+		// and check and fk constraints
+		if sql := d.stmtBuilder.duplicateColumn(c.column, c.asName, c.withoutNotNull, c.withType, d.withoutConstraint); sql != "" {
+			_, err := d.conn.ExecContext(ctx, sql)
+			if err != nil {
+				return err
+			}
+		}
+
+		// Duplicate the column's default value
+		if sql := d.stmtBuilder.duplicateDefault(c.column, c.asName); sql != "" {
+			_, err := d.conn.ExecContext(ctx, sql)
+			err = errorIgnoringErrorCode(err, dataTypeMismatchErrorCode)
+			if err != nil {
+				return err
+			}
+		}
+
+		if sql := d.stmtBuilder.duplicateComment(c.column, c.asName); sql != "" {
+			_, err := d.conn.ExecContext(ctx, sql)
+			if err != nil {
+				return err
+			}
+		}
+	}
+
+	// Generate SQL to duplicate any check constraints on the columns. This may faile
+	// if the check constraint is not valid for the new column type, in which case
+	// the error is ignored.
+	for _, sql := range d.stmtBuilder.duplicateCheckConstraints(d.withoutConstraint, colNames...) {
+		// Update the check constraint expression to use the new column names if any of the columns are duplicated
+		_, err := d.conn.ExecContext(ctx, sql)
+		err = errorIgnoringErrorCode(err, undefinedFunctionErrorCode)
+		if err != nil {
+			return err
+		}
+	}
+
+	// Generate SQL to duplicate any unique constraints on the columns
+	// The constraint is duplicated by adding a unique index on the column concurrently.
+	// The index is converted into a unique constraint on migration completion.
+	for _, sql := range d.stmtBuilder.duplicateUniqueConstraints(d.withoutConstraint, colNames...) {
+		// Update the unique constraint columns to use the new column names if any of the columns are duplicated
+		if _, err := d.conn.ExecContext(ctx, sql); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+func (d *duplicatorStmtBuilder) duplicateCheckConstraints(withoutConstraint []string, colNames ...string) []string {
+	stmts := make([]string, 0, len(d.table.CheckConstraints))
+	for _, cc := range d.table.CheckConstraints {
+		if slices.Contains(withoutConstraint, cc.Name) {
+			continue
+		}
+		if duplicatedConstraintColumns := d.duplicatedConstraintColumns(cc.Columns, colNames...); len(duplicatedConstraintColumns) > 0 {
+			stmts = append(stmts, fmt.Sprintf(cAlterTableAddCheckConstraintSQL,
+				pq.QuoteIdentifier(d.table.Name),
+				pq.QuoteIdentifier(DuplicationName(cc.Name)),
+				rewriteCheckExpression(cc.Definition, duplicatedConstraintColumns...),
+			))
+		}
+	}
+	return stmts
+}
+
+func (d *duplicatorStmtBuilder) duplicateUniqueConstraints(withoutConstraint []string, colNames ...string) []string {
+	stmts := make([]string, 0, len(d.table.UniqueConstraints))
+	for _, uc := range d.table.UniqueConstraints {
+		if slices.Contains(withoutConstraint, uc.Name) {
+			continue
+		}
+		if duplicatedMember, constraintColumns := d.allConstraintColumns(uc.Columns, colNames...); duplicatedMember {
+			stmts = append(stmts, fmt.Sprintf(cCreateUniqueIndexSQL,
+				pq.QuoteIdentifier(DuplicationName(uc.Name)),
+				pq.QuoteIdentifier(d.table.Name),
+				strings.Join(quoteColumnNames(constraintColumns), ", "),
+			))
+		}
+	}
+	return stmts
+}
+
+// duplicatedConstraintColumns returns a new slice of constraint columns with
+// the columns that are duplicated replaced with temporary names.
+func (d *duplicatorStmtBuilder) duplicatedConstraintColumns(constraintColumns []string, duplicatedColumns ...string) []string {
+	newConstraintColumns := make([]string, 0)
+	for _, column := range constraintColumns {
+		if slices.Contains(duplicatedColumns, column) {
+			newConstraintColumns = append(newConstraintColumns, column)
+		}
+	}
+	return newConstraintColumns
+}
+
+// allConstraintColumns returns a new slice of constraint columns with the columns
+// that are duplicated replaced with temporary names and a boolean indicating if
+// any of the columns are duplicated.
+func (d *duplicatorStmtBuilder) allConstraintColumns(constraintColumns []string, duplicatedColumns ...string) (bool, []string) {
+	duplicatedMember := false
+	newConstraintColumns := make([]string, len(constraintColumns))
+	for i, column := range constraintColumns {
+		if slices.Contains(duplicatedColumns, column) {
+			newConstraintColumns[i] = TemporaryName(column)
+			duplicatedMember = true
+		} else {
+			newConstraintColumns[i] = column
+		}
+	}
+	return duplicatedMember, newConstraintColumns
+}
+
+func (d *duplicatorStmtBuilder) duplicateColumn(
+	column *schema.Column,
+	asName string,
+	withoutNotNull bool,
+	withType string,
+	withoutConstraint []string,
+) string {
 	const (
-		cAlterTableSQL                   = `ALTER TABLE %s ADD COLUMN %s %s`
-		cAddForeignKeySQL                = `ADD CONSTRAINT %s FOREIGN KEY (%s) REFERENCES %s (%s) ON DELETE %s`
-		cAddCheckConstraintSQL           = `ADD CONSTRAINT %s %s NOT VALID`
-		cCreateUniqueIndexSQL            = `CREATE UNIQUE INDEX CONCURRENTLY %s ON %s (%s)`
-		cSetDefaultSQL                   = `ALTER TABLE %s ALTER COLUMN %s SET DEFAULT %s`
-		cAlterTableAddCheckConstraintSQL = `ALTER TABLE %s ADD CONSTRAINT %s %s NOT VALID`
-		cCommentOnColumnSQL              = `COMMENT ON COLUMN %s.%s IS %s`
+		cAlterTableSQL         = `ALTER TABLE %s ADD COLUMN %s %s`
+		cAddForeignKeySQL      = `ADD CONSTRAINT %s FOREIGN KEY (%s) REFERENCES %s (%s) ON DELETE %s`
+		cAddCheckConstraintSQL = `ADD CONSTRAINT %s %s NOT VALID`
 	)
 
 	// Generate SQL to duplicate the column's name and type
 	sql := fmt.Sprintf(cAlterTableSQL,
 		pq.QuoteIdentifier(d.table.Name),
-		pq.QuoteIdentifier(d.asName),
-		d.withType)
+		pq.QuoteIdentifier(asName),
+		withType)
 
 	// Generate SQL to add an unchecked NOT NULL constraint if the original column
 	// is NOT NULL. The constraint will be validated on migration completion.
-	if !d.column.Nullable && !d.withoutNotNull {
+	if !column.Nullable && !withoutNotNull {
 		sql += fmt.Sprintf(", "+cAddCheckConstraintSQL,
-			pq.QuoteIdentifier(DuplicationName(NotNullConstraintName(d.column.Name))),
-			fmt.Sprintf("CHECK (%s IS NOT NULL)", pq.QuoteIdentifier(d.asName)),
+			pq.QuoteIdentifier(DuplicationName(NotNullConstraintName(column.Name))),
+			fmt.Sprintf("CHECK (%s IS NOT NULL)", pq.QuoteIdentifier(asName)),
 		)
 	}
 
 	// Generate SQL to duplicate any foreign key constraints on the column
 	for _, fk := range d.table.ForeignKeys {
-		if fk.Name == d.withoutConstraint {
+		if slices.Contains(withoutConstraint, fk.Name) {
 			continue
 		}
 
-		if slices.Contains(fk.Columns, d.column.Name) {
+		if slices.Contains(fk.Columns, column.Name) {
 			sql += fmt.Sprintf(", "+cAddForeignKeySQL,
 				pq.QuoteIdentifier(DuplicationName(fk.Name)),
-				strings.Join(quoteColumnNames(copyAndReplace(fk.Columns, d.column.Name, d.asName)), ", "),
+				strings.Join(quoteColumnNames(copyAndReplace(fk.Columns, column.Name, asName)), ", "),
 				pq.QuoteIdentifier(fk.ReferencedTable),
 				strings.Join(quoteColumnNames(fk.ReferencedColumns), ", "),
 				fk.OnDelete,
 			)
 		}
 	}
 
-	_, err := d.conn.ExecContext(ctx, sql)
-	if err != nil {
-		return err
+	return sql
+}
+
+func (d *duplicatorStmtBuilder) duplicateDefault(column *schema.Column, asName string) string {
+	if column.Default == nil {
+		return ""
 	}
 
+	const cSetDefaultSQL = `ALTER TABLE %s ALTER COLUMN %s SET DEFAULT %s`
+
 	// Generate SQL to duplicate any default value on the column. This may fail
 	// if the default value is not valid for the new column type, in which case
 	// the error is ignored.
-	if d.column.Default != nil {
-		sql := fmt.Sprintf(cSetDefaultSQL, pq.QuoteIdentifier(d.table.Name), d.asName, *d.column.Default)
-
-		_, err := d.conn.ExecContext(ctx, sql)
+	return fmt.Sprintf(cSetDefaultSQL, pq.QuoteIdentifier(d.table.Name), asName, *column.Default)
+}
 
-		err = errorIgnoringErrorCode(err, dataTypeMismatchErrorCode)
-		if err != nil {
-			return err
-		}
+func (d *duplicatorStmtBuilder) duplicateComment(column *schema.Column, asName string) string {
+	if column.Comment == "" {
+		return ""
 	}
 
-	// Generate SQL to duplicate any check constraints on the column. This may faile
-	// if the check constraint is not valid for the new column type, in which case
-	// the error is ignored.
-	for _, cc := range d.table.CheckConstraints {
-		if cc.Name == d.withoutConstraint {
-			continue
-		}
-
-		if slices.Contains(cc.Columns, d.column.Name) {
-			sql := fmt.Sprintf(cAlterTableAddCheckConstraintSQL,
-				pq.QuoteIdentifier(d.table.Name),
-				pq.QuoteIdentifier(DuplicationName(cc.Name)),
-				rewriteCheckExpression(cc.Definition, d.column.Name, d.asName),
-			)
-
-			_, err := d.conn.ExecContext(ctx, sql)
-
-			err = errorIgnoringErrorCode(err, undefinedFunctionErrorCode)
-			if err != nil {
-				return err
-			}
-		}
-	}
+	const cCommentOnColumnSQL = `COMMENT ON COLUMN %s.%s IS %s`
 
 	// Generate SQL to duplicate the column's comment
-	if d.column.Comment != "" {
-		sql = fmt.Sprintf(cCommentOnColumnSQL,
-			pq.QuoteIdentifier(d.table.Name),
-			pq.QuoteIdentifier(d.asName),
-			pq.QuoteLiteral(d.column.Comment),
-		)
-
-		_, err = d.conn.ExecContext(ctx, sql)
-		if err != nil {
-			return err
-		}
-	}
-
-	// Generate SQL to duplicate any unique constraints on the column
-	// The constraint is duplicated by adding a unique index on the column concurrently.
-	// The index is converted into a unique constraint on migration completion.
-	for _, uc := range d.table.UniqueConstraints {
-		if uc.Name == d.withoutConstraint {
-			continue
-		}
-
-		if slices.Contains(uc.Columns, d.column.Name) {
-			sql = fmt.Sprintf(cCreateUniqueIndexSQL,
-				pq.QuoteIdentifier(DuplicationName(uc.Name)),
-				pq.QuoteIdentifier(d.table.Name),
-				strings.Join(quoteColumnNames(copyAndReplace(uc.Columns, d.column.Name, d.asName)), ", "),
-			)
-
-			_, err = d.conn.ExecContext(ctx, sql)
-			if err != nil {
-				return err
-			}
-		}
-	}
-
-	return nil
+	return fmt.Sprintf(cCommentOnColumnSQL,
+		pq.QuoteIdentifier(d.table.Name),
+		pq.QuoteIdentifier(asName),
+		pq.QuoteLiteral(column.Comment),
+	)
 }
 
 // DiplicationName returns the name of a duplicated column.