rkrtechnology/litestream
2
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Refactor Restore()

Browse Source 
This commit refactors out the complexity of downloading ordered WAL
files in parallel to a type called `WALDownloader`. This makes it
easier to test the restore separately from the download.
This commit is contained in:
Ben Johnson
2022-01-04 14:47:11 -07:00
parent 531e19ed6f
commit 3f0ec9fa9f
130 changed files with 2943 additions and 1254 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
cmd/litestream/databases.go
View File

@@ -10,12 +10,15 @@ import (
)
// DatabasesCommand is a command for listing managed databases.
type DatabasesCommand struct{}
type DatabasesCommand struct {
configPath string
noExpandEnv bool
}
// Run executes the command.
func (c *DatabasesCommand) Run(ctx context.Context, args []string) (err error) {
fs := flag.NewFlagSet("litestream-databases", flag.ContinueOnError)
configPath, noExpandEnv := registerConfigFlag(fs)
registerConfigFlag(fs, &c.configPath, &c.noExpandEnv)
fs.Usage = c.Usage
if err := fs.Parse(args); err != nil {
return err
@@ -24,10 +27,10 @@ func (c *DatabasesCommand) Run(ctx context.Context, args []string) (err error) {
}
// Load configuration.
if *configPath == "" {
*configPath = DefaultConfigPath()
if c.configPath == "" {
c.configPath = DefaultConfigPath()
}
config, err := ReadConfigFile(*configPath, !*noExpandEnv)
config, err := ReadConfigFile(c.configPath, !c.noExpandEnv)
if err != nil {
return err
}

17
cmd/litestream/generations.go
View File

@@ -13,12 +13,15 @@ import (
)
// GenerationsCommand represents a command to list all generations for a database.
type GenerationsCommand struct{}
type GenerationsCommand struct {
configPath string
noExpandEnv bool
}
// Run executes the command.
func (c *GenerationsCommand) Run(ctx context.Context, args []string) (err error) {
fs := flag.NewFlagSet("litestream-generations", flag.ContinueOnError)
configPath, noExpandEnv := registerConfigFlag(fs)
registerConfigFlag(fs, &c.configPath, &c.noExpandEnv)
replicaName := fs.String("replica", "", "replica name")
fs.Usage = c.Usage
if err := fs.Parse(args); err != nil {
@@ -33,19 +36,19 @@ func (c *GenerationsCommand) Run(ctx context.Context, args []string) (err error)
var r *litestream.Replica
dbUpdatedAt := time.Now()
if isURL(fs.Arg(0)) {
if *configPath != "" {
if c.configPath != "" {
return fmt.Errorf("cannot specify a replica URL and the -config flag")
}
if r, err = NewReplicaFromConfig(&ReplicaConfig{URL: fs.Arg(0)}, nil); err != nil {
return err
}
} else {
if *configPath == "" {
*configPath = DefaultConfigPath()
if c.configPath == "" {
c.configPath = DefaultConfigPath()
}
// Load configuration.
config, err := ReadConfigFile(*configPath, !*noExpandEnv)
config, err := ReadConfigFile(c.configPath, !c.noExpandEnv)
if err != nil {
return err
}
@@ -93,7 +96,7 @@ func (c *GenerationsCommand) Run(ctx context.Context, args []string) (err error)
// Iterate over each generation for the replica.
for _, generation := range generations {
createdAt, updatedAt, err := r.GenerationTimeBounds(ctx, generation)
createdAt, updatedAt, err := litestream.GenerationTimeBounds(ctx, r.Client, generation)
if err != nil {
log.Printf("%s: cannot determine generation time bounds: %s", r.Name(), err)
continue

17
cmd/litestream/main.go
View File

@@ -20,7 +20,6 @@ import (
"github.com/benbjohnson/litestream"
"github.com/benbjohnson/litestream/abs"
"github.com/benbjohnson/litestream/file"
"github.com/benbjohnson/litestream/gcs"
"github.com/benbjohnson/litestream/s3"
"github.com/benbjohnson/litestream/sftp"
@@ -126,7 +125,7 @@ func (m *Main) Run(ctx context.Context, args []string) (err error) {
return err
case "restore":
return (&RestoreCommand{}).Run(ctx, args)
return NewRestoreCommand().Run(ctx, args)
case "snapshots":
return (&SnapshotsCommand{}).Run(ctx, args)
case "version":
@@ -383,8 +382,8 @@ func NewReplicaFromConfig(c *ReplicaConfig, db *litestream.DB) (_ *litestream.Re
return r, nil
}
// newFileReplicaClientFromConfig returns a new instance of file.ReplicaClient built from config.
func newFileReplicaClientFromConfig(c *ReplicaConfig, r *litestream.Replica) (_ *file.ReplicaClient, err error) {
// newFileReplicaClientFromConfig returns a new instance of FileReplicaClient built from config.
func newFileReplicaClientFromConfig(c *ReplicaConfig, r *litestream.Replica) (_ *litestream.FileReplicaClient, err error) {
// Ensure URL & path are not both specified.
if c.URL != "" && c.Path != "" {
return nil, fmt.Errorf("cannot specify url & path for file replica")
@@ -409,9 +408,7 @@ func newFileReplicaClientFromConfig(c *ReplicaConfig, r *litestream.Replica) (_
}
// Instantiate replica and apply time fields, if set.
client := file.NewReplicaClient(path)
client.Replica = r
return client, nil
return litestream.NewFileReplicaClient(path), nil
}
// newS3ReplicaClientFromConfig returns a new instance of s3.ReplicaClient built from config.
@@ -669,9 +666,9 @@ func DefaultConfigPath() string {
return defaultConfigPath
}
func registerConfigFlag(fs *flag.FlagSet) (configPath *string, noExpandEnv *bool) {
return fs.String("config", "", "config path"),
fs.Bool("no-expand-env", false, "do not expand env vars in config")
func registerConfigFlag(fs *flag.FlagSet, configPath *string, noExpandEnv *bool) {
fs.StringVar(configPath, "config", "", "config path")
fs.BoolVar(noExpandEnv, "no-expand-env", false, "do not expand env vars in config")
}
// expand returns an absolute path for s.

3
cmd/litestream/main_test.go
View File

@@ -9,7 +9,6 @@ import (
"github.com/benbjohnson/litestream"
main "github.com/benbjohnson/litestream/cmd/litestream"
"github.com/benbjohnson/litestream/file"
"github.com/benbjohnson/litestream/gcs"
"github.com/benbjohnson/litestream/s3"
)
@@ -103,7 +102,7 @@ func TestNewFileReplicaFromConfig(t *testing.T) {
r, err := main.NewReplicaFromConfig(&main.ReplicaConfig{Path: "/foo"}, nil)
if err != nil {
t.Fatal(err)
} else if client, ok := r.Client.(*file.ReplicaClient); !ok {
} else if client, ok := r.Client.(*litestream.FileReplicaClient); !ok {
t.Fatal("unexpected replica type")
} else if got, want := client.Path(), "/foo"; got != want {
t.Fatalf("Path=%s, want %s", got, want)

16
cmd/litestream/replicate.go
View File

@@ -13,7 +13,6 @@ import (
"github.com/benbjohnson/litestream"
"github.com/benbjohnson/litestream/abs"
"github.com/benbjohnson/litestream/file"
"github.com/benbjohnson/litestream/gcs"
"github.com/benbjohnson/litestream/s3"
"github.com/benbjohnson/litestream/sftp"
@@ -23,6 +22,9 @@ import (
// ReplicateCommand represents a command that continuously replicates SQLite databases.
type ReplicateCommand struct {
configPath string
noExpandEnv bool
cmd *exec.Cmd // subcommand
execCh chan error // subcommand error channel
@@ -42,7 +44,7 @@ func NewReplicateCommand() *ReplicateCommand {
func (c *ReplicateCommand) ParseFlags(ctx context.Context, args []string) (err error) {
fs := flag.NewFlagSet("litestream-replicate", flag.ContinueOnError)
execFlag := fs.String("exec", "", "execute subcommand")
configPath, noExpandEnv := registerConfigFlag(fs)
registerConfigFlag(fs, &c.configPath, &c.noExpandEnv)
fs.Usage = c.Usage
if err := fs.Parse(args); err != nil {
return err
@@ -52,7 +54,7 @@ func (c *ReplicateCommand) ParseFlags(ctx context.Context, args []string) (err e
if fs.NArg() == 1 {
return fmt.Errorf("must specify at least one replica URL for %s", fs.Arg(0))
} else if fs.NArg() > 1 {
if *configPath != "" {
if c.configPath != "" {
return fmt.Errorf("cannot specify a replica URL and the -config flag")
}
@@ -66,10 +68,10 @@ func (c *ReplicateCommand) ParseFlags(ctx context.Context, args []string) (err e
}
c.Config.DBs = []*DBConfig{dbConfig}
} else {
if *configPath == "" {
*configPath = DefaultConfigPath()
if c.configPath == "" {
c.configPath = DefaultConfigPath()
}
if c.Config, err = ReadConfigFile(*configPath, !*noExpandEnv); err != nil {
if c.Config, err = ReadConfigFile(c.configPath, !c.noExpandEnv); err != nil {
return err
}
}
@@ -110,7 +112,7 @@ func (c *ReplicateCommand) Run(ctx context.Context) (err error) {
log.Printf("initialized db: %s", db.Path())
for _, r := range db.Replicas {
switch client := r.Client.(type) {
case *file.ReplicaClient:
case *litestream.FileReplicaClient:
log.Printf("replicating to: name=%q type=%q path=%q", r.Name(), client.Type(), client.Path())
case *s3.ReplicaClient:
log.Printf("replicating to: name=%q type=%q bucket=%q path=%q region=%q endpoint=%q sync-interval=%s", r.Name(), client.Type(), client.Bucket, client.Path, client.Region, client.Endpoint, r.SyncInterval)

206
cmd/litestream/restore.go
View File

@@ -7,31 +7,46 @@ import (
"fmt"
"log"
"os"
"path/filepath"
"strconv"
"time"
"github.com/benbjohnson/litestream"
)
// RestoreCommand represents a command to restore a database from a backup.
type RestoreCommand struct{}
type RestoreCommand struct {
snapshotIndex int // index of snapshot to start from
// CLI options
configPath string // path to config file
noExpandEnv bool // if true, do not expand env variables in config
outputPath string // path to restore database to
replicaName string // optional, name of replica to restore from
generation string // optional, generation to restore
targetIndex int // optional, last WAL index to replay
ifDBNotExists bool // if true, skips restore if output path already exists
ifReplicaExists bool // if true, skips if no backups exist
opt litestream.RestoreOptions
}
func NewRestoreCommand() *RestoreCommand {
return &RestoreCommand{
targetIndex: -1,
opt: litestream.NewRestoreOptions(),
}
}
// Run executes the command.
func (c *RestoreCommand) Run(ctx context.Context, args []string) (err error) {
opt := litestream.NewRestoreOptions()
opt.Verbose = true
fs := flag.NewFlagSet("litestream-restore", flag.ContinueOnError)
configPath, noExpandEnv := registerConfigFlag(fs)
fs.StringVar(&opt.OutputPath, "o", "", "output path")
fs.StringVar(&opt.ReplicaName, "replica", "", "replica name")
fs.StringVar(&opt.Generation, "generation", "", "generation name")
fs.Var((*indexVar)(&opt.Index), "index", "wal index")
fs.IntVar(&opt.Parallelism, "parallelism", opt.Parallelism, "parallelism")
ifDBNotExists := fs.Bool("if-db-not-exists", false, "")
ifReplicaExists := fs.Bool("if-replica-exists", false, "")
timestampStr := fs.String("timestamp", "", "timestamp")
verbose := fs.Bool("v", false, "verbose output")
registerConfigFlag(fs, &c.configPath, &c.noExpandEnv)
fs.StringVar(&c.outputPath, "o", "", "output path")
fs.StringVar(&c.replicaName, "replica", "", "replica name")
fs.StringVar(&c.generation, "generation", "", "generation name")
fs.Var((*indexVar)(&c.targetIndex), "index", "wal index")
fs.IntVar(&c.opt.Parallelism, "parallelism", c.opt.Parallelism, "parallelism")
fs.BoolVar(&c.ifDBNotExists, "if-db-not-exists", false, "")
fs.BoolVar(&c.ifReplicaExists, "if-replica-exists", false, "")
fs.Usage = c.Usage
if err := fs.Parse(args); err != nil {
return err
@@ -40,83 +55,100 @@ func (c *RestoreCommand) Run(ctx context.Context, args []string) (err error) {
} else if fs.NArg() > 1 {
return fmt.Errorf("too many arguments")
}
arg := fs.Arg(0)
// Parse timestamp, if specified.
if *timestampStr != "" {
if opt.Timestamp, err = time.Parse(time.RFC3339, *timestampStr); err != nil {
return errors.New("invalid -timestamp, must specify in ISO 8601 format (e.g. 2000-01-01T00:00:00Z)")
}
// Ensure a generation is specified if target index is specified.
if c.targetIndex != -1 && c.generation == "" {
return fmt.Errorf("must specify -generation when using -index flag")
}
// Instantiate logger if verbose output is enabled.
if *verbose {
opt.Logger = log.New(os.Stderr, "", log.LstdFlags|log.Lmicroseconds)
// Default to original database path if output path not specified.
if !isURL(arg) && c.outputPath == "" {
c.outputPath = arg
}
// Determine replica & generation to restore from.
var r *litestream.Replica
if isURL(fs.Arg(0)) {
if *configPath != "" {
return fmt.Errorf("cannot specify a replica URL and the -config flag")
}
if r, err = c.loadFromURL(ctx, fs.Arg(0), *ifDBNotExists, &opt); err == errSkipDBExists {
fmt.Println("database already exists, skipping")
return nil
// Exit successfully if the output file already exists and flag is set.
if _, err := os.Stat(c.outputPath); !os.IsNotExist(err) && c.ifDBNotExists {
fmt.Println("database already exists, skipping")
return nil
}
// Create parent directory if it doesn't already exist.
if err := os.MkdirAll(filepath.Dir(c.outputPath), 0700); err != nil {
return fmt.Errorf("cannot create parent directory: %w", err)
}
// Build replica from either a URL or config.
r, err := c.loadReplica(ctx, arg)
if err != nil {
return err
}
// Determine latest generation if one is not specified.
if c.generation == "" {
if c.generation, err = litestream.FindLatestGeneration(ctx, r.Client); err == litestream.ErrNoGeneration {
// Return an error if no matching targets found.
// If optional flag set, return success. Useful for automated recovery.
if c.ifReplicaExists {
fmt.Println("no matching backups found")
return nil
}
return fmt.Errorf("no matching backups found")
} else if err != nil {
return err
}
} else {
if *configPath == "" {
*configPath = DefaultConfigPath()
}
if r, err = c.loadFromConfig(ctx, fs.Arg(0), *configPath, !*noExpandEnv, *ifDBNotExists, &opt); err == errSkipDBExists {
fmt.Println("database already exists, skipping")
return nil
} else if err != nil {
return err
return fmt.Errorf("cannot determine latest generation: %w", err)
}
}
// Return an error if no matching targets found.
// If optional flag set, return success. Useful for automated recovery.
if opt.Generation == "" {
if *ifReplicaExists {
fmt.Println("no matching backups found")
return nil
// Determine the maximum available index for the generation if one is not specified.
if c.targetIndex == -1 {
if c.targetIndex, err = litestream.FindMaxIndexByGeneration(ctx, r.Client, c.generation); err != nil {
return fmt.Errorf("cannot determine latest index in generation %q: %w", c.generation, err)
}
return fmt.Errorf("no matching backups found")
}
return r.Restore(ctx, opt)
// Find lastest snapshot that occurs before the index.
// TODO: Optionally allow -snapshot-index
if c.snapshotIndex, err = litestream.FindSnapshotForIndex(ctx, r.Client, c.generation, c.targetIndex); err != nil {
return fmt.Errorf("cannot find snapshot index: %w", err)
}
c.opt.Logger = log.New(os.Stderr, "", log.LstdFlags|log.Lmicroseconds)
return litestream.Restore(ctx, r.Client, c.outputPath, c.generation, c.snapshotIndex, c.targetIndex, c.opt)
}
// loadFromURL creates a replica & updates the restore options from a replica URL.
func (c *RestoreCommand) loadFromURL(ctx context.Context, replicaURL string, ifDBNotExists bool, opt *litestream.RestoreOptions) (*litestream.Replica, error) {
if opt.OutputPath == "" {
func (c *RestoreCommand) loadReplica(ctx context.Context, arg string) (*litestream.Replica, error) {
if isURL(arg) {
return c.loadReplicaFromURL(ctx, arg)
}
return c.loadReplicaFromConfig(ctx, arg)
}
// loadReplicaFromURL creates a replica & updates the restore options from a replica URL.
func (c *RestoreCommand) loadReplicaFromURL(ctx context.Context, replicaURL string) (*litestream.Replica, error) {
if c.configPath != "" {
return nil, fmt.Errorf("cannot specify a replica URL and the -config flag")
} else if c.replicaName != "" {
return nil, fmt.Errorf("cannot specify a replica URL and the -replica flag")
} else if c.outputPath == "" {
return nil, fmt.Errorf("output path required")
}
// Exit successfully if the output file already exists.
if _, err := os.Stat(opt.OutputPath); !os.IsNotExist(err) && ifDBNotExists {
return nil, errSkipDBExists
}
syncInterval := litestream.DefaultSyncInterval
r, err := NewReplicaFromConfig(&ReplicaConfig{
return NewReplicaFromConfig(&ReplicaConfig{
URL: replicaURL,
SyncInterval: &syncInterval,
}, nil)
if err != nil {
return nil, err
}
opt.Generation, _, err = r.CalcRestoreTarget(ctx, *opt)
return r, err
}
// loadFromConfig returns a replica & updates the restore options from a DB reference.
func (c *RestoreCommand) loadFromConfig(ctx context.Context, dbPath, configPath string, expandEnv, ifDBNotExists bool, opt *litestream.RestoreOptions) (*litestream.Replica, error) {
// loadReplicaFromConfig returns replicas based on the specific config path.
func (c *RestoreCommand) loadReplicaFromConfig(ctx context.Context, dbPath string) (*litestream.Replica, error) {
if c.configPath == "" {
c.configPath = DefaultConfigPath()
}
// Load configuration.
config, err := ReadConfigFile(configPath, expandEnv)
config, err := ReadConfigFile(c.configPath, !c.noExpandEnv)
if err != nil {
return nil, err
}
@@ -132,25 +164,34 @@ func (c *RestoreCommand) loadFromConfig(ctx context.Context, dbPath, configPath
db, err := NewDBFromConfig(dbConfig)
if err != nil {
return nil, err
} else if len(db.Replicas) == 0 {
return nil, fmt.Errorf("database has no replicas: %s", dbPath)
}
// Restore into original database path if not specified.
if opt.OutputPath == "" {
opt.OutputPath = dbPath
// Filter by replica name if specified.
if c.replicaName != "" {
r := db.Replica(c.replicaName)
if r == nil {
return nil, fmt.Errorf("replica %q not found", c.replicaName)
}
return r, nil
}
// Exit successfully if the output file already exists.
if _, err := os.Stat(opt.OutputPath); !os.IsNotExist(err) && ifDBNotExists {
return nil, errSkipDBExists
// Choose only replica if only one available and no name is specified.
if len(db.Replicas) == 1 {
return db.Replicas[0], nil
}
// Determine the appropriate replica & generation to restore from,
r, generation, err := db.CalcRestoreTarget(ctx, *opt)
// A replica must be specified when restoring a specific generation with multiple replicas.
if c.generation != "" {
return nil, fmt.Errorf("must specify -replica when restoring from a specific generation")
}
// Determine latest replica to restore from.
r, err := litestream.LatestReplica(ctx, db.Replicas)
if err != nil {
return nil, err
return nil, fmt.Errorf("cannot determine latest replica: %w", err)
}
opt.Generation = generation
return r, nil
}
@@ -186,10 +227,6 @@ Arguments:
Restore up to a specific hex-encoded WAL index (inclusive).
Defaults to use the highest available index.
-timestamp TIMESTAMP
Restore to a specific point-in-time.
Defaults to use the latest available backup.
-o PATH
Output path of the restored database.
Defaults to original DB path.
@@ -213,9 +250,6 @@ Examples:
# Restore latest replica for database to original location.
$ litestream restore /path/to/db
# Restore replica for database to a given point in time.
$ litestream restore -timestamp 2020-01-01T00:00:00Z /path/to/db
# Restore latest replica for database to new /tmp directory
$ litestream restore -o /tmp/db /path/to/db

15
cmd/litestream/snapshots.go
View File

@@ -14,12 +14,15 @@ import (
)
// SnapshotsCommand represents a command to list snapshots for a command.
type SnapshotsCommand struct{}
type SnapshotsCommand struct {
configPath string
noExpandEnv bool
}
// Run executes the command.
func (c *SnapshotsCommand) Run(ctx context.Context, args []string) (err error) {
fs := flag.NewFlagSet("litestream-snapshots", flag.ContinueOnError)
configPath, noExpandEnv := registerConfigFlag(fs)
registerConfigFlag(fs, &c.configPath, &c.noExpandEnv)
replicaName := fs.String("replica", "", "replica name")
fs.Usage = c.Usage
if err := fs.Parse(args); err != nil {
@@ -33,19 +36,19 @@ func (c *SnapshotsCommand) Run(ctx context.Context, args []string) (err error) {
var db *litestream.DB
var r *litestream.Replica
if isURL(fs.Arg(0)) {
if *configPath != "" {
if c.configPath != "" {
return fmt.Errorf("cannot specify a replica URL and the -config flag")
}
if r, err = NewReplicaFromConfig(&ReplicaConfig{URL: fs.Arg(0)}, nil); err != nil {
return err
}
} else {
if *configPath == "" {
*configPath = DefaultConfigPath()
if c.configPath == "" {
c.configPath = DefaultConfigPath()
}
// Load configuration.
config, err := ReadConfigFile(*configPath, !*noExpandEnv)
config, err := ReadConfigFile(c.configPath, !c.noExpandEnv)
if err != nil {
return err
}

15
cmd/litestream/wal.go
View File

@@ -13,12 +13,15 @@ import (
)
// WALCommand represents a command to list WAL files for a database.
type WALCommand struct{}
type WALCommand struct {
configPath string
noExpandEnv bool
}
// Run executes the command.
func (c *WALCommand) Run(ctx context.Context, args []string) (err error) {
fs := flag.NewFlagSet("litestream-wal", flag.ContinueOnError)
configPath, noExpandEnv := registerConfigFlag(fs)
registerConfigFlag(fs, &c.configPath, &c.noExpandEnv)
replicaName := fs.String("replica", "", "replica name")
generation := fs.String("generation", "", "generation name")
fs.Usage = c.Usage
@@ -33,19 +36,19 @@ func (c *WALCommand) Run(ctx context.Context, args []string) (err error) {
var db *litestream.DB
var r *litestream.Replica
if isURL(fs.Arg(0)) {
if *configPath != "" {
if c.configPath != "" {
return fmt.Errorf("cannot specify a replica URL and the -config flag")
}
if r, err = NewReplicaFromConfig(&ReplicaConfig{URL: fs.Arg(0)}, nil); err != nil {
return err
}
} else {
if *configPath == "" {
*configPath = DefaultConfigPath()
if c.configPath == "" {
c.configPath = DefaultConfigPath()
}
// Load configuration.
config, err := ReadConfigFile(*configPath, !*noExpandEnv)
config, err := ReadConfigFile(c.configPath, !c.noExpandEnv)
if err != nil {
return err
}