Add optional progress information and other enhancements for restore-dump command

internal/cmd/database/restore.go

@@ -18,11 +18,18 @@ import (
 )
 
 type restoreFlags struct {
-	localAddr  string
-	remoteAddr string
-	dir        string
-	overwrite  bool
-	threads    int
+	localAddr                 string
+	remoteAddr                string
+	dir                       string
+	overwrite                 bool
+	schemaOnly                bool
+	dataOnly                  bool
+	showDetails               bool
+	startingTable             string
+	endingTable               string
+	allowDifferentDestination bool
+	maxQuerySize              int
+	threads                   int
 }
 
 // RestoreCmd encapsulates the commands for restore a database
@@ -42,7 +49,15 @@ func RestoreCmd(ch *cmdutil.Helper) *cobra.Command {
 	cmd.PersistentFlags().StringVar(&f.dir, "dir", "",
 		"Directory containing the files to be used for the restore (required)")
 	cmd.PersistentFlags().BoolVar(&f.overwrite, "overwrite-tables", false, "If true, will attempt to DROP TABLE before restoring.")
-
+	cmd.PersistentFlags().BoolVar(&f.schemaOnly, "schema-only", false, "If true, will only restore the schema files during the restore process.")
+	cmd.PersistentFlags().BoolVar(&f.dataOnly, "data-only", false, "If true, will only restore the data files during the restore process.")
+	cmd.PersistentFlags().BoolVar(&f.showDetails, "show-details", false, "If true, will add extra output during the restore process.")
+	cmd.PersistentFlags().StringVar(&f.startingTable, "starting-table", "",
+		"Table to start from for the restore (useful for restarting from a certain point)")
+	cmd.PersistentFlags().StringVar(&f.endingTable, "ending-table", "",
+		"Table to end at for the restore (useful for stopping restore at a certain point)")
+	cmd.PersistentFlags().BoolVar(&f.allowDifferentDestination, "allow-different-destination", false, "If true, will allow you to restore the files to a database with a different name without needing to rename the existing dump's files.")
+	cmd.PersistentFlags().IntVar(&f.maxQuerySize, "max-query-size", 16777216, "The maximum size allowed for each individual query processed by the command. Currently limited to 16777216 bytes (16 MiB).")
 	cmd.PersistentFlags().IntVar(&f.threads, "threads", 1, "Number of concurrent threads to use to restore the database.")
 	return cmd
 }
@@ -58,6 +73,11 @@ func restore(ch *cmdutil.Helper, cmd *cobra.Command, flags *restoreFlags, args [
 		return errors.New("--dir flag is missing, it's needed to restore the database")
 	}
 
+	if flags.endingTable != "" && flags.startingTable != "" && (flags.endingTable < flags.startingTable) {
+		return fmt.Errorf("provided ending table %s must come alphabetically after your provided starting table %s for the restore to continue",
+			printer.BoldBlue(flags.endingTable), printer.BoldBlue(flags.startingTable))
+	}
+
 	client, err := ch.Client()
 	if err != nil {
 		return err
@@ -148,19 +168,34 @@ func restore(ch *cmdutil.Helper, cmd *cobra.Command, flags *restoreFlags, args [
 	cfg.Password = "nobody"
 	cfg.Address = addr.String()
 	cfg.Debug = ch.Debug()
+	cfg.Printer = ch.Printer
 	cfg.IntervalMs = 10 * 1000
 	cfg.Outdir = flags.dir
 	cfg.OverwriteTables = flags.overwrite
+	cfg.SchemaOnly = flags.schemaOnly
+	cfg.DataOnly = flags.dataOnly
+	cfg.ShowDetails = flags.showDetails
+	cfg.AllowDifferentDestination = flags.allowDifferentDestination
+	cfg.Database = database // Needs to be passed in to allow for allowDifferentDestination flag to work
+	cfg.StartingTable = flags.startingTable
+	cfg.EndingTable = flags.endingTable
+	cfg.MaxQuerySize = flags.maxQuerySize
 
 	loader, err := dumper.NewLoader(cfg)
 	if err != nil {
 		return err
 	}
 
+	end := func() {}
+
 	ch.Printer.Printf("Starting to restore database %s from folder %s\n",
 		printer.BoldBlue(database), printer.BoldBlue(flags.dir))
 
-	end := ch.Printer.PrintProgress("Restoring database ...")
+	if flags.showDetails {
+		ch.Printer.Println("Restoring database ...")
+	} else {
+		end = ch.Printer.PrintProgress("Restoring database ...\n")
+	}
 	defer end()
 
 	start := time.Now()

internal/dumper/dumper.go

@@ -13,6 +13,7 @@ import (
 	"time"
 
 	"github.com/planetscale/cli/internal/cmdutil"
+	"github.com/planetscale/cli/internal/printer"
 	querypb "github.com/xelabs/go-mysqlstack/sqlparser/depends/query"
 
 	"go.uber.org/zap"
@@ -22,36 +23,44 @@ const VITESS_GHOST_TABLE_REGEX = "_vt_EVAC_.*|_vt_DROP_.*|_vt_PURGE_.*|_vt_HOLD_
 
 // Config describes the settings to dump from a database.
 type Config struct {
-	User                 string
-	Password             string
-	Address              string
-	ToUser               string
-	ToPassword           string
-	ToAddress            string
-	ToDatabase           string
-	ToEngine             string
-	Database             string
-	DatabaseRegexp       string
-	DatabaseInvertRegexp bool
-	Shard                string
-	Table                string
-	Outdir               string
-	SessionVars          []string
-	Threads              int
-	ChunksizeInMB        int
-	StmtSize             int
-	Allbytes             uint64
-	Allrows              uint64
-	OverwriteTables      bool
-	UseReplica           bool
-	UseRdonly            bool
-	Wheres               map[string]string
-	Selects              map[string]map[string]string
-	Filters              map[string]map[string]string
+	User                      string
+	Password                  string
+	Address                   string
+	ToUser                    string
+	ToPassword                string
+	ToAddress                 string
+	ToDatabase                string
+	ToEngine                  string
+	Database                  string
+	DatabaseRegexp            string
+	DatabaseInvertRegexp      bool
+	Shard                     string
+	Table                     string
+	Outdir                    string
+	SessionVars               []string
+	Threads                   int
+	ChunksizeInMB             int
+	StmtSize                  int
+	Allbytes                  uint64
+	Allrows                   uint64
+	OverwriteTables           bool
+	SchemaOnly                bool
+	DataOnly                  bool
+	ShowDetails               bool
+	StartingTable             string
+	EndingTable               string
+	AllowDifferentDestination bool
+	MaxQuerySize              int
+	UseReplica                bool
+	UseRdonly                 bool
+	Wheres                    map[string]string
+	Selects                   map[string]map[string]string
+	Filters                   map[string]map[string]string
 
 	// Interval in millisecond.
 	IntervalMs int
 	Debug      bool
+	Printer    *printer.Printer
 }
 
 func NewDefaultConfig() *Config {