accounts: add new UpdateAccountBalanceAndExpiry Store method

In this commit, we remove one call to the UpdateAccount store method and
replace it with a call to a new UpdateAccountBalanceAndExpiry method
which updates an accounts balance and/or expiry fields and finds the
account via the given ID. This method signature is more appropriate for
a SQL backend than the UpdateAccount method.

Author: ellemouton

accounts/interface.go

@@ -7,6 +7,7 @@ import (
 	"fmt"
 	"time"
 
+	"github.com/lightningnetwork/lnd/fn"
 	"github.com/lightningnetwork/lnd/lnrpc"
 	"github.com/lightningnetwork/lnd/lntypes"
 	"github.com/lightningnetwork/lnd/lnwire"
@@ -219,6 +220,12 @@ type Store interface {
 	// Accounts retrieves all accounts from the store and un-marshals them.
 	Accounts(ctx context.Context) ([]*OffChainBalanceAccount, error)
 
+	// UpdateAccountBalanceAndExpiry updates the balance and/or expiry of an
+	// account.
+	UpdateAccountBalanceAndExpiry(ctx context.Context, id AccountID,
+		newBalance fn.Option[lnwire.MilliSatoshi],
+		newExpiry fn.Option[time.Time]) error
+
 	// RemoveAccount finds an account by its ID and removes it from the¨
 	// store.
 	RemoveAccount(ctx context.Context, id AccountID) error

accounts/service.go

@@ -9,8 +9,8 @@ import (
 
 	"github.com/btcsuite/btcd/chaincfg"
 	"github.com/lightninglabs/lndclient"
-	"github.com/lightninglabs/taproot-assets/fn"
 	"github.com/lightningnetwork/lnd/channeldb"
+	"github.com/lightningnetwork/lnd/fn"
 	invpkg "github.com/lightningnetwork/lnd/invoices"
 	"github.com/lightningnetwork/lnd/lnrpc"
 	"github.com/lightningnetwork/lnd/lntypes"
@@ -313,36 +313,35 @@ func (s *InterceptorService) UpdateAccount(ctx context.Context,
 		return nil, ErrAccountServiceDisabled
 	}
 
-	account, err := s.store.Account(ctx, accountID)
-	if err != nil {
-		return nil, fmt.Errorf("error fetching account: %w", err)
-	}
-
 	// If the expiration date was set, parse it as a unix time stamp. A
 	// value of -1 signals "don't update the expiration date".
+	var expiry fn.Option[time.Time]
 	if expirationDate > 0 {
-		account.ExpirationDate = time.Unix(expirationDate, 0)
+		expiry = fn.Some(time.Unix(expirationDate, 0))
 	} else if expirationDate == 0 {
 		// Setting the expiration to 0 means don't expire in which case
 		// we use a zero time (zero unix time would still be 1970, so
 		// that doesn't work for us).
-		account.ExpirationDate = time.Time{}
+		expiry = fn.Some(time.Time{})
 	}
 
 	// If the new account balance was set, parse it as millisatoshis. A
 	// value of -1 signals "don't update the balance".
+	var balance fn.Option[lnwire.MilliSatoshi]
 	if accountBalance >= 0 {
 		// Convert from satoshis to millisatoshis for storage.
-		account.CurrentBalance = int64(accountBalance) * 1000
+		balance = fn.Some(lnwire.MilliSatoshi(accountBalance) * 1000)
 	}
 
 	// Create the actual account in the macaroon account store.
-	err = s.store.UpdateAccount(ctx, account)
+	err := s.store.UpdateAccountBalanceAndExpiry(
+		ctx, accountID, balance, expiry,
+	)
 	if err != nil {
 		return nil, fmt.Errorf("unable to update account: %w", err)
 	}
 
-	return account, nil
+	return s.store.Account(ctx, accountID)
 }
 
 // Account retrieves an account from the bolt DB and un-marshals it. If the

accounts/store_kvdb.go

@@ -11,6 +11,7 @@ import (
 	"time"
 
 	"github.com/btcsuite/btcwallet/walletdb"
+	"github.com/lightningnetwork/lnd/fn"
 	"github.com/lightningnetwork/lnd/kvdb"
 	"github.com/lightningnetwork/lnd/lnwire"
 	"go.etcd.io/bbolt"
@@ -194,6 +195,56 @@ func (s *BoltStore) UpdateAccount(_ context.Context,
 	}, func() {})
 }
 
+// UpdateAccountBalanceAndExpiry updates the balance and/or expiry of an
+// account.
+//
+// NOTE: This is part of the Store interface.
+func (s *BoltStore) UpdateAccountBalanceAndExpiry(_ context.Context,
+	id AccountID, newBalance fn.Option[lnwire.MilliSatoshi],
+	newExpiry fn.Option[time.Time]) error {
+
+	update := func(account *OffChainBalanceAccount) error {
+		newBalance.WhenSome(func(balance lnwire.MilliSatoshi) {
+			account.CurrentBalance = int64(balance)
+		})
+		newExpiry.WhenSome(func(expiry time.Time) {
+			account.ExpirationDate = expiry
+		})
+
+		return nil
+	}
+
+	return s.updateAccount(id, update)
+}
+
+func (s *BoltStore) updateAccount(id AccountID,
+	updateFn func(*OffChainBalanceAccount) error) error {
+
+	return s.db.Update(func(tx kvdb.RwTx) error {
+		bucket := tx.ReadWriteBucket(accountBucketName)
+		if bucket == nil {
+			return ErrAccountBucketNotFound
+		}
+
+		account, err := getAccount(bucket, id)
+		if err != nil {
+			return fmt.Errorf("error fetching account, %w", err)
+		}
+
+		err = updateFn(account)
+		if err != nil {
+			return fmt.Errorf("error updating account, %w", err)
+		}
+
+		err = storeAccount(bucket, account)
+		if err != nil {
+			return fmt.Errorf("error storing account, %w", err)
+		}
+
+		return nil
+	}, func() {})
+}
+
 // storeAccount serializes and writes the given account to the given account
 // bucket.
 func storeAccount(accountBucket kvdb.RwBucket,
@@ -209,6 +260,19 @@ func storeAccount(accountBucket kvdb.RwBucket,
 	return accountBucket.Put(account.ID[:], accountBinary)
 }
 
+// getAccount retrieves an account from the given account bucket and
+// deserializes it.
+func getAccount(accountBucket kvdb.RwBucket, id AccountID) (
+	*OffChainBalanceAccount, error) {
+
+	accountBinary := accountBucket.Get(id[:])
+	if len(accountBinary) == 0 {
+		return nil, ErrAccNotFound
+	}
+
+	return deserializeAccount(accountBinary)
+}
+
 // uniqueRandomAccountID generates a new random ID and makes sure it does not
 // yet exist in the DB.
 func uniqueRandomAccountID(accountBucket kvdb.RBucket) (AccountID, error) {

accounts/store_test.go

@@ -5,8 +5,10 @@ import (
 	"testing"
 	"time"
 
+	"github.com/lightningnetwork/lnd/fn"
 	"github.com/lightningnetwork/lnd/lnrpc"
 	"github.com/lightningnetwork/lnd/lntypes"
+	"github.com/lightningnetwork/lnd/lnwire"
 	"github.com/stretchr/testify/require"
 )
 
@@ -105,6 +107,78 @@ func assertEqualAccounts(t *testing.T, expected,
 	actual.LastUpdate = actualUpdate
 }
 
+// TestAccountUpdateMethods tests that all the Store methods that update an
+// account work correctly.
+func TestAccountUpdateMethods(t *testing.T) {
+	t.Parallel()
+	ctx := context.Background()
+
+	t.Run("UpdateAccountBalanceAndExpiry", func(t *testing.T) {
+		store := NewTestDB(t)
+
+		// Ensure that the function errors out if we try update an
+		// account that does not exist.
+		err := store.UpdateAccountBalanceAndExpiry(
+			ctx, AccountID{}, fn.None[lnwire.MilliSatoshi](),
+			fn.None[time.Time](),
+		)
+		require.ErrorIs(t, err, ErrAccNotFound)
+
+		acct, err := store.NewAccount(ctx, 0, time.Time{}, "foo")
+		require.NoError(t, err)
+
+		assertBalanceAndExpiry := func(balance lnwire.MilliSatoshi,
+			expiry time.Time) {
+
+			dbAcct, err := store.Account(ctx, acct.ID)
+			require.NoError(t, err)
+			require.EqualValues(t, balance, dbAcct.CurrentBalance)
+			require.WithinDuration(
+				t, expiry, dbAcct.ExpirationDate, 0,
+			)
+		}
+
+		// Get the account from the store and check to see what its
+		// initial balance and expiry fields are set to.
+		assertBalanceAndExpiry(0, time.Time{})
+
+		// Now, update just the balance of the account.
+		newBalance := lnwire.MilliSatoshi(123)
+		err = store.UpdateAccountBalanceAndExpiry(
+			ctx, acct.ID, fn.Some(newBalance), fn.None[time.Time](),
+		)
+		require.NoError(t, err)
+		assertBalanceAndExpiry(newBalance, time.Time{})
+
+		// Now update just the expiry of the account.
+		newExpiry := time.Now().Add(time.Hour)
+		err = store.UpdateAccountBalanceAndExpiry(
+			ctx, acct.ID, fn.None[lnwire.MilliSatoshi](),
+			fn.Some(newExpiry),
+		)
+		require.NoError(t, err)
+		assertBalanceAndExpiry(newBalance, newExpiry)
+
+		// Update both the balance and expiry of the account.
+		newBalance = 456
+		newExpiry = time.Now().Add(2 * time.Hour)
+		err = store.UpdateAccountBalanceAndExpiry(
+			ctx, acct.ID, fn.Some(newBalance), fn.Some(newExpiry),
+		)
+		require.NoError(t, err)
+		assertBalanceAndExpiry(newBalance, newExpiry)
+
+		// Finally, test an update that has no net changes to the
+		// balance or expiry.
+		err = store.UpdateAccountBalanceAndExpiry(
+			ctx, acct.ID, fn.None[lnwire.MilliSatoshi](),
+			fn.None[time.Time](),
+		)
+		require.NoError(t, err)
+		assertBalanceAndExpiry(newBalance, newExpiry)
+	})
+}
+
 // TestLastInvoiceIndexes makes sure the last known invoice indexes can be
 // stored and retrieved correctly.
 func TestLastInvoiceIndexes(t *testing.T) {