Merge pull request #484 from ahrtr/surgery_20230507

cmd: migrate 'surgery clear-page' command to cobra style command

cmd/bbolt/command_surgery.go

@@ -26,6 +26,7 @@ func newSurgeryCobraCommand() *cobra.Command {
 
 	surgeryCmd.AddCommand(newSurgeryRevertMetaPageCommand())
 	surgeryCmd.AddCommand(newSurgeryCopyPageCommand())
+	surgeryCmd.AddCommand(newSurgeryClearPageCommand())
 	surgeryCmd.AddCommand(newSurgeryClearPageElementsCommand())
 	surgeryCmd.AddCommand(newSurgeryFreelistCommand())
 
@@ -138,6 +139,10 @@ func newSurgeryCopyPageCommand() *cobra.Command {
 }
 
 func surgeryCopyPageFunc(srcDBPath string, cfg surgeryCopyPageOptions) error {
+	if _, err := checkSourceDBPath(srcDBPath); err != nil {
+		return err
+	}
+
 	if err := common.CopyFile(srcDBPath, cfg.outputDBFilePath); err != nil {
 		return fmt.Errorf("[copy-page] copy file failed: %w", err)
 	}
@@ -159,6 +164,74 @@ func surgeryCopyPageFunc(srcDBPath string, cfg surgeryCopyPageOptions) error {
 	return nil
 }
 
+type surgeryClearPageOptions struct {
+	surgeryBaseOptions
+	pageId uint64
+}
+
+func (o *surgeryClearPageOptions) AddFlags(fs *pflag.FlagSet) {
+	o.surgeryBaseOptions.AddFlags(fs)
+	fs.Uint64VarP(&o.pageId, "pageId", "", o.pageId, "page Id")
+}
+
+func (o *surgeryClearPageOptions) Validate() error {
+	if err := o.surgeryBaseOptions.Validate(); err != nil {
+		return err
+	}
+	if o.pageId < 2 {
+		return fmt.Errorf("the pageId must be at least 2, but got %d", o.pageId)
+	}
+	return nil
+}
+
+func newSurgeryClearPageCommand() *cobra.Command {
+	var o surgeryClearPageOptions
+	clearPageCmd := &cobra.Command{
+		Use:   "clear-page <bbolt-file> [options]",
+		Short: "Clears all elements from the given page, which can be a branch or leaf page",
+		Args: func(cmd *cobra.Command, args []string) error {
+			if len(args) == 0 {
+				return errors.New("db file path not provided")
+			}
+			if len(args) > 1 {
+				return errors.New("too many arguments")
+			}
+			return nil
+		},
+		RunE: func(cmd *cobra.Command, args []string) error {
+			if err := o.Validate(); err != nil {
+				return err
+			}
+			return surgeryClearPageFunc(args[0], o)
+		},
+	}
+	o.AddFlags(clearPageCmd.Flags())
+	return clearPageCmd
+}
+
+func surgeryClearPageFunc(srcDBPath string, cfg surgeryClearPageOptions) error {
+	if _, err := checkSourceDBPath(srcDBPath); err != nil {
+		return err
+	}
+
+	if err := common.CopyFile(srcDBPath, cfg.outputDBFilePath); err != nil {
+		return fmt.Errorf("[clear-page] copy file failed: %w", err)
+	}
+
+	needAbandonFreelist, err := surgeon.ClearPage(cfg.outputDBFilePath, common.Pgid(cfg.pageId))
+	if err != nil {
+		return fmt.Errorf("clear-page command failed: %w", err)
+	}
+
+	if needAbandonFreelist {
+		fmt.Fprintf(os.Stdout, "WARNING: The clearing has abandoned some pages that are not yet referenced from free list.\n")
+		fmt.Fprintf(os.Stdout, "Please consider executing `./bbolt surgery abandon-freelist ...`\n")
+	}
+
+	fmt.Fprintf(os.Stdout, "The page (%d) was cleared\n", cfg.pageId)
+	return nil
+}
+
 type surgeryClearPageElementsOptions struct {
 	surgeryBaseOptions
 	pageId          uint64
@@ -209,6 +282,10 @@ func newSurgeryClearPageElementsCommand() *cobra.Command {
 }
 
 func surgeryClearPageElementFunc(srcDBPath string, cfg surgeryClearPageElementsOptions) error {
+	if _, err := checkSourceDBPath(srcDBPath); err != nil {
+		return err
+	}
+
 	if err := common.CopyFile(srcDBPath, cfg.outputDBFilePath); err != nil {
 		return fmt.Errorf("[clear-page-element] copy file failed: %w", err)
 	}
@@ -227,9 +304,6 @@ func surgeryClearPageElementFunc(srcDBPath string, cfg surgeryClearPageElementsO
 	return nil
 }
 
-// TODO(ahrtr): add `bbolt surgery freelist rebuild/check ...` commands,
-// and move all `surgery freelist` commands into a separate file,
-// e.g command_surgery_freelist.go.
 func newSurgeryFreelistCommand() *cobra.Command {
 	cmd := &cobra.Command{
 		Use:   "freelist <subcommand>",
@@ -269,6 +343,10 @@ func newSurgeryFreelistAbandonCommand() *cobra.Command {
 }
 
 func surgeryFreelistAbandonFunc(srcDBPath string, cfg surgeryBaseOptions) error {
+	if _, err := checkSourceDBPath(srcDBPath); err != nil {
+		return err
+	}
+
 	if err := common.CopyFile(srcDBPath, cfg.outputDBFilePath); err != nil {
 		return fmt.Errorf("[freelist abandon] copy file failed: %w", err)
 	}

cmd/bbolt/command_surgery_test.go

@@ -99,6 +99,43 @@ func TestSurgery_CopyPage(t *testing.T) {
 	assert.Equal(t, pageDataWithoutPageId(srcPageId3Data), pageDataWithoutPageId(dstPageId2Data))
 }
 
+// TODO(ahrtr): add test case below for `surgery clear-page` command:
+//  1. The page is a branch page. All its children should become free pages.
+func TestSurgery_ClearPage(t *testing.T) {
+	pageSize := 4096
+	db := btesting.MustCreateDBWithOption(t, &bolt.Options{PageSize: pageSize})
+	srcPath := db.Path()
+
+	// Insert some sample data
+	t.Log("Insert some sample data")
+	err := db.Fill([]byte("data"), 1, 20,
+		func(tx int, k int) []byte { return []byte(fmt.Sprintf("%04d", k)) },
+		func(tx int, k int) []byte { return make([]byte, 10) },
+	)
+	require.NoError(t, err)
+
+	defer requireDBNoChange(t, dbData(t, srcPath), srcPath)
+
+	// clear page 3
+	t.Log("clear page 3")
+	rootCmd := main.NewRootCommand()
+	output := filepath.Join(t.TempDir(), "dstdb")
+	rootCmd.SetArgs([]string{
+		"surgery", "clear-page", srcPath,
+		"--output", output,
+		"--pageId", "3",
+	})
+	err = rootCmd.Execute()
+	require.NoError(t, err)
+
+	t.Log("Verify result")
+	dstPageId3Data := readPage(t, output, 3, pageSize)
+
+	p := common.LoadPage(dstPageId3Data)
+	assert.Equal(t, uint16(0), p.Count())
+	assert.Equal(t, uint32(0), p.Overflow())
+}
+
 func TestSurgery_ClearPageElements_Without_Overflow(t *testing.T) {
 	testCases := []struct {
 		name                 string

cmd/bbolt/main.go

@@ -140,8 +140,6 @@ func (m *Main) Run(args ...string) error {
 		return newPagesCommand(m).Run(args[1:]...)
 	case "stats":
 		return newStatsCommand(m).Run(args[1:]...)
-	case "surgery":
-		return newSurgeryCommand(m).Run(args[1:]...)
 	default:
 		return ErrUnknownCommand
 	}

cmd/bbolt/surgery_commands.go

@@ -1,146 +0,0 @@
-package main
-
-import (
-	"errors"
-	"flag"
-	"fmt"
-	"os"
-	"strconv"
-	"strings"
-
-	"go.etcd.io/bbolt/internal/common"
-	"go.etcd.io/bbolt/internal/surgeon"
-)
-
-// surgeryCommand represents the "surgery" command execution.
-type surgeryCommand struct {
-	baseCommand
-
-	srcPath string
-	dstPath string
-}
-
-// newSurgeryCommand returns a SurgeryCommand.
-func newSurgeryCommand(m *Main) *surgeryCommand {
-	c := &surgeryCommand{}
-	c.baseCommand = m.baseCommand
-	return c
-}
-
-// Run executes the `surgery` program.
-func (cmd *surgeryCommand) Run(args ...string) error {
-	// Require a command at the beginning.
-	if len(args) == 0 || strings.HasPrefix(args[0], "-") {
-		fmt.Fprintln(cmd.Stderr, cmd.Usage())
-		return ErrUsage
-	}
-
-	// Execute command.
-	switch args[0] {
-	case "help":
-		fmt.Fprintln(cmd.Stderr, cmd.Usage())
-		return ErrUsage
-	case "clear-page":
-		return newClearPageCommand(cmd).Run(args[1:]...)
-	default:
-		return ErrUnknownCommand
-	}
-}
-
-func (cmd *surgeryCommand) parsePathsAndCopyFile(fs *flag.FlagSet) error {
-	// Require database paths.
-	cmd.srcPath = fs.Arg(0)
-	if cmd.srcPath == "" {
-		return ErrPathRequired
-	}
-
-	cmd.dstPath = fs.Arg(1)
-	if cmd.dstPath == "" {
-		return errors.New("output file required")
-	}
-
-	// Copy database from SrcPath to DstPath
-	if err := common.CopyFile(cmd.srcPath, cmd.dstPath); err != nil {
-		return fmt.Errorf("failed to copy file: %w", err)
-	}
-
-	return nil
-}
-
-// Usage returns the help message.
-func (cmd *surgeryCommand) Usage() string {
-	return strings.TrimLeft(`
-Surgery is a command for performing low level update on bbolt databases.
-
-Usage:
-
-	bbolt surgery command [arguments]
-
-The commands are:
-    help                   print this screen
-    clear-page             clear all elements at the given pageId
-
-Use "bbolt surgery [command] -h" for more information about a command.
-`, "\n")
-}
-
-// clearPageCommand represents the "surgery clear-page" command execution.
-type clearPageCommand struct {
-	*surgeryCommand
-}
-
-// newClearPageCommand returns a clearPageCommand.
-func newClearPageCommand(m *surgeryCommand) *clearPageCommand {
-	c := &clearPageCommand{}
-	c.surgeryCommand = m
-	return c
-}
-
-// Run executes the command.
-func (cmd *clearPageCommand) Run(args ...string) error {
-	// Parse flags.
-	fs := flag.NewFlagSet("", flag.ContinueOnError)
-	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
-	}
-
-	if err := cmd.parsePathsAndCopyFile(fs); err != nil {
-		return fmt.Errorf("clearPageCommand failed to parse paths and copy file: %w", err)
-	}
-
-	// Read page id.
-	pageId, err := strconv.ParseUint(fs.Arg(2), 10, 64)
-	if err != nil {
-		return err
-	}
-
-	needAbandonFreelist, err := surgeon.ClearPage(cmd.dstPath, common.Pgid(pageId))
-	if err != nil {
-		return fmt.Errorf("clearPageCommand failed: %w", err)
-	}
-
-	if needAbandonFreelist {
-		fmt.Fprintf(os.Stdout, "WARNING: The clearing has abandoned some pages that are not yet referenced from free list.\n")
-		fmt.Fprintf(os.Stdout, "Please consider executing `./bbolt surgery abandon-freelist ...`\n")
-	}
-
-	fmt.Fprintf(cmd.Stdout, "Page (%d) was cleared\n", pageId)
-	return nil
-}
-
-// Usage returns the help message.
-func (cmd *clearPageCommand) Usage() string {
-	return strings.TrimLeft(`
-usage: bolt surgery clear-page SRC DST pageId
-
-ClearPage copies the database file at SRC to a newly created database
-file at DST. Afterwards, it clears all elements in the page at pageId
-in DST.
-
-The original database is left untouched.
-`, "\n")
-}

cmd/bbolt/surgery_commands_test.go

@@ -1,47 +0,0 @@
-package main_test
-
-import (
-	"fmt"
-	"path/filepath"
-	"testing"
-
-	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/require"
-
-	bolt "go.etcd.io/bbolt"
-	"go.etcd.io/bbolt/internal/btesting"
-	"go.etcd.io/bbolt/internal/common"
-)
-
-// TODO(ahrtr): add test case below for `surgery clear-page` command:
-//  1. The page is a branch page. All its children should become free pages.
-func TestSurgery_ClearPage(t *testing.T) {
-	pageSize := 4096
-	db := btesting.MustCreateDBWithOption(t, &bolt.Options{PageSize: pageSize})
-	srcPath := db.Path()
-
-	// Insert some sample data
-	t.Log("Insert some sample data")
-	err := db.Fill([]byte("data"), 1, 20,
-		func(tx int, k int) []byte { return []byte(fmt.Sprintf("%04d", k)) },
-		func(tx int, k int) []byte { return make([]byte, 10) },
-	)
-	require.NoError(t, err)
-
-	defer requireDBNoChange(t, dbData(t, srcPath), srcPath)
-
-	// clear page 3
-	t.Log("clear page 3")
-	dstPath := filepath.Join(t.TempDir(), "dstdb")
-	m := NewMain()
-	err = m.Run("surgery", "clear-page", srcPath, dstPath, "3")
-	require.NoError(t, err)
-
-	// The page 2 should have exactly the same data as page 3.
-	t.Log("Verify result")
-	dstPageId3Data := readPage(t, dstPath, 3, pageSize)
-
-	p := common.LoadPage(dstPageId3Data)
-	assert.Equal(t, uint16(0), p.Count())
-	assert.Equal(t, uint32(0), p.Overflow())
-}