Repository: ariga/atlas
Pull request #663 (2022-03-21 19:10:23)

sql/migrate: GlobStateReader will now try to get a lock before replay…
…ing the migration directory, check for the driver to be in an empty state and clean up after the replay is done


Patch:
diff --git a/cmd/action/migrate.go b/cmd/action/migrate.go
index 4d605b05..f243c6c6 100644
--- a/cmd/action/migrate.go
+++ b/cmd/action/migrate.go
@@ -6,7 +6,6 @@ package action
 
 import (
 	"context"
-	"errors"
 	"fmt"
 	"io"
 	"io/ioutil"
@@ -125,8 +124,14 @@ func CmdMigrateDiffRun(cmd *cobra.Command, args []string) error {
 		return err
 	}
 	defer dev.Close()
-	if err := checkClean(cmd.Context(), dev); err != nil {
-		return err
+	// Acquire a lock.
+	if l, ok := dev.Driver.(schema.Locker); ok {
+		unlock, err := l.Lock(cmd.Context(), "atlas_migrate_diff", 0)
+		if err != nil {
+			return err
+		}
+		// If unlocking fails notify the user about it.
+		defer cobra.CheckErr(unlock())
 	}
 	// Open the migration directory.
 	dir, err := dir()
@@ -147,7 +152,7 @@ func CmdMigrateDiffRun(cmd *cobra.Command, args []string) error {
 		return err
 	}
 	// Plan the changes and create a new migration file.
-	pl := migrate.NewPlanner(dev, dir, migrate.WithFormatter(tf))
+	pl := migrate.NewPlanner(dev.Driver, dir, migrate.WithFormatter(tf))
 	var name string
 	if len(args) > 0 {
 		name = args[0]
@@ -158,7 +163,6 @@ func CmdMigrateDiffRun(cmd *cobra.Command, args []string) error {
 	}
 	// Write the plan to a new file.
 	return pl.WritePlan(plan)
-	// TODO(masseelch): clean up dev after reading the state from migration dir.
 }
 
 // CmdMigrateHashRun is the command executed when running the CLI with 'migrate hash' args.
@@ -243,21 +247,6 @@ func to(ctx context.Context, d *Driver) (migrate.StateReader, error) {
 	}
 }
 
-func checkClean(ctx context.Context, drv *Driver) error {
-	realm, err := drv.InspectRealm(ctx, nil)
-	if err != nil {
-		return err
-	}
-	if len(realm.Schemas) == 0 {
-		return nil
-	}
-	// If this is an SQLite database it is considered clean if the "main" schema has no tables.
-	if strings.HasPrefix(MigrateFlags.DevURL, "sqlite") && realm.Schemas[0].Name == "main" && len(realm.Schemas[0].Tables) == 0 {
-		return nil
-	}
-	return errors.New("dev database must be clean")
-}
-
 var (
 	funcMap = template.FuncMap{
 		"now": func() string { return time.Now().Format("20060102150405") },
diff --git a/sql/migrate/migrate.go b/sql/migrate/migrate.go
index def3b642..d9402a6f 100644
--- a/sql/migrate/migrate.go
+++ b/sql/migrate/migrate.go
@@ -211,16 +211,76 @@ func DisableChecksum() PlannerOption {
 
 // GlobStateReader creates a StateReader that loads all files from Dir matching
 // glob in lexicographic order and uses the Driver to create a migration state.
+//
+// If the given Driver implements the Emptier interface the IsEmpty method will be used to determine if the Driver is
+// connected to an "empty" database. This behavior was added to support SQLite flavors.
 func GlobStateReader(dir Dir, drv Driver, glob string) StateReaderFunc {
-	return func(ctx context.Context) (*schema.Realm, error) {
+	var errNotClean = errors.New("sql/migrate: connected database is not clean")
+	return func(ctx context.Context) (realm *schema.Realm, err error) {
+		// Collect the migration files.
 		names, err := fs.Glob(dir, glob)
 		if err != nil {
 			return nil, err
 		}
+		// We need an empty database state to reliably replay the migration directory.
+		if c, ok := drv.(interface {
+			// The IsClean method can be added to a Driver to override how to
+			// determine if a connected database is in a clean state.
+			// This interface exists only to support SQLite favors.
+			IsClean(context.Context) (bool, error)
+		}); ok {
+			e, err := c.IsClean(ctx)
+			if err != nil {
+				return nil, fmt.Errorf("sql/migrate: checking database state: %w", err)
+			}
+			if !e {
+				return nil, errNotClean
+			}
+		} else {
+			realm, err = drv.InspectRealm(ctx, nil)
+			if err != nil {
+				return nil, err
+			}
+			if len(realm.Schemas) > 0 {
+				return nil, errNotClean
+			}
+		}
 		// Sort files lexicographically.
 		sort.Slice(names, func(i, j int) bool {
 			return names[i] < names[j]
 		})
+		// If the driver supports it, acquire a lock while replaying the migration changes.
+		if l, ok := drv.(schema.Locker); ok {
+			unlock, err := l.Lock(ctx, "atlas_migration_directory_state", 0)
+			if err != nil {
+				return nil, fmt.Errorf("sql/migrate: acquiring database lock: %w", err)
+			}
+			defer unlock()
+		}
+		// Clean up after ourselves.
+		defer func() {
+			if e, ok := drv.(interface {
+				// The Clean method can be added to a Driver to change how to clean a database.
+				// This interface exists only to support SQLite favors.
+				Clean(context.Context) error
+			}); ok {
+				if derr := e.Clean(ctx); derr != nil {
+					err = wrap(derr, err)
+				}
+				return
+			}
+			realm, derr := drv.InspectRealm(ctx, nil)
+			if derr != nil {
+				err = wrap(derr, err)
+			}
+			del := make([]schema.Change, len(realm.Schemas))
+			for i, s := range realm.Schemas {
+				del[i] = &schema.DropSchema{S: s, Extra: []schema.Clause{&schema.IfExists{}}}
+			}
+			if derr := drv.ApplyChanges(ctx, del); derr != nil {
+				err = wrap(derr, err)
+			}
+		}()
 		for _, n := range names {
 			f, err := dir.Open(n)
 			if err != nil {
@@ -241,7 +301,8 @@ func GlobStateReader(dir Dir, drv Driver, glob string) StateReaderFunc {
 				return nil, err
 			}
 		}
-		return drv.InspectRealm(ctx, nil)
+		realm, err = drv.InspectRealm(ctx, nil)
+		return
 	}
 }
 
@@ -549,3 +610,10 @@ func ensureSemicolonSuffix(s string) string {
 	}
 	return s
 }
+
+func wrap(err1, err2 error) error {
+	if err2 != nil {
+		return fmt.Errorf("sql/migrate: %w: %v", err2, err1)
+	}
+	return err1
+}
diff --git a/sql/migrate/testdata/sub/1.a_sub-up.sql b/sql/migrate/testdata/sub/1.a_sub.up.sql
similarity index 100%
rename from sql/migrate/testdata/sub/1.a_sub-up.sql
rename to sql/migrate/testdata/sub/1.a_sub.up.sql
diff --git a/sql/postgres/driver.go b/sql/postgres/driver.go
index 0235209d..c6987f2e 100644
--- a/sql/postgres/driver.go
+++ b/sql/postgres/driver.go
@@ -69,7 +69,7 @@ func Open(db schema.ExecQuerier) (*Driver, error) {
 }
 
 func (d *Driver) dev() *sqlx.DevDriver {
-	return &sqlx.DevDriver{Driver: d, MaxNameLen: 63, DropClause: []schema.Clause{&Cascade{}}}
+	return &sqlx.DevDriver{Driver: d, MaxNameLen: 63}
 }
 
 // NormalizeRealm returns the normal representation of the given database.
diff --git a/sql/postgres/inspect.go b/sql/postgres/inspect.go
index 600b6f23..a659090d 100644
--- a/sql/postgres/inspect.go
+++ b/sql/postgres/inspect.go
@@ -681,16 +681,6 @@ 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 0436a2a7..9bb58275 100644
--- a/sql/postgres/migrate.go
+++ b/sql/postgres/migrate.go
@@ -94,7 +94,7 @@ func (s *state) topLevel(changes []schema.Change) []schema.Change {
 			s.append(&migrate.Change{
 				Cmd:     b.String(),
 				Source:  c,
-				Reverse: Build("DROP SCHEMA").Ident(c.S.Name).String(),
+				Reverse: Build("DROP SCHEMA").Ident(c.S.Name).P("CASCADE").String(),
 				Comment: fmt.Sprintf("Add new schema named %q", c.S.Name),
 			})
 		case *schema.DropSchema:
@@ -102,10 +102,7 @@ func (s *state) topLevel(changes []schema.Change) []schema.Change {
 			if sqlx.Has(c.Extra, &schema.IfExists{}) {
 				b.P("IF EXISTS")
 			}
-			b.Ident(c.S.Name)
-			if sqlx.Has(c.Extra, &Cascade{}) {
-				b.P("CASCADE")
-			}
+			b.Ident(c.S.Name).P("CASCADE")
 			s.append(&migrate.Change{
 				Cmd:     b.String(),
 				Source:  c,
diff --git a/sql/sqlite/driver.go b/sql/sqlite/driver.go
index 9d86fede..50aa76bb 100644
--- a/sql/sqlite/driver.go
+++ b/sql/sqlite/driver.go
@@ -60,6 +60,30 @@ func Open(db schema.ExecQuerier) (*Driver, error) {
 	}, nil
 }
 
+// IsClean implements the inlined IsClean interface to override what to consider a clean database.
+func (d *Driver) IsClean(ctx context.Context) (bool, error) {
+	r, err := d.InspectRealm(ctx, nil)
+	if err != nil {
+		return false, err
+	}
+	return r == nil || len(r.Schemas) == 1 && r.Schemas[0].Name == "main" && len(r.Schemas[0].Tables) == 0, nil
+}
+
+// Clean implements the inlined migrate.Clean interface to override the "emptying" behavior.
+func (d *Driver) Clean(ctx context.Context) error {
+	for _, stmt := range []string{
+		"PRAGMA writable_schema = 1;",
+		"DELETE FROM sqlite_master WHERE type IN ('table', 'index', 'trigger');",
+		"PRAGMA writable_schema = 0;",
+		"VACUUM;",
+	} {
+		if _, err := d.ExecContext(ctx, stmt); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
 // SQLite standard data types as defined in its codebase and documentation.
 // https://www.sqlite.org/datatype3.html
 // https://github.com/sqlite/sqlite/blob/master/src/global.c


Test patch:
diff --git a/cmd/action/migrate_test.go b/cmd/action/migrate_test.go
index 0a59b923..763ecfef 100644
--- a/cmd/action/migrate_test.go
+++ b/cmd/action/migrate_test.go
@@ -5,6 +5,8 @@
 package action
 
 import (
+	"context"
+	"errors"
 	"fmt"
 	"io"
 	"io/ioutil"
@@ -16,6 +18,7 @@ import (
 	"time"
 
 	"ariga.io/atlas/sql/migrate"
+	"ariga.io/atlas/sql/schema"
 	"github.com/stretchr/testify/require"
 )
 
@@ -35,8 +38,8 @@ func TestMigrate_Diff(t *testing.T) {
 		"--dev-url", openSQLite(t, "create table t (c int);"),
 		"--to", hclURL(t),
 	)
-	require.True(t, strings.HasPrefix(s, "Error: dev database must be clean"))
-	require.EqualError(t, err, "dev database must be clean")
+	require.True(t, strings.HasPrefix(s, "Error: sql/migrate: connected database is not clean"))
+	require.EqualError(t, err, "sql/migrate: connected database is not clean")
 
 	// Works.
 	s, err = runCmd(
@@ -50,6 +53,28 @@ func TestMigrate_Diff(t *testing.T) {
 	require.NoError(t, err)
 	require.FileExists(t, filepath.Join(p, fmt.Sprintf("%s_name.sql", time.Now().Format("20060102150405"))))
 	require.FileExists(t, filepath.Join(p, "atlas.sum"))
+
+	// A lock will prevent diffing.
+	DefaultMux.RegisterProvider("sqlitelock", func(s string) (*Driver, error) {
+		drv, err := sqliteProvider(s)
+		if err != nil {
+			return nil, err
+		}
+		drv.Driver = &sqliteLockerDriver{Driver: drv.Driver}
+		return drv, nil
+	})
+	f, err := os.Create(filepath.Join(p, "test.db"))
+	require.NoError(t, err)
+	require.NoError(t, f.Close())
+	s, err = runCmd(
+		RootCmd, "migrate", "diff",
+		"name",
+		"--dir", "file://"+t.TempDir(),
+		"--dev-url", fmt.Sprintf("sqlitelock://file:%s?cache=shared&_fk=1", filepath.Join(p, "test.db")),
+		"--to", hclURL(t),
+	)
+	require.True(t, strings.HasPrefix(s, "Error: "+lockErr))
+	require.EqualError(t, err, lockErr)
 }
 
 func TestMigrate_Validate(t *testing.T) {
@@ -216,3 +241,11 @@ func copyFile(src, dst string) error {
 	_, err = io.Copy(df, sf)
 	return err
 }
+
+type sqliteLockerDriver struct{ migrate.Driver }
+
+const lockErr = "lockErr"
+
+func (d *sqliteLockerDriver) Lock(context.Context, string, time.Duration) (schema.UnlockFunc, error) {
+	return func() error { return nil }, errors.New(lockErr)
+}
diff --git a/sql/migrate/migrate_test.go b/sql/migrate/migrate_test.go
index 57c8faf7..b95babac 100644
--- a/sql/migrate/migrate_test.go
+++ b/sql/migrate/migrate_test.go
@@ -8,6 +8,8 @@ import (
 	"context"
 	"database/sql"
 	_ "embed"
+	"errors"
+	"fmt"
 	"io"
 	"io/fs"
 	"os"
@@ -133,6 +135,7 @@ func TestHash_MarshalText(t *testing.T) {
 	h, err := migrate.HashSum(d)
 	require.NoError(t, err)
 	ac, err := h.MarshalText()
+	fmt.Println(string(ac))
 	require.Equal(t, hash, ac)
 }
 
@@ -151,16 +154,45 @@ func TestGlobStateReader(t *testing.T) {
 		drv = &mockDriver{}
 		ctx = context.Background()
 	)
-	localFS, err := migrate.NewLocalDir("testdata")
+	d, err := migrate.NewLocalDir("testdata")
+	require.NoError(t, err)
+
+	_, err = migrate.GlobStateReader(d, drv, "*.up.sql").ReadState(ctx)
 	require.NoError(t, err)
+	require.Equal(t, []string{"CREATE TABLE t(c int);"}, drv.executed)
+	require.Equal(t, 1, drv.lockCounter)
+	require.Equal(t, 1, drv.unlockCounter)
 
-	_, err = migrate.GlobStateReader(localFS, drv, "*.up.sql").ReadState(ctx)
+	_, err = migrate.GlobStateReader(d, drv, "*.down.sql").ReadState(ctx)
 	require.NoError(t, err)
-	require.Equal(t, drv.executed, []string{"CREATE TABLE t(c int);"})
+	require.Equal(t, []string{"CREATE TABLE t(c int);", "DROP TABLE IF EXISTS t;"}, drv.executed)
+	require.Equal(t, 2, drv.lockCounter)
+	require.Equal(t, 2, drv.unlockCounter)
+
+	drv.locked = true
+	_, err = migrate.GlobStateReader(d, drv, "").ReadState(ctx)
+	require.EqualError(t, err, "sql/migrate: acquiring database lock: lockErr")
+	require.Equal(t, 2, drv.lockCounter)
+	require.Equal(t, 2, drv.unlockCounter)
+	drv.locked = false
+
+	drv.realm = schema.Realm{Schemas: []*schema.Schema{{Name: "schema"}}}
+	_, err = migrate.GlobStateReader(d, drv, "*.up.sql").ReadState(ctx)
+	require.EqualError(t, err, "sql/migrate: connected database is not clean")
+	require.Equal(t, 2, drv.lockCounter)
+	require.Equal(t, 2, drv.unlockCounter)
 
-	_, err = migrate.GlobStateReader(localFS, drv, "*.down.sql").ReadState(ctx)
+	edrv := &emptyMockDriver{drv}
+	_, err = migrate.GlobStateReader(d, edrv, "*.up.sql").ReadState(ctx)
 	require.NoError(t, err)
-	require.Equal(t, drv.executed, []string{"CREATE TABLE t(c int);", "DROP TABLE IF EXISTS t;"})
+	require.Equal(t, []schema.Change{
+		&schema.DropSchema{
+			S:     &schema.Schema{Name: "schema"},
+			Extra: []schema.Clause{&schema.IfExists{}},
+		},
+	}, edrv.applied)
+	require.Equal(t, 3, drv.lockCounter)
+	require.Equal(t, 3, drv.unlockCounter)
 }
 
 func TestLocalDir(t *testing.T) {
@@ -182,38 +214,54 @@ func TestLocalDir(t *testing.T) {
 	require.Equal(t, "content", string(c))
 }
 
-type mockHashFS struct{ hash []byte }
-
-func (h *mockHashFS) WriteHash(b []byte) error {
-	h.hash = b
-	return nil
-}
-
-func (h *mockHashFS) ReadHash() ([]byte, error) {
-	return h.hash, nil
-}
-
-type mockDriver struct {
-	migrate.Driver
-	plan     *migrate.Plan
-	changes  []schema.Change
-	executed []string
-}
+type (
+	mockDriver struct {
+		migrate.Driver
+		plan          *migrate.Plan
+		changes       []schema.Change
+		applied       []schema.Change
+		realm         schema.Realm
+		executed      []string
+		locked        bool
+		lockCounter   int
+		unlockCounter int
+	}
+	emptyMockDriver struct{ *mockDriver }
+)
 
 func (m *mockDriver) ExecContext(_ context.Context, query string, _ ...interface{}) (sql.Result, error) {
 	m.executed = append(m.executed, query)
 	return nil, nil
 }
 
-func (m mockDriver) InspectRealm(context.Context, *schema.InspectRealmOption) (*schema.Realm, error) {
-	return &schema.Realm{}, nil
+func (m *mockDriver) InspectRealm(context.Context, *schema.InspectRealmOption) (*schema.Realm, error) {
+	return &m.realm, nil
 }
-func (m mockDriver) RealmDiff(_, _ *schema.Realm) ([]schema.Change, error) {
+func (m *mockDriver) RealmDiff(_, _ *schema.Realm) ([]schema.Change, error) {
 	return m.changes, nil
 }
-func (m mockDriver) PlanChanges(context.Context, string, []schema.Change) (*migrate.Plan, error) {
+func (m *mockDriver) PlanChanges(context.Context, string, []schema.Change) (*migrate.Plan, error) {
 	return m.plan, nil
 }
+func (m *mockDriver) ApplyChanges(_ context.Context, changes []schema.Change) error {
+	m.applied = changes
+	return nil
+}
+func (m *mockDriver) Lock(context.Context, string, time.Duration) (schema.UnlockFunc, error) {
+	if m.locked {
+		return nil, errors.New("lockErr")
+	}
+	m.lockCounter++
+	m.locked = true
+	return func() error {
+		m.unlockCounter++
+		m.locked = false
+		return nil
+	}, nil
+}
+func (m *emptyMockDriver) IsClean(context.Context) (bool, error) {
+	return true, nil
+}
 
 func countFiles(t *testing.T, d migrate.Dir) int {
 	files, err := fs.ReadDir(d, "")
diff --git a/sql/migrate/testdata/atlas.sum b/sql/migrate/testdata/atlas.sum
index c513cf1d..c24d9663 100644
--- a/sql/migrate/testdata/atlas.sum
+++ b/sql/migrate/testdata/atlas.sum
@@ -1,4 +1,4 @@
 sum h1:DPwupu/FZJZYDnpuzD/l7aOg0jW2vhZxnajKPqhtKXI=
 1_initial.down.sql h1:EZiUHYOg5uOulzAVoV3ASKISA07Yh1ssDovA9hFGm1g=
 1_initial.up.sql h1:jDnx3o9bPJxVl9R0zpqnB+ETv1uyHR8JmvDw2G+kJEk=
-sub/1.a_sub-up.sql h1:XTo+9ZV78c5t2cuZCqcvTSwkQEf0B2L9Didd8p8N2gc=
+sub/1.a_sub.up.sql h1:XTo+9ZV78c5t2cuZCqcvTSwkQEf0B2L9Didd8p8N2gc=
diff --git a/sql/postgres/migrate_test.go b/sql/postgres/migrate_test.go
index 1288ceab..cce349eb 100644
--- a/sql/postgres/migrate_test.go
+++ b/sql/postgres/migrate_test.go
@@ -26,7 +26,7 @@ func TestPlanChanges(t *testing.T) {
 			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{}}},
+				&schema.DropSchema{S: schema.New("test"), Extra: []schema.Clause{}},
 			},
 			plan: &migrate.Plan{
 				Reversible:    false,
@@ -34,10 +34,10 @@ func TestPlanChanges(t *testing.T) {
 				Changes: []*migrate.Change{
 					{
 						Cmd:     `CREATE SCHEMA IF NOT EXISTS "test"`,
-						Reverse: `DROP SCHEMA "test"`,
+						Reverse: `DROP SCHEMA "test" CASCADE`,
 					},
 					{
-						Cmd: `DROP SCHEMA IF EXISTS "test"`,
+						Cmd: `DROP SCHEMA IF EXISTS "test" CASCADE`,
 					},
 					{
 						Cmd: `DROP SCHEMA "test" CASCADE`,
@@ -141,7 +141,7 @@ func TestPlanChanges(t *testing.T) {
 			plan: &migrate.Plan{
 				Reversible:    true,
 				Transactional: true,
-				Changes:       []*migrate.Change{{Cmd: `CREATE SCHEMA "test"`, Reverse: `DROP SCHEMA "test"`}}},
+				Changes:       []*migrate.Change{{Cmd: `CREATE SCHEMA "test"`, Reverse: `DROP SCHEMA "test" CASCADE`}}},
 		},
 		{
 			changes: []schema.Change{
@@ -149,7 +149,7 @@ func TestPlanChanges(t *testing.T) {
 			},
 			plan: &migrate.Plan{
 				Transactional: true,
-				Changes:       []*migrate.Change{{Cmd: `DROP SCHEMA "atlas"`}},
+				Changes:       []*migrate.Change{{Cmd: `DROP SCHEMA "atlas" CASCADE`}},
 			},
 		},
 		{