refactor(delegators): adopt samber/lo functional style for cleaner logic

Author: ccamel

pkg/delegators/enhancer.go

@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"strings"
 
+	"github.com/samber/lo"
 	"github.com/teambenny/goetl"
 	"github.com/teambenny/goetl/etldata"
 
@@ -21,16 +22,13 @@ type addressEnhancer struct {
 
 // NewAddressEnhancer returns a new processor that enrich the data with addresses with the given prefixes.
 func NewAddressEnhancer(prefixes []string, logger log.Logger) (goetl.Processor, error) {
-	keys := make([]string, len(prefixes))
-	for i, prefix := range prefixes {
-		keys[i] = fmt.Sprintf("delegator-%s-address", prefix)
-	}
-
 	return &addressEnhancer{
 		prefixes: prefixes,
 		logger:   logger,
-		keys:     keys,
-		name:     fmt.Sprintf("AddressEnhancer(%s)", strings.Join(prefixes, ",")),
+		keys: lo.Map(prefixes, func(prefix string, _ int) string {
+			return fmt.Sprintf("delegator_%s_address", prefix)
+		}),
+		name: fmt.Sprintf("AddressEnhancer(%s)", strings.Join(prefixes, ",")),
 	}, nil
 }
 

pkg/delegators/reader.go

@@ -9,16 +9,17 @@ import (
 
 	"github.com/axone-protocol/cosmos-extractor/pkg/keeper"
 	cmtproto "github.com/cometbft/cometbft/proto/tendermint/types"
+	"github.com/samber/lo"
 	"github.com/teambenny/goetl"
 	"github.com/teambenny/goetl/etldata"
-	"github.com/teambenny/goetl/etlutil"
 
 	"cosmossdk.io/collections"
 	"cosmossdk.io/log"
 	"cosmossdk.io/math"
 
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	bankkeeper "github.com/cosmos/cosmos-sdk/x/bank/keeper"
+	stakingkeeper "github.com/cosmos/cosmos-sdk/x/staking/keeper"
 	stakingtypes "github.com/cosmos/cosmos-sdk/x/staking/types"
 )
 
@@ -65,22 +66,11 @@ func (r *delegatorsReader) ProcessData(_ etldata.Payload, outputChan chan etldat
 
 	configureSdk(prefix)
 
-	err = IterateAllAddresses(ctx, keepers.Bank, func(addr sdk.AccAddress) (stop bool) {
-		for _, val := range validators {
-			valAddr, err := sdk.ValAddressFromBech32(val.OperatorAddress)
-			etlutil.KillPipelineIfErr(err, killChan)
-
-			delegation, err := keepers.Staking.GetDelegation(ctx, addr, valAddr)
-			if err != nil {
-				if errors.Is(err, stakingtypes.ErrNoDelegation) {
-					continue
-				}
-
-				r.logger.Error(err.Error())
-				killChan <- err
-				return true
-			}
+	err = iterateAllAddresses(ctx, keepers.Bank, func(addr sdk.AccAddress) (stop bool) {
+		delegations := lo.RejectMap(validators,
+			toDelegations(ctx, addr, r.logger, keepers.Staking, killChan))
 
+		for _, delegation := range delegations {
 			payload := Delegation{
 				ChainName:           r.chainName,
 				DelegatorNativeAddr: delegation.DelegatorAddress,
@@ -122,7 +112,7 @@ func (r *delegatorsReader) String() string {
 
 // IterateAllAddresses iterates over all the accounts that are provided to a callback.
 // If true is returned from the callback, iteration is halted.
-func IterateAllAddresses(ctx context.Context, bankKeeper bankkeeper.BaseKeeper, cb func(sdk.AccAddress) bool) error {
+func iterateAllAddresses(ctx context.Context, bankKeeper bankkeeper.BaseKeeper, cb func(sdk.AccAddress) bool) error {
 	lastSeenAddr := ""
 	err := bankKeeper.Balances.Walk(ctx, nil, func(key collections.Pair[sdk.AccAddress, string], _ math.Int) (stop bool, err error) {
 		addr := key.K1()
@@ -137,6 +127,31 @@ func IterateAllAddresses(ctx context.Context, bankKeeper bankkeeper.BaseKeeper,
 	return err
 }
 
+func toDelegations(
+	ctx context.Context, address sdk.AccAddress, logger log.Logger, stakingKeeper *stakingkeeper.Keeper, killChan chan error,
+) func(item stakingtypes.Validator, index int) (stakingtypes.Delegation, bool) {
+	return func(item stakingtypes.Validator, _ int) (stakingtypes.Delegation, bool) {
+		valAddr, err := sdk.ValAddressFromBech32(item.OperatorAddress)
+		if err != nil {
+			logger.Error(err.Error())
+			killChan <- err
+			return stakingtypes.Delegation{}, true
+		}
+
+		delegation, err := stakingKeeper.GetDelegation(ctx, address, valAddr)
+		if err != nil {
+			if errors.Is(err, stakingtypes.ErrNoDelegation) {
+				return stakingtypes.Delegation{}, true
+			}
+
+			logger.Error(err.Error())
+			killChan <- err
+			return stakingtypes.Delegation{}, true
+		}
+		return delegation, false
+	}
+}
+
 func guessPrefixFromValoper(valoper string) (string, error) {
 	if idx := strings.Index(valoper, "valoper"); idx != -1 {
 		return valoper[:idx], nil