Restrict generations command to single db

cmd/litestream/generations.go

@@ -12,53 +12,46 @@ import (
 	"time"
 )
 
-type GenerationsCommand struct {
+type GenerationsCommand struct{}
-	ConfigPath string
-	Config     Config
-
-	DBPath string
-}
 
 func NewGenerationsCommand() *GenerationsCommand {
 	return &GenerationsCommand{}
 }
 
 func (c *GenerationsCommand) Run(ctx context.Context, args []string) (err error) {
+	var configPath string
 	fs := flag.NewFlagSet("litestream-generations", flag.ContinueOnError)
-	registerConfigFlag(fs, &c.ConfigPath)
+	registerConfigFlag(fs, &configPath)
+	replicaName := fs.String("replica", "", "replica name")
 	fs.Usage = c.Usage
 	if err := fs.Parse(args); err != nil {
 		return err
+	} else if fs.NArg() == 0 || fs.Arg(0) == "" {
+		return fmt.Errorf("database path required")
 	} else if fs.NArg() > 1 {
 		return fmt.Errorf("too many arguments")
 	}
 
 	// Load configuration.
-	if c.ConfigPath == "" {
+	if configPath == "" {
 		return errors.New("-config required")
 	}
-	config, err := ReadConfigFile(c.ConfigPath)
+	config, err := ReadConfigFile(configPath)
 	if err != nil {
 		return err
 	}
 
-	// Determine absolute path for database, if specified.
+	// Determine absolute path for database.
-	if c.DBPath = fs.Arg(0); c.DBPath != "" {
+	dbPath, err := filepath.Abs(fs.Arg(0))
-		if c.DBPath, err = filepath.Abs(c.DBPath); err != nil {
+	if err != nil {
 		return err
 	}
-	}
-
-	// List each generation.
-	w := tabwriter.NewWriter(os.Stdout, 0, 8, 1, '\t', 0)
-	fmt.Fprintln(w, "db\tname\tgeneration\tlag\tstart\tend")
-	for _, dbConfig := range config.DBs {
-		// Filter database, if specified in the arguments.
-		if c.DBPath != "" && dbConfig.Path != c.DBPath {
-			continue
-		}
 
 	// Instantiate DB from from configuration.
+	dbConfig := config.DBConfig(dbPath)
+	if dbConfig == nil {
+		return fmt.Errorf("database not found in config: %s", dbPath)
+	}
 	db, err := newDBFromConfig(dbConfig)
 	if err != nil {
 		return err
@@ -70,8 +63,14 @@ func (c *GenerationsCommand) Run(ctx context.Context, args []string) (err error)
 		return err
 	}
 
-		// Iterate over each replica in the database.
+	// List each generation.
+	w := tabwriter.NewWriter(os.Stdout, 0, 8, 1, '\t', 0)
+	fmt.Fprintln(w, "name\tgeneration\tlag\tstart\tend")
 	for _, r := range db.Replicas {
+		if *replicaName != "" && r.Name() != *replicaName {
+			continue
+		}
+
 		generations, err := r.Generations(ctx)
 		if err != nil {
 			log.Printf("%s: cannot list generations: %s", r.Name(), err)
@@ -86,8 +85,7 @@ func (c *GenerationsCommand) Run(ctx context.Context, args []string) (err error)
 				continue
 			}
 
-				fmt.Fprintf(w, "%s\t%s\t%s\t%s\t%s\t%s\n",
+			fmt.Fprintf(w, "%s\t%s\t%s\t%s\t%s\n",
-					db.Path(),
 				r.Name(),
 				generation,
 				truncateDuration(stats.UpdatedAt.Sub(updatedAt)).String(),
@@ -96,7 +94,6 @@ func (c *GenerationsCommand) Run(ctx context.Context, args []string) (err error)
 			)
 		}
 	}
-	}
 	w.Flush()
 
 	return nil
@@ -104,7 +101,7 @@ func (c *GenerationsCommand) Run(ctx context.Context, args []string) (err error)
 
 func (c *GenerationsCommand) Usage() {
 	fmt.Printf(`
-The generations command lists all generations across all replicas along with
+The generations command lists all generations for a database. It also lists
 stats about their lag behind the primary database and the time range they cover.
 
 Usage:
@@ -116,6 +113,9 @@ Arguments:
 	-config PATH
 	    Specifies the configuration file. Defaults to %s
 
+	-replica NAME
+	    Optional, filters by replica.
+
 `[1:],
 		DefaultConfigPath,
 	)