New DBAction: CreateCheckConstraintAction (#837)

This PR adds a new `DBAction` that is responsible for adding new check
constraints.
It is used in adding check constraints to columns and tables.

Related #742

Author: kvch

pkg/migrations/dbactions.go

@@ -394,3 +394,50 @@ func (a *validateConstraintAction) Execute(ctx context.Context) error {
 		pq.QuoteIdentifier(a.constraint)))
 	return err
 }
+
+// CreateCheckConstraintAction creates a check constraint on a table.
+type CreateCheckConstraintAction struct {
+	conn           db.DB
+	table          string
+	columns        []string
+	constraint     string
+	check          string
+	noInherit      bool
+	skipValidation bool
+}
+
+func NewCreateCheckConstraintAction(conn db.DB, table, constraint, check string, columns []string, noInherit, skipValidation bool) *CreateCheckConstraintAction {
+	return &CreateCheckConstraintAction{
+		conn:           conn,
+		table:          table,
+		columns:        columns,
+		check:          check,
+		constraint:     constraint,
+		noInherit:      noInherit,
+		skipValidation: skipValidation,
+	}
+}
+
+func (a *CreateCheckConstraintAction) Execute(ctx context.Context) error {
+	sql := fmt.Sprintf("ALTER TABLE %s ADD ", pq.QuoteIdentifier(a.table))
+
+	writer := &ConstraintSQLWriter{
+		Name:           a.constraint,
+		SkipValidation: a.skipValidation,
+	}
+	sql += writer.WriteCheck(rewriteCheckExpression(a.check, a.columns...), a.noInherit)
+	_, err := a.conn.ExecContext(ctx, sql)
+	return err
+}
+
+// In order for the `check` expression to be easy to write, migration authors specify
+// the check expression as though it were being applied to the old column,
+// On migration start, however, the check is actually applied to the new (temporary)
+// column.
+// This function naively rewrites the check expression to apply to the new column.
+func rewriteCheckExpression(check string, columns ...string) string {
+	for _, col := range columns {
+		check = strings.ReplaceAll(check, col, TemporaryName(col))
+	}
+	return check
+}

pkg/migrations/op_create_constraint.go

@@ -93,7 +93,7 @@ func (o *OpCreateConstraint) Start(ctx context.Context, l Logger, conn db.DB, la
 	case OpCreateConstraintTypeUnique, OpCreateConstraintTypePrimaryKey:
 		return table, NewCreateUniqueIndexConcurrentlyAction(conn, s.Name, o.Name, table.Name, temporaryNames(o.Columns)...).Execute(ctx)
 	case OpCreateConstraintTypeCheck:
-		return table, o.addCheckConstraint(ctx, conn, table.Name)
+		return table, NewCreateCheckConstraintAction(conn, table.Name, o.Name, *o.Check, o.Columns, o.NoInherit, true).Execute(ctx)
 	case OpCreateConstraintTypeForeignKey:
 		return table, o.addForeignKeyConstraint(ctx, conn, table)
 	}
@@ -284,18 +284,6 @@ func (o *OpCreateConstraint) Validate(ctx context.Context, s *schema.Schema) err
 	return nil
 }
 
-func (o *OpCreateConstraint) addCheckConstraint(ctx context.Context, conn db.DB, tableName string) error {
-	sql := fmt.Sprintf("ALTER TABLE %s ADD ", pq.QuoteIdentifier(tableName))
-
-	writer := &ConstraintSQLWriter{
-		Name:           o.Name,
-		SkipValidation: true,
-	}
-	sql += writer.WriteCheck(rewriteCheckExpression(*o.Check, o.Columns...), o.NoInherit)
-	_, err := conn.ExecContext(ctx, sql)
-	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))
 

pkg/migrations/op_set_check.go

@@ -5,9 +5,6 @@ package migrations
 import (
 	"context"
 	"fmt"
-	"strings"
-
-	"github.com/lib/pq"
 
 	"github.com/xataio/pgroll/pkg/db"
 	"github.com/xataio/pgroll/pkg/schema"
@@ -32,7 +29,7 @@ func (o *OpSetCheckConstraint) Start(ctx context.Context, l Logger, conn db.DB,
 	}
 
 	// Add the check constraint to the new column as NOT VALID.
-	if err := o.addCheckConstraint(ctx, conn, s); err != nil {
+	if err := NewCreateCheckConstraintAction(conn, table.Name, o.Check.Name, o.Check.Constraint, []string{o.Column}, o.Check.NoInherit, true).Execute(ctx); err != nil {
 		return nil, fmt.Errorf("failed to add check constraint: %w", err)
 	}
 
@@ -88,29 +85,3 @@ func (o *OpSetCheckConstraint) Validate(ctx context.Context, s *schema.Schema) e
 
 	return nil
 }
-
-func (o *OpSetCheckConstraint) addCheckConstraint(ctx context.Context, conn db.DB, s *schema.Schema) error {
-	table := s.GetTable(o.Table)
-	sql := fmt.Sprintf("ALTER TABLE %s ADD ", pq.QuoteIdentifier(table.Name))
-
-	writer := &ConstraintSQLWriter{
-		Name:           o.Check.Name,
-		SkipValidation: true,
-	}
-	sql += writer.WriteCheck(rewriteCheckExpression(o.Check.Constraint, o.Column), o.Check.NoInherit)
-	_, err := conn.ExecContext(ctx, sql)
-
-	return err
-}
-
-// In order for the `check` expression to be easy to write, migration authors specify
-// the check expression as though it were being applied to the old column,
-// On migration start, however, the check is actually applied to the new (temporary)
-// column.
-// This function naively rewrites the check expression to apply to the new column.
-func rewriteCheckExpression(check string, columns ...string) string {
-	for _, col := range columns {
-		check = strings.ReplaceAll(check, col, TemporaryName(col))
-	}
-	return check
-}