Repository: ariga/atlas
Pull request #601 (2022-02-24 16:13:51)

sql/postgres: fix statements generation when IfExists/IfNotExists is used
Also, add support for `DROP ... CASCADE` statements.


Patch:
diff --git a/sql/postgres/inspect.go b/sql/postgres/inspect.go
index 489ece46..4bb687e0 100644
--- a/sql/postgres/inspect.go
+++ b/sql/postgres/inspect.go
@@ -681,6 +681,16 @@ type (
 		schema.Attr
 		Columns []string
 	}
+
+	// Cascade clause for appending to the different DROP changes.
+	// For example, A schema.DropSchema with this clause will be
+	// applied as follows:
+	//
+	//	DROP SCHEMA <name> CASCADE
+	//
+	Cascade struct {
+		schema.Clause
+	}
 )
 
 const (
diff --git a/sql/postgres/migrate.go b/sql/postgres/migrate.go
index ee9119ca..7a7f258b 100644
--- a/sql/postgres/migrate.go
+++ b/sql/postgres/migrate.go
@@ -86,10 +86,11 @@ func (s *state) topLevel(changes []schema.Change) []schema.Change {
 	for _, c := range changes {
 		switch c := c.(type) {
 		case *schema.AddSchema:
-			b := Build("CREATE SCHEMA").Ident(c.S.Name)
+			b := Build("CREATE SCHEMA")
 			if sqlx.Has(c.Extra, &schema.IfNotExists{}) {
 				b.P("IF NOT EXISTS")
 			}
+			b.Ident(c.S.Name)
 			s.append(&migrate.Change{
 				Cmd:     b.String(),
 				Source:  c,
@@ -97,9 +98,13 @@ func (s *state) topLevel(changes []schema.Change) []schema.Change {
 				Comment: fmt.Sprintf("Add new schema named %q", c.S.Name),
 			})
 		case *schema.DropSchema:
-			b := Build("DROP SCHEMA").Ident(c.S.Name)
+			b := Build("DROP SCHEMA")
 			if sqlx.Has(c.Extra, &schema.IfExists{}) {
-				b.P("IF NOT EXISTS")
+				b.P("IF EXISTS")
+			}
+			b.Ident(c.S.Name)
+			if sqlx.Has(c.Extra, &Cascade{}) {
+				b.P("CASCADE")
 			}
 			s.append(&migrate.Change{
 				Cmd:     b.String(),
@@ -119,10 +124,11 @@ func (s *state) addTable(ctx context.Context, add *schema.AddTable) error {
 	if err := s.addTypes(ctx, add.T.Columns...); err != nil {
 		return err
 	}
-	b := Build("CREATE TABLE").Table(add.T)
+	b := Build("CREATE TABLE")
 	if sqlx.Has(add.Extra, &schema.IfNotExists{}) {
 		b.P("IF NOT EXISTS")
 	}
+	b.Table(add.T)
 	b.Wrap(func(b *sqlx.Builder) {
 		b.MapComma(add.T.Columns, func(i int, b *sqlx.Builder) {
 			s.column(b, add.T.Columns[i])
@@ -155,10 +161,11 @@ func (s *state) addTable(ctx context.Context, add *schema.AddTable) error {
 
 // dropTable builds and executes the query for dropping a table from a schema.
 func (s *state) dropTable(drop *schema.DropTable) {
-	b := Build("DROP TABLE").Table(drop.T)
+	b := Build("DROP TABLE")
 	if sqlx.Has(drop.Extra, &schema.IfExists{}) {
 		b.P("IF EXISTS")
 	}
+	b.Table(drop.T)
 	s.append(&migrate.Change{
 		Cmd:     b.String(),
 		Source:  drop,


Test patch:
diff --git a/sql/postgres/migrate_test.go b/sql/postgres/migrate_test.go
index 9a7a4d8f..346ef00e 100644
--- a/sql/postgres/migrate_test.go
+++ b/sql/postgres/migrate_test.go
@@ -22,6 +22,29 @@ func TestPlanChanges(t *testing.T) {
 		mock    func(mock)
 		plan    *migrate.Plan
 	}{
+		{
+			changes: []schema.Change{
+				&schema.AddSchema{S: schema.New("test"), Extra: []schema.Clause{&schema.IfNotExists{}}},
+				&schema.DropSchema{S: schema.New("test"), Extra: []schema.Clause{&schema.IfExists{}}},
+				&schema.DropSchema{S: schema.New("test"), Extra: []schema.Clause{&Cascade{}}},
+			},
+			plan: &migrate.Plan{
+				Reversible:    false,
+				Transactional: true,
+				Changes: []*migrate.Change{
+					{
+						Cmd:     `CREATE SCHEMA IF NOT EXISTS "test"`,
+						Reverse: `DROP SCHEMA "test"`,
+					},
+					{
+						Cmd: `DROP SCHEMA IF EXISTS "test"`,
+					},
+					{
+						Cmd: `DROP SCHEMA "test" CASCADE`,
+					},
+				},
+			},
+		},
 		{
 			changes: []schema.Change{
 				func() schema.Change {