Introduce DBAction interface and refactoring of column dropping

pkg/migrations/dbactions.go

@@ -0,0 +1,49 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package migrations
+
+import (
+	"context"
+	"fmt"
+	"strings"
+
+	"github.com/lib/pq"
+	"github.com/xataio/pgroll/pkg/db"
+)
+
+// DBAction is an interface for common database actions
+// pgroll runs during migrations.
+type DBAction interface {
+	Execute(context.Context) error
+}
+
+// dropColumnAction is a DBAction that drops one or more columns from a table.
+type dropColumnAction struct {
+	conn db.DB
+
+	table   string
+	columns []string
+}
+
+func NewDropColumnAction(conn db.DB, table string, columns ...string) *dropColumnAction {
+	return &dropColumnAction{
+		conn:    conn,
+		table:   table,
+		columns: columns,
+	}
+}
+
+func (a *dropColumnAction) Execute(ctx context.Context) error {
+	_, err := a.conn.ExecContext(ctx, fmt.Sprintf("ALTER TABLE IF EXISTS %s %s",
+		pq.QuoteIdentifier(a.table),
+		a.dropMultipleColumns()))
+	return err
+}
+
+func (a *dropColumnAction) dropMultipleColumns() string {
+	cols := make([]string, len(a.columns))
+	for i, col := range a.columns {
+		cols[i] = "DROP COLUMN IF EXISTS " + pq.QuoteIdentifier(col)
+	}
+	return strings.Join(cols, ", ")
+}

pkg/migrations/op_add_column.go

@@ -117,10 +117,8 @@ func (o *OpAddColumn) Complete(ctx context.Context, conn db.DB, s *schema.Schema
 		return err
 	}
 
-	// Remove the needs backfill column
-	_, err = conn.ExecContext(ctx, fmt.Sprintf("ALTER TABLE IF EXISTS %s DROP COLUMN IF EXISTS %s",
-		pq.QuoteIdentifier(o.Table),
-		pq.QuoteIdentifier(CNeedsBackfillColumn)))
+	removeBackfillColumn := NewDropColumnAction(conn, o.Table, CNeedsBackfillColumn)
+	err = removeBackfillColumn.Execute(ctx)
 	if err != nil {
 		return err
 	}
@@ -187,9 +185,8 @@ func (o *OpAddColumn) Rollback(ctx context.Context, conn db.DB, s *schema.Schema
 		return ColumnDoesNotExistError{Table: o.Table, Name: o.Column.Name}
 	}
 
-	_, err := conn.ExecContext(ctx, fmt.Sprintf("ALTER TABLE IF EXISTS %s DROP COLUMN IF EXISTS %s",
-		pq.QuoteIdentifier(table.Name),
-		pq.QuoteIdentifier(column.Name)))
+	rollbackAddColumn := NewDropColumnAction(conn, table.Name, column.Name)
+	err := rollbackAddColumn.Execute(ctx)
 	if err != nil {
 		return err
 	}
@@ -200,10 +197,8 @@ func (o *OpAddColumn) Rollback(ctx context.Context, conn db.DB, s *schema.Schema
 		return err
 	}
 
-	// Remove the needs backfill column
-	_, err = conn.ExecContext(ctx, fmt.Sprintf("ALTER TABLE IF EXISTS %s DROP COLUMN IF EXISTS %s",
-		pq.QuoteIdentifier(table.Name),
-		pq.QuoteIdentifier(CNeedsBackfillColumn)))
+	removeBackfillColumn := NewDropColumnAction(conn, table.Name, CNeedsBackfillColumn)
+	err = removeBackfillColumn.Execute(ctx)
 
 	return err
 }

pkg/migrations/op_alter_column.go

@@ -96,10 +96,8 @@ func (o *OpAlterColumn) Complete(ctx context.Context, conn db.DB, s *schema.Sche
 		return err
 	}
 
-	// Drop the old column
-	_, err := conn.ExecContext(ctx, fmt.Sprintf("ALTER TABLE IF EXISTS %s DROP COLUMN IF EXISTS %s",
-		pq.QuoteIdentifier(o.Table),
-		pq.QuoteIdentifier(o.Column)))
+	removeOldColumn := NewDropColumnAction(conn, o.Table, o.Column)
+	err := removeOldColumn.Execute(ctx)
 	if err != nil {
 		return err
 	}
@@ -118,10 +116,8 @@ func (o *OpAlterColumn) Complete(ctx context.Context, conn db.DB, s *schema.Sche
 		return err
 	}
 
-	// Remove the needs backfill column
-	_, err = conn.ExecContext(ctx, fmt.Sprintf("ALTER TABLE IF EXISTS %s DROP COLUMN IF EXISTS %s",
-		pq.QuoteIdentifier(o.Table),
-		pq.QuoteIdentifier(CNeedsBackfillColumn)))
+	removeBackfillColumn := NewDropColumnAction(conn, o.Table, CNeedsBackfillColumn)
+	err = removeBackfillColumn.Execute(ctx)
 	if err != nil {
 		return err
 	}
@@ -160,11 +156,8 @@ func (o *OpAlterColumn) Rollback(ctx context.Context, conn db.DB, s *schema.Sche
 		}
 	}
 
-	// Drop the new column
-	_, err := conn.ExecContext(ctx, fmt.Sprintf("ALTER TABLE %s DROP COLUMN IF EXISTS %s",
-		pq.QuoteIdentifier(table.Name),
-		pq.QuoteIdentifier(column.Name),
-	))
+	rollbackAddColumn := NewDropColumnAction(conn, table.Name, column.Name)
+	err := rollbackAddColumn.Execute(ctx)
 	if err != nil {
 		return err
 	}
@@ -185,12 +178,8 @@ func (o *OpAlterColumn) Rollback(ctx context.Context, conn db.DB, s *schema.Sche
 		return err
 	}
 
-	// Remove the needs backfill column
-	_, err = conn.ExecContext(ctx, fmt.Sprintf("ALTER TABLE IF EXISTS %s DROP COLUMN IF EXISTS %s",
-		pq.QuoteIdentifier(table.Name),
-		pq.QuoteIdentifier(CNeedsBackfillColumn)))
-
-	return err
+	removeBackfillColumn := NewDropColumnAction(conn, table.Name, CNeedsBackfillColumn)
+	return removeBackfillColumn.Execute(ctx)
 }
 
 func (o *OpAlterColumn) Validate(ctx context.Context, s *schema.Schema) error {

pkg/migrations/op_create_constraint.go

@@ -133,11 +133,8 @@ func (o *OpCreateConstraint) Complete(ctx context.Context, conn db.DB, s *schema
 		}
 	}
 
-	// remove old columns
-	_, err := conn.ExecContext(ctx, fmt.Sprintf("ALTER TABLE %s %s",
-		pq.QuoteIdentifier(o.Table),
-		dropMultipleColumns(quoteColumnNames(o.Columns)),
-	))
+	removeOldColumns := NewDropColumnAction(conn, o.Table, o.Columns...)
+	err := removeOldColumns.Execute(ctx)
 	if err != nil {
 		return err
 	}
@@ -161,10 +158,8 @@ func (o *OpCreateConstraint) Complete(ctx context.Context, conn db.DB, s *schema
 		return err
 	}
 
-	// Remove the needs backfill column
-	_, err = conn.ExecContext(ctx, fmt.Sprintf("ALTER TABLE IF EXISTS %s DROP COLUMN IF EXISTS %s",
-		pq.QuoteIdentifier(o.Table),
-		pq.QuoteIdentifier(CNeedsBackfillColumn)))
+	removeBackfillColumn := NewDropColumnAction(conn, o.Table, CNeedsBackfillColumn)
+	err = removeBackfillColumn.Execute(ctx)
 
 	return err
 }
@@ -175,10 +170,8 @@ func (o *OpCreateConstraint) Rollback(ctx context.Context, conn db.DB, s *schema
 		return TableDoesNotExistError{Name: o.Table}
 	}
 
-	_, err := conn.ExecContext(ctx, fmt.Sprintf("ALTER TABLE %s %s",
-		pq.QuoteIdentifier(table.Name),
-		dropMultipleColumns(quotedTemporaryNames(o.Columns)),
-	))
+	removeDuplicatedColumns := NewDropColumnAction(conn, table.Name, temporaryNames(o.Columns)...)
+	err := removeDuplicatedColumns.Execute(ctx)
 	if err != nil {
 		return err
 	}
@@ -187,10 +180,8 @@ func (o *OpCreateConstraint) Rollback(ctx context.Context, conn db.DB, s *schema
 		return err
 	}
 
-	// Remove the needs backfill column
-	_, err = conn.ExecContext(ctx, fmt.Sprintf("ALTER TABLE IF EXISTS %s DROP COLUMN IF EXISTS %s",
-		pq.QuoteIdentifier(table.Name),
-		pq.QuoteIdentifier(CNeedsBackfillColumn)))
+	removeBackfillColumn := NewDropColumnAction(conn, table.Name, CNeedsBackfillColumn)
+	err = removeBackfillColumn.Execute(ctx)
 
 	return err
 }
@@ -207,13 +198,6 @@ func (o *OpCreateConstraint) removeTriggers(ctx context.Context, conn db.DB) err
 	return err
 }
 
-func dropMultipleColumns(columns []string) string {
-	for i, col := range columns {
-		columns[i] = "DROP COLUMN IF EXISTS " + col
-	}
-	return strings.Join(columns, ", ")
-}
-
 func (o *OpCreateConstraint) Validate(ctx context.Context, s *schema.Schema) error {
 	table := s.GetTable(o.Table)
 	if table == nil {
@@ -323,11 +307,3 @@ func temporaryNames(columns []string) []string {
 	}
 	return names
 }
-
-func quotedTemporaryNames(columns []string) []string {
-	names := make([]string, len(columns))
-	for i, col := range columns {
-		names[i] = pq.QuoteIdentifier(TemporaryName(col))
-	}
-	return names
-}

pkg/migrations/op_drop_column.go

@@ -44,9 +44,8 @@ func (o *OpDropColumn) Start(ctx context.Context, conn db.DB, latestSchema strin
 }
 
 func (o *OpDropColumn) Complete(ctx context.Context, conn db.DB, s *schema.Schema) error {
-	_, err := conn.ExecContext(ctx, fmt.Sprintf("ALTER TABLE %s DROP COLUMN %s",
-		pq.QuoteIdentifier(o.Table),
-		pq.QuoteIdentifier(o.Column)))
+	removeColumn := NewDropColumnAction(conn, o.Table, o.Column)
+	err := removeColumn.Execute(ctx)
 	if err != nil {
 		return err
 	}
@@ -57,10 +56,8 @@ func (o *OpDropColumn) Complete(ctx context.Context, conn db.DB, s *schema.Schem
 		return err
 	}
 
-	// Remove the needs backfill column
-	_, err = conn.ExecContext(ctx, fmt.Sprintf("ALTER TABLE IF EXISTS %s DROP COLUMN IF EXISTS %s",
-		pq.QuoteIdentifier(o.Table),
-		pq.QuoteIdentifier(CNeedsBackfillColumn)))
+	removeBackfillColumn := NewDropColumnAction(conn, o.Table, CNeedsBackfillColumn)
+	err = removeBackfillColumn.Execute(ctx)
 	if err != nil {
 		return err
 	}
@@ -77,10 +74,8 @@ func (o *OpDropColumn) Rollback(ctx context.Context, conn db.DB, s *schema.Schem
 		return err
 	}
 
-	// Remove the needs backfill column
-	_, err = conn.ExecContext(ctx, fmt.Sprintf("ALTER TABLE IF EXISTS %s DROP COLUMN IF EXISTS %s",
-		pq.QuoteIdentifier(table.Name),
-		pq.QuoteIdentifier(CNeedsBackfillColumn)))
+	removeBackfillColumn := NewDropColumnAction(conn, table.Name, CNeedsBackfillColumn)
+	err = removeBackfillColumn.Execute(ctx)
 	if err != nil {
 		return err
 	}

pkg/migrations/op_drop_constraint.go

@@ -95,18 +95,16 @@ func (o *OpDropConstraint) Complete(ctx context.Context, conn db.DB, s *schema.S
 		return err
 	}
 
-	// Remove the needs backfill column
-	_, err = conn.ExecContext(ctx, fmt.Sprintf("ALTER TABLE IF EXISTS %s DROP COLUMN IF EXISTS %s",
-		pq.QuoteIdentifier(table.Name),
-		pq.QuoteIdentifier(CNeedsBackfillColumn)))
+	removeBackfillColumn := NewDropColumnAction(conn, table.Name, CNeedsBackfillColumn)
+	err = removeBackfillColumn.Execute(ctx)
 	if err != nil {
 		return err
 	}
 
-	// Drop the old column
-	_, err = conn.ExecContext(ctx, fmt.Sprintf("ALTER TABLE IF EXISTS %s DROP COLUMN IF EXISTS %s",
-		pq.QuoteIdentifier(o.Table),
-		pq.QuoteIdentifier(column.Name)))
+	removeOldColumn := NewDropColumnAction(conn,
+		o.Table,
+		column.Name)
+	err = removeOldColumn.Execute(ctx)
 	if err != nil {
 		return err
 	}
@@ -124,11 +122,8 @@ func (o *OpDropConstraint) Rollback(ctx context.Context, conn db.DB, s *schema.S
 	table := s.GetTable(o.Table)
 	columnName := table.GetConstraintColumns(o.Name)[0]
 
-	// Drop the new column
-	_, err := conn.ExecContext(ctx, fmt.Sprintf("ALTER TABLE %s DROP COLUMN IF EXISTS %s",
-		pq.QuoteIdentifier(o.Table),
-		pq.QuoteIdentifier(TemporaryName(columnName)),
-	))
+	removeNewColumn := NewDropColumnAction(conn, o.Table, TemporaryName(columnName))
+	err := removeNewColumn.Execute(ctx)
 	if err != nil {
 		return err
 	}
@@ -149,10 +144,8 @@ func (o *OpDropConstraint) Rollback(ctx context.Context, conn db.DB, s *schema.S
 		return err
 	}
 
-	// Remove the needs backfill column
-	_, err = conn.ExecContext(ctx, fmt.Sprintf("ALTER TABLE IF EXISTS %s DROP COLUMN IF EXISTS %s",
-		pq.QuoteIdentifier(table.Name),
-		pq.QuoteIdentifier(CNeedsBackfillColumn)))
+	removeBackfillColumn := NewDropColumnAction(conn, table.Name, CNeedsBackfillColumn)
+	err = removeBackfillColumn.Execute(ctx)
 
 	return err
 }

pkg/migrations/op_drop_multicolumn_constraint.go

@@ -108,18 +108,14 @@ func (o *OpDropMultiColumnConstraint) Complete(ctx context.Context, conn db.DB,
 			return err
 		}
 
-		// Remove the needs backfill column
-		_, err = conn.ExecContext(ctx, fmt.Sprintf("ALTER TABLE IF EXISTS %s DROP COLUMN IF EXISTS %s",
-			pq.QuoteIdentifier(o.Table),
-			pq.QuoteIdentifier(CNeedsBackfillColumn)))
+		removeBackfillColumn := NewDropColumnAction(conn, o.Table, CNeedsBackfillColumn)
+		err = removeBackfillColumn.Execute(ctx)
 		if err != nil {
 			return err
 		}
 
-		// Drop the old column
-		_, err = conn.ExecContext(ctx, fmt.Sprintf("ALTER TABLE IF EXISTS %s DROP COLUMN IF EXISTS %s",
-			pq.QuoteIdentifier(o.Table),
-			pq.QuoteIdentifier(columnName)))
+		removeOldColumn := NewDropColumnAction(conn, o.Table, columnName)
+		err = removeOldColumn.Execute(ctx)
 		if err != nil {
 			return err
 		}
@@ -138,11 +134,8 @@ func (o *OpDropMultiColumnConstraint) Rollback(ctx context.Context, conn db.DB,
 	table := s.GetTable(o.Table)
 
 	for _, columnName := range table.GetConstraintColumns(o.Name) {
-		// Drop the new column
-		_, err := conn.ExecContext(ctx, fmt.Sprintf("ALTER TABLE %s DROP COLUMN IF EXISTS %s",
-			pq.QuoteIdentifier(table.Name),
-			pq.QuoteIdentifier(TemporaryName(columnName)),
-		))
+		removeNewColumn := NewDropColumnAction(conn, table.Name, TemporaryName(columnName))
+		err := removeNewColumn.Execute(ctx)
 		if err != nil {
 			return err
 		}
@@ -163,10 +156,8 @@ func (o *OpDropMultiColumnConstraint) Rollback(ctx context.Context, conn db.DB,
 			return err
 		}
 
-		// Remove the needs backfill column
-		_, err = conn.ExecContext(ctx, fmt.Sprintf("ALTER TABLE IF EXISTS %s DROP COLUMN IF EXISTS %s",
-			pq.QuoteIdentifier(table.Name),
-			pq.QuoteIdentifier(CNeedsBackfillColumn)))
+		removeBackfillColumn := NewDropColumnAction(conn, table.Name, CNeedsBackfillColumn)
+		err = removeBackfillColumn.Execute(ctx)
 		if err != nil {
 			return err
 		}