Re-add point-in-time restore

cmd/litestream/restore.go

@@ -9,6 +9,7 @@ import (
 	"os"
 	"path/filepath"
 	"strconv"
+	"time"
 
 	"github.com/benbjohnson/litestream"
 )
@@ -28,6 +29,7 @@ type RestoreCommand struct {
 	replicaName     string    // optional, name of replica to restore from
 	generation      string    // optional, generation to restore
 	targetIndex     int       // optional, last WAL index to replay
+	timestamp       time.Time // optional, restore to point-in-time (ISO 8601)
 	ifDBNotExists   bool      // if true, skips restore if output path already exists
 	ifReplicaExists bool      // if true, skips if no backups exist
 	opt             litestream.RestoreOptions
@@ -53,6 +55,7 @@ func (c *RestoreCommand) Run(ctx context.Context, args []string) (err error) {
 	fs.StringVar(&c.replicaName, "replica", "", "replica name")
 	fs.StringVar(&c.generation, "generation", "", "generation name")
 	fs.Var((*indexVar)(&c.targetIndex), "index", "wal index")
+	timestampStr := fs.String("timestamp", "", "point-in-time restore (ISO 8601)")
 	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, "")
@@ -66,9 +69,20 @@ func (c *RestoreCommand) Run(ctx context.Context, args []string) (err error) {
 	}
 	pathOrURL := fs.Arg(0)
 
+	// Parse timestamp.
+	if *timestampStr != "" {
+		if c.timestamp, err = time.Parse(time.RFC3339Nano, *timestampStr); err != nil {
+			return fmt.Errorf("invalid -timestamp, expected ISO 8601: %w", err)
+		}
+	}
+
 	// Ensure a generation is specified if target index is specified.
-	if c.targetIndex != -1 && c.generation == "" {
+	if c.targetIndex != -1 && !c.timestamp.IsZero() {
+		return fmt.Errorf("cannot specify both -index flag and -timestamp flag")
+	} else if c.targetIndex != -1 && c.generation == "" {
 		return fmt.Errorf("must specify -generation flag when using -index flag")
+	} else if !c.timestamp.IsZero() && c.generation == "" {
+		return fmt.Errorf("must specify -generation flag when using -timestamp flag")
 	}
 
 	// Default to original database path if output path not specified.
@@ -117,7 +131,11 @@ func (c *RestoreCommand) Run(ctx context.Context, args []string) (err error) {
 	}
 
 	// Determine the maximum available index for the generation if one is not specified.
-	if c.targetIndex == -1 {
+	if !c.timestamp.IsZero() {
+		if c.targetIndex, err = litestream.FindIndexByTimestamp(ctx, r.Client(), c.generation, c.timestamp); err != nil {
+			return fmt.Errorf("cannot find index for timestamp in generation %q: %w", c.generation, err)
+		}
+	} else 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)
 		}
@@ -239,6 +257,10 @@ Arguments:
 	    Restore up to a specific hex-encoded WAL index (inclusive).
 	    Defaults to use the highest available index.
 
+	-timestamp DATETIME
+	    Restore up to a specific point-in-time. Must be ISO 8601.
+	    Cannot be specified with -index flag.
+
 	-o PATH
 	    Output path of the restored database.
 	    Defaults to original DB path.
@@ -271,6 +293,9 @@ Examples:
 	# Restore database from specific generation on S3.
 	$ litestream restore -replica s3 -generation xxxxxxxx /path/to/db
 
+	# Restore database to a specific point in time.
+	$ litestream restore -generation xxxxxxxx -timestamp 2000-01-01T00:00:00Z /path/to/db
+
 `[1:],
 		DefaultConfigPath(),
 	)

replica_client.go

@@ -234,6 +234,88 @@ func ReplicaClientTimeBounds(ctx context.Context, client ReplicaClient) (min, ma
 	return min, max, nil
 }
 
+// FindIndexByTimestamp returns the highest index before a given point-in-time
+// within a generation. Returns ErrNoSnapshots if no index exists on the replica
+// for the generation.
+func FindIndexByTimestamp(ctx context.Context, client ReplicaClient, generation string, timestamp time.Time) (index int, err error) {
+	snapshotIndex, err := FindSnapshotIndexByTimestamp(ctx, client, generation, timestamp)
+	if err == ErrNoSnapshots {
+		return 0, err
+	} else if err != nil {
+		return 0, fmt.Errorf("max snapshot index: %w", err)
+	}
+
+	// Determine the highest available WAL index.
+	walIndex, err := FindWALIndexByTimestamp(ctx, client, generation, timestamp)
+	if err != nil && err != ErrNoWALSegments {
+		return 0, fmt.Errorf("max wal index: %w", err)
+	}
+
+	// Use snapshot index if it's after the last WAL index.
+	if snapshotIndex > walIndex {
+		return snapshotIndex, nil
+	}
+	return walIndex, nil
+}
+
+// FindSnapshotIndexByTimestamp returns the highest snapshot index before timestamp.
+// Returns ErrNoSnapshots if no snapshots exist for the generation on the replica.
+func FindSnapshotIndexByTimestamp(ctx context.Context, client ReplicaClient, generation string, timestamp time.Time) (index int, err error) {
+	itr, err := client.Snapshots(ctx, generation)
+	if err != nil {
+		return 0, fmt.Errorf("snapshots: %w", err)
+	}
+	defer func() { _ = itr.Close() }()
+
+	// Iterate over snapshots to find the highest index.
+	var n int
+	for ; itr.Next(); n++ {
+		if info := itr.Snapshot(); info.CreatedAt.After(timestamp) {
+			continue
+		} else if info.Index > index {
+			index = info.Index
+		}
+	}
+	if err := itr.Close(); err != nil {
+		return 0, fmt.Errorf("snapshot iteration: %w", err)
+	}
+
+	// Return an error if no snapshots were found.
+	if n == 0 {
+		return 0, ErrNoSnapshots
+	}
+	return index, nil
+}
+
+// FindWALIndexByTimestamp returns the highest WAL index before timestamp.
+// Returns ErrNoWALSegments if no segments exist for the generation on the replica.
+func FindWALIndexByTimestamp(ctx context.Context, client ReplicaClient, generation string, timestamp time.Time) (index int, err error) {
+	itr, err := client.WALSegments(ctx, generation)
+	if err != nil {
+		return 0, fmt.Errorf("wal segments: %w", err)
+	}
+	defer func() { _ = itr.Close() }()
+
+	// Iterate over WAL segments to find the highest index.
+	var n int
+	for ; itr.Next(); n++ {
+		if info := itr.WALSegment(); info.CreatedAt.After(timestamp) {
+			continue
+		} else if info.Index > index {
+			index = info.Index
+		}
+	}
+	if err := itr.Close(); err != nil {
+		return 0, fmt.Errorf("wal segment iteration: %w", err)
+	}
+
+	// Return an error if no WAL segments were found.
+	if n == 0 {
+		return 0, ErrNoWALSegments
+	}
+	return index, nil
+}
+
 // FindMaxIndexByGeneration returns the last index within a generation.
 // Returns ErrNoSnapshots if no index exists on the replica for the generation.
 func FindMaxIndexByGeneration(ctx context.Context, client ReplicaClient, generation string) (index int, err error) {

replica_client_test.go

@@ -489,6 +489,167 @@ func TestFindMaxIndexByGeneration(t *testing.T) {
 	})
 }
 
+func TestFindSnapshotIndexByTimestamp(t *testing.T) {
+	t.Run("OK", func(t *testing.T) {
+		client := litestream.NewFileReplicaClient(filepath.Join("testdata", "snapshot-index-by-timestamp", "ok"))
+		if index, err := litestream.FindSnapshotIndexByTimestamp(context.Background(), client, "0000000000000000", time.Date(2000, 1, 3, 0, 0, 0, 0, time.UTC)); err != nil {
+			t.Fatal(err)
+		} else if got, want := index, 0x000007d0; got != want {
+			t.Fatalf("index=%d, want %d", got, want)
+		}
+	})
+
+	t.Run("ErrNoSnapshots", func(t *testing.T) {
+		client := litestream.NewFileReplicaClient(filepath.Join("testdata", "snapshot-index-by-timestamp", "no-snapshots"))
+
+		_, err := litestream.FindSnapshotIndexByTimestamp(context.Background(), client, "0000000000000000", time.Date(2000, 1, 1, 0, 0, 0, 0, time.UTC))
+		if err != litestream.ErrNoSnapshots {
+			t.Fatalf("unexpected error: %s", err)
+		}
+	})
+
+	t.Run("ErrSnapshots", func(t *testing.T) {
+		var client mock.ReplicaClient
+		client.SnapshotsFunc = func(ctx context.Context, generation string) (litestream.SnapshotIterator, error) {
+			return nil, fmt.Errorf("marker")
+		}
+
+		_, err := litestream.FindSnapshotIndexByTimestamp(context.Background(), &client, "0000000000000000", time.Date(2000, 1, 1, 0, 0, 0, 0, time.UTC))
+		if err == nil || err.Error() != `snapshots: marker` {
+			t.Fatalf("unexpected error: %s", err)
+		}
+	})
+
+	t.Run("ErrSnapshotIteration", func(t *testing.T) {
+		var itr mock.SnapshotIterator
+		itr.NextFunc = func() bool { return false }
+		itr.CloseFunc = func() error { return fmt.Errorf("marker") }
+
+		var client mock.ReplicaClient
+		client.SnapshotsFunc = func(ctx context.Context, generation string) (litestream.SnapshotIterator, error) {
+			return &itr, nil
+		}
+
+		_, err := litestream.FindSnapshotIndexByTimestamp(context.Background(), &client, "0000000000000000", time.Date(2000, 1, 1, 0, 0, 0, 0, time.UTC))
+		if err == nil || err.Error() != `snapshot iteration: marker` {
+			t.Fatalf("unexpected error: %s", err)
+		}
+	})
+}
+
+func TestFindWALIndexByTimestamp(t *testing.T) {
+	t.Run("OK", func(t *testing.T) {
+		client := litestream.NewFileReplicaClient(filepath.Join("testdata", "wal-index-by-timestamp", "ok"))
+		if index, err := litestream.FindWALIndexByTimestamp(context.Background(), client, "0000000000000000", time.Date(2000, 1, 3, 0, 0, 0, 0, time.UTC)); err != nil {
+			t.Fatal(err)
+		} else if got, want := index, 1; got != want {
+			t.Fatalf("index=%d, want %d", got, want)
+		}
+	})
+
+	t.Run("ErrNoWALSegments", func(t *testing.T) {
+		client := litestream.NewFileReplicaClient(filepath.Join("testdata", "wal-index-by-timestamp", "no-wal"))
+
+		_, err := litestream.FindWALIndexByTimestamp(context.Background(), client, "0000000000000000", time.Date(2000, 1, 1, 0, 0, 0, 0, time.UTC))
+		if err != litestream.ErrNoWALSegments {
+			t.Fatalf("unexpected error: %s", err)
+		}
+	})
+
+	t.Run("ErrWALSegments", func(t *testing.T) {
+		var client mock.ReplicaClient
+		client.WALSegmentsFunc = func(ctx context.Context, generation string) (litestream.WALSegmentIterator, error) {
+			return nil, fmt.Errorf("marker")
+		}
+
+		_, err := litestream.FindWALIndexByTimestamp(context.Background(), &client, "0000000000000000", time.Date(2000, 1, 1, 0, 0, 0, 0, time.UTC))
+		if err == nil || err.Error() != `wal segments: marker` {
+			t.Fatalf("unexpected error: %s", err)
+		}
+	})
+
+	t.Run("ErrWALSegmentIteration", func(t *testing.T) {
+		var itr mock.WALSegmentIterator
+		itr.NextFunc = func() bool { return false }
+		itr.CloseFunc = func() error { return fmt.Errorf("marker") }
+
+		var client mock.ReplicaClient
+		client.WALSegmentsFunc = func(ctx context.Context, generation string) (litestream.WALSegmentIterator, error) {
+			return &itr, nil
+		}
+
+		_, err := litestream.FindWALIndexByTimestamp(context.Background(), &client, "0000000000000000", time.Date(2000, 1, 1, 0, 0, 0, 0, time.UTC))
+		if err == nil || err.Error() != `wal segment iteration: marker` {
+			t.Fatalf("unexpected error: %s", err)
+		}
+	})
+}
+
+func TestFindIndexByTimestamp(t *testing.T) {
+	t.Run("OK", func(t *testing.T) {
+		client := litestream.NewFileReplicaClient(filepath.Join("testdata", "index-by-timestamp", "ok"))
+		if index, err := litestream.FindIndexByTimestamp(context.Background(), client, "0000000000000000", time.Date(2000, 1, 4, 0, 0, 0, 0, time.UTC)); err != nil {
+			t.Fatal(err)
+		} else if got, want := index, 0x00000002; got != want {
+			t.Fatalf("index=%d, want %d", got, want)
+		}
+	})
+
+	t.Run("NoWAL", func(t *testing.T) {
+		client := litestream.NewFileReplicaClient(filepath.Join("testdata", "index-by-timestamp", "no-wal"))
+		if index, err := litestream.FindIndexByTimestamp(context.Background(), client, "0000000000000000", time.Date(2000, 1, 2, 0, 0, 0, 0, time.UTC)); err != nil {
+			t.Fatal(err)
+		} else if got, want := index, 0x00000001; got != want {
+			t.Fatalf("index=%d, want %d", got, want)
+		}
+	})
+
+	t.Run("SnapshotLaterThanWAL", func(t *testing.T) {
+		client := litestream.NewFileReplicaClient(filepath.Join("testdata", "index-by-timestamp", "snapshot-later-than-wal"))
+		if index, err := litestream.FindIndexByTimestamp(context.Background(), client, "0000000000000000", time.Date(2000, 1, 3, 0, 0, 0, 0, time.UTC)); err != nil {
+			t.Fatal(err)
+		} else if got, want := index, 0x00000001; got != want {
+			t.Fatalf("index=%d, want %d", got, want)
+		}
+	})
+
+	t.Run("ErrNoSnapshots", func(t *testing.T) {
+		client := litestream.NewFileReplicaClient(filepath.Join("testdata", "index-by-timestamp", "no-snapshots"))
+
+		_, err := litestream.FindIndexByTimestamp(context.Background(), client, "0000000000000000", time.Date(2000, 1, 1, 0, 0, 0, 0, time.UTC))
+		if err != litestream.ErrNoSnapshots {
+			t.Fatalf("unexpected error: %s", err)
+		}
+	})
+
+	t.Run("ErrSnapshots", func(t *testing.T) {
+		var client mock.ReplicaClient
+		client.SnapshotsFunc = func(ctx context.Context, generation string) (litestream.SnapshotIterator, error) {
+			return nil, fmt.Errorf("marker")
+		}
+
+		_, err := litestream.FindIndexByTimestamp(context.Background(), &client, "0000000000000000", time.Date(2000, 1, 1, 0, 0, 0, 0, time.UTC))
+		if err == nil || err.Error() != `max snapshot index: snapshots: marker` {
+			t.Fatalf("unexpected error: %s", err)
+		}
+	})
+
+	t.Run("ErrWALSegments", func(t *testing.T) {
+		var client mock.ReplicaClient
+		client.SnapshotsFunc = func(ctx context.Context, generation string) (litestream.SnapshotIterator, error) {
+			return litestream.NewSnapshotInfoSliceIterator([]litestream.SnapshotInfo{{Index: 0x00000001}}), nil
+		}
+		client.WALSegmentsFunc = func(ctx context.Context, generation string) (litestream.WALSegmentIterator, error) {
+			return nil, fmt.Errorf("marker")
+		}
+
+		_, err := litestream.FindIndexByTimestamp(context.Background(), &client, "0000000000000000", time.Date(2000, 1, 1, 0, 0, 0, 0, time.UTC))
+		if err == nil || err.Error() != `max wal index: wal segments: marker` {
+			t.Fatalf("unexpected error: %s", err)
+		}
+	})
+}
+
 func TestRestore(t *testing.T) {
 	t.Run("OK", func(t *testing.T) {
 		testDir := filepath.Join("testdata", "restore", "ok")

testdata/Makefile

@@ -1,8 +1,13 @@
 .PHONY: default
 default:
 	make -C find-latest-generation/ok
+	make -C index-by-timestamp/no-wal
+	make -C index-by-timestamp/ok
+	make -C index-by-timestamp/snapshot-later-than-wal
 	make -C generation-time-bounds/ok
 	make -C generation-time-bounds/snapshots-only
 	make -C replica-client-time-bounds/ok
 	make -C snapshot-time-bounds/ok
+	make -C snapshot-index-by-timestamp/ok
 	make -C wal-time-bounds/ok
+	make -C wal-index-by-timestamp/ok

testdata/index-by-timestamp/no-wal/Makefile

@@ -0,0 +1,6 @@
+.PHONY: default
+default:
+	TZ=UTC touch -ct 200001010000 generations/0000000000000000/snapshots/0000000000000000.snapshot.lz4
+	TZ=UTC touch -ct 200001020000 generations/0000000000000000/snapshots/0000000000000001.snapshot.lz4
+	TZ=UTC touch -ct 200001030000 generations/0000000000000000/snapshots/0000000000000002.snapshot.lz4
+

testdata/index-by-timestamp/ok/Makefile

@@ -0,0 +1,11 @@
+.PHONY: default
+default:
+	TZ=UTC touch -ct 200001010000 generations/0000000000000000/snapshots/0000000000000000.snapshot.lz4
+	TZ=UTC touch -ct 200001030000 generations/0000000000000000/snapshots/0000000000000001.snapshot.lz4
+
+	TZ=UTC touch -ct 200001010000 generations/0000000000000000/wal/0000000000000000/0000000000000000.wal.lz4
+	TZ=UTC touch -ct 200001020000 generations/0000000000000000/wal/0000000000000000/0000000000001234.wal.lz4
+	TZ=UTC touch -ct 200001030000 generations/0000000000000000/wal/0000000000000001/0000000000000000.wal.lz4
+	TZ=UTC touch -ct 200001040000 generations/0000000000000000/wal/0000000000000002/0000000000000000.wal.lz4
+	TZ=UTC touch -ct 200001050000 generations/0000000000000000/wal/0000000000000003/0000000000000000.wal.lz4
+

testdata/index-by-timestamp/snapshot-later-than-wal/Makefile

@@ -0,0 +1,7 @@
+.PHONY: default
+default:
+	TZ=UTC touch -ct 200001010000 generations/0000000000000000/snapshots/0000000000000000.snapshot.lz4
+	TZ=UTC touch -ct 200001030000 generations/0000000000000000/snapshots/0000000000000001.snapshot.lz4
+
+	TZ=UTC touch -ct 200001020000 generations/0000000000000000/wal/0000000000000000/0000000000000000.wal.lz4
+	TZ=UTC touch -ct 200001020000 generations/0000000000000000/wal/0000000000000000/0000000000001234.wal.lz4

testdata/snapshot-index-by-timestamp/ok/Makefile

@@ -0,0 +1,5 @@
+.PHONY: default
+default:
+	TZ=UTC touch -ct 200001010000 generations/0000000000000000/snapshots/0000000000000000.snapshot.lz4
+	TZ=UTC touch -ct 200001020000 generations/0000000000000000/snapshots/00000000000003e8.snapshot.lz4
+	TZ=UTC touch -ct 200001030000 generations/0000000000000000/snapshots/00000000000007d0.snapshot.lz4

testdata/wal-index-by-timestamp/ok/Makefile

@@ -0,0 +1,6 @@
+.PHONY: default
+default:
+	TZ=UTC touch -ct 200001010000 generations/0000000000000000/wal/0000000000000000/0000000000000000.wal.lz4
+	TZ=UTC touch -ct 200001020000 generations/0000000000000000/wal/0000000000000000/0000000000001234.wal.lz4
+	TZ=UTC touch -ct 200001030000 generations/0000000000000000/wal/0000000000000001/0000000000000000.wal.lz4
+