addresed review comments

Signed-off-by: Benjamin Wang <wachao@vmware.com>

cmd/bbolt/surgery_commands.go

@@ -80,15 +80,12 @@ func newRevertMetaPageCommand(m *SurgeryCommand) *RevertMetaPageCommand {
 func (cmd *RevertMetaPageCommand) Run(args ...string) error {
 	// Parse flags.
 	fs := flag.NewFlagSet("", flag.ContinueOnError)
-	fs.SetOutput(io.Discard)
-	fs.StringVar(&cmd.DstPath, "o", "", "")
-	if err := fs.Parse(args); err == flag.ErrHelp {
+	help := fs.Bool("h", false, "")
+	if err := fs.Parse(args); err != nil {
+		return err
+	} else if *help {
 		fmt.Fprintln(cmd.Stderr, cmd.Usage())
 		return ErrUsage
-	} else if err != nil {
-		return err
-	} else if cmd.DstPath == "" {
-		return errors.New("output file required")
 	}
 
 	// Require database paths.
@@ -97,14 +94,18 @@ func (cmd *RevertMetaPageCommand) Run(args ...string) error {
 		return ErrPathRequired
 	}
 
+	cmd.DstPath = fs.Arg(1)
+	if cmd.DstPath == "" {
+		return errors.New("output file required")
+	}
+
 	// Ensure source file exists.
-	fi, err := os.Stat(cmd.SrcPath)
+	_, err := os.Stat(cmd.SrcPath)
 	if os.IsNotExist(err) {
 		return ErrFileNotFound
 	} else if err != nil {
 		return err
 	}
-	initialSize := fi.Size()
 
 	// Ensure output file not exist.
 	_, err = os.Stat(cmd.DstPath)
@@ -115,30 +116,44 @@ func (cmd *RevertMetaPageCommand) Run(args ...string) error {
 	}
 
 	// Copy database from SrcPath to DstPath
-	srcDB, err := os.Open(cmd.SrcPath)
+	if err := copyFile(cmd.SrcPath, cmd.DstPath); err != nil {
+		return fmt.Errorf("failed to copy file: %w", err)
+	}
+
+	// revert the meta page
+	if err = surgeon.RevertMetaPage(cmd.DstPath); err != nil {
+		return err
+	}
+
+	fmt.Fprintln(cmd.Stdout, "The meta page is reverted.")
+	return nil
+}
+
+func copyFile(srcPath, dstPath string) error {
+	srcDB, err := os.Open(srcPath)
 	if err != nil {
-		return fmt.Errorf("failed to open source file %q: %w", cmd.SrcPath, err)
+		return fmt.Errorf("failed to open source file %q: %w", srcPath, err)
 	}
 	defer srcDB.Close()
-	dstDB, err := os.Create(cmd.DstPath)
+	dstDB, err := os.Create(dstPath)
 	if err != nil {
-		return fmt.Errorf("failed to create output file %q: %w", cmd.DstPath, err)
+		return fmt.Errorf("failed to create output file %q: %w", dstPath, err)
 	}
 	defer dstDB.Close()
 	written, err := io.Copy(dstDB, srcDB)
 	if err != nil {
-		return fmt.Errorf("failed to copy database file from %q to %q: %w", cmd.SrcPath, cmd.DstPath, err)
-	}
-	if initialSize != written {
-		return fmt.Errorf("the byte copied (%q: %d) isn't equal to the initial db size (%q: %d)", cmd.DstPath, written, cmd.SrcPath, initialSize)
+		return fmt.Errorf("failed to copy database file from %q to %q: %w", srcPath, dstPath, err)
 	}
 
-	// revert the meta page
-	if err = surgeon.RevertMetaPage(cmd.DstPath); err != nil {
-		return err
+	srcFi, err := srcDB.Stat()
+	if err != nil {
+		return fmt.Errorf("failed to get source file info %q: %w", srcPath, err)
+	}
+	initialSize := srcFi.Size()
+	if initialSize != written {
+		return fmt.Errorf("the byte copied (%q: %d) isn't equal to the initial db size (%q: %d)", dstPath, written, srcPath, initialSize)
 	}
 
-	fmt.Fprintln(cmd.Stdout, "The meta page is reverted.")
 	return nil
 }
 

cmd/bbolt/surgery_commands_test.go

@@ -1,6 +1,7 @@
 package main_test
 
 import (
+	bolt "go.etcd.io/bbolt"
 	"os"
 	"path/filepath"
 	"testing"
@@ -13,8 +14,8 @@ import (
 )
 
 func TestSurgery_RevertMetaPage(t *testing.T) {
-	pageSize := os.Getpagesize()
-	db := btesting.MustCreateDB(t)
+	pageSize := 4096
+	db := btesting.MustCreateDBWithOption(t, &bolt.Options{PageSize: pageSize})
 	srcPath := db.Path()
 
 	srcFile, err := os.Open(srcPath)
@@ -38,15 +39,19 @@ func TestSurgery_RevertMetaPage(t *testing.T) {
 	// revert the meta page
 	dstPath := filepath.Join(t.TempDir(), "dstdb")
 	m := NewMain()
-	err = m.Run("surgery", "revert-meta-page", "-o", dstPath, srcPath)
+	err = m.Run("surgery", "revert-meta-page", srcPath, dstPath)
 	require.NoError(t, err)
 
 	// read both meta0 and meta1 from dst file
 	dstBuf0, dstBuf1 := readBothMetaPages(t, dstPath, pageSize)
 
 	// check result. Note we should skip the page ID
-	assert.Equal(t, nonActiveSrcBuf[8:], dstBuf0[8:])
-	assert.Equal(t, nonActiveSrcBuf[8:], dstBuf1[8:])
+	assert.Equal(t, pageDataWithoutPageId(nonActiveSrcBuf), pageDataWithoutPageId(dstBuf0))
+	assert.Equal(t, pageDataWithoutPageId(nonActiveSrcBuf), pageDataWithoutPageId(dstBuf1))
+}
+
+func pageDataWithoutPageId(buf []byte) []byte {
+	return buf[8:]
 }
 
 func readBothMetaPages(t *testing.T, filePath string, pageSize int) ([]byte, []byte) {