Repository: ariga/atlas
Pull request #1514 (2023-03-07 20:47:52)

cmd/atlas/internal: add sumfile to remotedir memdir


Patch:
diff --git a/cmd/atlas/internal/cmdext/cmdext.go b/cmd/atlas/internal/cmdext/cmdext.go
index 9b54a29b..8da0d207 100644
--- a/cmd/atlas/internal/cmdext/cmdext.go
+++ b/cmd/atlas/internal/cmdext/cmdext.go
@@ -11,6 +11,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"io"
 	"io/fs"
 	"net/url"
 	"os"
@@ -403,6 +404,15 @@ func memdir(client *cloudapi.Client, dirName string, tag string) (string, error)
 			return "", err
 		}
 	}
+	if hf, err := dir.Open(migrate.HashFileName); err == nil {
+		b, err := io.ReadAll(hf)
+		if err != nil {
+			return "", err
+		}
+		if err := md.WriteFile(migrate.HashFileName, b); err != nil {
+			return "", err
+		}
+	}
 	return "mem://" + dirName, nil
 }
 
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/cmd/atlas/internal/cmdext/cmdext_test.go b/cmd/atlas/internal/cmdext/cmdext_test.go
index 70d1c3c0..f3b838d6 100644
--- a/cmd/atlas/internal/cmdext/cmdext_test.go
+++ b/cmd/atlas/internal/cmdext/cmdext_test.go
@@ -258,6 +258,13 @@ func TestRemoteDir(t *testing.T) {
 			if err := d.WriteFile("1.sql", []byte("create table t(c int);")); err != nil {
 				http.Error(w, err.Error(), http.StatusInternalServerError)
 			}
+			sum, err := d.Checksum()
+			if err != nil {
+				http.Error(w, err.Error(), http.StatusInternalServerError)
+			}
+			if err := migrate.WriteSumFile(&d, sum); err != nil {
+				http.Error(w, err.Error(), http.StatusInternalServerError)
+			}
 			arch, err := migrate.ArchiveDir(&d)
 			if err != nil {
 				http.Error(w, err.Error(), http.StatusInternalServerError)
@@ -300,7 +307,6 @@ dir = data.remote_dir.hello.url
 	require.NoError(t, err)
 	require.Equal(t, "Bearer token", token)
 	require.Equal(t, "xyz", tag)
-
 	md := migrate.OpenMemDir(strings.TrimPrefix(v.Dir, "mem://"))
 	defer md.Close()
 	files, err := md.Files()
@@ -308,4 +314,7 @@ dir = data.remote_dir.hello.url
 	require.Len(t, files, 1)
 	require.Equal(t, "1.sql", files[0].Name())
 	require.Equal(t, "create table t(c int);", string(files[0].Bytes()))
+	require.NoError(t, migrate.Validate(md))
+	_, err = md.Open(migrate.HashFileName)
+	require.NoError(t, err)
 }
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)`,
+					},
 				},
 			},
 		},