Repository: ariga/atlas
Pull request #362 (2022-01-03 07:22:58)

sql/postgres: support modifying identity column attributes


Patch:
diff --git a/sql/postgres/migrate.go b/sql/postgres/migrate.go
index 00020513..abfa33fd 100644
--- a/sql/postgres/migrate.go
+++ b/sql/postgres/migrate.go
@@ -247,8 +247,12 @@ func (s *state) alterTable(t *schema.Table, changes []schema.Change) error {
 			b.P("DROP COLUMN").Ident(change.C.Name)
 			reversible = false
 		case *schema.ModifyColumn:
-			s.alterColumn(b, change.Change, change.To)
-			s.alterColumn(reverse, change.Change, change.From)
+			if err := s.alterColumn(b, change.Change, change.To); err != nil {
+				errors = append(errors, err.Error())
+			}
+			if err := s.alterColumn(reverse, change.Change, change.From); err != nil {
+				errors = append(errors, err.Error())
+			}
 		case *schema.AddForeignKey:
 			b.P("ADD")
 			s.fks(b, change.F)
@@ -481,7 +485,7 @@ func (s *state) columnDefault(b *sqlx.Builder, c *schema.Column) {
 	}
 }
 
-func (s *state) alterColumn(b *sqlx.Builder, k schema.ChangeKind, c *schema.Column) {
+func (s *state) alterColumn(b *sqlx.Builder, k schema.ChangeKind, c *schema.Column) error {
 	for !k.Is(schema.NoChange) {
 		b.P("ALTER COLUMN").Ident(c.Name)
 		switch {
@@ -501,19 +505,25 @@ func (s *state) alterColumn(b *sqlx.Builder, k schema.ChangeKind, c *schema.Colu
 			b.P("DROP DEFAULT")
 			k &= ^schema.ChangeDefault
 		case k.Is(schema.ChangeDefault) && c.Default != nil:
-			x, ok := c.Default.(*schema.RawExpr)
+			s.columnDefault(b.P("SET"), c)
+			k &= ^schema.ChangeDefault
+		case k.Is(schema.ChangeAttr):
+			id, ok := identity(c.Attrs)
 			if !ok {
-				panic(fmt.Sprintf("unexpected column default: %T", c.Default))
+				return fmt.Errorf("unexpected attribute change: %v", c.Attrs)
 			}
-			b.P("SET DEFAULT", x.X)
-			k &= ^schema.ChangeDefault
+			// The syntax for altering identity columns is identical to sequence_options.
+			// https://www.postgresql.org/docs/current/sql-altersequence.html
+			b.P("SET START WITH", strconv.FormatInt(id.Sequence.Start, 10), "SET INCREMENT BY", strconv.FormatInt(id.Sequence.Increment, 10), "RESTART")
+			k &= ^schema.ChangeAttr
 		default:
-			panic(fmt.Sprintf("unexpected column change: %d", k))
+			return fmt.Errorf("unexpected column change: %d", k)
 		}
 		if !k.Is(schema.NoChange) {
 			b.Comma()
 		}
 	}
+	return nil
 }
 
 func (s *state) indexParts(b *sqlx.Builder, parts []*schema.IndexPart) {


Test patch:
diff --git a/sql/postgres/migrate_test.go b/sql/postgres/migrate_test.go
index 20c1062d..ade42a3c 100644
--- a/sql/postgres/migrate_test.go
+++ b/sql/postgres/migrate_test.go
@@ -192,6 +192,11 @@ func TestPlanChanges(t *testing.T) {
 							&schema.DropColumn{
 								C: &schema.Column{Name: "name", Type: &schema.ColumnType{Type: &schema.StringType{T: "varchar(255)"}}},
 							},
+							&schema.ModifyColumn{
+								From:   &schema.Column{Name: "id", Type: &schema.ColumnType{Type: &schema.IntegerType{T: "bigint"}}, Attrs: []schema.Attr{&Identity{}}},
+								To:     &schema.Column{Name: "id", Type: &schema.ColumnType{Type: &schema.IntegerType{T: "bigint"}}, Attrs: []schema.Attr{&Identity{Sequence: &Sequence{Start: 1024}}}},
+								Change: schema.ChangeAttr,
+							},
 						},
 					}
 				}(),
@@ -200,7 +205,7 @@ func TestPlanChanges(t *testing.T) {
 				Reversible:    false,
 				Transactional: true,
 				Changes: []*migrate.Change{
-					{Cmd: `ALTER TABLE "users" DROP COLUMN "name"`},
+					{Cmd: `ALTER TABLE "users" DROP COLUMN "name", ALTER COLUMN "id" SET START WITH 1024 SET INCREMENT BY 1 RESTART`},
 				},
 			},
 		},