Merge pull request #723 from guggero/drop-history

Export dropwtxmgr functionality as utility function

cmd/dropwtxmgr/main.go

@@ -10,12 +10,10 @@ import (
 	"os"
 	"path/filepath"
 
-	"github.com/btcsuite/btcd/chaincfg/chainhash"
 	"github.com/btcsuite/btcutil"
-	"github.com/btcsuite/btcwallet/waddrmgr"
+	"github.com/btcsuite/btcwallet/wallet"
 	"github.com/btcsuite/btcwallet/walletdb"
 	_ "github.com/btcsuite/btcwallet/walletdb/bdb"
-	"github.com/btcsuite/btcwallet/wtxmgr"
 	"github.com/jessevdk/go-flags"
 )
 
@@ -40,15 +38,6 @@ func init() {
 	}
 }
 
-var (
-	// Namespace keys.
-	waddrmgrNamespace = []byte("waddrmgr")
-	wtxmgrNamespace   = []byte("wtxmgr")
-
-	// Bucket names.
-	bucketTxLabels = []byte("l")
-)
-
 func yes(s string) bool {
 	switch s {
 	case "y", "Y", "yes", "Yes":
@@ -113,62 +102,7 @@ func mainInt() int {
 
 	fmt.Println("Dropping btcwallet transaction history")
 
-	err = walletdb.Update(db, func(tx walletdb.ReadWriteTx) error {
+	err = wallet.DropTransactionHistory(db, !opts.DropLabels)
-		// If we want to keep our tx labels, we read them out so we
-		// can re-add them after we have deleted our wtxmgr.
-		var labels map[chainhash.Hash]string
-		if !opts.DropLabels {
-			labels, err = fetchAllLabels(tx)
-			if err != nil {
-				return err
-			}
-		}
-
-		err := tx.DeleteTopLevelBucket(wtxmgrNamespace)
-		if err != nil && err != walletdb.ErrBucketNotFound {
-			return err
-		}
-		ns, err := tx.CreateTopLevelBucket(wtxmgrNamespace)
-		if err != nil {
-			return err
-		}
-		err = wtxmgr.Create(ns)
-		if err != nil {
-			return err
-		}
-
-		// If we want to re-add our labels, we do so now.
-		if !opts.DropLabels {
-			if err := putTxLabels(ns, labels); err != nil {
-				return err
-			}
-		}
-
-		ns = tx.ReadWriteBucket(waddrmgrNamespace)
-		birthdayBlock, err := waddrmgr.FetchBirthdayBlock(ns)
-		if err != nil {
-			fmt.Println("Wallet does not have a birthday block " +
-				"set, falling back to rescan from genesis")
-
-			startBlock, err := waddrmgr.FetchStartBlock(ns)
-			if err != nil {
-				return err
-			}
-			return waddrmgr.PutSyncedTo(ns, startBlock)
-		}
-
-		// We'll need to remove our birthday block first because it
-		// serves as a barrier when updating our state to detect reorgs
-		// due to the wallet not storing all block hashes of the chain.
-		if err := waddrmgr.DeleteBirthdayBlock(ns); err != nil {
-			return err
-		}
-
-		if err := waddrmgr.PutSyncedTo(ns, &birthdayBlock); err != nil {
-			return err
-		}
-		return waddrmgr.PutBirthdayBlock(ns, birthdayBlock)
-	})
 	if err != nil {
 		fmt.Println("Failed to drop and re-create namespace:", err)
 		return 1
@@ -176,69 +110,3 @@ func mainInt() int {
 
 	return 0
 }
-
-// fetchAllLabels returns a map of hex-encoded txid to label.
-func fetchAllLabels(tx walletdb.ReadWriteTx) (map[chainhash.Hash]string,
-	error) {
-
-	// Get our top level bucket, if it does not exist we just exit.
-	txBucket := tx.ReadBucket(wtxmgrNamespace)
-	if txBucket == nil {
-		return nil, nil
-	}
-
-	// If we do not have a labels bucket, there are no labels so we exit.
-	labelsBucket := txBucket.NestedReadBucket(bucketTxLabels)
-	if labelsBucket == nil {
-		return nil, nil
-	}
-
-	labels := make(map[chainhash.Hash]string)
-	if err := labelsBucket.ForEach(func(k, v []byte) error {
-		txid, err := chainhash.NewHash(k)
-		if err != nil {
-			return err
-		}
-
-		label, err := wtxmgr.DeserializeLabel(v)
-		if err != nil {
-			return err
-		}
-
-		// Add an entry to our map of labels.
-		labels[*txid] = label
-
-		return nil
-	}); err != nil {
-		return nil, err
-	}
-
-	return labels, nil
-}
-
-// putTxLabels re-adds a nested labels bucket and entries to the bucket provided
-// if there are any labels present.
-func putTxLabels(ns walletdb.ReadWriteBucket,
-	labels map[chainhash.Hash]string) error {
-
-	// If there are no labels, exit early.
-	if len(labels) == 0 {
-		return nil
-	}
-
-	// First, we create a labels bucket which we will add all labels to.
-	labelBucket, err := ns.CreateBucketIfNotExists(bucketTxLabels)
-	if err != nil {
-		return err
-	}
-
-	// Next, we re-add every label to the bucket.
-	for txid, label := range labels {
-		err := wtxmgr.PutTxLabel(labelBucket, txid, label)
-		if err != nil {
-			return err
-		}
-	}
-
-	return nil
-}

wallet/history.go

@@ -0,0 +1,161 @@
+// Copyright (c) 2015-2020 The btcsuite developers
+// Use of this source code is governed by an ISC
+// license that can be found in the LICENSE file.
+
+package wallet
+
+import (
+	"fmt"
+
+	"github.com/btcsuite/btcd/chaincfg/chainhash"
+	"github.com/btcsuite/btcwallet/waddrmgr"
+	"github.com/btcsuite/btcwallet/walletdb"
+	"github.com/btcsuite/btcwallet/wtxmgr"
+)
+
+var (
+	// bucketTxLabels is the name of the label sub bucket of the wtxmgr
+	// top level bucket that stores the mapping between a txid and a
+	// user-defined transaction label.
+	bucketTxLabels = []byte("l")
+)
+
+// DropTransactionHistory completely removes and re-creates the transaction
+// manager namespace from the given wallet database. This can be used to force
+// a full chain rescan of all wallet transaction and UTXO data. User-defined
+// transaction labels can optionally be kept by setting keepLabels to true.
+func DropTransactionHistory(db walletdb.DB, keepLabels bool) error {
+	log.Infof("Dropping btcwallet transaction history")
+
+	err := walletdb.Update(db, func(tx walletdb.ReadWriteTx) error {
+		// If we want to keep our tx labels, we read them out so we
+		// can re-add them after we have deleted our wtxmgr.
+		var (
+			labels map[chainhash.Hash]string
+			err    error
+		)
+		if keepLabels {
+			labels, err = fetchAllLabels(tx)
+			if err != nil {
+				return err
+			}
+		}
+
+		err = tx.DeleteTopLevelBucket(wtxmgrNamespaceKey)
+		if err != nil && err != walletdb.ErrBucketNotFound {
+			return err
+		}
+		ns, err := tx.CreateTopLevelBucket(wtxmgrNamespaceKey)
+		if err != nil {
+			return err
+		}
+		err = wtxmgr.Create(ns)
+		if err != nil {
+			return err
+		}
+
+		// If we want to re-add our labels, we do so now.
+		if keepLabels {
+			if err := putTxLabels(ns, labels); err != nil {
+				return err
+			}
+		}
+
+		ns = tx.ReadWriteBucket(waddrmgrNamespaceKey)
+		birthdayBlock, err := waddrmgr.FetchBirthdayBlock(ns)
+		if err != nil {
+			log.Warnf("Wallet does not have a birthday block " +
+				"set, falling back to rescan from genesis")
+
+			startBlock, err := waddrmgr.FetchStartBlock(ns)
+			if err != nil {
+				return err
+			}
+			return waddrmgr.PutSyncedTo(ns, startBlock)
+		}
+
+		// We'll need to remove our birthday block first because it
+		// serves as a barrier when updating our state to detect reorgs
+		// due to the wallet not storing all block hashes of the chain.
+		if err := waddrmgr.DeleteBirthdayBlock(ns); err != nil {
+			return err
+		}
+
+		if err := waddrmgr.PutSyncedTo(ns, &birthdayBlock); err != nil {
+			return err
+		}
+		return waddrmgr.PutBirthdayBlock(ns, birthdayBlock)
+	})
+	if err != nil {
+		return fmt.Errorf("failed to drop and re-create namespace: %v",
+			err)
+	}
+
+	return nil
+}
+
+// fetchAllLabels returns a map of hex-encoded txid to label.
+func fetchAllLabels(tx walletdb.ReadWriteTx) (map[chainhash.Hash]string,
+	error) {
+
+	// Get our top level bucket, if it does not exist we just exit.
+	txBucket := tx.ReadBucket(wtxmgrNamespaceKey)
+	if txBucket == nil {
+		return nil, nil
+	}
+
+	// If we do not have a labels bucket, there are no labels so we exit.
+	labelsBucket := txBucket.NestedReadBucket(bucketTxLabels)
+	if labelsBucket == nil {
+		return nil, nil
+	}
+
+	labels := make(map[chainhash.Hash]string)
+	if err := labelsBucket.ForEach(func(k, v []byte) error {
+		txid, err := chainhash.NewHash(k)
+		if err != nil {
+			return err
+		}
+
+		label, err := wtxmgr.DeserializeLabel(v)
+		if err != nil {
+			return err
+		}
+
+		// Add an entry to our map of labels.
+		labels[*txid] = label
+
+		return nil
+	}); err != nil {
+		return nil, err
+	}
+
+	return labels, nil
+}
+
+// putTxLabels re-adds a nested labels bucket and entries to the bucket provided
+// if there are any labels present.
+func putTxLabels(ns walletdb.ReadWriteBucket,
+	labels map[chainhash.Hash]string) error {
+
+	// If there are no labels, exit early.
+	if len(labels) == 0 {
+		return nil
+	}
+
+	// First, we create a labels bucket which we will add all labels to.
+	labelBucket, err := ns.CreateBucketIfNotExists(bucketTxLabels)
+	if err != nil {
+		return err
+	}
+
+	// Next, we re-add every label to the bucket.
+	for txid, label := range labels {
+		err := wtxmgr.PutTxLabel(labelBucket, txid, label)
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+}