cmd/atlas: add support for 'migrate checkpoint'

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
 		}

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")
+}

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 . "  " }}'
 ```

sql/migrate/dir.go

@@ -172,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())
 }
 
@@ -426,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())
 }