From acfa0868c56a7eb7c33685d3ab46f35c148d712a Mon Sep 17 00:00:00 2001
From: Mustafa Elbehery <melbeher@redhat.com>
Date: Mon, 22 Apr 2024 20:43:25 +0200
Subject: [PATCH] add pageId flag to check cmd

Signed-off-by: Mustafa Elbehery <melbeher@redhat.com>
---
 cmd/bbolt/command_check.go      | 21 ++++++++--
 cmd/bbolt/command_check_test.go | 69 ++++++++++++++++++++++++---------
 2 files changed, 69 insertions(+), 21 deletions(-)

diff --git a/cmd/bbolt/command_check.go b/cmd/bbolt/command_check.go
index 237e02f91..cb6e3b47d 100644
--- a/cmd/bbolt/command_check.go
+++ b/cmd/bbolt/command_check.go
@@ -4,25 +4,36 @@ import (
 	"fmt"
 
 	"github.com/spf13/cobra"
+	"github.com/spf13/pflag"
 
 	bolt "go.etcd.io/bbolt"
 	"go.etcd.io/bbolt/internal/guts_cli"
 )
 
+type checkOptions struct {
+	fromPageID uint64
+}
+
+func (o *checkOptions) AddFlags(fs *pflag.FlagSet) {
+	fs.Uint64VarP(&o.fromPageID, "from-page", "", o.fromPageID, "check db integrity starting from the given page ID")
+}
+
 func newCheckCommand() *cobra.Command {
+	var o checkOptions
 	checkCmd := &cobra.Command{
 		Use:   "check <bbolt-file>",
 		Short: "verify integrity of bbolt database data",
 		Args:  cobra.ExactArgs(1),
 		RunE: func(cmd *cobra.Command, args []string) error {
-			return checkFunc(cmd, args[0])
+			return checkFunc(cmd, args[0], o)
 		},
 	}
 
+	o.AddFlags(checkCmd.Flags())
 	return checkCmd
 }
 
-func checkFunc(cmd *cobra.Command, dbPath string) error {
+func checkFunc(cmd *cobra.Command, dbPath string, cfg checkOptions) error {
 	if _, err := checkSourceDBPath(dbPath); err != nil {
 		return err
 	}
@@ -37,10 +48,14 @@ func checkFunc(cmd *cobra.Command, dbPath string) error {
 	}
 	defer db.Close()
 
+	opts := []bolt.CheckOption{bolt.WithKVStringer(CmdKvStringer())}
+	if cfg.fromPageID != 0 {
+		opts = append(opts, bolt.WithPageId(cfg.fromPageID))
+	}
 	// Perform consistency check.
 	return db.View(func(tx *bolt.Tx) error {
 		var count int
-		for err := range tx.Check(bolt.WithKVStringer(CmdKvStringer())) {
+		for err := range tx.Check(opts...) {
 			fmt.Fprintln(cmd.OutOrStdout(), err)
 			count++
 		}
diff --git a/cmd/bbolt/command_check_test.go b/cmd/bbolt/command_check_test.go
index 02500745c..a2cdc6716 100644
--- a/cmd/bbolt/command_check_test.go
+++ b/cmd/bbolt/command_check_test.go
@@ -9,25 +9,58 @@ import (
 
 	main "go.etcd.io/bbolt/cmd/bbolt"
 	"go.etcd.io/bbolt/internal/btesting"
+	"go.etcd.io/bbolt/internal/guts_cli"
 )
 
 func TestCheckCommand_Run(t *testing.T) {
-	db := btesting.MustCreateDB(t)
-	db.Close()
-	defer requireDBNoChange(t, dbData(t, db.Path()), db.Path())
-
-	rootCmd := main.NewRootCommand()
-	// capture output for assertion
-	outputBuf := bytes.NewBufferString("")
-	rootCmd.SetOut(outputBuf)
-
-	rootCmd.SetArgs([]string{
-		"check", db.Path(),
-	})
-	err := rootCmd.Execute()
-	require.NoError(t, err)
-
-	output, err := io.ReadAll(outputBuf)
-	require.NoError(t, err)
-	require.Equalf(t, "OK\n", string(output), "unexpected stdout:\n\n%s", string(output))
+	testCases := []struct {
+		name      string
+		args      []string
+		expErr    error
+		expOutput string
+	}{
+		{
+			name:      "check whole db",
+			args:      []string{"check", "path"},
+			expErr:    nil,
+			expOutput: "OK\n",
+		},
+		{
+			name:      "check valid pageId",
+			args:      []string{"check", "path", "--from-page", "3"},
+			expErr:    nil,
+			expOutput: "OK\n",
+		},
+		{
+			name:      "check invalid pageId",
+			args:      []string{"check", "path", "--from-page", "1"},
+			expErr:    guts_cli.ErrCorrupt,
+			expOutput: "page ID (1) out of range [2, 4)",
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+
+			t.Log("Creating sample DB")
+			db := btesting.MustCreateDB(t)
+			db.Close()
+			defer requireDBNoChange(t, dbData(t, db.Path()), db.Path())
+
+			t.Log("Running check cmd")
+			rootCmd := main.NewRootCommand()
+			outputBuf := bytes.NewBufferString("") // capture output for assertion
+			rootCmd.SetOut(outputBuf)
+
+			tc.args[1] = db.Path() // path to be replaced with db.Path()
+			rootCmd.SetArgs(tc.args)
+			err := rootCmd.Execute()
+			require.Equal(t, tc.expErr, err)
+
+			t.Log("Checking output")
+			output, err := io.ReadAll(outputBuf)
+			require.NoError(t, err)
+			require.Containsf(t, string(output), tc.expOutput, "unexpected stdout:\n\n%s", string(output))
+		})
+	}
 }