improve token balances endpoint performance (#157)

### TL;DR
Added balance aggregation support for token balances API endpoint and simplified response model

### What changed?
- Modified token balance query to support aggregated balances with GROUP BY functionality
- Simplified BalanceModel response by removing redundant fields (chainId, tokenType, owner)
- Added serialization functions to transform token balances into the response model
- Updated Swagger documentation for the balance endpoint path
- Added support for custom column selection in token balance queries
- Implemented conditional balance filtering based on aggregation status

### Why make this change?
To provide more efficient token balance querying by allowing balance aggregation and reducing redundant data in the response. This improves API performance and simplifies the response structure for consumers.

Author: iuwqyir

internal/handlers/token_handlers.go

@@ -8,15 +8,13 @@ import (
 	"github.com/gin-gonic/gin"
 	"github.com/rs/zerolog/log"
 	"github.com/thirdweb-dev/indexer/api"
+	"github.com/thirdweb-dev/indexer/internal/common"
 	"github.com/thirdweb-dev/indexer/internal/storage"
 )
 
 // BalanceModel return type for Swagger documentation
 type BalanceModel struct {
-	ChainId      string   `json:"chain_id" ch:"chain_id"`
-	TokenType    string   `json:"token_type" ch:"token_type"`
 	TokenAddress string   `json:"token_address" ch:"address"`
-	Owner        string   `json:"owner" ch:"owner"`
 	TokenId      string   `json:"token_id" ch:"token_id"`
 	Balance      *big.Int `json:"balance" ch:"balance"`
 }
@@ -37,7 +35,7 @@ type BalanceModel struct {
 // @Failure 400 {object} api.Error
 // @Failure 401 {object} api.Error
 // @Failure 500 {object} api.Error
-// @Router /{chainId}/events [get]
+// @Router /{chainId}/balances/{owner}/{type} [get]
 func GetTokenBalancesByType(c *gin.Context) {
 	chainId, err := api.GetChainId(c)
 	if err != nil {
@@ -60,12 +58,21 @@ func GetTokenBalancesByType(c *gin.Context) {
 		return
 	}
 	hideZeroBalances := c.Query("hide_zero_balances") != "false"
+
+	columns := []string{"address", "sum(balance) as balance"}
+	groupBy := []string{"address"}
+	if tokenType != "erc20" {
+		columns = []string{"address", "token_id", "sum(balance) as balance"}
+		groupBy = []string{"address", "token_id"}
+	}
+
 	qf := storage.BalancesQueryFilter{
 		ChainId:      chainId,
 		Owner:        owner,
 		TokenType:    tokenType,
 		TokenAddress: tokenAddress,
 		ZeroBalance:  hideZeroBalances,
+		GroupBy:      groupBy,
 		SortBy:       c.Query("sort_by"),
 		SortOrder:    c.Query("sort_order"),
 		Page:         api.ParseIntQueryParam(c.Query("page"), 0),
@@ -87,13 +94,34 @@ func GetTokenBalancesByType(c *gin.Context) {
 		return
 	}
 
-	balancesResult, err := mainStorage.GetTokenBalances(qf)
+	balancesResult, err := mainStorage.GetTokenBalances(qf, columns...)
 	if err != nil {
 		log.Error().Err(err).Msg("Error querying balances")
 		// TODO: might want to choose BadRequestError if it's due to not-allowed functions
 		api.InternalErrorHandler(c)
 		return
 	}
-	queryResult.Data = balancesResult.Data
+	queryResult.Data = serializeBalances(balancesResult.Data)
 	sendJSONResponse(c, queryResult)
 }
+
+func serializeBalances(balances []common.TokenBalance) []BalanceModel {
+	balanceModels := make([]BalanceModel, len(balances))
+	for i, balance := range balances {
+		balanceModels[i] = serializeBalance(balance)
+	}
+	return balanceModels
+}
+
+func serializeBalance(balance common.TokenBalance) BalanceModel {
+	return BalanceModel{
+		TokenAddress: balance.TokenAddress,
+		Balance:      balance.Balance,
+		TokenId: func() string {
+			if balance.TokenId != nil {
+				return balance.TokenId.String()
+			}
+			return ""
+		}(),
+	}
+}

internal/storage/clickhouse.go

@@ -1373,16 +1373,38 @@ func (c *ClickHouseConnector) getTableName(chainId *big.Int, defaultTable string
 	return defaultTable
 }
 
-func (c *ClickHouseConnector) GetTokenBalances(qf BalancesQueryFilter) (QueryResult[common.TokenBalance], error) {
+func (c *ClickHouseConnector) GetTokenBalances(qf BalancesQueryFilter, fields ...string) (QueryResult[common.TokenBalance], error) {
 	columns := "chain_id, token_type, address, owner, token_id, balance"
+	if len(fields) > 0 {
+		columns = strings.Join(fields, ", ")
+	}
+	query := fmt.Sprintf("SELECT %s FROM %s.token_balances WHERE chain_id = ? AND token_type = ? AND owner = ?", columns, c.cfg.Database)
+
+	if qf.TokenAddress != "" {
+		query += fmt.Sprintf(" AND address = '%s'", qf.TokenAddress)
+	}
+
+	isBalanceAggregated := false
+	for _, field := range fields {
+		if strings.Contains(field, "balance") && strings.TrimSpace(field) != "balance" {
+			isBalanceAggregated = true
+			break
+		}
+	}
 	balanceCondition := ">="
 	if qf.ZeroBalance {
 		balanceCondition = ">"
 	}
-	query := fmt.Sprintf("SELECT %s FROM %s.token_balances FINAL WHERE chain_id = ? AND token_type = ? AND owner = ? AND balance %s 0", columns, c.cfg.Database, balanceCondition)
+	if !isBalanceAggregated {
+		query += fmt.Sprintf(" AND balance %s 0", balanceCondition)
+	}
 
-	if qf.TokenAddress != "" {
-		query += fmt.Sprintf(" AND address = '%s'", qf.TokenAddress)
+	if len(qf.GroupBy) > 0 {
+		query += fmt.Sprintf(" GROUP BY %s", strings.Join(qf.GroupBy, ", "))
+
+		if isBalanceAggregated {
+			query += fmt.Sprintf(" HAVING balance %s 0", balanceCondition)
+		}
 	}
 
 	// Add ORDER BY clause

internal/storage/connector.go

@@ -31,6 +31,7 @@ type BalancesQueryFilter struct {
 	TokenAddress string
 	Owner        string
 	ZeroBalance  bool
+	GroupBy      []string
 	SortBy       string
 	SortOrder    string
 	Page         int
@@ -80,7 +81,7 @@ type IMainStorage interface {
 	GetBlockHeadersDescending(chainId *big.Int, from *big.Int, to *big.Int) (blockHeaders []common.BlockHeader, err error)
 	DeleteBlockData(chainId *big.Int, blockNumbers []*big.Int) error
 
-	GetTokenBalances(qf BalancesQueryFilter) (QueryResult[common.TokenBalance], error)
+	GetTokenBalances(qf BalancesQueryFilter, fields ...string) (QueryResult[common.TokenBalance], error)
 }
 
 func NewStorageConnector(cfg *config.StorageConfig) (IStorage, error) {