xataio · agedemenli · Feb 11, 2025 · Feb 10, 2025 · Feb 11, 2025 · Feb 11, 2025
diff --git a/pkg/migrations/duplicate.go b/pkg/migrations/duplicate.go
@@ -353,3 +353,36 @@ func errorIgnoringErrorCode(err error, code pq.ErrorCode) error {
 
 	return err
 }
+
+func alterSequenceOwnerToDuplicatedColumn(ctx context.Context, conn db.DB, tableName, columnName string) error {
+	sequenceName := getSequenceNameForColumn(ctx, conn, tableName, columnName)
+	if sequenceName == "" {
+		// No sequence for the column
+		return nil
+	}
+	_, err := conn.ExecContext(ctx, fmt.Sprintf("ALTER SEQUENCE IF EXISTS %s OWNED BY %s.%s",
+		sequenceName,
+		pq.QuoteIdentifier(tableName),
+		pq.QuoteIdentifier(TemporaryName(columnName)),
+	))
+
+	return err
+}
+
+func getSequenceNameForColumn(ctx context.Context, conn db.DB, tableName, columnName string) string {
+	var sequenceName string
+	query := fmt.Sprintf(`
+		SELECT pg_get_serial_sequence('%s', '%s')
+	`, pq.QuoteIdentifier(tableName), columnName)
+	rows, err := conn.QueryContext(ctx, query)
+	if err != nil {
+		return ""
+	}
+	defer rows.Close()
+
+	if err := db.ScanFirstValue(rows, &sequenceName); err != nil {
+		return ""
+	}
+
+	return sequenceName
+}
diff --git a/pkg/migrations/op_alter_column.go b/pkg/migrations/op_alter_column.go
@@ -86,6 +86,10 @@ func (o *OpAlterColumn) Complete(ctx context.Context, conn db.DB, tr SQLTransfor
 		}
 	}
 
+	if err := alterSequenceOwnerToDuplicatedColumn(ctx, conn, o.Table, o.Column); err != nil {
+		return err
+	}
+
 	// Drop the old column
 	_, err := conn.ExecContext(ctx, fmt.Sprintf("ALTER TABLE IF EXISTS %s DROP COLUMN IF EXISTS %s",
 		pq.QuoteIdentifier(o.Table),

diff --git a/pkg/migrations/op_create_constraint.go b/pkg/migrations/op_create_constraint.go
@@ -113,6 +113,12 @@ func (o *OpCreateConstraint) Complete(ctx context.Context, conn db.DB, tr SQLTra
 		}
 	}
 
+	for _, col := range o.Columns {
+		if err := alterSequenceOwnerToDuplicatedColumn(ctx, conn, o.Table, col); err != nil {
+			return err
+		}
+	}
+
 	// remove old columns
 	_, err := conn.ExecContext(ctx, fmt.Sprintf("ALTER TABLE %s %s",
 		pq.QuoteIdentifier(o.Table),

diff --git a/pkg/migrations/op_create_constraint_test.go b/pkg/migrations/op_create_constraint_test.go
@@ -931,5 +931,66 @@ func TestCreateConstraintValidation(t *testing.T) {
 			afterRollback: func(t *testing.T, db *sql.DB, schema string) {},
 			afterComplete: func(t *testing.T, db *sql.DB, schema string) {},
 		},
+		{
+			name: "create unique constraint on serial column",
+			migrations: []migrations.Migration{
+				{
+					Name: "01_add_table",
+					Operations: migrations.Operations{
+						&migrations.OpCreateTable{
+							Name: "users",
+							Columns: []migrations.Column{
+								{
+									Name: "email",
+									Type: "text",
+									Pk:   true,
+								},
+								{
+									Name: "id",
+									Type: "serial",
+								},
+							},
+						},
+					},
+				},
+				{
+					Name: "02_create_constraint",
+					Operations: migrations.Operations{
+						&migrations.OpCreateConstraint{
+							Name:    "unique_id",
+							Table:   "users",
+							Type:    "unique",
+							Columns: []string{"id"},
+							Up: map[string]string{
+								"id": "id",
+							},
+							Down: map[string]string{
+								"id": "id",
+							},
+						},
+					},
+				},
+			},
+			afterStart: func(t *testing.T, db *sql.DB, schema string) {
+				// The index has been created on the underlying table.
+				IndexMustExist(t, db, schema, "users", "unique_id")
+			},
+			afterRollback: func(t *testing.T, db *sql.DB, schema string) {
+				// The index has been dropped from the the underlying table.
+				IndexMustNotExist(t, db, schema, "users", "unique_id")
+			},
+			afterComplete: func(t *testing.T, db *sql.DB, schema string) {
+				// Functions, triggers and temporary columns are dropped.
+				TableMustBeCleanedUp(t, db, schema, "users", "id")
+
+				// Inserting values into the new schema that violate uniqueness should fail.
+				MustInsert(t, db, schema, "02_create_constraint", "users", map[string]string{
+					"email": "alice@xata.io", "id": "1",
+				})
+				MustNotInsert(t, db, schema, "02_create_constraint", "users", map[string]string{
+					"email": "bob@xata.io", "id": "1",
+				}, testutils.UniqueViolationErrorCode)
+			},
+		},
 	})
 }
diff --git a/pkg/migrations/op_drop_constraint.go b/pkg/migrations/op_drop_constraint.go
@@ -86,6 +86,10 @@ func (o *OpDropConstraint) Complete(ctx context.Context, conn db.DB, tr SQLTrans
 		return err
 	}
 
+	if err := alterSequenceOwnerToDuplicatedColumn(ctx, conn, o.Table, column.Name); err != nil {
+		return err
+	}
+
 	// Drop the old column
 	_, err = conn.ExecContext(ctx, fmt.Sprintf("ALTER TABLE IF EXISTS %s DROP COLUMN IF EXISTS %s",
 		pq.QuoteIdentifier(o.Table),

diff --git a/pkg/migrations/op_drop_constraint_test.go b/pkg/migrations/op_drop_constraint_test.go
@@ -409,6 +409,69 @@ func TestDropConstraint(t *testing.T) {
 				})
 			},
 		},
+		{
+			name: "drop unique constraint from serial column",
+			migrations: []migrations.Migration{
+				{
+					Name: "01_add_tables",
+					Operations: migrations.Operations{
+						&migrations.OpCreateTable{
+							Name: "users",
+							Columns: []migrations.Column{
+								{
+									Name: "id",
+									Type: "serial",
+									Pk:   true,
+								},
+								{
+									Name:   "secondary_id",
+									Type:   "serial",
+									Unique: true,
+								},
+							},
+						},
+					},
+				},
+				{
+					Name: "02_drop_unique_constraint",
+					Operations: migrations.Operations{
+						&migrations.OpDropConstraint{
+							Table: "users",
+							Name:  "users_secondary_id_key",
+							Up:    "secondary_id",
+							Down:  "secondary_id",
+						},
+					},
+				},
+			},
+			afterStart: func(t *testing.T, db *sql.DB, schema string) {
+				// The new (temporary) `secondary_id` column should exist on the underlying table.
+				ColumnMustExist(t, db, schema, "users", migrations.TemporaryName("secondary_id"))
+
+				// Inserting a row that meets the unique constraint into the old view works.
+				MustInsert(t, db, schema, "01_add_tables", "users", map[string]string{
+					"secondary_id": "1",
+				})
+
+				// Inserting a row that does not meet the unique constraint into the old view fails.
+				MustNotInsert(t, db, schema, "01_add_tables", "users", map[string]string{
+					"secondary_id": "1",
+				}, testutils.UniqueViolationErrorCode)
+			},
+			afterRollback: func(t *testing.T, db *sql.DB, schema string) {
+				// The table is cleaned up; temporary columns, trigger functions and triggers no longer exist.
+				TableMustBeCleanedUp(t, db, schema, "users", "secondary_id")
+			},
+			afterComplete: func(t *testing.T, db *sql.DB, schema string) {
+				// The table is cleaned up; temporary columns, trigger functions and triggers no longer exist.
+				TableMustBeCleanedUp(t, db, schema, "users", "secondary_id")
+
+				// Inserting a row that does not meet the unique constraint into the new view works.
+				MustInsert(t, db, schema, "02_drop_unique_constraint", "users", map[string]string{
+					"secondary_id": "1",
+				})
+			},
+		},
 		{
 			name: "dropping a unique constraint preserves the column's default value",
 			migrations: []migrations.Migration{

diff --git a/pkg/migrations/op_drop_multicolumn_constraint.go b/pkg/migrations/op_drop_multicolumn_constraint.go
@@ -92,6 +92,10 @@ func (o *OpDropMultiColumnConstraint) Complete(ctx context.Context, conn db.DB,
 			return err
 		}
 
+		if err := alterSequenceOwnerToDuplicatedColumn(ctx, conn, o.Table, columnName); err != nil {
+			return err
+		}
+
 		// Drop the old column
 		_, err = conn.ExecContext(ctx, fmt.Sprintf("ALTER TABLE IF EXISTS %s DROP COLUMN IF EXISTS %s",
 			pq.QuoteIdentifier(o.Table),

diff --git a/pkg/migrations/op_drop_multicolumn_constraint_test.go b/pkg/migrations/op_drop_multicolumn_constraint_test.go
@@ -269,7 +269,7 @@ func TestDropMultiColumnConstraint(t *testing.T) {
 							Columns: []migrations.Column{
 								{
 									Name: "id",
-									Type: "integer",
+									Type: "serial",
 									Pk:   true,
 								},
 								{

diff --git a/pkg/migrations/op_rename_column_test.go b/pkg/migrations/op_rename_column_test.go
@@ -85,6 +85,77 @@ func TestOpRenameColumn(t *testing.T) {
 				}, rows)
 			},
 		},
+		{
+			name: "rename serial column",
+			migrations: []migrations.Migration{
+				{
+					Name: "01_create_table",
+					Operations: migrations.Operations{
+						&migrations.OpCreateTable{
+							Name: "orders",
+							Columns: []migrations.Column{
+								{
+									Name: "order_id",
+									Type: "serial",
+									Pk:   true,
+								},
+								{
+									Name: "description",
+									Type: "varchar(255)",
+								},
+							},
+						},
+					},
+				},
+				{
+					Name: "02_rename_serial_column",
+					Operations: migrations.Operations{
+						&migrations.OpRenameColumn{
+							Table: "orders",
+							From:  "order_id",
+							To:    "id",
+						},
+					},
+				},
+			},
+			afterStart: func(t *testing.T, db *sql.DB, schema string) {
+				// The column in the underlying table has not been renamed.
+				ColumnMustExist(t, db, schema, "orders", "order_id")
+
+				// Insertions to the new column name in the new version schema should work.
+				MustInsert(t, db, schema, "02_rename_serial_column", "orders", map[string]string{
+					"id":          "1",
+					"description": "first order",
+				})
+
+				// Insertions to the old column name in the old version schema should work.
+				MustInsert(t, db, schema, "01_create_table", "orders", map[string]string{
+					"order_id":    "2",
+					"description": "second order",
+				})
+
+				// Data can be read from the view in the new version schema.
+				rows := MustSelect(t, db, schema, "02_rename_serial_column", "orders")
+				assert.Equal(t, []map[string]any{
+					{"id": 1, "description": "first order"},
+					{"id": 2, "description": "second order"},
+				}, rows)
+			},
+			afterRollback: func(t *testing.T, db *sql.DB, schema string) {
+				// no-op
+			},
+			afterComplete: func(t *testing.T, db *sql.DB, schema string) {
+				// The column in the underlying table has been renamed.
+				ColumnMustExist(t, db, schema, "orders", "id")
+
+				// Data can be read from the view in the new version schema.
+				rows := MustSelect(t, db, schema, "02_rename_serial_column", "orders")
+				assert.Equal(t, []map[string]any{
+					{"id": 1, "description": "first order"},
+					{"id": 2, "description": "second order"},
+				}, rows)
+			},
+		},
 	})
 }