From 178152bcd036ea378c34900576535cebb11f5b00 Mon Sep 17 00:00:00 2001
From: Oliver Gugger <gugger@gmail.com>
Date: Mon, 19 Jul 2021 13:30:40 +0200
Subject: [PATCH] wallet: refactor ImportAccountDryRun to use walletdb.Update

To make sure we don't create any manual DB transactions, we refactor the
ImportAccountDryRun method to use walletdb.Update and the new
walletdb.ErrDryRunRollBack error for making sure a rollback is issued.
---
 wallet/import.go | 103 ++++++++++++++++++++++++++---------------------
 1 file changed, 58 insertions(+), 45 deletions(-)

diff --git a/wallet/import.go b/wallet/import.go
index af2c3f0..9d2dcfa 100644
--- a/wallet/import.go
+++ b/wallet/import.go
@@ -245,57 +245,70 @@ func (w *Wallet) ImportAccountDryRun(name string,
 	*waddrmgr.AccountProperties, []waddrmgr.ManagedAddress,
 	[]waddrmgr.ManagedAddress, error) {
 
+	var (
+		accountProps  *waddrmgr.AccountProperties
+		externalAddrs []waddrmgr.ManagedAddress
+		internalAddrs []waddrmgr.ManagedAddress
+	)
+
 	// Start a database transaction that we'll never commit and always
-	// rollback.
-	tx, err := w.db.BeginReadWriteTx()
-	if err != nil {
-		return nil, nil, nil, err
-	}
-	defer func() {
-		_ = tx.Rollback()
-	}()
-	ns := tx.ReadWriteBucket(waddrmgrNamespaceKey)
+	// rollback because we'll return a specific error in the end.
+	err := walletdb.Update(w.db, func(tx walletdb.ReadWriteTx) error {
+		ns := tx.ReadWriteBucket(waddrmgrNamespaceKey)
 
-	// Import the account as usual.
-	accountProps, err := w.importAccount(
-		ns, name, accountPubKey, masterKeyFingerprint, addrType,
-	)
-	if err != nil {
-		return nil, nil, nil, err
-	}
+		// Import the account as usual.
+		var err error
+		accountProps, err = w.importAccount(
+			ns, name, accountPubKey, masterKeyFingerprint, addrType,
+		)
+		if err != nil {
+			return err
+		}
 
-	// Derive the external and internal addresses. Note that we could do
-	// this based on the provided accountPubKey alone, but we go through the
-	// ScopedKeyManager instead to ensure addresses will be derived as
-	// expected from the wallet's point-of-view.
-	manager, err := w.Manager.FetchScopedKeyManager(accountProps.KeyScope)
-	if err != nil {
-		return nil, nil, nil, err
-	}
+		// Derive the external and internal addresses. Note that we
+		// could do this based on the provided accountPubKey alone, but
+		// we go through the ScopedKeyManager instead to ensure
+		// addresses will be derived as expected from the wallet's
+		// point-of-view.
+		manager, err := w.Manager.FetchScopedKeyManager(
+			accountProps.KeyScope,
+		)
+		if err != nil {
+			return err
+		}
 
-	// The importAccount method above will cache the imported account within
-	// the scoped manager. Since this is a dry-run attempt, we'll want to
-	// invalidate the cache for it.
-	defer manager.InvalidateAccountCache(accountProps.AccountNumber)
+		// The importAccount method above will cache the imported
+		// account within the scoped manager. Since this is a dry-run
+		// attempt, we'll want to invalidate the cache for it.
+		defer manager.InvalidateAccountCache(accountProps.AccountNumber)
 
-	externalAddrs, err := manager.NextExternalAddresses(
-		ns, accountProps.AccountNumber, numAddrs,
-	)
-	if err != nil {
-		return nil, nil, nil, err
-	}
-	internalAddrs, err := manager.NextInternalAddresses(
-		ns, accountProps.AccountNumber, numAddrs,
-	)
-	if err != nil {
-		return nil, nil, nil, err
-	}
+		externalAddrs, err = manager.NextExternalAddresses(
+			ns, accountProps.AccountNumber, numAddrs,
+		)
+		if err != nil {
+			return err
+		}
+		internalAddrs, err = manager.NextInternalAddresses(
+			ns, accountProps.AccountNumber, numAddrs,
+		)
+		if err != nil {
+			return err
+		}
 
-	// Refresh the account's properties after generating the addresses.
-	accountProps, err = manager.AccountProperties(
-		ns, accountProps.AccountNumber,
-	)
-	if err != nil {
+		// Refresh the account's properties after generating the
+		// addresses.
+		accountProps, err = manager.AccountProperties(
+			ns, accountProps.AccountNumber,
+		)
+		if err != nil {
+			return err
+		}
+
+		// Make sure we always roll back the dry-run transaction by
+		// returning an error here.
+		return walletdb.ErrDryRunRollBack
+	})
+	if err != nil && err != walletdb.ErrDryRunRollBack {
 		return nil, nil, nil, err
 	}