dropwtxmgr: keep tx labels, with an optional DropLabels flag

cmd/dropwtxmgr/main.go

@@ -10,6 +10,7 @@ import (
 	"os"
 	"path/filepath"
 
+	"github.com/btcsuite/btcd/chaincfg/chainhash"
 	"github.com/btcsuite/btcutil"
 	"github.com/btcsuite/btcwallet/waddrmgr"
 	"github.com/btcsuite/btcwallet/walletdb"
@@ -24,8 +25,9 @@ var datadir = btcutil.AppDataDir("btcwallet", false)
 
 // Flags.
 var opts = struct {
-	Force  bool   `short:"f" description:"Force removal without prompt"`
-	DbPath string `long:"db" description:"Path to wallet database"`
+	Force      bool   `short:"f" description:"Force removal without prompt"`
+	DbPath     string `long:"db" description:"Path to wallet database"`
+	DropLabels bool   `long:"droplabels" description:"Drop transaction labels"`
 }{
 	Force:  false,
 	DbPath: filepath.Join(datadir, defaultNet, "wallet.db"),
@@ -42,6 +44,9 @@ var (
 	// Namespace keys.
 	waddrmgrNamespace = []byte("waddrmgr")
 	wtxmgrNamespace   = []byte("wtxmgr")
+
+	// Bucket names.
+	bucketTxLabels = []byte("l")
 )
 
 func yes(s string) bool {
@@ -109,6 +114,16 @@ func mainInt() int {
 	fmt.Println("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
+		if !opts.DropLabels {
+			labels, err = fetchAllLabels(tx)
+			if err != nil {
+				return err
+			}
+		}
+
 		err := tx.DeleteTopLevelBucket(wtxmgrNamespace)
 		if err != nil && err != walletdb.ErrBucketNotFound {
 			return err
@@ -122,6 +137,13 @@ func mainInt() int {
 			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 {
@@ -154,3 +176,69 @@ 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
+}