Add new DBActions for setting and dropping default values (#892)

Related to #742

Author: kvch

pkg/migrations/dbactions.go

@@ -628,3 +628,48 @@ func (a *setNotNullAction) Execute(ctx context.Context) error {
 		pq.QuoteIdentifier(a.column)))
 	return err
 }
+
+type setDefaultAction struct {
+	conn         db.DB
+	table        string
+	column       string
+	defaultValue string
+}
+
+func NewSetDefaultValueAction(conn db.DB, table, column, defaultValue string) *setDefaultAction {
+	return &setDefaultAction{
+		conn:         conn,
+		table:        table,
+		column:       column,
+		defaultValue: defaultValue,
+	}
+}
+
+func (a *setDefaultAction) Execute(ctx context.Context) error {
+	_, err := a.conn.ExecContext(ctx, fmt.Sprintf("ALTER TABLE IF EXISTS %s ALTER COLUMN %s SET DEFAULT %s",
+		pq.QuoteIdentifier(a.table),
+		pq.QuoteIdentifier(a.column),
+		a.defaultValue))
+	return err
+}
+
+type dropDefaultAction struct {
+	conn   db.DB
+	table  string
+	column string
+}
+
+func NewDropDefaultValueAction(conn db.DB, table, column string) *dropDefaultAction {
+	return &dropDefaultAction{
+		conn:   conn,
+		table:  table,
+		column: column,
+	}
+}
+
+func (a *dropDefaultAction) Execute(ctx context.Context) error {
+	_, err := a.conn.ExecContext(ctx, fmt.Sprintf("ALTER TABLE IF EXISTS %s ALTER COLUMN %s DROP DEFAULT",
+		pq.QuoteIdentifier(a.table),
+		pq.QuoteIdentifier(a.column)))
+	return err
+}

pkg/migrations/op_add_column.go

@@ -191,11 +191,7 @@ func (o *OpAddColumn) Complete(ctx context.Context, l Logger, conn db.DB, s *sch
 	// optimization, set it here.
 	column := s.GetTable(o.Table).GetColumn(TemporaryName(o.Column.Name))
 	if o.Column.HasDefault() && column.Default == nil {
-		_, err := conn.ExecContext(ctx, fmt.Sprintf("ALTER TABLE IF EXISTS %s ALTER COLUMN %s SET DEFAULT %s",
-			pq.QuoteIdentifier(o.Table),
-			pq.QuoteIdentifier(o.Column.Name),
-			*o.Column.Default,
-		))
+		err := NewSetDefaultValueAction(conn, o.Table, o.Column.Name, *o.Column.Default).Execute(ctx)
 		if err != nil {
 			return err
 		}

pkg/migrations/op_set_default.go

@@ -4,9 +4,6 @@ package migrations
 
 import (
 	"context"
-	"fmt"
-
-	"github.com/lib/pq"
 
 	"github.com/xataio/pgroll/pkg/db"
 	"github.com/xataio/pgroll/pkg/schema"
@@ -36,17 +33,10 @@ func (o *OpSetDefault) Start(ctx context.Context, l Logger, conn db.DB, latestSc
 
 	var err error
 	if o.Default == nil {
-		_, err = conn.ExecContext(ctx, fmt.Sprintf(`ALTER TABLE %s ALTER COLUMN %s DROP DEFAULT`,
-			pq.QuoteIdentifier(table.Name),
-			pq.QuoteIdentifier(column.Name)))
-
+		err = NewDropDefaultValueAction(conn, table.Name, column.Name).Execute(ctx)
 		column.Default = nil
 	} else {
-		_, err = conn.ExecContext(ctx, fmt.Sprintf(`ALTER TABLE %s ALTER COLUMN %s SET DEFAULT %s`,
-			pq.QuoteIdentifier(table.Name),
-			pq.QuoteIdentifier(column.Name),
-			*o.Default))
-
+		err = NewSetDefaultValueAction(conn, table.Name, column.Name, *o.Default).Execute(ctx)
 		column.Default = o.Default
 	}
 	if err != nil {