New DBAction for setting replica identity (#938)

Relates to #742

Author: kvch

pkg/migrations/dbactions.go

@@ -736,3 +736,33 @@ func (a *rawSQLAction) Execute(ctx context.Context) error {
 	_, err := a.conn.ExecContext(ctx, a.sql)
 	return err
 }
+
+type setReplicaIdentityAction struct {
+	conn     db.DB
+	table    string
+	identity string
+	index    string
+}
+
+func NewSetReplicaIdentityAction(conn db.DB, table string, identityType, index string) *setReplicaIdentityAction {
+	return &setReplicaIdentityAction{
+		conn:     conn,
+		table:    table,
+		identity: strings.ToUpper(identityType),
+		index:    index,
+	}
+}
+
+func (a *setReplicaIdentityAction) Execute(ctx context.Context) error {
+	// build the correct form of the `SET REPLICA IDENTITY` statement based on the`identity type
+	identitySQL := a.identity
+	if identitySQL == "INDEX" {
+		identitySQL = fmt.Sprintf("USING INDEX %s", pq.QuoteIdentifier(a.index))
+	}
+
+	// set the replica identity on the underlying table
+	_, err := a.conn.ExecContext(ctx, fmt.Sprintf("ALTER TABLE %s REPLICA IDENTITY %s",
+		pq.QuoteIdentifier(a.table),
+		identitySQL))
+	return err
+}

pkg/migrations/op_set_replica_identity.go

@@ -4,12 +4,9 @@ package migrations
 
 import (
 	"context"
-	"fmt"
 	"slices"
 	"strings"
 
-	"github.com/lib/pq"
-
 	"github.com/xataio/pgroll/pkg/backfill"
 	"github.com/xataio/pgroll/pkg/db"
 	"github.com/xataio/pgroll/pkg/schema"
@@ -20,17 +17,7 @@ var _ Operation = (*OpSetReplicaIdentity)(nil)
 func (o *OpSetReplicaIdentity) Start(ctx context.Context, l Logger, conn db.DB, latestSchema string, s *schema.Schema) (*backfill.Task, error) {
 	l.LogOperationStart(o)
 
-	// build the correct form of the `SET REPLICA IDENTITY` statement based on the`identity type
-	identitySQL := strings.ToUpper(o.Identity.Type)
-	if identitySQL == "INDEX" {
-		identitySQL = fmt.Sprintf("USING INDEX %s", pq.QuoteIdentifier(o.Identity.Index))
-	}
-
-	// set the replica identity on the underlying table
-	_, err := conn.ExecContext(ctx, fmt.Sprintf("ALTER TABLE %s REPLICA IDENTITY %s",
-		pq.QuoteIdentifier(o.Table),
-		identitySQL))
-	return nil, err
+	return nil, NewSetReplicaIdentityAction(conn, o.Table, o.Identity.Type, o.Identity.Index).Execute(ctx)
 }
 
 func (o *OpSetReplicaIdentity) Complete(l Logger, conn db.DB, s *schema.Schema) ([]DBAction, error) {