Add db revert cmd (#2216)

cmd/juno/dbcmd.go

@@ -15,6 +15,10 @@ import (
 	"github.com/spf13/cobra"
 )
 
+const (
+	dbRevertToBlockF = "to-block"
+)
+
 type DBInfo struct {
 	Network         string     `json:"network"`
 	ChainHeight     uint64     `json:"chain_height"`
@@ -33,7 +37,7 @@ func DBCmd(defaultDBPath string) *cobra.Command {
 	}
 
 	dbCmd.PersistentFlags().String(dbPathF, defaultDBPath, dbPathUsage)
-	dbCmd.AddCommand(DBInfoCmd(), DBSizeCmd())
+	dbCmd.AddCommand(DBInfoCmd(), DBSizeCmd(), DBRevertCmd())
 	return dbCmd
 }
 
@@ -55,21 +59,29 @@ func DBSizeCmd() *cobra.Command {
 	}
 }
 
+func DBRevertCmd() *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   "revert",
+		Short: "Revert current head to given position",
+		Long:  `This subcommand revert all data related to all blocks till given so it becomes new head.`,
+		RunE:  dbRevert,
+	}
+	cmd.Flags().Uint64(dbRevertToBlockF, 0, "New head (this block won't be reverted)")
+
+	return cmd
+}
+
 func dbInfo(cmd *cobra.Command, args []string) error {
 	dbPath, err := cmd.Flags().GetString(dbPathF)
 	if err != nil {
 		return err
 	}
 
-	if _, err = os.Stat(dbPath); os.IsNotExist(err) {
-		fmt.Fprintln(cmd.OutOrStdout(), "Database path does not exist")
-		return nil
-	}
-
-	database, err := pebble.New(dbPath)
+	database, err := openDB(dbPath)
 	if err != nil {
-		return fmt.Errorf("open DB: %w", err)
+		return err
 	}
+	defer database.Close()
 
 	chain := blockchain.New(database, nil)
 	info := DBInfo{}
@@ -110,6 +122,50 @@ func dbInfo(cmd *cobra.Command, args []string) error {
 	return nil
 }
 
+func dbRevert(cmd *cobra.Command, args []string) error {
+	dbPath, err := cmd.Flags().GetString(dbPathF)
+	if err != nil {
+		return err
+	}
+
+	revertToBlock, err := cmd.Flags().GetUint64(dbRevertToBlockF)
+	if err != nil {
+		return err
+	}
+
+	if revertToBlock == 0 {
+		return fmt.Errorf("--%v cannot be 0", dbRevertToBlockF)
+	}
+
+	database, err := openDB(dbPath)
+	if err != nil {
+		return err
+	}
+	defer database.Close()
+
+	for {
+		chain := blockchain.New(database, nil)
+		head, err := chain.Head()
+		if err != nil {
+			return fmt.Errorf("failed to get the latest block information: %v", err)
+		}
+
+		if head.Number == revertToBlock {
+			fmt.Fprintf(cmd.OutOrStdout(), "Successfully reverted all blocks to %d\n", revertToBlock)
+			break
+		}
+
+		err = chain.RevertHead()
+		if err != nil {
+			return fmt.Errorf("failed to revert head at block %d: %v", head.Number, err)
+		}
+
+		fmt.Fprintf(cmd.OutOrStdout(), "Reverted head at block %d\n", head.Number)
+	}
+
+	return nil
+}
+
 func dbSize(cmd *cobra.Command, args []string) error {
 	dbPath, err := cmd.Flags().GetString(dbPathF)
 	if err != nil {
@@ -120,15 +176,11 @@ func dbSize(cmd *cobra.Command, args []string) error {
 		return fmt.Errorf("--%v cannot be empty", dbPathF)
 	}
 
-	if _, err = os.Stat(dbPath); os.IsNotExist(err) {
-		fmt.Fprintln(cmd.OutOrStdout(), "Database path does not exist")
-		return nil
-	}
-
-	pebbleDB, err := pebble.New(dbPath)
+	pebbleDB, err := openDB(dbPath)
 	if err != nil {
 		return err
 	}
+	defer pebbleDB.Close()
 
 	var (
 		totalSize  utils.DataSize
@@ -201,3 +253,17 @@ func getNetwork(head *core.Block, stateDiff *core.StateDiff) string {
 
 	return "unknown"
 }
+
+func openDB(path string) (db.DB, error) {
+	_, err := os.Stat(path)
+	if os.IsNotExist(err) {
+		return nil, fmt.Errorf("database path does not exist")
+	}
+
+	database, err := pebble.New(path)
+	if err != nil {
+		return nil, fmt.Errorf("failed to open db: %w", err)
+	}
+
+	return database, nil
+}

cmd/juno/dbcmd_test.go

@@ -2,6 +2,7 @@ package main_test
 
 import (
 	"context"
+	"strconv"
 	"testing"
 
 	"github.com/NethermindEth/juno/blockchain"
@@ -12,6 +13,7 @@ import (
 	adaptfeeder "github.com/NethermindEth/juno/starknetdata/feeder"
 	"github.com/NethermindEth/juno/utils"
 	"github.com/spf13/cobra"
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
@@ -27,27 +29,69 @@ func TestDBCmd(t *testing.T) {
 		cmd := juno.DBSizeCmd()
 		executeCmdInDB(t, cmd)
 	})
+
+	t.Run("revert db by 1 block", func(t *testing.T) {
+		network := utils.Mainnet
+
+		const (
+			syncToBlock   = uint64(2)
+			revertToBlock = syncToBlock - 1
+		)
+
+		cmd := juno.DBRevertCmd()
+		cmd.Flags().String("db-path", "", "")
+
+		dbPath := prepareDB(t, &network, syncToBlock)
+
+		require.NoError(t, cmd.Flags().Set("db-path", dbPath))
+		require.NoError(t, cmd.Flags().Set("to-block", strconv.Itoa(int(revertToBlock))))
+		require.NoError(t, cmd.Execute())
+
+		// unfortunately we cannot use blockchain from prepareDB because
+		// inside revert cmd another pebble instance is used which will panic if there are other instances
+		// that use the same db path
+		db, err := pebble.New(dbPath)
+		require.NoError(t, err)
+		t.Cleanup(func() {
+			require.NoError(t, db.Close())
+		})
+
+		chain := blockchain.New(db, &network)
+		block, err := chain.Head()
+		require.NoError(t, err)
+		assert.Equal(t, revertToBlock, block.Number)
+	})
 }
 
 func executeCmdInDB(t *testing.T, cmd *cobra.Command) {
 	cmd.Flags().String("db-path", "", "")
 
-	client := feeder.NewTestClient(t, &utils.Mainnet)
-	gw := adaptfeeder.New(client)
-	block0, err := gw.BlockByNumber(context.Background(), 0)
-	require.NoError(t, err)
+	dbPath := prepareDB(t, &utils.Mainnet, 0)
 
-	stateUpdate0, err := gw.StateUpdate(context.Background(), 0)
-	require.NoError(t, err)
+	require.NoError(t, cmd.Flags().Set("db-path", dbPath))
+	require.NoError(t, cmd.Execute())
+}
+
+func prepareDB(t *testing.T, network *utils.Network, syncToBlock uint64) string {
+	client := feeder.NewTestClient(t, network)
+	gw := adaptfeeder.New(client)
 
 	dbPath := t.TempDir()
 	testDB, err := pebble.New(dbPath)
 	require.NoError(t, err)
 
-	chain := blockchain.New(testDB, &utils.Mainnet)
-	require.NoError(t, chain.Store(block0, &emptyCommitments, stateUpdate0, nil))
-	testDB.Close()
+	chain := blockchain.New(testDB, network)
 
-	require.NoError(t, cmd.Flags().Set("db-path", dbPath))
-	require.NoError(t, cmd.Execute())
+	for blockNumber := uint64(0); blockNumber <= syncToBlock; blockNumber++ {
+		block, err := gw.BlockByNumber(context.Background(), blockNumber)
+		require.NoError(t, err)
+
+		stateUpdate, err := gw.StateUpdate(context.Background(), blockNumber)
+		require.NoError(t, err)
+
+		require.NoError(t, chain.Store(block, &emptyCommitments, stateUpdate, nil))
+	}
+	require.NoError(t, testDB.Close())
+
+	return dbPath
 }

docs/docs/configuring.md

@@ -109,6 +109,7 @@ Juno provides several subcommands to perform specific tasks or operations. Here
 - `db`: Perform database-related operations
   - `db info`: Retrieve information about the database.
   - `db size`: Calculate database size information for each data type.
+  - `db revert`: Reverts the database to a specific block number.
 
 To use a subcommand, append it when running Juno: