loop/staticaddr/withdraw/sql_store.go

package withdraw

import (
	"bytes"
	"context"
	"database/sql"

	"github.com/btcsuite/btcd/btcutil/v2"
	"github.com/btcsuite/btcd/chainhash/v2"
	"github.com/btcsuite/btcd/wire/v2"
	"github.com/lightninglabs/loop/loopdb"
	"github.com/lightninglabs/loop/loopdb/sqlc"
	"github.com/lightninglabs/loop/staticaddr/deposit"
	"github.com/lightningnetwork/lnd/clock"
)

type Querier interface {
	// CreateWithdrawal inserts a new withdrawal.
	CreateWithdrawal(ctx context.Context,
		arg sqlc.CreateWithdrawalParams) error

	// UpdateWithdrawal updates a withdrawal with confirmation parameters.
	UpdateWithdrawal(ctx context.Context,
		arg sqlc.UpdateWithdrawalParams) error

	// GetWithdrawalIDByDepositID retrieves the withdrawal ID associated
	// with a given deposit ID.
	GetWithdrawalIDByDepositID(ctx context.Context, depositID []byte) (
		[]byte, error)

	// CreateWithdrawalDeposit links withdrawal to deposits.
	CreateWithdrawalDeposit(ctx context.Context,
		arg sqlc.CreateWithdrawalDepositParams) error

	// GetWithdrawalDeposits retrieves the deposit IDs associated with a
	// withdrawal.
	GetWithdrawalDeposits(ctx context.Context, withdrawalID []byte) (
		[][]byte, error)

	// GetAllWithdrawals retrieves all withdrawals from the database.
	GetAllWithdrawals(ctx context.Context) ([]sqlc.Withdrawal, error)
}

// BaseDB is the interface that contains all the queries generated by sqlc for
// the static_address_swaps table and transaction functionality.
type BaseDB interface {
	Querier

	// ExecTx allows for executing a function in the context of a database
	// transaction.
	ExecTx(ctx context.Context, txOptions loopdb.TxOptions,
		txBody func(Querier) error) error
}

// SqlStore is the backing store for static address withdrawals.
type SqlStore struct {
	baseDB       BaseDB
	depositStore deposit.Store
	clock        clock.Clock
}

// NewSqlStore constructs a new SQLStore from a BaseDB. The BaseDB is agnostic
// to the underlying driver which can be postgres or sqlite.
func NewSqlStore(db BaseDB, depositStore deposit.Store) *SqlStore {
	return &SqlStore{
		baseDB:       db,
		depositStore: depositStore,
		clock:        clock.NewDefaultClock(),
	}
}

// CreateWithdrawal creates a static address withdrawal record in the database.
func (s *SqlStore) CreateWithdrawal(ctx context.Context,
	deposits []*deposit.Deposit) error {

	id, err := GetRandomWithdrawalID()
	if err != nil {
		return err
	}

	var totalAmount btcutil.Amount
	for _, deposit := range deposits {
		totalAmount += deposit.Value
	}

	createArgs := sqlc.CreateWithdrawalParams{
		WithdrawalID:       id[:],
		TotalDepositAmount: int64(totalAmount),
		InitiationTime:     s.clock.Now().UTC(),
	}
	return s.baseDB.ExecTx(ctx, &loopdb.SqliteTxOptions{},
		func(q Querier) error {
			err := q.CreateWithdrawal(ctx, createArgs)
			if err != nil {
				return err
			}

			for _, deposit := range deposits {
				err = q.CreateWithdrawalDeposit(
					ctx, sqlc.CreateWithdrawalDepositParams{
						WithdrawalID: id[:],
						DepositID:    deposit.ID[:],
					})
				if err != nil {
					return err
				}
			}

			return nil
		})
}

// UpdateWithdrawal updates a withdrawal record with the transaction
// information, including the withdrawn amount, change amount, and
// confirmation height. It is expected that the withdrawal has already been
// created with CreateWithdrawal, and that the deposits slice contains the
// deposits associated with the withdrawal.
func (s *SqlStore) UpdateWithdrawal(ctx context.Context,
	deposits []*deposit.Deposit, tx *wire.MsgTx, confirmationHeight uint32,
	changePkScript []byte) error {

	// Populate the optional change amount.
	withdrawnAmount, changeAmount := int64(0), int64(0)
	if len(tx.TxOut) == 1 {
		withdrawnAmount = tx.TxOut[0].Value
	} else if len(tx.TxOut) == 2 {
		withdrawnAmount, changeAmount = tx.TxOut[0].Value, tx.TxOut[1].Value
		if bytes.Equal(changePkScript, tx.TxOut[0].PkScript) {
			changeAmount = tx.TxOut[0].Value
			withdrawnAmount = tx.TxOut[1].Value
		}
	}

	updateArgs := sqlc.UpdateWithdrawalParams{
		WithdrawalTxID: sql.NullString{
			String: tx.TxHash().String(),
			Valid:  true,
		},
		WithdrawnAmount: sql.NullInt64{
			Int64: withdrawnAmount,
			Valid: withdrawnAmount > 0,
		},
		ChangeAmount: sql.NullInt64{
			Int64: changeAmount,
			Valid: changeAmount > 0,
		},
		ConfirmationHeight: sql.NullInt64{
			Int64: int64(confirmationHeight),
			Valid: confirmationHeight > 0,
		},
	}

	return s.baseDB.ExecTx(ctx, &loopdb.SqliteTxOptions{},
		func(q Querier) error {
			withdrawalID, err := q.GetWithdrawalIDByDepositID(
				ctx, deposits[0].ID[:],
			)
			if err != nil {
				return err
			}

			updateArgs.WithdrawalID = withdrawalID
			err = q.UpdateWithdrawal(ctx, updateArgs)
			if err != nil {
				return err
			}

			return nil
		})
}

// GetAllWithdrawals retrieves all static address withdrawals from the
// database. It returns a slice of Withdrawal structs, each containing a list
// of associated deposits.
func (s *SqlStore) GetAllWithdrawals(ctx context.Context) ([]Withdrawal,
	error) {

	withdrawals, err := s.baseDB.GetAllWithdrawals(ctx)
	if err != nil {
		return nil, err
	}

	result := make([]Withdrawal, 0, len(withdrawals))
	for _, w := range withdrawals {
		depositIDs, err := s.baseDB.GetWithdrawalDeposits(ctx,
			w.WithdrawalID)

		if err != nil {
			return nil, err
		}

		deposits := make([]*deposit.Deposit, 0, len(depositIDs))
		for _, dID := range depositIDs {
			deposit, err := s.depositStore.GetDeposit(
				ctx, deposit.ID(dID),
			)
			if err != nil {
				return nil, err
			}
			deposits = append(deposits, deposit)
		}

		txID, err := chainhash.NewHashFromStr(w.WithdrawalTxID.String)
		if err != nil {
			return nil, err
		}

		result = append(result, Withdrawal{
			ID:                 ID(w.WithdrawalID),
			TxID:               *txID,
			Deposits:           deposits,
			TotalDepositAmount: btcutil.Amount(w.TotalDepositAmount),
			WithdrawnAmount:    btcutil.Amount(w.WithdrawnAmount.Int64),
			ChangeAmount:       btcutil.Amount(w.ChangeAmount.Int64),
			InitiationTime:     w.InitiationTime,
			ConfirmationHeight: w.ConfirmationHeight.Int64,
		})
	}

	return result, nil
}