Make GetVtxos accept multiple addresses

Author: altafan

api-spec/openapi/swagger/ark/v1/indexer.swagger.json

@@ -310,7 +310,7 @@
         ]
       }
     },
-    "/v1/getVtxos/{address}": {
+    "/v1/getVtxos/{addresses}": {
       "get": {
         "operationId": "IndexerService_GetVtxos",
         "responses": {
@@ -329,10 +329,15 @@
         },
         "parameters": [
           {
-            "name": "address",
+            "name": "addresses",
             "in": "path",
             "required": true,
-            "type": "string"
+            "type": "array",
+            "items": {
+              "type": "string"
+            },
+            "collectionFormat": "csv",
+            "minItems": 1
           },
           {
             "name": "spendableOnly",

api-spec/protobuf/ark/v1/indexer.proto

@@ -37,7 +37,7 @@ service IndexerService {
   };
   rpc GetVtxos(GetVtxosRequest) returns (GetVtxosResponse) {
     option (google.api.http) = {
-      get: "/v1/getVtxos/{address}"
+      get: "/v1/getVtxos/{addresses}"
     };
   };
   rpc GetTransactionHistory(GetTransactionHistoryRequest) returns (GetTransactionHistoryResponse) {
@@ -121,7 +121,7 @@ message GetVtxoTreeLeavesResponse {
 }
 
 message GetVtxosRequest {
-  string address = 1;
+  repeated string addresses = 1;
   bool spendable_only = 2;
   bool spent_only = 3;
   IndexerPageRequest page = 4;

api-spec/protobuf/gen/ark/v1/indexer.pb.go

@@ -32,7 +32,7 @@ type IndexerService interface {
 	GetVtxoTreeLeaves(ctx context.Context, batchOutpoint Outpoint, page *Page) (*VtxoTreeLeavesResp, error)
 	GetForfeitTxs(ctx context.Context, txid string, page *Page) (*ForfeitTxsResp, error)
 	GetConnectors(ctx context.Context, txid string, page *Page) (*ConnectorResp, error)
-	GetVtxos(ctx context.Context, pubkey string, spendableOnly, spendOnly bool, page *Page) (*SpendableVtxosResp, error)
+	GetVtxos(ctx context.Context, pubkeys []string, spendableOnly, spendOnly bool, page *Page) (*SpendableVtxosResp, error)
 	GetTransactionHistory(ctx context.Context, pubkey string, start, end int64, page *Page) (*TxHistoryResp, error)
 	GetVtxoChain(ctx context.Context, vtxoKey Outpoint, page *Page) (*VtxoChainResp, error)
 	GetVirtualTxs(ctx context.Context, txids []string, page *Page) (*VirtualTxsResp, error)
@@ -160,28 +160,23 @@ func (i *indexerService) GetConnectors(ctx context.Context, txid string, page *P
 }
 
 func (i *indexerService) GetVtxos(
-	ctx context.Context, pubkey string, spendableOnly, spentOnly bool, page *Page,
+	ctx context.Context, pubkeys []string, spendableOnly, spentOnly bool, page *Page,
 ) (*SpendableVtxosResp, error) {
 	if spendableOnly && spentOnly {
 		return nil, fmt.Errorf("spendable and spent only can't be true at the same time")
 	}
-	spendable, spent, err := i.repoManager.Vtxos().GetAllVtxosWithPubKey(ctx, pubkey)
+
+	vtxos, err := i.repoManager.Vtxos().GetAllVtxosWithPubKeys(
+		ctx, pubkeys, spendableOnly, spentOnly,
+	)
 	if err != nil {
 		return nil, err
 	}
 
-	vtxos := append(spendable, spent...)
-	if spendableOnly {
-		vtxos = spendable
-	}
-	if spentOnly {
-		vtxos = spent
-	}
-
-	spendableVtxosPaged, pageResp := paginate(vtxos, page, maxPageSizeSpendableVtxos)
+	pagedVtxos, pageResp := paginate(vtxos, page, maxPageSizeSpendableVtxos)
 
 	return &SpendableVtxosResp{
-		Vtxos: spendableVtxosPaged,
+		Vtxos: pagedVtxos,
 		Page:  pageResp,
 	}, nil
 }

api-spec/protobuf/gen/ark/v1/indexer.pb.gw.go

@@ -41,6 +41,7 @@ type VtxoRepository interface {
 	GetSpendableVtxosWithPubKey(ctx context.Context, pubkey string) ([]Vtxo, error)
 	GetAll(ctx context.Context) ([]Vtxo, error)
 	GetAllVtxosWithPubKey(ctx context.Context, pubkey string) ([]Vtxo, []Vtxo, error)
+	GetAllVtxosWithPubKeys(ctx context.Context, pubkeys []string, spendableOnly, spentOnly bool) ([]Vtxo, error)
 	UpdateExpireAt(ctx context.Context, vtxos []VtxoKey, expireAt int64) error
 	GetLeafVtxosForRound(ctx context.Context, txid string) ([]Vtxo, error)
 	Close()

server/internal/core/application/indexer.go

@@ -5,6 +5,7 @@ import (
 	"errors"
 	"fmt"
 	"path/filepath"
+	"sort"
 	"strings"
 	"time"
 
@@ -209,6 +210,47 @@ func (r *vtxoRepository) GetAllVtxosWithPubKey(
 	return unspentVtxos, spentVtxos, nil
 }
 
+func (r *vtxoRepository) GetAllVtxosWithPubKeys(
+	ctx context.Context, pubkeys []string, spendableOnly, spentOnly bool,
+) ([]domain.Vtxo, error) {
+	if spendableOnly && spendableOnly == spentOnly {
+		return nil, fmt.Errorf("spendable and spent only can't be true at the same time")
+	}
+
+	allVtxos := make([]domain.Vtxo, 0)
+	for _, pubkey := range pubkeys {
+		query := badgerhold.Where("PubKey").Eq(pubkey)
+		vtxos, err := r.findVtxos(ctx, query)
+		if err != nil {
+			return nil, err
+		}
+		sort.SliceStable(vtxos, func(i, j int) bool {
+			return vtxos[i].CreatedAt > vtxos[j].CreatedAt
+		})
+
+		if spendableOnly {
+			spendableVtxos := make([]domain.Vtxo, 0, len(vtxos))
+			for _, vtxo := range vtxos {
+				if !vtxo.Spent && !vtxo.Swept && !vtxo.Redeemed {
+					spendableVtxos = append(spendableVtxos, vtxo)
+				}
+			}
+			vtxos = spendableVtxos
+		}
+		if spentOnly {
+			spentVtxos := make([]domain.Vtxo, 0, len(vtxos))
+			for _, vtxo := range vtxos {
+				if vtxo.Spent || vtxo.Swept || vtxo.Redeemed {
+					spentVtxos = append(spentVtxos, vtxo)
+				}
+			}
+			vtxos = spentVtxos
+		}
+		allVtxos = append(allVtxos, vtxos...)
+	}
+	return allVtxos, nil
+}
+
 func (r *vtxoRepository) Close() {
 	r.store.Close()
 }

server/internal/core/domain/round_repo.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"database/sql"
 	"fmt"
+	"sort"
 
 	"github.com/ark-network/ark/server/internal/core/domain"
 	"github.com/ark-network/ark/server/internal/infrastructure/db/sqlite/sqlc/queries"
@@ -330,6 +331,58 @@ func (v *vtxoRepository) GetAllVtxosWithPubKey(
 	return unspentVtxos, spentVtxos, nil
 }
 
+func (v *vtxoRepository) GetAllVtxosWithPubKeys(
+	ctx context.Context, pubkeys []string, spendableOnly, spentOnly bool,
+) ([]domain.Vtxo, error) {
+	if spendableOnly && spendableOnly == spentOnly {
+		return nil, fmt.Errorf("spendable and spent only can't be true at the same time")
+	}
+
+	allVtxos := make([]domain.Vtxo, 0)
+	// TODO: make this a proper sql query
+	for _, pubkey := range pubkeys {
+		res, err := v.querier.SelectVtxosWithPubkey(ctx, pubkey)
+		if err != nil {
+			return nil, err
+		}
+		rows := make([]queries.Vtxo, 0, len(res))
+		for _, row := range res {
+			rows = append(rows, row.Vtxo)
+		}
+
+		vtxos, err := readRows(rows)
+		if err != nil {
+			return nil, err
+		}
+		sort.SliceStable(vtxos, func(i, j int) bool {
+			return vtxos[i].CreatedAt > vtxos[j].CreatedAt
+		})
+
+		if spendableOnly {
+			spendableVtxos := make([]domain.Vtxo, 0, len(vtxos))
+			for _, vtxo := range vtxos {
+				if !vtxo.Spent && !vtxo.Swept && !vtxo.Redeemed {
+					spendableVtxos = append(spendableVtxos, vtxo)
+				}
+			}
+			vtxos = spendableVtxos
+		}
+		if spentOnly {
+			spentVtxos := make([]domain.Vtxo, 0, len(vtxos))
+			for _, vtxo := range vtxos {
+				if vtxo.Spent || vtxo.Swept || vtxo.Redeemed {
+					spentVtxos = append(spentVtxos, vtxo)
+				}
+			}
+			vtxos = spentVtxos
+		}
+
+		allVtxos = append(allVtxos, vtxos...)
+	}
+
+	return allVtxos, nil
+}
+
 func rowToVtxo(row queries.Vtxo) domain.Vtxo {
 	return domain.Vtxo{
 		VtxoKey: domain.VtxoKey{

server/internal/infrastructure/db/badger/vtxo_repo.go

@@ -202,7 +202,7 @@ func (e indexerService) GetConnectors(ctx context.Context, request *arkv1.GetCon
 }
 
 func (e indexerService) GetVtxos(ctx context.Context, request *arkv1.GetVtxosRequest) (*arkv1.GetVtxosResponse, error) {
-	pubkey, err := parseArkAddress(request.GetAddress())
+	pubkeys, err := parseArkAddresses(request.GetAddresses())
 	if err != nil {
 		return nil, status.Error(codes.InvalidArgument, err.Error())
 	}
@@ -215,7 +215,7 @@ func (e indexerService) GetVtxos(ctx context.Context, request *arkv1.GetVtxosReq
 	}
 
 	resp, err := e.indexerSvc.GetVtxos(
-		ctx, pubkey, request.GetSpendableOnly(), request.GetSpentOnly(), page,
+		ctx, pubkeys, request.GetSpendableOnly(), request.GetSpentOnly(), page,
 	)
 	if err != nil {
 		return nil, status.Errorf(codes.Internal, "failed to get vtxos: %v", err)
@@ -458,3 +458,18 @@ func protoPage(page application.PageResp) *arkv1.IndexerPageResponse {
 		Total:   page.Total,
 	}
 }
+
+func parseArkAddresses(addresses []string) ([]string, error) {
+	if len(addresses) == 0 {
+		return nil, fmt.Errorf("missing addresses")
+	}
+	pubkeys := make([]string, 0, len(addresses))
+	for _, address := range addresses {
+		pubkey, err := parseArkAddress(address)
+		if err != nil {
+			return nil, err
+		}
+		pubkeys = append(pubkeys, pubkey)
+	}
+	return pubkeys, nil
+}