Add custom query handling for unbonding balances

app/app.go

@@ -560,6 +560,7 @@ func New(
 			app.TransferKeeper,
 			app.AccessControlKeeper,
 			&app.EvmKeeper,
+			app.StakingKeeper,
 		),
 		wasmOpts...,
 	)

wasmbinding/queries.go

@@ -3,9 +3,12 @@ package wasmbinding
 import (
 	"encoding/json"
 	"errors"
+	"fmt"
+	"math"
 
 	sdk "github.com/cosmos/cosmos-sdk/types"
 
+	stakingkeeper "github.com/cosmos/cosmos-sdk/x/staking/keeper"
 	"github.com/sei-protocol/sei-chain/utils/metrics"
 	epochwasm "github.com/sei-protocol/sei-chain/x/epoch/client/wasm"
 	epochbindings "github.com/sei-protocol/sei-chain/x/epoch/client/wasm/bindings"
@@ -25,15 +28,17 @@ type QueryPlugin struct {
 	epochHandler        epochwasm.EpochWasmQueryHandler
 	tokenfactoryHandler tokenfactorywasm.TokenFactoryWasmQueryHandler
 	evmHandler          evmwasm.EVMQueryHandler
+	stakingKeeper       stakingkeeper.Keeper
 }
 
 // NewQueryPlugin returns a reference to a new QueryPlugin.
-func NewQueryPlugin(oh *oraclewasm.OracleWasmQueryHandler, eh *epochwasm.EpochWasmQueryHandler, th *tokenfactorywasm.TokenFactoryWasmQueryHandler, evmh *evmwasm.EVMQueryHandler) *QueryPlugin {
+func NewQueryPlugin(oh *oraclewasm.OracleWasmQueryHandler, eh *epochwasm.EpochWasmQueryHandler, th *tokenfactorywasm.TokenFactoryWasmQueryHandler, evmh *evmwasm.EVMQueryHandler, sk stakingkeeper.Keeper) *QueryPlugin {
 	return &QueryPlugin{
 		oracleHandler:       *oh,
 		epochHandler:        *eh,
 		tokenfactoryHandler: *th,
 		evmHandler:          *evmh,
+		stakingKeeper:       sk,
 	}
 }
 
@@ -238,3 +243,72 @@ func (qp QueryPlugin) HandleEVMQuery(ctx sdk.Context, queryData json.RawMessage)
 		return nil, errors.New("unknown EVM query")
 	}
 }
+
+type StakingExtQueryType string
+
+const (
+	UnbondingDelegationsType StakingExtQueryType = "staking_ext_unbonding_delegations"
+)
+
+type StakingExtQuery struct {
+	UnbondingDelegations *UnbondingDelegationsRequest `json:"unbonding_delegations,omitempty"`
+}
+
+func (seq *StakingExtQuery) GetQueryType() StakingExtQueryType {
+	if seq.UnbondingDelegations != nil {
+		return UnbondingDelegationsType
+	}
+	return ""
+}
+
+type UnbondingDelegationsRequest struct {
+	Delegator string `json:"delegator,omitempty"`
+}
+
+type UnbondingDelegationsResponse struct {
+	Entries []UnbondingDelegationEntry `json:"entries"`
+}
+
+type UnbondingDelegationEntry struct {
+	CreationHeight int64   `json:"creation_height"`
+	CompletionTime string  `json:"completion_time"`
+	InitialBalance sdk.Int `json:"initial_balance"`
+	Balance        sdk.Int `json:"balance"`
+}
+
+func (qp QueryPlugin) HandleStakingExtQuery(ctx sdk.Context, queryData json.RawMessage) (res []byte, err error) {
+	var queryType StakingExtQueryType
+	var parsedQuery StakingExtQuery
+	if err := json.Unmarshal(queryData, &parsedQuery); err != nil {
+		return nil, errors.New("invalid Staking extension query")
+	}
+	queryType = parsedQuery.GetQueryType()
+
+	defer func() {
+		metrics.IncrementErrorMetrics(string(queryType), err)
+	}()
+
+	switch queryType {
+	case UnbondingDelegationsType:
+		c := parsedQuery.UnbondingDelegations
+		delegator, err := sdk.AccAddressFromBech32(c.Delegator)
+		if err != nil {
+			return nil, fmt.Errorf("invalid delegator string: %s", c.Delegator)
+		}
+		unbondingDelegations := qp.stakingKeeper.GetUnbondingDelegations(ctx, delegator, math.MaxUint16)
+		response := UnbondingDelegationsResponse{}
+		for _, ud := range unbondingDelegations {
+			for _, entry := range ud.Entries {
+				response.Entries = append(response.Entries, UnbondingDelegationEntry{
+					CreationHeight: entry.CreationHeight,
+					CompletionTime: entry.CompletionTime.String(),
+					InitialBalance: entry.Balance,
+					Balance:        entry.Balance,
+				})
+			}
+		}
+		return json.Marshal(response)
+	default:
+		return nil, errors.New("unknown Staking extension query")
+	}
+}

wasmbinding/query_plugin.go

@@ -13,6 +13,7 @@ const (
 	EpochRoute        = "epoch"
 	TokenFactoryRoute = "tokenfactory"
 	EVMRoute          = "evm"
+	StakingExtRoute   = "stakingext"
 )
 
 type SeiQueryWrapper struct {
@@ -37,6 +38,8 @@ func CustomQuerier(qp *QueryPlugin) func(ctx sdk.Context, request json.RawMessag
 			return qp.HandleTokenFactoryQuery(ctx, contractQuery.QueryData)
 		case EVMRoute:
 			return qp.HandleEVMQuery(ctx, contractQuery.QueryData)
+		case StakingExtRoute:
+			return qp.HandleStakingExtQuery(ctx, contractQuery.QueryData)
 		default:
 			return nil, wasmvmtypes.UnsupportedRequest{Kind: "Unknown Sei Query Route"}
 		}

wasmbinding/test/query_test.go

@@ -36,7 +36,7 @@ func SetupWasmbindingTest(t *testing.T) (*app.TestWrapper, func(ctx sdk.Context,
 	eh := epochwasm.NewEpochWasmQueryHandler(&testWrapper.App.EpochKeeper)
 	th := tokenfactorywasm.NewTokenFactoryWasmQueryHandler(&testWrapper.App.TokenFactoryKeeper)
 	evmh := evmwasm.NewEVMQueryHandler(&testWrapper.App.EvmKeeper)
-	qp := wasmbinding.NewQueryPlugin(oh, eh, th, evmh)
+	qp := wasmbinding.NewQueryPlugin(oh, eh, th, evmh, testWrapper.App.StakingKeeper)
 	return testWrapper, wasmbinding.CustomQuerier(qp)
 }
 

wasmbinding/wasm.go

@@ -7,6 +7,7 @@ import (
 	codectypes "github.com/cosmos/cosmos-sdk/codec/types"
 	aclkeeper "github.com/cosmos/cosmos-sdk/x/accesscontrol/keeper"
 	authkeeper "github.com/cosmos/cosmos-sdk/x/auth/keeper"
+	stakingkeeper "github.com/cosmos/cosmos-sdk/x/staking/keeper"
 	epochwasm "github.com/sei-protocol/sei-chain/x/epoch/client/wasm"
 	epochkeeper "github.com/sei-protocol/sei-chain/x/epoch/keeper"
 	evmwasm "github.com/sei-protocol/sei-chain/x/evm/client/wasm"
@@ -30,12 +31,13 @@ func RegisterCustomPlugins(
 	portSource wasmtypes.ICS20TransferPortSource,
 	aclKeeper aclkeeper.Keeper,
 	evmKeeper *evmkeeper.Keeper,
+	stakingKeeper stakingkeeper.Keeper,
 ) []wasmkeeper.Option {
 	oracleHandler := oraclewasm.NewOracleWasmQueryHandler(oracle)
 	epochHandler := epochwasm.NewEpochWasmQueryHandler(epoch)
 	tokenfactoryHandler := tokenfactorywasm.NewTokenFactoryWasmQueryHandler(tokenfactory)
 	evmHandler := evmwasm.NewEVMQueryHandler(evmKeeper)
-	wasmQueryPlugin := NewQueryPlugin(oracleHandler, epochHandler, tokenfactoryHandler, evmHandler)
+	wasmQueryPlugin := NewQueryPlugin(oracleHandler, epochHandler, tokenfactoryHandler, evmHandler, stakingKeeper)
 
 	queryPluginOpt := wasmkeeper.WithQueryPlugins(&wasmkeeper.QueryPlugins{
 		Custom: CustomQuerier(wasmQueryPlugin),