address review notes

Author: kvch

pkg/migrations/dbactions.go

@@ -10,6 +10,7 @@ import (
 
 	"github.com/lib/pq"
 	"github.com/xataio/pgroll/pkg/db"
+	"github.com/xataio/pgroll/pkg/schema"
 )
 
 // DBAction is an interface for common database actions
@@ -842,75 +843,46 @@ func (a *setReplicaIdentityAction) Execute(ctx context.Context) error {
 }
 
 type alterReferencesAction struct {
-	conn   db.DB
-	table  string
-	column string
+	conn         db.DB
+	referencedBy map[string][]*schema.ReferencedBy
+	table        string
+	column       string
 }
 
-func NewAlterReferencesAction(conn db.DB, table, column string) *alterReferencesAction {
+func NewAlterReferencesAction(conn db.DB, referencedBy map[string][]*schema.ReferencedBy, table, column string) *alterReferencesAction {
 	return &alterReferencesAction{
-		conn:   conn,
-		table:  table,
-		column: column,
+		conn:         conn,
+		referencedBy: referencedBy,
+		table:        table,
+		column:       column,
 	}
 }
 
 func (a *alterReferencesAction) Execute(ctx context.Context) error {
-	definitions, err := a.constraintDefinitions(ctx)
-	if err != nil {
-		return err
-	}
-	for _, def := range definitions {
-		// Drop the existing constraint
-		_, err := a.conn.ExecContext(ctx, fmt.Sprintf("ALTER TABLE %s DROP CONSTRAINT %s",
-			pq.QuoteIdentifier(def.table),
-			pq.QuoteIdentifier(def.name),
-		))
-		if err != nil {
-			return fmt.Errorf("dropping constraint %s on %s: %w", def.name, def.table, err)
-		}
-
-		// Recreate the constraint with the table and new column
-		newDef := strings.ReplaceAll(def.def, a.column, pq.QuoteIdentifier(TemporaryName(a.column)))
-		newDef = strings.ReplaceAll(newDef, a.table, pq.QuoteIdentifier(a.table))
-		_, err = a.conn.ExecContext(ctx, fmt.Sprintf("ALTER TABLE %s ADD CONSTRAINT %s %s",
-			pq.QuoteIdentifier(def.table),
-			pq.QuoteIdentifier(def.name),
-			newDef,
-		))
-		if err != nil {
-			return fmt.Errorf("altering references for %s.%s: %w", a.table, a.column, err)
-		}
-	}
-	return nil
-}
-
-type constraintDefinition struct {
-	name  string
-	table string
-	def   string
-}
+	for table, constraints := range a.referencedBy {
+		for _, constraint := range constraints {
+			// Drop the existing constraint
+			_, err := a.conn.ExecContext(ctx, fmt.Sprintf("ALTER TABLE %s DROP CONSTRAINT %s",
+				pq.QuoteIdentifier(table),
+				pq.QuoteIdentifier(constraint.Name),
+			))
+			if err != nil {
+				return fmt.Errorf("dropping constraint %s on %s: %w", constraint.Name, table, err)
+			}
 
-func (a *alterReferencesAction) constraintDefinitions(ctx context.Context) ([]constraintDefinition, error) {
-	rows, err := a.conn.QueryContext(ctx, fmt.Sprintf(`
-SELECT conname, r.conrelid::regclass, pg_catalog.pg_get_constraintdef(r.oid, true) as condef
-FROM pg_catalog.pg_constraint r
-WHERE confrelid = %s::regclass AND r.contype = 'f'`,
-		pq.QuoteIdentifier(a.table),
-	))
-	// No FK constraint for table
-	if err != nil {
-		return nil, nil
-	}
-	defer rows.Close()
+			// Recreate the constraint with the table and new column
+			newDef := strings.ReplaceAll(constraint.Definition, a.column, pq.QuoteIdentifier(TemporaryName(a.column)))
+			newDef = strings.ReplaceAll(newDef, a.table, pq.QuoteIdentifier(a.table))
+			_, err = a.conn.ExecContext(ctx, fmt.Sprintf("ALTER TABLE %s ADD CONSTRAINT %s %s",
+				pq.QuoteIdentifier(table),
+				pq.QuoteIdentifier(constraint.Name),
+				newDef,
+			))
+			if err != nil {
+				return fmt.Errorf("altering references for %s.%s: %w", a.table, a.column, err)
+			}
 
-	defs := make([]constraintDefinition, 0)
-	for rows.Next() {
-		var def constraintDefinition
-		if err := rows.Scan(&def.name, &def.table, &def.def); err != nil {
-			return nil, fmt.Errorf("scanning referencing constraints for %s.%s: %w", a.table, a.column, err)
 		}
-		defs = append(defs, def)
 	}
-	return defs, rows.Err()
+	return nil
 }

pkg/migrations/op_alter_column.go

@@ -118,7 +118,7 @@ func (o *OpAlterColumn) Complete(l Logger, conn db.DB, s *schema.Schema) ([]DBAc
 
 	return append(dbActions, []DBAction{
 		NewAlterSequenceOwnerAction(conn, o.Table, o.Column, TemporaryName(o.Column)),
-		NewAlterReferencesAction(conn, o.Table, o.Column),
+		NewAlterReferencesAction(conn, table.ReferencedBy, o.Table, o.Column),
 		NewDropColumnAction(conn, o.Table, o.Column),
 		NewDropFunctionAction(conn,
 			backfill.TriggerFunctionName(o.Table, o.Column),

pkg/migrations/op_alter_column_test.go

@@ -726,7 +726,7 @@ func TestAlterPrimaryKeyColumns(t *testing.T) {
 					Name: "02_alter_column",
 					Operations: migrations.Operations{
 						&migrations.OpAlterColumn{
-							Table:  "people",
+							Table:  "events",
 							Column: "id",
 							Type:   ptr("bigint"),
 							Up:     "CAST(id AS bigint)",
@@ -740,29 +740,32 @@ func TestAlterPrimaryKeyColumns(t *testing.T) {
 				ValidatedForeignKeyMustExist(t, db, schema, "people", "person_manages_event_fk")
 
 				bigint := "31474836471" // A value larger than int can hold
-				integer := "100"
 
 				// Inserting a row with integer id into the old schema should succeed
 				MustInsert(t, db, schema, "01_create_table", "events", map[string]string{
+					"id":   "1",
 					"name": "pgroll v1 release party",
 				})
 				MustInsert(t, db, schema, "01_create_table", "people", map[string]string{
-					"id":      integer,
+					"id":      "1",
 					"name":    "alice",
 					"manages": "1",
 				})
 				// Inserting a row with integer bigint id into the old schema should fail
-				MustNotInsert(t, db, schema, "01_create_table", "people", map[string]string{
-					"id":      bigint,
-					"name":    "alice",
-					"manages": "1",
+				MustNotInsert(t, db, schema, "01_create_table", "events", map[string]string{
+					"id":   bigint,
+					"name": "pgroll v2 release party",
 				}, testutils.NumericValueOutOfRangeErrorCode)
 
 				// Inserting a row with bigint id into the new schema should succeed
+				MustInsert(t, db, schema, "02_alter_column", "events", map[string]string{
+					"id":   bigint,
+					"name": "pgroll v2 release party",
+				})
 				MustInsert(t, db, schema, "02_alter_column", "people", map[string]string{
-					"id":      bigint,
-					"name":    "alice",
-					"manages": "1",
+					"id":      "2",
+					"name":    "bob",
+					"manages": bigint,
 				})
 				// Inserting a row into the `people` table with a `manages` field that
 				// violates the FK constraint fails
@@ -781,16 +784,15 @@ func TestAlterPrimaryKeyColumns(t *testing.T) {
 				ValidatedForeignKeyMustExist(t, db, schema, "people", "person_manages_event_fk")
 
 				// Inserting a row with integer bigint into the new schema should succeed
-				MustInsert(t, db, schema, "02_alter_column", "people", map[string]string{
-					"id":      "31474836472",
-					"name":    "bob",
-					"manages": "1",
+				MustInsert(t, db, schema, "02_alter_column", "events", map[string]string{
+					"id":   "31474836472",
+					"name": "pgroll v3 release party",
 				})
 				// Inserting a row into the `people` table with a `manages` field that
 				// violates the FK constraint fails
 				MustNotInsert(t, db, schema, "02_alter_column", "people", map[string]string{
-					"id":      "31474836474",
-					"name":    "alice",
+					"id":      "3",
+					"name":    "carol",
 					"manages": "2",
 				}, testutils.FKViolationErrorCode)
 			},

pkg/schema/schema.go

@@ -62,10 +62,26 @@ type Table struct {
 	// ExcludeConstraints is a map of all exclude constraints defined on the table
 	ExcludeConstraints map[string]*ExcludeConstraint `json:"excludeConstraints"`
 
+	// ReferencedBy is a map of table names that reference this table by foreign key
+	// The key is the name of the referencing table, and the value is the name and
+	// the definition of the foreign key constraint
+	ReferencedBy map[string][]*ReferencedBy `json:"referencedBy"`
+
 	// Whether or not the table has been deleted in the virtual schema
 	Deleted bool `json:"-"`
 }
 
+type ReferencedBy struct {
+	// Name is the name of the foreign key constraint
+	Name string `json:"name"`
+
+	// Table is the name of the table that is referenced
+	Table string `json:"table"`
+
+	// Definition is the definition of the foreign key constraint
+	Definition string `json:"definition"`
+}
+
 // Column represents a column in a table
 type Column struct {
 	// Name is the actual name in postgres