Skip to content

Commit f3432da

Browse files
agedemenliagedemenli
authored
Alter sequence owner before dropping the column (#668)
Check if there's any sequence for the column before dropping it. If the column type is serial, `pg_get_serial_sequence` will return the sequence name. With this PR, we will be changing the owner column of the seqeunce to the new column so that we can safely drop the old one. Fixes: #633
1 parent 1558ef7 commit f3432da

9 files changed

+247
-1
lines changed

‎pkg/migrations/duplicate.go

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -358,3 +358,36 @@ func errorIgnoringErrorCode(err error, code pq.ErrorCode) error {
358358

359359
return err
360360
}
361+
362+
func alterSequenceOwnerToDuplicatedColumn(ctx context.Context, conn db.DB, tableName, columnName string) error {
363+
sequenceName := getSequenceNameForColumn(ctx, conn, tableName, columnName)
364+
if sequenceName == "" {
365+
// No sequence for the column
366+
return nil
367+
}
368+
_, err := conn.ExecContext(ctx, fmt.Sprintf("ALTER SEQUENCE IF EXISTS %s OWNED BY %s.%s",
369+
sequenceName,
370+
pq.QuoteIdentifier(tableName),
371+
pq.QuoteIdentifier(TemporaryName(columnName)),
372+
))
373+
374+
return err
375+
}
376+
377+
func getSequenceNameForColumn(ctx context.Context, conn db.DB, tableName, columnName string) string {
378+
var sequenceName string
379+
query := fmt.Sprintf(`
380+
SELECT pg_get_serial_sequence('%s', '%s')
381+
`, pq.QuoteIdentifier(tableName), columnName)
382+
rows, err := conn.QueryContext(ctx, query)
383+
if err != nil {
384+
return ""
385+
}
386+
defer rows.Close()
387+
388+
if err := db.ScanFirstValue(rows, &sequenceName); err != nil {
389+
return ""
390+
}
391+
392+
return sequenceName
393+
}

‎pkg/migrations/op_alter_column.go

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -86,6 +86,10 @@ func (o *OpAlterColumn) Complete(ctx context.Context, conn db.DB, tr SQLTransfor
8686
}
8787
}
8888

89+
if err := alterSequenceOwnerToDuplicatedColumn(ctx, conn, o.Table, o.Column); err != nil {
90+
return err
91+
}
92+
8993
// Drop the old column
9094
_, err := conn.ExecContext(ctx, fmt.Sprintf("ALTER TABLE IF EXISTS %s DROP COLUMN IF EXISTS %s",
9195
pq.QuoteIdentifier(o.Table),

‎pkg/migrations/op_create_constraint.go

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -120,6 +120,12 @@ func (o *OpCreateConstraint) Complete(ctx context.Context, conn db.DB, tr SQLTra
120120
}
121121
}
122122

123+
for _, col := range o.Columns {
124+
if err := alterSequenceOwnerToDuplicatedColumn(ctx, conn, o.Table, col); err != nil {
125+
return err
126+
}
127+
}
128+
123129
// remove old columns
124130
_, err := conn.ExecContext(ctx, fmt.Sprintf("ALTER TABLE %s %s",
125131
pq.QuoteIdentifier(o.Table),

‎pkg/migrations/op_create_constraint_test.go

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1159,5 +1159,66 @@ func TestCreateConstraintValidation(t *testing.T) {
11591159
afterRollback: func(t *testing.T, db *sql.DB, schema string) {},
11601160
afterComplete: func(t *testing.T, db *sql.DB, schema string) {},
11611161
},
1162+
{
1163+
name: "create unique constraint on serial column",
1164+
migrations: []migrations.Migration{
1165+
{
1166+
Name: "01_add_table",
1167+
Operations: migrations.Operations{
1168+
&migrations.OpCreateTable{
1169+
Name: "users",
1170+
Columns: []migrations.Column{
1171+
{
1172+
Name: "email",
1173+
Type: "text",
1174+
Pk: true,
1175+
},
1176+
{
1177+
Name: "id",
1178+
Type: "serial",
1179+
},
1180+
},
1181+
},
1182+
},
1183+
},
1184+
{
1185+
Name: "02_create_constraint",
1186+
Operations: migrations.Operations{
1187+
&migrations.OpCreateConstraint{
1188+
Name: "unique_id",
1189+
Table: "users",
1190+
Type: "unique",
1191+
Columns: []string{"id"},
1192+
Up: map[string]string{
1193+
"id": "id",
1194+
},
1195+
Down: map[string]string{
1196+
"id": "id",
1197+
},
1198+
},
1199+
},
1200+
},
1201+
},
1202+
afterStart: func(t *testing.T, db *sql.DB, schema string) {
1203+
// The index has been created on the underlying table.
1204+
IndexMustExist(t, db, schema, "users", "unique_id")
1205+
},
1206+
afterRollback: func(t *testing.T, db *sql.DB, schema string) {
1207+
// The index has been dropped from the the underlying table.
1208+
IndexMustNotExist(t, db, schema, "users", "unique_id")
1209+
},
1210+
afterComplete: func(t *testing.T, db *sql.DB, schema string) {
1211+
// Functions, triggers and temporary columns are dropped.
1212+
TableMustBeCleanedUp(t, db, schema, "users", "id")
1213+
1214+
// Inserting values into the new schema that violate uniqueness should fail.
1215+
MustInsert(t, db, schema, "02_create_constraint", "users", map[string]string{
1216+
"email": "alice@xata.io", "id": "1",
1217+
})
1218+
MustNotInsert(t, db, schema, "02_create_constraint", "users", map[string]string{
1219+
"email": "bob@xata.io", "id": "1",
1220+
}, testutils.UniqueViolationErrorCode)
1221+
},
1222+
},
11621223
})
11631224
}

‎pkg/migrations/op_drop_constraint.go

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -85,6 +85,10 @@ func (o *OpDropConstraint) Complete(ctx context.Context, conn db.DB, tr SQLTrans
8585
return err
8686
}
8787

88+
if err := alterSequenceOwnerToDuplicatedColumn(ctx, conn, o.Table, column.Name); err != nil {
89+
return err
90+
}
91+
8892
// Drop the old column
8993
_, err = conn.ExecContext(ctx, fmt.Sprintf("ALTER TABLE IF EXISTS %s DROP COLUMN IF EXISTS %s",
9094
pq.QuoteIdentifier(o.Table),

‎pkg/migrations/op_drop_constraint_test.go

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -409,6 +409,69 @@ func TestDropConstraint(t *testing.T) {
409409
})
410410
},
411411
},
412+
{
413+
name: "drop unique constraint from serial column",
414+
migrations: []migrations.Migration{
415+
{
416+
Name: "01_add_tables",
417+
Operations: migrations.Operations{
418+
&migrations.OpCreateTable{
419+
Name: "users",
420+
Columns: []migrations.Column{
421+
{
422+
Name: "id",
423+
Type: "serial",
424+
Pk: true,
425+
},
426+
{
427+
Name: "secondary_id",
428+
Type: "serial",
429+
Unique: true,
430+
},
431+
},
432+
},
433+
},
434+
},
435+
{
436+
Name: "02_drop_unique_constraint",
437+
Operations: migrations.Operations{
438+
&migrations.OpDropConstraint{
439+
Table: "users",
440+
Name: "users_secondary_id_key",
441+
Up: "secondary_id",
442+
Down: "secondary_id",
443+
},
444+
},
445+
},
446+
},
447+
afterStart: func(t *testing.T, db *sql.DB, schema string) {
448+
// The new (temporary) `secondary_id` column should exist on the underlying table.
449+
ColumnMustExist(t, db, schema, "users", migrations.TemporaryName("secondary_id"))
450+
451+
// Inserting a row that meets the unique constraint into the old view works.
452+
MustInsert(t, db, schema, "01_add_tables", "users", map[string]string{
453+
"secondary_id": "1",
454+
})
455+
456+
// Inserting a row that does not meet the unique constraint into the old view fails.
457+
MustNotInsert(t, db, schema, "01_add_tables", "users", map[string]string{
458+
"secondary_id": "1",
459+
}, testutils.UniqueViolationErrorCode)
460+
},
461+
afterRollback: func(t *testing.T, db *sql.DB, schema string) {
462+
// The table is cleaned up; temporary columns, trigger functions and triggers no longer exist.
463+
TableMustBeCleanedUp(t, db, schema, "users", "secondary_id")
464+
},
465+
afterComplete: func(t *testing.T, db *sql.DB, schema string) {
466+
// The table is cleaned up; temporary columns, trigger functions and triggers no longer exist.
467+
TableMustBeCleanedUp(t, db, schema, "users", "secondary_id")
468+
469+
// Inserting a row that does not meet the unique constraint into the new view works.
470+
MustInsert(t, db, schema, "02_drop_unique_constraint", "users", map[string]string{
471+
"secondary_id": "1",
472+
})
473+
},
474+
},
412475
{
413476
name: "dropping a unique constraint preserves the column's default value",
414477
migrations: []migrations.Migration{

‎pkg/migrations/op_drop_multicolumn_constraint.go

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -91,6 +91,10 @@ func (o *OpDropMultiColumnConstraint) Complete(ctx context.Context, conn db.DB,
9191
return err
9292
}
9393

94+
if err := alterSequenceOwnerToDuplicatedColumn(ctx, conn, o.Table, columnName); err != nil {
95+
return err
96+
}
97+
9498
// Drop the old column
9599
_, err = conn.ExecContext(ctx, fmt.Sprintf("ALTER TABLE IF EXISTS %s DROP COLUMN IF EXISTS %s",
96100
pq.QuoteIdentifier(o.Table),

‎pkg/migrations/op_drop_multicolumn_constraint_test.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -269,7 +269,7 @@ func TestDropMultiColumnConstraint(t *testing.T) {
269269
Columns: []migrations.Column{
270270
{
271271
Name: "id",
272-
Type: "integer",
272+
Type: "serial",
273273
Pk: true,
274274
},
275275
{

‎pkg/migrations/op_rename_column_test.go

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -85,6 +85,77 @@ func TestOpRenameColumn(t *testing.T) {
8585
}, rows)
8686
},
8787
},
88+
{
89+
name: "rename serial column",
90+
migrations: []migrations.Migration{
91+
{
92+
Name: "01_create_table",
93+
Operations: migrations.Operations{
94+
&migrations.OpCreateTable{
95+
Name: "orders",
96+
Columns: []migrations.Column{
97+
{
98+
Name: "order_id",
99+
Type: "serial",
100+
Pk: true,
101+
},
102+
{
103+
Name: "description",
104+
Type: "varchar(255)",
105+
},
106+
},
107+
},
108+
},
109+
},
110+
{
111+
Name: "02_rename_serial_column",
112+
Operations: migrations.Operations{
113+
&migrations.OpRenameColumn{
114+
Table: "orders",
115+
From: "order_id",
116+
To: "id",
117+
},
118+
},
119+
},
120+
},
121+
afterStart: func(t *testing.T, db *sql.DB, schema string) {
122+
// The column in the underlying table has not been renamed.
123+
ColumnMustExist(t, db, schema, "orders", "order_id")
124+
125+
// Insertions to the new column name in the new version schema should work.
126+
MustInsert(t, db, schema, "02_rename_serial_column", "orders", map[string]string{
127+
"id": "1",
128+
"description": "first order",
129+
})
130+
131+
// Insertions to the old column name in the old version schema should work.
132+
MustInsert(t, db, schema, "01_create_table", "orders", map[string]string{
133+
"order_id": "2",
134+
"description": "second order",
135+
})
136+
137+
// Data can be read from the view in the new version schema.
138+
rows := MustSelect(t, db, schema, "02_rename_serial_column", "orders")
139+
assert.Equal(t, []map[string]any{
140+
{"id": 1, "description": "first order"},
141+
{"id": 2, "description": "second order"},
142+
}, rows)
143+
},
144+
afterRollback: func(t *testing.T, db *sql.DB, schema string) {
145+
// no-op
146+
},
147+
afterComplete: func(t *testing.T, db *sql.DB, schema string) {
148+
// The column in the underlying table has been renamed.
149+
ColumnMustExist(t, db, schema, "orders", "id")
150+
151+
// Data can be read from the view in the new version schema.
152+
rows := MustSelect(t, db, schema, "02_rename_serial_column", "orders")
153+
assert.Equal(t, []map[string]any{
154+
{"id": 1, "description": "first order"},
155+
{"id": 2, "description": "second order"},
156+
}, rows)
157+
},
158+
},
88159
})
89160
}
90161

0 commit comments

Comments
 (0)