Feature: filtering balances by contract (#558)

baking-bad · Mar 23, 2021 · b6a79cf · b6a79cf
1 parent 6e18c08
commit b6a79cf
Show file tree

Hide file tree

Showing 6 changed files with 28 additions and 16 deletions.
diff --git a/cmd/api/handlers/account.go b/cmd/api/handlers/account.go
@@ -74,6 +74,7 @@ func (ctx *Context) GetInfo(c *gin.Context) {
 // @Param address path string true "Address" minlength(36) maxlength(36)
 // @Param offset query integer false "Offset"
 // @Param size query integer false "Requested count" minimum(0) maximum(10)
+// @Param contract query string false "Contract address"
 // @Accept  json
 // @Produce  json
 // @Success 200 {object} TokenBalances
@@ -85,19 +86,19 @@ func (ctx *Context) GetAccountTokenBalances(c *gin.Context) {
 	if err := c.BindUri(&req); ctx.handleError(c, err, http.StatusBadRequest) {
 		return
 	}
-	var pageReq pageableRequest
-	if err := c.BindQuery(&pageReq); ctx.handleError(c, err, http.StatusBadRequest) {
+	var queryParams tokenBalanceRequest
+	if err := c.BindQuery(&queryParams); ctx.handleError(c, err, http.StatusBadRequest) {
 		return
 	}
-	balances, err := ctx.getAccountBalances(req.Network, req.Address, pageReq.Size, pageReq.Offset)
+	balances, err := ctx.getAccountBalances(req.Network, req.Address, queryParams)
 	if ctx.handleError(c, err, 0) {
 		return
 	}
 	c.JSON(http.StatusOK, balances)
 }
 
-func (ctx *Context) getAccountBalances(network, address string, size, offset int64) (*TokenBalances, error) {
-	tokenBalances, total, err := ctx.TokenBalances.GetAccountBalances(network, address, size, offset)
+func (ctx *Context) getAccountBalances(network, address string, req tokenBalanceRequest) (*TokenBalances, error) {
+	tokenBalances, total, err := ctx.TokenBalances.GetAccountBalances(network, address, req.Contract, req.Size, req.Offset)
 	if err != nil {
 		return nil, err
 	}

diff --git a/cmd/api/handlers/requests.go b/cmd/api/handlers/requests.go
@@ -367,3 +367,8 @@ type tokenRequest struct {
 	MaxLevel int64 `form:"max_level,omitempty" binding:"omitempty,gt_int64_ptr=MinLevel"`
 	MinLevel int64 `form:"min_level,omitempty" binding:"omitempty"`
 }
+
+type tokenBalanceRequest struct {
+	pageableRequest
+	Contract string `form:"contract" binding:"omitempty,address"`
+}
diff --git a/internal/elastic/tokenbalance/storage.go b/internal/elastic/tokenbalance/storage.go
@@ -85,12 +85,19 @@ func (storage *Storage) GetHolders(network, contract string, tokenID int64) ([]t
 }
 
 // GetAccountBalances -
-func (storage *Storage) GetAccountBalances(network, address string, size, offset int64) ([]tokenbalance.TokenBalance, int64, error) {
+func (storage *Storage) GetAccountBalances(network, address, contract string, size, offset int64) ([]tokenbalance.TokenBalance, int64, error) {
+	filters := []core.Item{
+		core.MatchPhrase("address", address),
+		core.Match("network", network),
+	}
+
+	if contract != "" {
+		filters = append(filters, core.MatchPhrase("contract", contract))
+	}
 	query := core.NewQuery().Query(
 		core.Bool(
 			core.Filter(
-				core.MatchPhrase("address", address),
-				core.Match("network", network),
+				filters...,
 			),
 			core.MustNot(
 				core.Term("balance", "0"),
@@ -112,8 +119,7 @@ func (storage *Storage) GetAccountBalances(network, address string, size, offset
 	countQuery := core.NewQuery().Query(
 		core.Bool(
 			core.Filter(
-				core.MatchPhrase("address", address),
-				core.Match("network", network),
+				filters...,
 			),
 			core.MustNot(
 				core.Term("balance", "0"),

diff --git a/internal/models/mock/tokenbalance/mock.go b/internal/models/mock/tokenbalance/mock.go
diff --git a/internal/models/tokenbalance/repository.go b/internal/models/tokenbalance/repository.go
@@ -2,7 +2,7 @@ package tokenbalance
 
 // Repository -
 type Repository interface {
-	GetAccountBalances(network string, address string, size, offset int64) ([]TokenBalance, int64, error)
+	GetAccountBalances(network, address, contract string, size, offset int64) ([]TokenBalance, int64, error)
 	Update(updates []*TokenBalance) error
 	GetHolders(network, contract string, tokenID int64) ([]TokenBalance, error)
 	BurnNft(network, contract string, tokenID int64) error

diff --git a/internal/reindexer/tokenbalance/storage.go b/internal/reindexer/tokenbalance/storage.go
@@ -84,7 +84,7 @@ func (storage *Storage) GetHolders(network, contract string, tokenID int64) (bal
 }
 
 // GetAccountBalances -
-func (storage *Storage) GetAccountBalances(network, address string, size, offset int64) (tokenBalances []tokenbalance.TokenBalance, total int64, err error) {
+func (storage *Storage) GetAccountBalances(network, address, contract string, size, offset int64) (tokenBalances []tokenbalance.TokenBalance, total int64, err error) {
 	query := storage.db.Query(models.DocTokenBalances).
 		Match("network", network).
 		Match("address", address)