Skip to content

Commit

Permalink
CNS-930: change get all store return types
Browse files Browse the repository at this point in the history
  • Loading branch information
oren-lava committed Mar 27, 2024
1 parent e861f77 commit ebb7dc0
Show file tree
Hide file tree
Showing 4 changed files with 83 additions and 95 deletions.
55 changes: 3 additions & 52 deletions x/pairing/genesis.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,6 @@ package pairing

import (
sdk "github.com/cosmos/cosmos-sdk/types"
"github.com/lavanet/lava/utils"
"github.com/lavanet/lava/x/pairing/keeper"
"github.com/lavanet/lava/x/pairing/types"
)
Expand Down Expand Up @@ -37,57 +36,9 @@ func InitGenesis(ctx sdk.Context, k keeper.Keeper, genState types.GenesisState)
func ExportGenesis(ctx sdk.Context, k keeper.Keeper) *types.GenesisState {
genesis := types.DefaultGenesis()
genesis.Params = k.GetParams(ctx)

epochs, uniqueEpochSessions := k.GetAllUniqueEpochSessionStore(ctx)
for i := range epochs {
provider, project, chainID, sessionID, err := types.DecodeUniqueEpochSessionKey(uniqueEpochSessions[i])
if err != nil {
utils.LavaFormatError("could not decode UniqueEpochSessionKey", err, utils.LogAttr("key", uniqueEpochSessions[i]))
continue
}
ues := types.UniqueEpochSessionGenesis{
Epoch: epochs[i],
Provider: provider,
Project: project,
ChainId: chainID,
SessionId: sessionID,
}
genesis.UniqueEpochSessions = append(genesis.UniqueEpochSessions, ues)
}

epochs, keys, providerEpochCus := k.GetAllProviderEpochCuStore(ctx)
for i := range epochs {
provider, chainID, err := types.DecodeProviderEpochCuKey(keys[i])
if err != nil {
utils.LavaFormatError("could not decode ProviderEpochCu key", err, utils.LogAttr("key", keys[i]))
continue
}
pec := types.ProviderEpochCuGenesis{
Epoch: epochs[i],
Provider: provider,
ChainId: chainID,
ProviderEpochCu: providerEpochCus[i],
}
genesis.ProviderEpochCus = append(genesis.ProviderEpochCus, pec)
}

epochs, keys, providerConsumerEpochCus := k.GetAllProviderConsumerEpochCuStore(ctx)
for i := range epochs {
provider, project, chainID, err := types.DecodeProviderConsumerEpochCuKey(keys[i])
if err != nil {
utils.LavaFormatError("could not decode ProviderConsumerEpochCu key", err, utils.LogAttr("key", keys[i]))
continue
}
pcec := types.ProviderConsumerEpochCuGenesis{
Epoch: epochs[i],
Provider: provider,
Project: project,
ChainId: chainID,
ProviderConsumerEpochCu: providerConsumerEpochCus[i],
}
genesis.ProviderConsumerEpochCus = append(genesis.ProviderConsumerEpochCus, pcec)
}

genesis.UniqueEpochSessions = k.GetAllUniqueEpochSessionStore(ctx)
genesis.ProviderEpochCus = k.GetAllProviderEpochCuStore(ctx)
genesis.ProviderConsumerEpochCus = k.GetAllProviderConsumerEpochCuStore(ctx)
genesis.BadgeUsedCuList = k.GetAllBadgeUsedCu(ctx)
genesis.BadgesTS = k.ExportBadgesTimers(ctx)
genesis.ProviderQosFS = k.ExportProviderQoS(ctx)
Expand Down
63 changes: 46 additions & 17 deletions x/pairing/keeper/epoch_cu.go
Original file line number Diff line number Diff line change
Expand Up @@ -64,7 +64,8 @@ func (k Keeper) GetAllUniqueEpochSessionForEpoch(ctx sdk.Context, epoch uint64)
}

// GetAllUniqueEpochSessionStore gets all the UniqueEpochSession objects from the store (used for genesis)
func (k Keeper) GetAllUniqueEpochSessionStore(ctx sdk.Context) (epochs []uint64, keys []string) {
func (k Keeper) GetAllUniqueEpochSessionStore(ctx sdk.Context) []types.UniqueEpochSessionGenesis {
info := []types.UniqueEpochSessionGenesis{}
store := prefix.NewStore(ctx.KVStore(k.storeKey), []byte(types.UniqueEpochSessionPrefix))
iterator := sdk.KVStorePrefixIterator(store, []byte{}) // Get an iterator with no prefix to iterate over all keys
defer iterator.Close()
Expand All @@ -73,14 +74,23 @@ func (k Keeper) GetAllUniqueEpochSessionStore(ctx sdk.Context) (epochs []uint64,
split := strings.Split(key, "/") // key structure: epoch/unique_epoch_session_key
epoch, err := strconv.ParseUint(split[0], 10, 64)
if err != nil {
utils.LavaFormatError("could not decode UniqueEpochSessionKey", err, utils.LogAttr("key", string(iterator.Key())))
utils.LavaFormatError("could not decode UniqueEpochSessionKey with epoch", err, utils.LogAttr("key", string(iterator.Key())))
continue
}
epochs = append(epochs, epoch)
keys = append(keys, split[1])
provider, project, chainID, sessionID, err := types.DecodeUniqueEpochSessionKey(split[1])
if err != nil {
utils.LavaFormatError("could not decode UniqueEpochSessionKey", err, utils.LogAttr("key", split[1]))
}
info = append(info, types.UniqueEpochSessionGenesis{
Epoch: epoch,
Provider: provider,
Project: project,
ChainId: chainID,
SessionId: sessionID,
})
}

return epochs, keys
return info
}

/* ########## ProviderEpochCu ############ */
Expand Down Expand Up @@ -115,25 +125,34 @@ func (k Keeper) RemoveAllProviderEpochCu(ctx sdk.Context, epoch uint64) {
}

// GetAllProviderEpochCuStore returns all the ProviderEpochCu from the store (used for genesis)
func (k Keeper) GetAllProviderEpochCuStore(ctx sdk.Context) (epochs []uint64, keys []string, providerEpochCus []types.ProviderEpochCu) {
func (k Keeper) GetAllProviderEpochCuStore(ctx sdk.Context) []types.ProviderEpochCuGenesis {
info := []types.ProviderEpochCuGenesis{}
store := prefix.NewStore(ctx.KVStore(k.storeKey), []byte(types.ProviderEpochCuPrefix))
iterator := sdk.KVStorePrefixIterator(store, []byte{}) // Get an iterator with no prefix to iterate over all keys
defer iterator.Close()
for ; iterator.Valid(); iterator.Next() {
split := strings.Split(string(iterator.Key()), "/") // key structure: epoch/provider_epoch_cu_key
epoch, err := strconv.ParseUint(split[0], 10, 64)
if err != nil {
utils.LavaFormatError("could not decode ProviderEpochCuKey", err, utils.LogAttr("key", string(iterator.Key())))
utils.LavaFormatError("could not decode ProviderEpochCuKey with epoch", err, utils.LogAttr("key", string(iterator.Key())))
continue
}
provider, chainID, err := types.DecodeProviderEpochCuKey(split[1])
if err != nil {
utils.LavaFormatError("could not decode ProviderEpochCuKey with epoch", err, utils.LogAttr("key", split[1]))
continue
}
epochs = append(epochs, epoch)
keys = append(keys, split[1])
var pec types.ProviderEpochCu
k.cdc.MustUnmarshal(iterator.Value(), &pec)
providerEpochCus = append(providerEpochCus, pec)
info = append(info, types.ProviderEpochCuGenesis{
Epoch: epoch,
Provider: provider,
ChainId: chainID,
ProviderEpochCu: pec,
})
}

return epochs, keys, providerEpochCus
return info
}

/* ########## ProviderConsumerEpochCu ############ */
Expand Down Expand Up @@ -179,23 +198,33 @@ func (k Keeper) GetAllProviderConsumerEpochCu(ctx sdk.Context, epoch uint64) (ke
}

// GetAllProviderConsumerEpochCuStore returns all the ProviderConsumerEpochCu from the store (used for genesis)
func (k Keeper) GetAllProviderConsumerEpochCuStore(ctx sdk.Context) (epochs []uint64, keys []string, providerConsumerEpochCus []types.ProviderConsumerEpochCu) {
func (k Keeper) GetAllProviderConsumerEpochCuStore(ctx sdk.Context) []types.ProviderConsumerEpochCuGenesis {
info := []types.ProviderConsumerEpochCuGenesis{}
store := prefix.NewStore(ctx.KVStore(k.storeKey), []byte(types.ProviderConsumerEpochCuPrefix))
iterator := sdk.KVStorePrefixIterator(store, []byte{}) // Get an iterator with no prefix to iterate over all keys
defer iterator.Close()
for ; iterator.Valid(); iterator.Next() {
split := strings.Split(string(iterator.Key()), "/") // key structure: epoch/provider_consumer_epoch_key
epoch, err := strconv.ParseUint(split[0], 10, 64)
if err != nil {
utils.LavaFormatError("could not decode ProviderConsumerEpochCuKey", err, utils.LogAttr("key", string(iterator.Key())))
utils.LavaFormatError("could not decode ProviderConsumerEpochCuKey with epoch", err, utils.LogAttr("key", string(iterator.Key())))
continue
}
provider, project, chainID, err := types.DecodeProviderConsumerEpochCuKey(split[1])
if err != nil {
utils.LavaFormatError("could not decode ProviderConsumerEpochCuKey", err, utils.LogAttr("key", split[1]))
continue
}
epochs = append(epochs, epoch)
keys = append(keys, split[1])
var pcec types.ProviderConsumerEpochCu
k.cdc.MustUnmarshal(iterator.Value(), &pcec)
providerConsumerEpochCus = append(providerConsumerEpochCus, pcec)
info = append(info, types.ProviderConsumerEpochCuGenesis{
Epoch: epoch,
Provider: provider,
Project: project,
ChainId: chainID,
ProviderConsumerEpochCu: pcec,
})
}

return epochs, keys, providerConsumerEpochCus
return info
}
46 changes: 29 additions & 17 deletions x/pairing/keeper/epoch_cu_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -63,13 +63,18 @@ func TestUniqueEpochSessionGetAll(t *testing.T) {
func TestUniqueEpochSessionGetAllStore(t *testing.T) {
keeper, ctx := keepertest.PairingKeeper(t)
items := createNUniqueEpochSession(keeper, ctx, 10)
expectedKeys := []string{}
expectedInfo := []types.UniqueEpochSessionGenesis{}
for i, item := range items {
key := string(types.UniqueEpochSessionKey(item, item, item, uint64(i)))
expectedKeys = append(expectedKeys, key)
expectedInfo = append(expectedInfo, types.UniqueEpochSessionGenesis{
Epoch: uint64(i),
Provider: item,
Project: item,
ChainId: item,
SessionId: uint64(i),
})
}
_, keys := keeper.GetAllUniqueEpochSessionStore(ctx)
require.ElementsMatch(t, nullify.Fill(expectedKeys), nullify.Fill(keys))
info := keeper.GetAllUniqueEpochSessionStore(ctx)
require.ElementsMatch(t, nullify.Fill(expectedInfo), nullify.Fill(info))
}

/* ########## ProviderEpochCu ############ */
Expand Down Expand Up @@ -109,15 +114,18 @@ func TestProviderEpochCuRemove(t *testing.T) {
func TestProviderEpochCuGetAllStore(t *testing.T) {
keeper, ctx := keepertest.PairingKeeper(t)
items := createNProviderEpochCu(keeper, ctx, 10)
expectedKeys := []string{}
expectedInfo := []types.ProviderEpochCuGenesis{}
for i := range items {
name := strconv.Itoa(i)
key := string(types.ProviderEpochCuKey(name, name))
expectedKeys = append(expectedKeys, key)
expectedInfo = append(expectedInfo, types.ProviderEpochCuGenesis{
Epoch: uint64(i),
Provider: name,
ChainId: name,
ProviderEpochCu: types.ProviderEpochCu{ServicedCu: uint64(i)},
})
}
_, keys, pecs := keeper.GetAllProviderEpochCuStore(ctx)
require.ElementsMatch(t, nullify.Fill(expectedKeys), nullify.Fill(keys))
require.ElementsMatch(t, items, pecs)
info := keeper.GetAllProviderEpochCuStore(ctx)
require.ElementsMatch(t, nullify.Fill(expectedInfo), nullify.Fill(info))
}

/* ########## ProviderConsumerEpochCu ############ */
Expand Down Expand Up @@ -176,13 +184,17 @@ func TestProviderConsumerEpochCuGetAll(t *testing.T) {
func TestProviderConsumerEpochCuGetAllStore(t *testing.T) {
keeper, ctx := keepertest.PairingKeeper(t)
items := createNProviderConsumerEpochCu(keeper, ctx, 10)
expectedKeys := []string{}
expectedInfo := []types.ProviderConsumerEpochCuGenesis{}
for i := range items {
name := strconv.Itoa(i)
key := string(types.ProviderConsumerEpochCuKey(name, name, name))
expectedKeys = append(expectedKeys, key)
expectedInfo = append(expectedInfo, types.ProviderConsumerEpochCuGenesis{
Epoch: uint64(i),
Provider: name,
Project: name,
ChainId: name,
ProviderConsumerEpochCu: types.ProviderConsumerEpochCu{Cu: uint64(i)},
})
}
_, keys, pecs := keeper.GetAllProviderConsumerEpochCuStore(ctx)
require.ElementsMatch(t, nullify.Fill(expectedKeys), nullify.Fill(keys))
require.ElementsMatch(t, items, pecs)
info := keeper.GetAllProviderConsumerEpochCuStore(ctx)
require.ElementsMatch(t, nullify.Fill(expectedInfo), nullify.Fill(info))
}
14 changes: 5 additions & 9 deletions x/pairing/keeper/grpc_query_providers_epoch_cu.go
Original file line number Diff line number Diff line change
Expand Up @@ -18,16 +18,12 @@ func (k Keeper) ProvidersEpochCu(goCtx context.Context, req *types.QueryProvider
ctx := sdk.UnwrapSDKContext(goCtx)
infoMap := map[string]uint64{}

_, keys, pecs := k.GetAllProviderEpochCuStore(ctx)
for i := range pecs {
provider, _, err := types.DecodeProviderEpochCuKey(keys[i])
if err != nil {
continue
}
if _, ok := infoMap[provider]; !ok {
infoMap[provider] = pecs[i].ServicedCu
infos := k.GetAllProviderEpochCuStore(ctx)
for _, info := range infos {
if _, ok := infoMap[info.Provider]; !ok {
infoMap[info.Provider] = info.ProviderEpochCu.ServicedCu
} else {
infoMap[provider] += pecs[i].ServicedCu
infoMap[info.Provider] += info.ProviderEpochCu.ServicedCu
}
}

Expand Down

0 comments on commit ebb7dc0

Please sign in to comment.