Make fundRawTransactionInternal use DBIOActions (#4575)

* Make fundRawTransactionInternal use DBIOActions

* Fix to correctly use callback

* Move reservedUTXOsCallbackF to FundRawTxHelper

core/src/main/scala/org/bitcoins/core/wallet/builder/FundRawTxHelper.scala

@@ -4,10 +4,13 @@ import org.bitcoins.core.protocol.transaction.Transaction
 import org.bitcoins.core.wallet.fee.FeeUnit
 import org.bitcoins.core.wallet.utxo.{InputInfo, ScriptSignatureParams}
 
+import scala.concurrent.Future
+
 case class FundRawTxHelper[T <: RawTxFinalizer](
     txBuilderWithFinalizer: RawTxBuilderWithFinalizer[T],
     scriptSigParams: Vector[ScriptSignatureParams[InputInfo]],
-    feeRate: FeeUnit) {
+    feeRate: FeeUnit,
+    reservedUTXOsCallbackF: Future[Unit]) {
 
   /** Produces the unsigned transaction built by fundrawtransaction */
   def unsignedTx: Transaction = txBuilderWithFinalizer.buildTx()

dlc-wallet/src/main/scala/org/bitcoins/dlc/wallet/DLCWallet.scala

@@ -103,7 +103,7 @@ abstract class DLCWallet
     DLCActionBuilder(dlcWalletDAOs)
   }
 
-  private lazy val safeDatabase: SafeDatabase = dlcDAO.safeDatabase
+  private lazy val safeDLCDatabase: SafeDatabase = dlcDAO.safeDatabase
   private lazy val walletDatabase: SafeDatabase = addressDAO.safeDatabase
 
   /** Updates the contract Id in the wallet database for the given offer and accept */
@@ -224,7 +224,7 @@ abstract class DLCWallet
     val updateOracleSigsA =
       actionBuilder.updateDLCOracleSigsAction(outcomeAndSigByNonce)
     for {
-      updates <- safeDatabase.runVec(updateOracleSigsA)
+      updates <- safeDLCDatabase.runVec(updateOracleSigsA)
     } yield updates
   }
 
@@ -269,7 +269,7 @@ abstract class DLCWallet
       // allow this to fail in the case they have already been unreserved
       _ <- unmarkUTXOsAsReserved(dbs).recover { case _: Throwable => () }
       action = actionBuilder.deleteDLCAction(dlcId)
-      _ <- safeDatabase.run(action)
+      _ <- safeDLCDatabase.run(action)
     } yield ()
   }
 
@@ -484,7 +484,7 @@ abstract class DLCWallet
         dlcInputs = dlcInputs,
         dlcOfferDb = dlcOfferDb)
 
-      _ <- safeDatabase.run(offerActions)
+      _ <- safeDLCDatabase.run(offerActions)
       status <- findDLC(dlcId)
       _ <- dlcConfig.walletCallbacks.executeOnDLCStateChange(logger, status.get)
     } yield offer
@@ -624,7 +624,7 @@ abstract class DLCWallet
            acceptDb,
            inputsDb,
            contractDataDb) <-
-            safeDatabase.run(zipped)
+            safeDLCDatabase.run(zipped)
           announcementDataDbs =
             createdDbs ++ groupedAnnouncements.existingAnnouncements
 
@@ -652,7 +652,7 @@ abstract class DLCWallet
           createAnnouncementAction = dlcAnnouncementDAO.createAllAction(
             dlcAnnouncementDbs)
 
-          _ <- safeDatabase.run(
+          _ <- safeDLCDatabase.run(
             DBIOAction.seq(createNonceAction, createAnnouncementAction))
         } yield {
           InitializedAccept(
@@ -867,7 +867,7 @@ abstract class DLCWallet
           cetSigsDb = sigsDbs,
           refundSigsDb = refundSigsDb
         )
-        _ <- safeDatabase.run(actions)
+        _ <- safeDLCDatabase.run(actions)
         dlcDb <- updateDLCContractIds(offer, accept)
         _ = logger.info(
           s"Created DLCAccept for tempContractId ${offer.tempContractId.hex} with contract Id ${contractId.toHex}")
@@ -968,7 +968,7 @@ abstract class DLCWallet
                    sigsAction,
                    refundSigAction,
                    acceptDbAction))
-          _ <- safeDatabase.run(actions)
+          _ <- safeDLCDatabase.run(actions)
 
           // .get is safe here because we must have an offer if we have a dlcDAO
           offerDb <- dlcOfferDAO.findByDLCId(dlc.dlcId).map(_.head)
@@ -1733,7 +1733,7 @@ abstract class DLCWallet
         dlc
       }
     }
-    safeDatabase.run(dlcAction).flatMap { intermediaries =>
+    safeDLCDatabase.run(dlcAction).flatMap { intermediaries =>
       val actions = intermediaries.map { intermediary =>
         getWalletDLCDbsAction(intermediary).map {
           case (closingTxOpt, payoutAddrOpt) =>
@@ -1778,7 +1778,7 @@ abstract class DLCWallet
   }
 
   override def findDLC(dlcId: Sha256Digest): Future[Option[DLCStatus]] = {
-    val intermediaryF = safeDatabase.run(findDLCAction(dlcId))
+    val intermediaryF = safeDLCDatabase.run(findDLCAction(dlcId))
 
     intermediaryF.flatMap {
       case None => Future.successful(None)

dlc-wallet/src/main/scala/org/bitcoins/dlc/wallet/internal/DLCTransactionProcessing.scala

@@ -32,7 +32,7 @@ import scala.concurrent._
   */
 private[bitcoins] trait DLCTransactionProcessing extends TransactionProcessing {
   self: DLCWallet =>
-  private lazy val safeDatabase: SafeDatabase = dlcDAO.safeDatabase
+  private lazy val safeDLCDatabase: SafeDatabase = dlcDAO.safeDatabase
 
   /** Calculates the new state of the DLCDb based on the closing transaction,
     * will delete old CET sigs that are no longer needed after execution
@@ -224,7 +224,7 @@ private[bitcoins] trait DLCTransactionProcessing extends TransactionProcessing {
             _ <- updateAnnouncementA
           } yield updatedDlcDb
         }
-        updatedDlcDb <- safeDatabase.run(actions)
+        updatedDlcDb <- safeDLCDatabase.run(actions)
       } yield {
         logger.info(
           s"Done calculating RemoteClaimed outcome for dlcId=${dlcId.hex}")

wallet-test/src/test/scala/org/bitcoins/wallet/WalletSendingTest.scala

@@ -525,7 +525,8 @@ class WalletSendingTest extends BitcoinSWalletTest {
         recoverToExceptionIf[RuntimeException](failedTx)
 
       exnF.map(err =>
-        assert(err.getMessage.contains("Failed to reserve all utxos")))
+        assert(err.getMessage.contains(
+          "Not enough value in given outputs to make transaction spending 599500000 sats plus fees")))
 
   }
 }

wallet/src/main/scala/org/bitcoins/wallet/Wallet.scala

@@ -94,7 +94,7 @@ abstract class Wallet
   private[bitcoins] val stateDescriptorDAO: WalletStateDescriptorDAO =
     WalletStateDescriptorDAO()
 
-  private val safeDatabase: SafeDatabase = spendingInfoDAO.safeDatabase
+  protected lazy val safeDatabase: SafeDatabase = spendingInfoDAO.safeDatabase
   val nodeApi: NodeApi
   val chainQueryApi: ChainQueryApi
   val creationTime: Instant = keyManager.creationTime
@@ -447,7 +447,7 @@ abstract class Wallet
       _ = require(
         tmp.outputs.size == 1,
         s"Created tx is not as expected, does not have 1 output, got $tmp")
-      rawTxHelper = FundRawTxHelper(withFinalizer, utxos, feeRate)
+      rawTxHelper = FundRawTxHelper(withFinalizer, utxos, feeRate, Future.unit)
       tx <- finishSend(rawTxHelper,
                        tmp.outputs.head.value,
                        feeRate,
@@ -495,7 +495,7 @@ abstract class Wallet
         utxos,
         feeRate,
         changeAddr.scriptPubKey)
-      rawTxHelper = FundRawTxHelper(txBuilder, utxos, feeRate)
+      rawTxHelper = FundRawTxHelper(txBuilder, utxos, feeRate, Future.unit)
       tx <- finishSend(rawTxHelper, amount, feeRate, newTags)
     } yield tx
   }
@@ -596,7 +596,10 @@ abstract class Wallet
                                                                 sequence)
 
       amount = outputs.foldLeft(CurrencyUnits.zero)(_ + _.value)
-      rawTxHelper = FundRawTxHelper(txBuilder, spendingInfos, newFeeRate)
+      rawTxHelper = FundRawTxHelper(txBuilder,
+                                    spendingInfos,
+                                    newFeeRate,
+                                    Future.unit)
       tx <-
         finishSend(rawTxHelper, amount, newFeeRate, Vector.empty)
     } yield tx

wallet/src/main/scala/org/bitcoins/wallet/internal/AddressHandling.scala

@@ -19,7 +19,6 @@ import org.bitcoins.core.wallet.utxo.{
   AddressTagType
 }
 import org.bitcoins.crypto.ECPublicKey
-import org.bitcoins.db.SafeDatabase
 import org.bitcoins.wallet._
 import slick.dbio.{DBIOAction, Effect, NoStream}
 
@@ -32,8 +31,6 @@ import scala.util.{Failure, Success}
 private[wallet] trait AddressHandling extends WalletLogger {
   self: Wallet =>
 
-  private lazy val safeDatabase: SafeDatabase = addressDAO.safeDatabase
-
   def contains(
       address: BitcoinAddress,
       accountOpt: Option[HDAccount]): Future[Boolean] = {
@@ -295,6 +292,13 @@ private[wallet] trait AddressHandling extends WalletLogger {
     getNewAddressHelperAction(account, HDChainType.External)
   }
 
+  def getNewChangeAddressAction(account: AccountDb): DBIOAction[
+    BitcoinAddress,
+    NoStream,
+    Effect.Read with Effect.Write with Effect.Transactional] = {
+    getNewAddressHelperAction(account, HDChainType.Change)
+  }
+
   def getNewAddress(account: AccountDb): Future[BitcoinAddress] = {
     safeDatabase.run(getNewAddressAction(account))
   }

wallet/src/main/scala/org/bitcoins/wallet/internal/FundTransactionHandling.scala

@@ -1,7 +1,8 @@
 package org.bitcoins.wallet.internal
 
-import org.bitcoins.core.api.wallet.db.{AccountDb, SpendingInfoDb}
 import org.bitcoins.core.api.wallet._
+import org.bitcoins.core.api.wallet.db.AccountDb
+import org.bitcoins.core.hd.HDAccount
 import org.bitcoins.core.policy.Policy
 import org.bitcoins.core.protocol.transaction._
 import org.bitcoins.core.wallet.builder._
@@ -10,9 +11,9 @@ import org.bitcoins.core.wallet.utxo._
 import org.bitcoins.wallet.{Wallet, WalletLogger}
 
 import scala.concurrent.Future
-import scala.util.control.NonFatal
 
 trait FundTransactionHandling extends WalletLogger { self: Wallet =>
+  import walletConfig.profile.api._
 
   def fundRawTransaction(
       destinations: Vector[TransactionOutput],
@@ -71,24 +72,52 @@ trait FundTransactionHandling extends WalletLogger { self: Wallet =>
       fromTagOpt: Option[AddressTag],
       markAsReserved: Boolean): Future[
     FundRawTxHelper[ShufflingNonInteractiveFinalizer]] = {
+    val action = fundRawTransactionInternalAction(destinations,
+                                                  feeRate,
+                                                  fromAccount,
+                                                  coinSelectionAlgo,
+                                                  fromTagOpt,
+                                                  markAsReserved)
+
+    for {
+      txHelper <- safeDatabase.run(action)
+      _ <- txHelper.reservedUTXOsCallbackF
+    } yield txHelper
+  }
+
+  private[bitcoins] def fundRawTransactionInternalAction(
+      destinations: Vector[TransactionOutput],
+      feeRate: FeeUnit,
+      fromAccount: AccountDb,
+      coinSelectionAlgo: CoinSelectionAlgo = CoinSelectionAlgo.LeastWaste,
+      fromTagOpt: Option[AddressTag],
+      markAsReserved: Boolean): DBIOAction[
+    FundRawTxHelper[ShufflingNonInteractiveFinalizer],
+    NoStream,
+    Effect.Read with Effect.Write with Effect.Transactional] = {
     val amts = destinations.map(_.value)
     //need to allow 0 for OP_RETURN outputs
     require(amts.forall(_.satoshis.toBigInt >= 0),
             s"Cannot fund a transaction for a negative amount, got=$amts")
     val amt = amts.sum
-    logger.info(s"Attempting to fund a tx for amt=${amt} with feeRate=$feeRate")
+    logger.info(s"Attempting to fund a tx for amt=$amt with feeRate=$feeRate")
-    val utxosF: Future[Vector[(SpendingInfoDb, Transaction)]] =
+    val utxosA =
       for {
         utxos <- fromTagOpt match {
           case None =>
-            listUtxos(fromAccount.hdAccount)
+            spendingInfoDAO.findAllUnspentForAccountAction(
+              fromAccount.hdAccount)
           case Some(tag) =>
-            listUtxos(fromAccount.hdAccount, tag)
+            spendingInfoDAO.findAllUnspentForTagAction(tag).map { utxos =>
+              utxos.filter(utxo =>
+                HDAccount.isSameAccount(bip32Path = utxo.privKeyPath,
+                                        account = fromAccount.hdAccount))
             }
-        utxoWithTxs <- Future.sequence {
+        }
+        utxoWithTxs <- DBIO.sequence {
           utxos.map { utxo =>
             transactionDAO
-              .findByOutPoint(utxo.outPoint)
+              .findByTxIdAction(utxo.outPoint.txIdBE)
               .map(tx => (utxo, tx.get.transaction))
           }
         }
@@ -99,9 +128,9 @@ trait FundTransactionHandling extends WalletLogger { self: Wallet =>
       } yield utxoWithTxs.filter(utxo =>
         !immatureCoinbases.exists(_._1 == utxo._1))
 
-    val selectedUtxosF: Future[Vector[(SpendingInfoDb, Transaction)]] =
+    val selectedUtxosA =
       for {
-        walletUtxos <- utxosF
+        walletUtxos <- utxosA
 
         // filter out dust
         selectableUtxos = walletUtxos
@@ -118,14 +147,14 @@ trait FundTransactionHandling extends WalletLogger { self: Wallet =>
         )
         filtered = walletUtxos.filter(utxo =>
           utxos.exists(_.outPoint == utxo._1.outPoint))
-        _ <-
+        (_, callbackF) <-
-          if (markAsReserved) markUTXOsAsReserved(filtered.map(_._1))
+          if (markAsReserved) markUTXOsAsReservedAction(filtered.map(_._1))
-          else Future.unit
+          else DBIO.successful((Vector.empty, Future.unit))
-      } yield filtered
+      } yield (filtered, callbackF)
 
-    val resultF = for {
+    for {
-      selectedUtxos <- selectedUtxosF
+      (selectedUtxos, callbackF) <- selectedUtxosA
-      change <- getNewChangeAddress(fromAccount)
+      change <- getNewChangeAddressAction(fromAccount)
       utxoSpendingInfos = {
         selectedUtxos.map { case (utxo, prevTx) =>
           utxo.toUTXOInfo(keyManager = self.keyManager, prevTx)
@@ -147,23 +176,12 @@ trait FundTransactionHandling extends WalletLogger { self: Wallet =>
                                                        feeRate,
                                                        change.scriptPubKey)
 
-      FundRawTxHelper(txBuilderWithFinalizer = txBuilder,
+      val fundTxHelper = FundRawTxHelper(txBuilderWithFinalizer = txBuilder,
                                          scriptSigParams = utxoSpendingInfos,
-                      feeRate)
+                                         feeRate = feeRate,
-    }
+                                         reservedUTXOsCallbackF = callbackF)
 
-    resultF.recoverWith { case NonFatal(error) =>
+      fundTxHelper
-      logger.error(
+    }
-        s"Failed to reserve utxos for amount=${amt} feeRate=$feeRate, unreserving the selected utxos")
-      // un-reserve utxos since we failed to create valid spending infos
-      if (markAsReserved) {
-        for {
-          utxos <- selectedUtxosF
-          _ <- unmarkUTXOsAsReserved(utxos.map(_._1))
-        } yield error
-      } else Future.failed(error)
-    }
-
-    resultF
   }
 }

wallet/src/main/scala/org/bitcoins/wallet/internal/RescanHandling.scala

@@ -15,7 +15,6 @@ import org.bitcoins.core.protocol.script.ScriptPubKey
 import org.bitcoins.core.protocol.{BitcoinAddress, BlockStamp}
 import org.bitcoins.core.wallet.rescan.RescanState
 import org.bitcoins.crypto.DoubleSha256Digest
-import org.bitcoins.db.SafeDatabase
 import org.bitcoins.wallet.{Wallet, WalletLogger}
 import slick.dbio.{DBIOAction, Effect, NoStream}
 
@@ -28,7 +27,6 @@ private[wallet] trait RescanHandling extends WalletLogger {
   /////////////////////
   // Public facing API
 
-  private lazy val safeDatabase: SafeDatabase = addressDAO.safeDatabase
   override def isRescanning(): Future[Boolean] = stateDescriptorDAO.isRescanning
 
   /** @inheritdoc */

wallet/src/main/scala/org/bitcoins/wallet/internal/TransactionProcessing.scala

@@ -17,7 +17,6 @@ import org.bitcoins.core.wallet.fee.FeeUnit
 import org.bitcoins.core.wallet.utxo.TxoState._
 import org.bitcoins.core.wallet.utxo.{AddressTag, TxoState}
 import org.bitcoins.crypto.{DoubleSha256Digest, DoubleSha256DigestBE}
-import org.bitcoins.db.SafeDatabase
 import org.bitcoins.wallet._
 
 import scala.concurrent.{Future, Promise}
@@ -33,8 +32,6 @@ private[bitcoins] trait TransactionProcessing extends WalletLogger {
 
   import walletConfig.profile.api._
 
-  private lazy val safeDatabase: SafeDatabase = transactionDAO.safeDatabase
-
   /////////////////////
   // Public facing API
 

wallet/src/main/scala/org/bitcoins/wallet/internal/UtxoHandling.scala

@@ -4,12 +4,7 @@ import org.bitcoins.core.api.wallet.db._
 import org.bitcoins.core.consensus.Consensus
 import org.bitcoins.core.hd.HDAccount
 import org.bitcoins.core.protocol.script.{P2WPKHWitnessSPKV0, P2WPKHWitnessV0}
-import org.bitcoins.core.protocol.transaction.{
+import org.bitcoins.core.protocol.transaction._
-  CoinbaseInput,
-  Transaction,
-  TransactionOutPoint,
-  TransactionOutput
-}
 import org.bitcoins.core.util.BlockHashWithConfs
 import org.bitcoins.core.wallet.utxo.TxoState._
 import org.bitcoins.core.wallet.utxo._
@@ -25,6 +20,7 @@ import scala.concurrent.Future
   */
 private[wallet] trait UtxoHandling extends WalletLogger {
   self: Wallet =>
+  import walletConfig.profile.api._
 
   /** @inheritdoc */
   def listDefaultAccountUtxos(): Future[Vector[SpendingInfoDb]] =
@@ -296,13 +292,24 @@ private[wallet] trait UtxoHandling extends WalletLogger {
 
   override def markUTXOsAsReserved(
       utxos: Vector[SpendingInfoDb]): Future[Vector[SpendingInfoDb]] = {
+    for {
+      (utxos, callbackF) <- safeDatabase.run(markUTXOsAsReservedAction(utxos))
+      _ <- callbackF
+    } yield utxos
+  }
+
+  protected def markUTXOsAsReservedAction(
+      utxos: Vector[SpendingInfoDb]): DBIOAction[
+    (Vector[SpendingInfoDb], Future[Unit]),
+    NoStream,
+    Effect.Read with Effect.Write] = {
     val outPoints = utxos.map(_.outPoint)
     logger.info(s"Reserving utxos=$outPoints")
     val updated = utxos.map(_.copyWithState(TxoState.Reserved))
-    for {
+    spendingInfoDAO.markAsReservedAction(updated).map { utxos =>
-      utxos <- spendingInfoDAO.markAsReserved(updated)
+      val callbackF = walletCallbacks.executeOnReservedUtxos(logger, utxos)
-      _ <- walletCallbacks.executeOnReservedUtxos(logger, utxos)
+      (utxos, callbackF)
-    } yield utxos
+    }
   }
 
   /** @inheritdoc */

wallet/src/main/scala/org/bitcoins/wallet/models/SpendingInfoDAO.scala

@@ -457,9 +457,16 @@ case class SpendingInfoDAO()(implicit
       UTXORecord.fromSpendingInfoDb(utxo, spks(utxo.output.scriptPubKey)))
 
   def findAllUnspent(): Future[Vector[SpendingInfoDb]] = {
+    safeDatabase.run(findAllUnspentAction())
+  }
+
+  def findAllUnspentAction(): DBIOAction[
+    Vector[SpendingInfoDb],
+    NoStream,
+    Effect.Read] = {
     for {
-      utxos <- _findAllUnspent()
+      utxos <- _findAllUnspentAction()
-      infos <- utxoToInfo(utxos)
+      infos <- utxoToInfoAction(utxos)
     } yield infos
   }
 
@@ -510,6 +517,14 @@ case class SpendingInfoDAO()(implicit
     allUtxosF.map(filterUtxosByAccount(_, hdAccount))
   }
 
+  def findAllUnspentForAccountAction(hdAccount: HDAccount): DBIOAction[
+    Vector[SpendingInfoDb],
+    NoStream,
+    Effect.Read] = {
+    val allUtxosA = findAllUnspentAction()
+    allUtxosA.map(filterUtxosByAccount(_, hdAccount))
+  }
+
   def findAllForAccountAction(hdAccount: HDAccount): DBIOAction[
     Vector[SpendingInfoDb],
     NoStream,
@@ -597,8 +612,11 @@ case class SpendingInfoDAO()(implicit
         })
   }
 
-  def findAllUnspentForTag(tag: AddressTag): Future[Vector[SpendingInfoDb]] = {
+  def findAllUnspentForTagAction(tag: AddressTag): DBIOAction[
-    val query = table
+    Vector[SpendingInfoDb],
+    NoStream,
+    Effect.Read] = {
+    table
       .join(spkTable)
       .on(_.scriptPubKeyId === _.id)
       .filter(_._1.state.inSet(TxoState.receivedStates))
@@ -608,24 +626,26 @@ case class SpendingInfoDAO()(implicit
       .on(_._2.address === _.address)
       .filter(_._2.tagName === tag.tagName)
       .filter(_._2.tagType === tag.tagType)
-
+      .result
-    safeDatabase
+      .map(_.toVector)
-      .runVec(query.result)
       .map(_.map { case (((utxoRecord, spkDb), _), _) =>
         utxoRecord.toSpendingInfoDb(spkDb.scriptPubKey)
       })
   }
 
-  def markAsReserved(
+  def findAllUnspentForTag(tag: AddressTag): Future[Vector[SpendingInfoDb]] = {
-      ts: Vector[SpendingInfoDb]): Future[Vector[SpendingInfoDb]] = {
+    safeDatabase.run(findAllUnspentForTagAction(tag))
+  }
+
+  def markAsReservedAction(ts: Vector[SpendingInfoDb]): DBIOAction[
+    Vector[SpendingInfoDb],
+    NoStream,
+    Effect.Read with Effect.Write] = {
     //1. Check if any are reserved already
     //2. if not, reserve them
     //3. if they are reserved, throw an exception?
     val outPoints = ts.map(_.outPoint)
-    val action: DBIOAction[
+    table
-      Int,
-      NoStream,
-      Effect.Write with Effect.Transactional] = table
       .filter(_.outPoint.inSet(outPoints))
       .filter(
         _.state.inSet(TxoState.receivedStates)
@@ -638,16 +658,17 @@ case class SpendingInfoDAO()(implicit
             s"Failed to reserve all utxos, expected=${ts.length} actual=$count")
           DBIO.failed(exn)
         } else {
-
           DBIO.successful(count)
         }
       }
-
-    safeDatabase
-      .run(action)
       .map(_ => ts.map(_.copyWithState(TxoState.Reserved)))
   }
 
+  def markAsReserved(
+      ts: Vector[SpendingInfoDb]): Future[Vector[SpendingInfoDb]] = {
+    safeDatabase.run(markAsReservedAction(ts))
+  }
+
   def createOutPointsIndexIfNeeded(): Future[Unit] = Future {
     withStatement(
       s"CREATE UNIQUE INDEX IF NOT EXISTS utxo_outpoints ON $fullTableName (tx_outpoint)") {