func runBenchmarkBackup(b *testing.B, count int) {
cleanupFn, _, kvDB, dir := setupBackupRestoreDB(b, count)
defer cleanupFn()
b.ResetTimer()
for i := 0; i < b.N; i++ {
b.StopTimer()
if err := os.RemoveAll(dir); err != nil {
b.Fatal(err)
}
b.StartTimer()
if _, err := sql.Backup(context.Background(), *kvDB, dir); err != nil {
b.Fatal(err)
}
}
}
func runBackup(cmd *cobra.Command, args []string) error {
if len(args) != 1 {
return errors.New("output basepath argument is required")
}
base := args[0]
ctx := context.Background()
kvDB, stopper := makeDBClient()
defer stopper.Stop()
desc, err := sql.Backup(ctx, *kvDB, base)
if err != nil {
return err
}
fmt.Printf("Backed up %d ranges to %s\n", len(desc.Ranges), base)
return nil
}
func runBenchmarkRestore(b *testing.B, count int) {
cleanupFn, sqlDB, kvDB, dir := setupBackupRestoreDB(b, count)
defer cleanupFn()
if _, err := sql.Backup(context.Background(), *kvDB, dir); err != nil {
b.Fatal(err)
}
b.ResetTimer()
for i := 0; i < b.N; i++ {
b.StopTimer()
if _, err := sqlDB.Exec(`TRUNCATE d.foo`); err != nil {
b.Fatal(err)
}
b.StartTimer()
if err := sql.Restore(context.Background(), *kvDB, dir, "foo", true); err != nil {
b.Fatal(err)
}
}
}
func TestBackupRestore(t *testing.T) {
defer leaktest.AfterTest(t)()
count := 1000
cleanupFn, sqlDB, kvDB, dir := setupBackupRestoreDB(t, count)
defer cleanupFn()
if _, err := sql.Backup(context.Background(), *kvDB, dir); err != nil {
t.Fatal(err)
}
if _, err := sqlDB.Exec(`TRUNCATE d.foo`); err != nil {
t.Fatal(err)
}
var rowCount int
if err := sqlDB.QueryRow(`SELECT COUNT(*) FROM d.foo`).Scan(&rowCount); err != nil {
t.Fatal(err)
}
if rowCount != 0 {
t.Fatalf("expected 0 rows but found %d", rowCount)
}
// TODO(dan): Shut down the cluster and restore into a fresh one.
if err := sql.Restore(context.Background(), *kvDB, dir, "foo", true); err != nil {
t.Fatal(err)
}
if err := sqlDB.QueryRow(`SELECT COUNT(*) FROM d.foo`).Scan(&rowCount); err != nil {
t.Fatal(err)
}
if rowCount != count {
t.Fatalf("expected %d rows but found %d", count, rowCount)
}
}