Skip to content

Commit be1a7cf

Browse files
kvchkvch
authored
Support altering primary key columns that are not included in FK constraints (xataio#930)
In the past `pgroll` could not preserve primary key constraint on altered columns. Instead the primary key constraint was down-graded to a unique, not null constraint. Now the PK constraint is not removed during completing the migration. Futhermore, if the table had rows, `pgroll` could not backfill rows. `pgroll` "updates" the primary keys of the backfilled table to trigger backfilling existing columns. Previously, finding the identity columns returned the duplicated column `__pgroll_new_id` instead of the correct `id`. In this case we need the current PK, so we can copy the correct values to the new column. I opened follow-up issue to support altering mutiple columns of the same composite key in the same migration: xataio#930 Support for alter PK columns that are part of FK constraints is in a follow-up PR: xataio#935 Closes xataio#888
1 parent bc40891 commit be1a7cf

4 files changed

Lines changed: 190 additions & 14 deletions

File tree

‎internal/testutils/error_codes.go‎

Lines changed: 8 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -3,11 +3,12 @@
33
package testutils
44

55
const (
6-
CheckViolationErrorCode string = "check_violation"
7-
ExclusionViolationErrorCode string = "exclusion_violation"
8-
FKViolationErrorCode string = "foreign_key_violation"
9-
NotNullViolationErrorCode string = "not_null_violation"
10-
UndefinedColumnErrorCode string = "undefined_column"
11-
UndefinedTableErrorCode string = "undefined_table"
12-
UniqueViolationErrorCode string = "unique_violation"
6+
CheckViolationErrorCode string = "check_violation"
7+
ExclusionViolationErrorCode string = "exclusion_violation"
8+
FKViolationErrorCode string = "foreign_key_violation"
9+
NotNullViolationErrorCode string = "not_null_violation"
10+
UndefinedColumnErrorCode string = "undefined_column"
11+
UndefinedTableErrorCode string = "undefined_table"
12+
UniqueViolationErrorCode string = "unique_violation"
13+
NumericValueOutOfRangeErrorCode string = "numeric_value_out_of_range"
1314
)

‎pkg/backfill/backfill.go‎

Lines changed: 2 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -165,13 +165,8 @@ func getRowCount(ctx context.Context, conn db.DB, tableName string) (int64, erro
165165

166166
// getIdentityColumn will return a column suitable for use in a backfill operation.
167167
func getIdentityColumns(table *schema.Table) []string {
168-
pks := table.GetPrimaryKey()
169-
if len(pks) != 0 {
170-
pkNames := make([]string, len(pks))
171-
for i, pk := range pks {
172-
pkNames[i] = pk.Name
173-
}
174-
return pkNames
168+
if len(table.PrimaryKey) != 0 {
169+
return table.PrimaryKey
175170
}
176171

177172
// If there is no primary key, look for a unique not null column

‎pkg/migrations/op_alter_column_test.go‎

Lines changed: 169 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -508,6 +508,175 @@ func TestAlterColumnMultipleSubOperations(t *testing.T) {
508508
})
509509
}
510510

511+
func TestAlterPrimaryKeyColumns(t *testing.T) {
512+
t.Parallel()
513+
514+
ExecuteTests(t, TestCases{
515+
{
516+
name: "alter a single primary key column",
517+
migrations: []migrations.Migration{
518+
{
519+
Name: "01_create_table",
520+
Operations: migrations.Operations{
521+
&migrations.OpCreateTable{
522+
Name: "people",
523+
Columns: []migrations.Column{
524+
{
525+
Name: "id",
526+
Type: "int",
527+
Pk: true,
528+
},
529+
{
530+
Name: "name",
531+
Type: "varchar(255)",
532+
Nullable: true,
533+
},
534+
},
535+
},
536+
},
537+
},
538+
{
539+
Name: "02_alter_column",
540+
Operations: migrations.Operations{
541+
&migrations.OpAlterColumn{
542+
Table: "people",
543+
Column: "id",
544+
Type: ptr("bigint"),
545+
Up: "CAST(id AS bigint)",
546+
Down: "SELECT CASE WHEN id < 2147483647 THEN CAST(id AS int) ELSE 0 END",
547+
},
548+
},
549+
},
550+
},
551+
afterStart: func(t *testing.T, db *sql.DB, schema string) {
552+
PrimaryKeyConstraintMustExist(t, db, schema, "people", "people_pkey")
553+
ColumnMustBePK(t, db, schema, "people", "id")
554+
555+
bigint := "31474836471" // A value larger than int can hold
556+
integer := "100"
557+
558+
// Inserting a row with integer id into the old schema should succeed
559+
MustInsert(t, db, schema, "01_create_table", "people", map[string]string{
560+
"id": integer,
561+
"name": "alice",
562+
})
563+
// Inserting a row with integer bigint id into the old schema should fail
564+
MustNotInsert(t, db, schema, "01_create_table", "people", map[string]string{
565+
"id": bigint,
566+
"name": "alice",
567+
}, testutils.NumericValueOutOfRangeErrorCode)
568+
569+
// Inserting a row with bigint id into the new schema should succeed
570+
MustInsert(t, db, schema, "02_alter_column", "people", map[string]string{
571+
"id": bigint,
572+
"name": "alice",
573+
})
574+
},
575+
afterRollback: func(t *testing.T, db *sql.DB, schema string) {
576+
PrimaryKeyConstraintMustExist(t, db, schema, "people", "people_pkey")
577+
ColumnMustBePK(t, db, schema, "people", "id")
578+
},
579+
afterComplete: func(t *testing.T, db *sql.DB, schema string) {
580+
PrimaryKeyConstraintMustExist(t, db, schema, "people", "people_pkey")
581+
ColumnMustBePK(t, db, schema, "people", "id")
582+
583+
// Inserting a row with integer bigint into the new schema should succeed
584+
MustInsert(t, db, schema, "02_alter_column", "people", map[string]string{
585+
"id": "31474836472",
586+
"name": "bob",
587+
})
588+
},
589+
},
590+
{
591+
name: "alter a composite primary key: one column changes type, the other is unchanged",
592+
migrations: []migrations.Migration{
593+
{
594+
Name: "01_create_table",
595+
Operations: migrations.Operations{
596+
&migrations.OpCreateTable{
597+
Name: "people",
598+
Columns: []migrations.Column{
599+
{
600+
Name: "id1",
601+
Type: "int",
602+
Pk: true,
603+
},
604+
{
605+
Name: "id2",
606+
Type: "int",
607+
Pk: true,
608+
},
609+
{
610+
Name: "name",
611+
Type: "varchar(255)",
612+
Nullable: true,
613+
},
614+
},
615+
},
616+
},
617+
},
618+
{
619+
Name: "02_alter_column",
620+
Operations: migrations.Operations{
621+
&migrations.OpAlterColumn{
622+
Table: "people",
623+
Column: "id1",
624+
Type: ptr("bigint"),
625+
Up: "CAST(id1 AS bigint)",
626+
Down: "SELECT CASE WHEN id1 < 2147483647 THEN CAST(id1 AS int) ELSE 0 END",
627+
},
628+
},
629+
},
630+
},
631+
afterStart: func(t *testing.T, db *sql.DB, schema string) {
632+
PrimaryKeyConstraintMustExist(t, db, schema, "people", "people_pkey")
633+
ColumnMustBePK(t, db, schema, "people", "id1")
634+
ColumnMustBePK(t, db, schema, "people", "id2")
635+
636+
bigint := "31474836471" // A value larger than int can hold
637+
integer := "100"
638+
639+
// Inserting a row with integer id into the old schema should succeed
640+
MustInsert(t, db, schema, "01_create_table", "people", map[string]string{
641+
"id1": integer,
642+
"id2": integer,
643+
"name": "alice",
644+
})
645+
// Inserting a row with integer bigint id into the old schema should fail
646+
MustNotInsert(t, db, schema, "01_create_table", "people", map[string]string{
647+
"id1": bigint,
648+
"id2": integer,
649+
"name": "alice",
650+
}, testutils.NumericValueOutOfRangeErrorCode)
651+
652+
// Inserting a row with bigint id into the new schema should succeed
653+
MustInsert(t, db, schema, "02_alter_column", "people", map[string]string{
654+
"id1": bigint,
655+
"id2": integer,
656+
"name": "alice",
657+
})
658+
},
659+
afterRollback: func(t *testing.T, db *sql.DB, schema string) {
660+
PrimaryKeyConstraintMustExist(t, db, schema, "people", "people_pkey")
661+
ColumnMustBePK(t, db, schema, "people", "id1")
662+
ColumnMustBePK(t, db, schema, "people", "id2")
663+
},
664+
afterComplete: func(t *testing.T, db *sql.DB, schema string) {
665+
PrimaryKeyConstraintMustExist(t, db, schema, "people", "people_pkey")
666+
ColumnMustBePK(t, db, schema, "people", "id1")
667+
ColumnMustBePK(t, db, schema, "people", "id2")
668+
669+
// Inserting a row with integer bigint into the new schema should succeed
670+
MustInsert(t, db, schema, "02_alter_column", "people", map[string]string{
671+
"id1": "31474836472",
672+
"id2": "72",
673+
"name": "bob",
674+
})
675+
},
676+
},
677+
})
678+
}
679+
511680
func TestAlterColumnValidation(t *testing.T) {
512681
t.Parallel()
513682

‎pkg/migrations/rename.go‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -130,5 +130,16 @@ func (a *renameDuplicatedColumnAction) Execute(ctx context.Context) error {
130130
}
131131
}
132132

133+
if slices.Contains(a.table.PrimaryKey, a.to) {
134+
err := NewAddPrimaryKeyAction(a.conn, a.table.Name, primaryKeyName(a.table.Name)).Execute(ctx)
135+
if err != nil {
136+
return fmt.Errorf("failed to re-add primary key constraint: %w", err)
137+
}
138+
}
139+
133140
return nil
134141
}
142+
143+
func primaryKeyName(tableName string) string {
144+
return fmt.Sprintf("%s_pkey", tableName)
145+
}

0 commit comments

Comments
 (0)