Repository: ariga/atlas
Pull request #1982 (2023-08-13 13:10:17)

cmd/atlas: add support for 'migrate checkpoint'


Patch:
diff --git a/cmd/atlas/internal/cmdapi/migrate.go b/cmd/atlas/internal/cmdapi/migrate.go
index 48d4461b..11f6f08c 100644
--- a/cmd/atlas/internal/cmdapi/migrate.go
+++ b/cmd/atlas/internal/cmdapi/migrate.go
@@ -44,6 +44,7 @@ func init() {
 	migrateCmd := migrateCmd()
 	migrateCmd.AddCommand(
 		migrateApplyCmd(),
+		migrateCheckpointCmd(),
 		migrateDiffCmd(),
 		migrateHashCmd(),
 		migrateImportCmd(),
@@ -524,11 +525,12 @@ func migrateDiffCmd() *cobra.Command {
 		cmd   = &cobra.Command{
 			Use:   "diff [flags] [name]",
 			Short: "Compute the diff between the migration directory and a desired state and create a new migration file.",
-			Long: `'atlas migrate diff' uses the dev-database to re-run all migration files in the migration directory, compares
-it to a given desired state and create a new migration file containing SQL statements to migrate the migration
-directory state to the desired schema. The desired state can be another connected database or an HCL file.`,
-			Example: `  atlas migrate diff --dev-url mysql://user:pass@localhost:3306/dev --to file://schema.hcl
-  atlas migrate diff --dev-url mysql://user:pass@localhost:3306/dev --to file://atlas.hcl add_users_table
+			Long: `The 'atlas migrate diff' command uses the dev-database to calculate the current state of the migration directory
+by executing its files. It then compares its state to the desired state and create a new migration file containing
+SQL statements for moving from the current to the desired state. The desired state can be another another database,
+an HCL, SQL, or ORM schema. See: https://atlasgo.io/versioned/diff`,
+			Example: `  atlas migrate diff --dev-url docker://mysql/8/dev --to file://schema.hcl
+  atlas migrate diff --dev-url "docker://postgres/15/dev?search_path=public" --to file://atlas.hcl add_users_table
   atlas migrate diff --dev-url mysql://user:pass@localhost:3306/dev --to mysql://user:pass@localhost:3306/dbname
   atlas migrate diff --env dev --format '{{ sql . "  " }}'`,
 			Args: cobra.MaximumNArgs(1),
@@ -658,6 +660,135 @@ func migrateDiffRun(cmd *cobra.Command, args []string, flags migrateDiffFlags, e
 	}
 }
 
+type migrateCheckpointFlags struct {
+	edit              bool
+	dirURL, dirFormat string
+	devURL            string
+	schemas           []string
+	lockTimeout       time.Duration
+	format            string
+	qualifier         string // optional table qualifier
+}
+
+// migrateCheckpointCmd represents the 'atlas migrate checkpoint' subcommand.
+func migrateCheckpointCmd() *cobra.Command {
+	var (
+		flags migrateCheckpointFlags
+		cmd   = &cobra.Command{
+			Use:   "checkpoint [flags] [tag]",
+			Short: "Generate a checkpoint file representing the state of the migration directory.",
+			Long: `The 'atlas migrate checkpoint' command uses the dev-database to calculate the current state of the migration directory
+by executing its files. It then creates a checkpoint file that represents this state, enabling new environments to bypass
+previous files and immediately skip to this checkpoint when executing the 'atlas migrate apply' command.`,
+			Example: `  atlas migrate checkpoint --dev-url docker://mysql/8/dev
+  atlas migrate checkpoint --dev-url "docker://postgres/15/dev?search_path=public"
+  atlas migrate checkpoint --dev-url "sqlite://dev?mode=memory"
+  atlas migrate checkpoint --env dev --format '{{ sql . "  " }}'`,
+			Args: cobra.MaximumNArgs(1),
+			PreRunE: func(cmd *cobra.Command, args []string) error {
+				if err := migrateFlagsFromConfig(cmd); err != nil {
+					return err
+				}
+				if err := dirFormatBC(flags.dirFormat, &flags.dirURL); err != nil {
+					return err
+				}
+				return checkDir(cmd, flags.dirURL, false)
+			},
+			RunE: func(cmd *cobra.Command, args []string) error {
+				env, err := selectEnv(cmd)
+				if err != nil {
+					return err
+				}
+				return migrateCheckpointRun(cmd, args, flags, env)
+			},
+		}
+	)
+	cmd.Flags().SortFlags = false
+	addFlagDevURL(cmd.Flags(), &flags.devURL)
+	addFlagDirURL(cmd.Flags(), &flags.dirURL)
+	addFlagDirFormat(cmd.Flags(), &flags.dirFormat)
+	addFlagSchemas(cmd.Flags(), &flags.schemas)
+	addFlagLockTimeout(cmd.Flags(), &flags.lockTimeout)
+	addFlagFormat(cmd.Flags(), &flags.format)
+	cmd.Flags().StringVar(&flags.qualifier, flagQualifier, "", "qualify tables with custom qualifier when working on a single schema")
+	cmd.Flags().BoolVarP(&flags.edit, flagEdit, "", false, "edit the generated migration file(s)")
+	cobra.CheckErr(cmd.MarkFlagRequired(flagDevURL))
+	return cmd
+}
+
+func migrateCheckpointRun(cmd *cobra.Command, args []string, flags migrateCheckpointFlags, env *Env) error {
+	ctx := cmd.Context()
+	dev, err := sqlclient.Open(ctx, flags.devURL)
+	if err != nil {
+		return err
+	}
+	defer dev.Close()
+	if l, ok := dev.Driver.(schema.Locker); ok {
+		unlock, err := l.Lock(ctx, "atlas_migrate_diff", flags.lockTimeout)
+		if err != nil {
+			return fmt.Errorf("acquiring database lock: %w", err)
+		}
+		defer func() { cobra.CheckErr(unlock()) }()
+	}
+	u, err := url.Parse(flags.dirURL)
+	if err != nil {
+		return err
+	}
+	dir, err := cmdmigrate.DirURL(u, false)
+	if err != nil {
+		return err
+	}
+	dir, ok := dir.(migrate.CheckpointDir)
+	if !ok {
+		return fmt.Errorf("migration directory %s does not support checkpoint files", u)
+	}
+	files, err := dir.Files()
+	if err != nil {
+		return err
+	}
+	if len(files) == 0 {
+		cmd.Println("The migration directory is empty, no checkpoint to be made")
+		return nil
+	}
+	if flags.edit {
+		dir = &editDir{dir}
+	}
+	var (
+		tag, indent string
+		name        = "checkpoint"
+	)
+	if len(args) > 0 {
+		tag = args[0]
+	}
+	f, err := cmdmigrate.Formatter(u)
+	if err != nil {
+		return err
+	}
+	if f, indent, err = mayIndent(u, f, flags.format); err != nil {
+		return err
+	}
+	opts := []migrate.PlannerOption{
+		migrate.PlanFormat(f),
+		migrate.PlanWithIndent(indent),
+		migrate.PlanWithDiffOptions(env.DiffOptions()...),
+	}
+	if dev.URL.Schema != "" {
+		// Disable tables qualifier in schema-mode.
+		opts = append(opts, migrate.PlanWithSchemaQualifier(flags.qualifier))
+	}
+	pl := migrate.NewPlanner(dev.Driver, dir, opts...)
+	plan, err := func() (*migrate.Plan, error) {
+		if s := dev.URL.Schema; s != "" {
+			return pl.CheckpointSchema(ctx, name)
+		}
+		return pl.Checkpoint(ctx, name)
+	}()
+	if err != nil {
+		return err
+	}
+	return pl.WriteCheckpoint(plan, tag)
+}
+
 func mayIndent(dir *url.URL, f migrate.Formatter, format string) (migrate.Formatter, string, error) {
 	if format == "" {
 		return f, "", nil
@@ -1700,7 +1831,7 @@ func setMigrateEnvFlags(cmd *cobra.Command, env *Env) error {
 		if err := maySetFlag(cmd, flagLockTimeout, env.Migration.LockTimeout); err != nil {
 			return err
 		}
-	case "diff":
+	case "diff", "checkpoint":
 		if err := maySetFlag(cmd, flagLockTimeout, env.Migration.LockTimeout); err != nil {
 			return err
 		}
diff --git a/doc/md/reference.md b/doc/md/reference.md
index baabfbc9..277c89da 100644
--- a/doc/md/reference.md
+++ b/doc/md/reference.md
@@ -99,6 +99,42 @@ If run with the "--dry-run" flag, atlas will not execute any SQL.
 ```
 
 
+### atlas migrate checkpoint
+
+Generate a checkpoint file representing the state of the migration directory.
+
+#### Usage
+```
+atlas migrate checkpoint [flags] [tag]
+```
+
+#### Details
+The 'atlas migrate checkpoint' command uses the dev-database to calculate the current state of the migration directory
+by executing its files. It then creates a checkpoint file that represents this state, enabling new environments to bypass
+previous files and immediately skip to this checkpoint when executing the 'atlas migrate apply' command.
+
+#### Example
+
+```
+  atlas migrate checkpoint --dev-url docker://mysql/8/dev
+  atlas migrate checkpoint --dev-url "docker://postgres/15/dev?search_path=public"
+  atlas migrate checkpoint --dev-url "sqlite://dev?mode=memory"
+  atlas migrate checkpoint --env dev --format '{{ sql . "  " }}'
+```
+#### Flags
+```
+      --dev-url string          [driver://username:password@address/dbname?param=value] select a dev database using the URL format
+      --dir string              select migration directory using URL format (default "file://migrations")
+      --dir-format string       select migration file format (default "atlas")
+  -s, --schema strings          set schema names
+      --lock-timeout duration   set how long to wait for the database lock (default 10s)
+      --format string           Go template to use to format the output
+      --qualifier string        qualify tables with custom qualifier when working on a single schema
+      --edit                    edit the generated migration file(s)
+
+```
+
+
 ### atlas migrate diff
 
 Compute the diff between the migration directory and a desired state and create a new migration file.
@@ -109,15 +145,16 @@ atlas migrate diff [flags] [name]
 ```
 
 #### Details
-'atlas migrate diff' uses the dev-database to re-run all migration files in the migration directory, compares
-it to a given desired state and create a new migration file containing SQL statements to migrate the migration
-directory state to the desired schema. The desired state can be another connected database or an HCL file.
+The 'atlas migrate diff' command uses the dev-database to calculate the current state of the migration directory
+by executing its files. It then compares its state to the desired state and create a new migration file containing
+SQL statements for moving from the current to the desired state. The desired state can be another another database,
+an HCL, SQL, or ORM schema. See: https://atlasgo.io/versioned/diff
 
 #### Example
 
 ```
-  atlas migrate diff --dev-url mysql://user:pass@localhost:3306/dev --to file://schema.hcl
-  atlas migrate diff --dev-url mysql://user:pass@localhost:3306/dev --to file://atlas.hcl add_users_table
+  atlas migrate diff --dev-url docker://mysql/8/dev --to file://schema.hcl
+  atlas migrate diff --dev-url "docker://postgres/15/dev?search_path=public" --to file://atlas.hcl add_users_table
   atlas migrate diff --dev-url mysql://user:pass@localhost:3306/dev --to mysql://user:pass@localhost:3306/dbname
   atlas migrate diff --env dev --format '{{ sql . "  " }}'
 ```
diff --git a/sql/migrate/dir.go b/sql/migrate/dir.go
index 6192b892..96789e73 100644
--- a/sql/migrate/dir.go
+++ b/sql/migrate/dir.go
@@ -65,6 +65,7 @@ type (
 	// CheckpointDir wraps the functionality used to interact
 	// with a migration directory that support checkpoints.
 	CheckpointDir interface {
+		Dir
 		// WriteCheckpoint writes the given checkpoint file to the migration directory.
 		WriteCheckpoint(name, tag string, content []byte) error
 
@@ -81,6 +82,7 @@ type (
 	// CheckpointFile wraps the functionality used to interact with files
 	// returned from a CheckpointDir.
 	CheckpointFile interface {
+		File
 		// IsCheckpoint returns true if the file is a checkpoint file.
 		IsCheckpoint() bool
 
@@ -105,10 +107,7 @@ type LocalDir struct {
 	path string
 }
 
-var _ interface {
-	Dir
-	CheckpointDir
-} = (*LocalDir)(nil)
+var _ CheckpointDir = (*LocalDir)(nil)
 
 // NewLocalDir returns a new the Dir used by a Planner to work on the given local path.
 func NewLocalDir(path string) (*LocalDir, error) {
@@ -173,8 +172,14 @@ func (d *LocalDir) Checksum() (HashFile, error) {
 
 // WriteCheckpoint is like WriteFile, but marks the file as a checkpoint file.
 func (d *LocalDir) WriteCheckpoint(name, tag string, b []byte) error {
-	f := NewLocalFile(name, b)
-	f.AddDirective(directiveCheckpoint, tag)
+	var (
+		args []string
+		f    = NewLocalFile(name, b)
+	)
+	if tag != "" {
+		args = append(args, tag)
+	}
+	f.AddDirective(directiveCheckpoint, args...)
 	return d.WriteFile(name, f.Bytes())
 }
 
@@ -194,10 +199,7 @@ type LocalFile struct {
 	b []byte
 }
 
-var _ interface {
-	File
-	CheckpointFile
-} = (*LocalFile)(nil)
+var _ CheckpointFile = (*LocalFile)(nil)
 
 // NewLocalFile returns a new local file.
 func NewLocalFile(name string, data []byte) *LocalFile {
@@ -357,10 +359,7 @@ var (
 		sync.Mutex
 		opened map[string]*openedMem
 	}
-	_ interface {
-		Dir
-		CheckpointDir
-	} = (*MemDir)(nil)
+	_ CheckpointDir = (*MemDir)(nil)
 )
 
 // OpenMemDir opens an in-memory directory and registers it in the process namespace
@@ -433,8 +432,14 @@ func (d *MemDir) WriteFile(name string, data []byte) error {
 
 // WriteCheckpoint is like WriteFile, but marks the file as a checkpoint file.
 func (d *MemDir) WriteCheckpoint(name, tag string, b []byte) error {
-	f := NewLocalFile(name, b)
-	f.AddDirective(directiveCheckpoint, tag)
+	var (
+		args []string
+		f    = NewLocalFile(name, b)
+	)
+	if tag != "" {
+		args = append(args, tag)
+	}
+	f.AddDirective(directiveCheckpoint, args...)
 	return d.WriteFile(name, f.Bytes())
 }
 
diff --git a/sql/migrate/migrate.go b/sql/migrate/migrate.go
index 1c577766..bdd07025 100644
--- a/sql/migrate/migrate.go
+++ b/sql/migrate/migrate.go
@@ -185,7 +185,7 @@ func SchemaConn(drv Driver, name string, opts *schema.InspectOptions) StateReade
 }
 
 type (
-	// Planner can plan the steps to take to migrate from one state to another. It uses the enclosed Dir to write
+	// Planner can plan the steps to take to migrate from one state to another. It uses the enclosed Dir to
 	// those changes to versioned migration files.
 	Planner struct {
 		drv      Driver              // driver to use
@@ -382,18 +382,7 @@ func (p *Planner) PlanSchema(ctx context.Context, name string, to StateReader) (
 }
 
 func (p *Planner) plan(ctx context.Context, name string, to StateReader, realmScope bool) (*Plan, error) {
-	from, err := NewExecutor(p.drv, p.dir, NopRevisionReadWriter{})
-	if err != nil {
-		return nil, err
-	}
-	current, err := from.Replay(ctx, func() StateReader {
-		if realmScope {
-			return RealmConn(p.drv, nil)
-		}
-		// In case the scope is the schema connection,
-		// inspect it and return its connected realm.
-		return SchemaConn(p.drv, "", nil)
-	}())
+	current, err := p.current(ctx, realmScope)
 	if err != nil {
 		return nil, err
 	}
@@ -434,6 +423,63 @@ func (p *Planner) plan(ctx context.Context, name string, to StateReader, realmSc
 	return p.drv.PlanChanges(ctx, name, changes, p.planOpts...)
 }
 
+// Checkpoint calculate the current state of the migration directory by executing its files,
+// and return a migration (checkpoint) Plan that represents its states.
+func (p *Planner) Checkpoint(ctx context.Context, name string) (*Plan, error) {
+	return p.checkpoint(ctx, name, true)
+}
+
+// CheckpointSchema is like Checkpoint but limits its scope to the schema connection.
+// Note, the operation fails in case the connection was not set to a schema.
+func (p *Planner) CheckpointSchema(ctx context.Context, name string) (*Plan, error) {
+	return p.checkpoint(ctx, name, false)
+}
+
+func (p *Planner) checkpoint(ctx context.Context, name string, realmScope bool) (*Plan, error) {
+	current, err := p.current(ctx, realmScope)
+	if err != nil {
+		return nil, err
+	}
+	var changes []schema.Change
+	switch {
+	case realmScope:
+		changes, err = p.drv.RealmDiff(schema.NewRealm(), current, p.diffOpts...)
+	default:
+		switch n := len(current.Schemas); {
+		case n == 0:
+			return nil, errors.New("no schema was found in current state after replaying migration directory")
+		case n > 1:
+			return nil, fmt.Errorf("%d schemas were found in current state after replaying migration directory", len(current.Schemas))
+		default:
+			changes, err = p.drv.SchemaDiff(schema.New(current.Schemas[0].Name), current.Schemas[0], p.diffOpts...)
+		}
+	}
+	if err != nil {
+		return nil, err
+	}
+	// No changes mean an empty checkpoint.
+	if len(changes) == 0 {
+		return &Plan{Name: name}, nil
+	}
+	return p.drv.PlanChanges(ctx, name, changes, p.planOpts...)
+}
+
+// current returns the current realm state.
+func (p *Planner) current(ctx context.Context, realmScope bool) (*schema.Realm, error) {
+	from, err := NewExecutor(p.drv, p.dir, NopRevisionReadWriter{})
+	if err != nil {
+		return nil, err
+	}
+	return from.Replay(ctx, func() StateReader {
+		if realmScope {
+			return RealmConn(p.drv, nil)
+		}
+		// In case the scope is the schema connection,
+		// inspect it and return its connected realm.
+		return SchemaConn(p.drv, "", nil)
+	}())
+}
+
 // WritePlan writes the given Plan to the Dir based on the configured Formatter.
 func (p *Planner) WritePlan(plan *Plan) error {
 	// Format the plan into files.
@@ -447,15 +493,39 @@ func (p *Planner) WritePlan(plan *Plan) error {
 			return err
 		}
 	}
-	// If enabled, update the sum file.
-	if p.sum {
-		sum, err := p.dir.Checksum()
-		if err != nil {
-			return err
-		}
-		return WriteSumFile(p.dir, sum)
+	return p.writeSum()
+}
+
+// WriteCheckpoint writes the given Plan as a checkpoint file to the Dir based on the configured Formatter.
+func (p *Planner) WriteCheckpoint(plan *Plan, tag string) error {
+	ck, ok := p.dir.(CheckpointDir)
+	if !ok {
+		return fmt.Errorf("checkpoint is not supported by %T", p.dir)
 	}
-	return nil
+	// Format the plan into files.
+	files, err := p.fmt.Format(plan)
+	if err != nil {
+		return err
+	}
+	if len(files) != 1 {
+		return fmt.Errorf("expected one checkpoint file, got %d", len(files))
+	}
+	if err := ck.WriteCheckpoint(files[0].Name(), tag, files[0].Bytes()); err != nil {
+		return err
+	}
+	return p.writeSum()
+}
+
+// writeSum writes the sum file to the Dir, if enabled.
+func (p *Planner) writeSum() error {
+	if !p.sum {
+		return nil
+	}
+	sum, err := p.dir.Checksum()
+	if err != nil {
+		return err
+	}
+	return WriteSumFile(p.dir, sum)
 }
 
 var (


Test patch:
diff --git a/cmd/atlas/internal/cmdapi/migrate_test.go b/cmd/atlas/internal/cmdapi/migrate_test.go
index 1b9351b5..e2ac78f1 100644
--- a/cmd/atlas/internal/cmdapi/migrate_test.go
+++ b/cmd/atlas/internal/cmdapi/migrate_test.go
@@ -1496,6 +1496,107 @@ table "users" {
 	})
 }
 
+func TestMigrate_Checkpoint(t *testing.T) {
+	p := filepath.Join(t.TempDir(), "migrations")
+	_, err := runCmd(
+		migrateCheckpointCmd(),
+		"--dir", "file://"+p,
+		"--dev-url", openSQLite(t, ""),
+	)
+	require.Error(t, err)
+
+	require.NoError(t, os.MkdirAll(p, 0755))
+	s, err := runCmd(
+		migrateCheckpointCmd(),
+		"--dir", "file://"+p,
+		"--dev-url", openSQLite(t, ""),
+	)
+	require.NoError(t, err)
+	require.Equal(t, "The migration directory is empty, no checkpoint to be made\n", s)
+
+	// Add t1.
+	_, err = runCmd(
+		migrateDiffCmd(),
+		"add_t1",
+		"--dir", "file://"+p,
+		"--dev-url", openSQLite(t, ""),
+		"--to", openSQLite(t, "CREATE TABLE t1(c int);"),
+	)
+	require.NoError(t, err)
+	_, err = runCmd(
+		migrateCheckpointCmd(),
+		"--dir", "file://"+p,
+		"--dev-url", openSQLite(t, ""),
+	)
+	require.NoError(t, err)
+	dir, err := migrate.NewLocalDir(p)
+	require.NoError(t, err)
+	files, err := dir.Files()
+	require.NoError(t, err)
+	require.Equal(t, []string{
+		`-- Create "t1" table`,
+		"CREATE TABLE `t1` (`c` int NULL);",
+	}, lines(files[0]))
+	require.Equal(t, []string{
+		"-- atlas:checkpoint",
+		"",
+		`-- Create "t1" table`,
+		"CREATE TABLE `t1` (`c` int NULL);",
+	}, lines(files[1]))
+
+	// Rename the checkpoint file to avoid timestamp conflicts.
+	ck := files[1].Name()
+	require.NoError(t, os.Rename(filepath.Join(p, ck), filepath.Join(p, strings.Replace(ck, "checkpoint", "add_t1_checkpoint", -1))))
+	_, err = runCmd(migrateHashCmd(), "--dir", "file://"+p)
+	require.NoError(t, err)
+
+	// Add t2 and t3 in two different files (after checkpoint).
+	_, err = runCmd(
+		migrateDiffCmd(),
+		"add_t2",
+		"--dir", "file://"+p,
+		"--dev-url", openSQLite(t, ""),
+		"--to", openSQLite(t, "CREATE TABLE t1(c int); CREATE TABLE t2(c int);"),
+	)
+	require.NoError(t, err)
+	_, err = runCmd(
+		migrateDiffCmd(),
+		"add_t3",
+		"--dir", "file://"+p,
+		"--dev-url", openSQLite(t, ""),
+		"--to", openSQLite(t, "CREATE TABLE t1(c int); CREATE TABLE t2(c int); CREATE TABLE t3(c int);"),
+	)
+	require.NoError(t, err)
+	_, err = runCmd(
+		migrateCheckpointCmd(),
+		"v2",
+		"--dir", "file://"+p,
+		"--dev-url", openSQLite(t, ""),
+	)
+	files2, err := dir.Files()
+	require.NoError(t, err)
+	require.Equal(t, files2[0].Bytes(), files[0].Bytes())
+	require.Equal(t, files2[1].Bytes(), files[1].Bytes())
+	require.Equal(t, []string{
+		`-- Create "t2" table`,
+		"CREATE TABLE `t2` (`c` int NULL);",
+	}, lines(files2[2]))
+	require.Equal(t, []string{
+		`-- Create "t3" table`,
+		"CREATE TABLE `t3` (`c` int NULL);",
+	}, lines(files2[3]))
+	require.Equal(t, []string{
+		"-- atlas:checkpoint v2",
+		"",
+		`-- Create "t1" table`,
+		"CREATE TABLE `t1` (`c` int NULL);",
+		`-- Create "t2" table`,
+		"CREATE TABLE `t2` (`c` int NULL);",
+		`-- Create "t3" table`,
+		"CREATE TABLE `t3` (`c` int NULL);",
+	}, lines(files2[4]))
+}
+
 func TestMigrate_StatusJSON(t *testing.T) {
 	p := t.TempDir()
 	s, err := runCmd(
@@ -1979,3 +2080,7 @@ func sed(t *testing.T, r, p string) {
 	buf, err := exec.Command("sed", append(args, r, p)...).CombinedOutput()
 	require.NoError(t, err, string(buf))
 }
+
+func lines(f migrate.File) []string {
+	return strings.Split(strings.TrimSpace(string(f.Bytes())), "\n")
+}
diff --git a/sql/migrate/dir_test.go b/sql/migrate/dir_test.go
index 4942b452..3055c8a6 100644
--- a/sql/migrate/dir_test.go
+++ b/sql/migrate/dir_test.go
@@ -202,10 +202,7 @@ func TestLocalDir(t *testing.T) {
 func TestCheckpointDir(t *testing.T) {
 	local, err := migrate.NewLocalDir(t.TempDir())
 	require.NoError(t, err)
-	for _, d := range []interface {
-		migrate.Dir
-		migrate.CheckpointDir
-	}{&migrate.MemDir{}, local} {
+	for _, d := range []migrate.CheckpointDir{&migrate.MemDir{}, local} {
 		files, err := d.Files()
 		require.NoError(t, err)
 		require.Empty(t, files)
diff --git a/sql/migrate/migrate_test.go b/sql/migrate/migrate_test.go
index 95e54c64..744ee46a 100644
--- a/sql/migrate/migrate_test.go
+++ b/sql/migrate/migrate_test.go
@@ -76,6 +76,32 @@ func TestPlanner_WritePlan(t *testing.T) {
 	requireFileEqual(t, d, "add_t1_and_t2.down.sql", "DROP TABLE t1 IF EXISTS\nDROP TABLE t2\n")
 }
 
+func TestPlanner_WriteCheckpoint(t *testing.T) {
+	p := t.TempDir()
+	d, err := migrate.NewLocalDir(p)
+	require.NoError(t, err)
+	plan := &migrate.Plan{
+		Name: "checkpoint",
+		Changes: []*migrate.Change{
+			{Cmd: "CREATE TABLE t1(c int)", Reverse: "DROP TABLE t1 IF EXISTS"},
+			{Cmd: "CREATE TABLE t2(c int)", Reverse: "DROP TABLE t2"},
+		},
+	}
+
+	// DefaultFormatter
+	pl := migrate.NewPlanner(nil, d)
+	require.NotNil(t, pl)
+	require.NoError(t, pl.WriteCheckpoint(plan, "v1"))
+	files, err := d.Files()
+	require.NoError(t, err)
+	require.Len(t, files, 1)
+	require.Equal(t, `-- atlas:checkpoint v1
+
+CREATE TABLE t1(c int);
+CREATE TABLE t2(c int);
+`, string(files[0].Bytes()))
+}
+
 func TestPlanner_Plan(t *testing.T) {
 	var (
 		drv = &mockDriver{}
@@ -139,6 +165,56 @@ func TestPlanner_PlanSchema(t *testing.T) {
 	require.Nil(t, plan)
 }
 
+func TestPlanner_Checkpoint(t *testing.T) {
+	var (
+		drv = &mockDriver{}
+		ctx = context.Background()
+	)
+	d, err := migrate.NewLocalDir(t.TempDir())
+	require.NoError(t, err)
+
+	// Nothing to do.
+	pl := migrate.NewPlanner(drv, d)
+	plan, err := pl.Checkpoint(ctx, "empty")
+	require.NoError(t, err)
+	require.Equal(t, &migrate.Plan{Name: "empty"}, plan)
+
+	// There are changes.
+	drv.changes = []schema.Change{
+		&schema.AddTable{T: schema.NewTable("t1").AddColumns(schema.NewIntColumn("c", "int"))},
+		&schema.AddTable{T: schema.NewTable("t2").AddColumns(schema.NewIntColumn("c", "int"))},
+	}
+	drv.plan = &migrate.Plan{
+		Changes: []*migrate.Change{
+			{Cmd: "CREATE TABLE t1(c int);"},
+			{Cmd: "CREATE TABLE t2(c int);"},
+		},
+	}
+	plan, err = pl.Checkpoint(ctx, "checkpoint")
+	require.NoError(t, err)
+	require.Equal(t, drv.plan, plan)
+}
+
+func TestPlanner_CheckpointSchema(t *testing.T) {
+	var (
+		drv = &mockDriver{}
+		ctx = context.Background()
+	)
+	d, err := migrate.NewLocalDir(t.TempDir())
+	require.NoError(t, err)
+
+	// Schema is missing in dev connection.
+	pl := migrate.NewPlanner(drv, d)
+	plan, err := pl.CheckpointSchema(ctx, "empty")
+	require.EqualError(t, err, `not found`)
+	require.Nil(t, plan)
+
+	drv.realm = *schema.NewRealm(schema.New("test"))
+	pl = migrate.NewPlanner(drv, d)
+	plan, err = pl.CheckpointSchema(ctx, "empty")
+	require.Equal(t, &migrate.Plan{Name: "empty"}, plan)
+}
+
 func TestExecutor_Replay(t *testing.T) {
 	ctx := context.Background()
 	d, err := migrate.NewLocalDir(filepath.FromSlash("testdata/migrate"))