Repository: ariga/atlas
Pull request #1515 (2023-03-08 08:09:57)

sql/postgres: drop tables with CASCADE on cleanup
In case other resources like views were created in the migration, replay will fail


Patch:
diff --git a/sql/postgres/driver.go b/sql/postgres/driver.go
index 6a47e95b..a1194245 100644
--- a/sql/postgres/driver.go
+++ b/sql/postgres/driver.go
@@ -196,7 +196,7 @@ func (d *Driver) Snapshot(ctx context.Context) (migrate.RestoreFunc, error) {
 			if err != nil {
 				return err
 			}
-			return d.ApplyChanges(ctx, changes)
+			return d.ApplyChanges(ctx, withCascade(changes))
 		}, nil
 	}
 	// Not bound to a schema.
@@ -213,7 +213,7 @@ func (d *Driver) Snapshot(ctx context.Context) (migrate.RestoreFunc, error) {
 		if err != nil {
 			return err
 		}
-		return d.ApplyChanges(ctx, changes)
+		return d.ApplyChanges(ctx, withCascade(changes))
 	}
 	// Postgres is considered clean, if there are no schemas or the public schema has no tables.
 	if len(realm.Schemas) == 0 {
@@ -228,6 +228,15 @@ func (d *Driver) Snapshot(ctx context.Context) (migrate.RestoreFunc, error) {
 	return nil, &migrate.NotCleanError{Reason: fmt.Sprintf("found schema %q", realm.Schemas[0].Name)}
 }
 
+func withCascade(changes schema.Changes) schema.Changes {
+	for _, c := range changes {
+		if d, ok := c.(*schema.DropTable); ok {
+			d.Extra = append(d.Extra, &Cascade{})
+		}
+	}
+	return changes
+}
+
 // CheckClean implements migrate.CleanChecker.
 func (d *Driver) CheckClean(ctx context.Context, revT *migrate.TableIdent) error {
 	if revT == nil { // accept nil values
diff --git a/sql/postgres/inspect.go b/sql/postgres/inspect.go
index ca193e38..184fbc23 100644
--- a/sql/postgres/inspect.go
+++ b/sql/postgres/inspect.go
@@ -871,6 +871,12 @@ type (
 		C     *schema.Column
 		Attrs []schema.Attr
 	}
+
+	// Cascade describes that a CASCADE clause should be added to the DROP [TABLE|SCHEMA]
+	// operation. Note, this clause is automatically added to DROP SCHEMA by the planner.
+	Cascade struct {
+		schema.Clause
+	}
 )
 
 // IsUnique reports if the type is unique constraint.
diff --git a/sql/postgres/migrate.go b/sql/postgres/migrate.go
index 3454f130..d695bb44 100644
--- a/sql/postgres/migrate.go
+++ b/sql/postgres/migrate.go
@@ -225,6 +225,9 @@ func (s *state) dropTable(ctx context.Context, drop *schema.DropTable) error {
 		b.P("IF EXISTS")
 	}
 	b.Table(drop.T)
+	if sqlx.Has(drop.Extra, &Cascade{}) {
+		b.P("CASCADE")
+	}
 	cmd.main = &migrate.Change{
 		Cmd:     b.String(),
 		Source:  drop,


Test patch:
diff --git a/sql/postgres/migrate_test.go b/sql/postgres/migrate_test.go
index 9f5c2972..8729c1f8 100644
--- a/sql/postgres/migrate_test.go
+++ b/sql/postgres/migrate_test.go
@@ -328,6 +328,7 @@ func TestPlanChanges(t *testing.T) {
 		{
 			changes: []schema.Change{
 				&schema.DropTable{T: schema.NewTable("posts").AddColumns(schema.NewIntColumn("id", "int"))},
+				&schema.DropTable{T: schema.NewTable("users").AddColumns(schema.NewIntColumn("id", "int")), Extra: []schema.Clause{&Cascade{}}},
 			},
 			wantPlan: &migrate.Plan{
 				Reversible:    true,
@@ -337,6 +338,10 @@ func TestPlanChanges(t *testing.T) {
 						Cmd:     `DROP TABLE "posts"`,
 						Reverse: `CREATE TABLE "posts" ("id" integer NOT NULL)`,
 					},
+					{
+						Cmd:     `DROP TABLE "users" CASCADE`,
+						Reverse: `CREATE TABLE "users" ("id" integer NOT NULL)`,
+					},
 				},
 			},
 		},