tapdb: add includeLeased parameter

Author: gijswijs

rpcserver.go

@@ -922,7 +922,7 @@ func (r *rpcServer) checkBalanceOverflow(ctx context.Context,
 	case assetID != nil:
 		// Retrieve the current asset balance.
 		balances, err := r.cfg.AssetStore.QueryBalancesByAsset(
-			ctx, assetID,
+			ctx, assetID, true,
 		)
 		if err != nil {
 			return fmt.Errorf("unable to query asset balance: %w",
@@ -938,7 +938,7 @@ func (r *rpcServer) checkBalanceOverflow(ctx context.Context,
 	case groupPubKey != nil:
 		// Retrieve the current balance of the group.
 		balances, err := r.cfg.AssetStore.QueryAssetBalancesByGroup(
-			ctx, groupPubKey,
+			ctx, groupPubKey, true,
 		)
 		if err != nil {
 			return fmt.Errorf("unable to query group balance: %w",
@@ -1106,9 +1106,12 @@ func (r *rpcServer) MarshalChainAsset(ctx context.Context, a *asset.ChainAsset,
 }
 
 func (r *rpcServer) listBalancesByAsset(ctx context.Context,
-	assetID *asset.ID) (*taprpc.ListBalancesResponse, error) {
+	assetID *asset.ID, includeLeased bool) (*taprpc.ListBalancesResponse,
+	error) {
 
-	balances, err := r.cfg.AssetStore.QueryBalancesByAsset(ctx, assetID)
+	balances, err := r.cfg.AssetStore.QueryBalancesByAsset(
+		ctx, assetID, includeLeased,
+	)
 	if err != nil {
 		return nil, fmt.Errorf("unable to list balances: %w", err)
 	}
@@ -1138,10 +1141,11 @@ func (r *rpcServer) listBalancesByAsset(ctx context.Context,
 }
 
 func (r *rpcServer) listBalancesByGroupKey(ctx context.Context,
-	groupKey *btcec.PublicKey) (*taprpc.ListBalancesResponse, error) {
+	groupKey *btcec.PublicKey,
+	includeLeased bool) (*taprpc.ListBalancesResponse, error) {
 
 	balances, err := r.cfg.AssetStore.QueryAssetBalancesByGroup(
-		ctx, groupKey,
+		ctx, groupKey, includeLeased,
 	)
 	if err != nil {
 		return nil, fmt.Errorf("unable to list balances: %w", err)
@@ -1293,7 +1297,7 @@ func (r *rpcServer) ListBalances(ctx context.Context,
 			copy(assetID[:], req.AssetFilter)
 		}
 
-		return r.listBalancesByAsset(ctx, assetID)
+		return r.listBalancesByAsset(ctx, assetID, req.IncludeLeased)
 
 	case *taprpc.ListBalancesRequest_GroupKey:
 		if !groupBy.GroupKey {
@@ -1310,7 +1314,9 @@ func (r *rpcServer) ListBalances(ctx context.Context,
 			}
 		}
 
-		return r.listBalancesByGroupKey(ctx, groupKey)
+		return r.listBalancesByGroupKey(
+			ctx, groupKey, req.IncludeLeased,
+		)
 
 	default:
 		return nil, fmt.Errorf("invalid group_by")

tapdb/asset_minting_test.go

@@ -1131,7 +1131,7 @@ func TestCommitBatchChainActions(t *testing.T) {
 
 	// We'll now query for the set of balances to ensure they all line up
 	// with the assets we just created, including the group genesis asset.
-	assetBalances, err := confAssets.QueryBalancesByAsset(ctx, nil)
+	assetBalances, err := confAssets.QueryBalancesByAsset(ctx, nil, false)
 	require.NoError(t, err)
 	require.Equal(t, numSeedlings+1, len(assetBalances))
 
@@ -1153,7 +1153,7 @@ func TestCommitBatchChainActions(t *testing.T) {
 	}
 	numKeyGroups := fn.Reduce(mintedAssets, keyGroupSumReducer)
 	assetBalancesByGroup, err := confAssets.QueryAssetBalancesByGroup(
-		ctx, nil,
+		ctx, nil, false,
 	)
 	require.NoError(t, err)
 	require.Equal(t, numKeyGroups, len(assetBalancesByGroup))

tapdb/assets_store.go

@@ -76,6 +76,16 @@ type (
 	// or for things like coin selection.
 	QueryAssetFilters = sqlc.QueryAssetsParams
 
+	// QueryAssetBalancesByGroupFilters lets us query the asset balances for
+	// asset groups or alternatively for a selected one that matches the
+	// passed filter.
+	QueryAssetBalancesByGroupFilters = sqlc.QueryAssetBalancesByGroupParams
+
+	// QueryAssetBalancesByAssetFilters lets us query the asset balances for
+	// assets or alternatively for a selected one that matches the passed
+	// filter.
+	QueryAssetBalancesByAssetFilters = sqlc.QueryAssetBalancesByAssetParams
+
 	// UtxoQuery lets us query a managed UTXO by either the transaction it
 	// references, or the outpoint.
 	UtxoQuery = sqlc.FetchManagedUTXOParams
@@ -184,14 +194,15 @@ type ActiveAssetsStore interface {
 	// QueryAssetBalancesByAsset queries the balances for assets or
 	// alternatively for a selected one that matches the passed asset ID
 	// filter.
-	QueryAssetBalancesByAsset(context.Context, []byte) ([]RawAssetBalance,
-		error)
+	QueryAssetBalancesByAsset(context.Context,
+		QueryAssetBalancesByAssetFilters) ([]RawAssetBalance, error)
 
 	// QueryAssetBalancesByGroup queries the asset balances for asset
 	// groups or alternatively for a selected one that matches the passed
 	// filter.
 	QueryAssetBalancesByGroup(context.Context,
-		[]byte) ([]RawAssetGroupBalance, error)
+		QueryAssetBalancesByGroupFilters) ([]RawAssetGroupBalance,
+		error)
 
 	// FetchGroupedAssets fetches all assets with non-nil group keys.
 	FetchGroupedAssets(context.Context) ([]RawGroupedAsset, error)
@@ -960,18 +971,35 @@ type AssetQueryFilters struct {
 // QueryBalancesByAsset queries the balances for assets or alternatively
 // for a selected one that matches the passed asset ID filter.
 func (a *AssetStore) QueryBalancesByAsset(ctx context.Context,
-	assetID *asset.ID) (map[asset.ID]AssetBalance, error) {
+	assetID *asset.ID,
+	includeLeased bool) (map[asset.ID]AssetBalance, error) {
+
+	// We'll now map the application level filtering to the type of
+	// filtering our database query understands.
+	assetBalancesFilter := QueryAssetBalancesByAssetFilters{
+		Now: sql.NullTime{
+			Time:  a.clock.Now().UTC(),
+			Valid: true,
+		},
+	}
 
-	var assetFilter []byte
+	// By default, we only show assets that are not leased.
+	if !includeLeased {
+		assetBalancesFilter.Leased = sqlBool(false)
+	}
+
+	// Only show assets that match the filter that has been passed
 	if assetID != nil {
-		assetFilter = assetID[:]
+		assetBalancesFilter.AssetIDFilter = assetID[:]
 	}
 
 	balances := make(map[asset.ID]AssetBalance)
 
 	readOpts := NewAssetStoreReadTx()
 	dbErr := a.db.ExecTx(ctx, &readOpts, func(q ActiveAssetsStore) error {
-		dbBalances, err := q.QueryAssetBalancesByAsset(ctx, assetFilter)
+		dbBalances, err := q.QueryAssetBalancesByAsset(
+			ctx, assetBalancesFilter,
+		)
 		if err != nil {
 			return fmt.Errorf("unable to query asset "+
 				"balances by asset: %w", err)
@@ -1014,20 +1042,37 @@ func (a *AssetStore) QueryBalancesByAsset(ctx context.Context,
 // QueryAssetBalancesByGroup queries the asset balances for asset groups or
 // alternatively for a selected one that matches the passed filter.
 func (a *AssetStore) QueryAssetBalancesByGroup(ctx context.Context,
-	groupKey *btcec.PublicKey) (map[asset.SerializedKey]AssetGroupBalance,
+	groupKey *btcec.PublicKey,
+	includeLeased bool) (map[asset.SerializedKey]AssetGroupBalance,
 	error) {
 
-	var groupFilter []byte
+	// We'll now map the application level filtering to the type of
+	// filtering our database query understands.
+	assetBalancesFilter := QueryAssetBalancesByGroupFilters{
+		Now: sql.NullTime{
+			Time:  a.clock.Now().UTC(),
+			Valid: true,
+		},
+	}
+
+	// By default, we only show assets that are not leased.
+	if !includeLeased {
+		assetBalancesFilter.Leased = sqlBool(false)
+	}
+
+	// Only show specific group if a groupKey has been passed.
 	if groupKey != nil {
 		groupKeySerialized := groupKey.SerializeCompressed()
-		groupFilter = groupKeySerialized[:]
+		assetBalancesFilter.KeyGroupFilter = groupKeySerialized[:]
 	}
 
 	balances := make(map[asset.SerializedKey]AssetGroupBalance)
 
 	readOpts := NewAssetStoreReadTx()
 	dbErr := a.db.ExecTx(ctx, &readOpts, func(q ActiveAssetsStore) error {
-		dbBalances, err := q.QueryAssetBalancesByGroup(ctx, groupFilter)
+		dbBalances, err := q.QueryAssetBalancesByGroup(
+			ctx, assetBalancesFilter,
+		)
 		if err != nil {
 			return fmt.Errorf("unable to query asset "+
 				"balances by asset: %w", err)