romans suggestion

app/upgrades/v17/upgrades.go

@@ -152,27 +152,28 @@ func CreateUpgradeHandler(
 }
 
 func FlipTwapSpotPriceRecords(ctx sdk.Context, poolIds []uint64, keepers *keepers.AppKeepers) error {
-	twapRecordHistoricalTimeIndexed, err := keepers.TwapKeeper.GetAllHistoricalTimeIndexedTWAPs(ctx)
-	if err != nil {
-		return err
-	}
-
-	for _, historicalTwapRecord := range twapRecordHistoricalTimeIndexed {
-		oldRecord := historicalTwapRecord
-		historicalTwapRecord.Asset0Denom, historicalTwapRecord.Asset1Denom = oldRecord.Asset1Denom, oldRecord.Asset0Denom
-		historicalTwapRecord.P0LastSpotPrice, historicalTwapRecord.P1LastSpotPrice = oldRecord.P1LastSpotPrice, oldRecord.P0LastSpotPrice
-
-		keepers.TwapKeeper.StoreHistoricalTWAP(ctx, historicalTwapRecord)
-		keepers.TwapKeeper.DeleteHistoricalRecord(ctx, oldRecord)
-	}
-
 	for _, poolId := range poolIds {
 		// check that this is a cl pool
 		_, err := keepers.ConcentratedLiquidityKeeper.GetConcentratedPoolById(ctx, poolId)
 		if err != nil {
 			return err
 		}
 
+		// check if the pool is CL osmo/dai pool
+		twapRecordHistoricalPoolIndexed, err := keepers.TwapKeeper.GetAllHistoricalPoolIndexedTWAPsForPoolId(ctx, poolId)
+		if err != nil {
+			return err
+		}
+
+		for _, historicalTwapRecord := range twapRecordHistoricalPoolIndexed {
+			oldRecord := historicalTwapRecord
+			historicalTwapRecord.Asset0Denom, historicalTwapRecord.Asset1Denom = oldRecord.Asset1Denom, oldRecord.Asset0Denom
+			historicalTwapRecord.P0LastSpotPrice, historicalTwapRecord.P1LastSpotPrice = oldRecord.P1LastSpotPrice, oldRecord.P0LastSpotPrice
+
+			keepers.TwapKeeper.StoreHistoricalTWAP(ctx, historicalTwapRecord)
+			keepers.TwapKeeper.DeleteHistoricalRecord(ctx, oldRecord)
+		}
+
 		// check that the twap record exists
 		clPoolTwapRecords, err := keepers.TwapKeeper.GetAllMostRecentRecordsForPool(ctx, poolId)
 		if err != nil {

app/upgrades/v17/upgrades_test.go

@@ -19,9 +19,8 @@ import (
 	"github.com/osmosis-labs/osmosis/v17/app/keepers"
 	v17 "github.com/osmosis-labs/osmosis/v17/app/upgrades/v17"
 	cltypes "github.com/osmosis-labs/osmosis/v17/x/concentrated-liquidity/types"
-	"github.com/osmosis-labs/osmosis/v17/x/twap/types"
-
 	poolmanagertypes "github.com/osmosis-labs/osmosis/v17/x/poolmanager/types"
+	"github.com/osmosis-labs/osmosis/v17/x/twap/types"
 )
 
 type UpgradeTestSuite struct {
@@ -88,12 +87,12 @@ func (suite *UpgradeTestSuite) TestUpgrade() {
 
 	testCases := []struct {
 		name        string
-		pre_upgrade func(sdk.Context, *keepers.AppKeepers) (sdk.Coins, uint64, []types.TwapRecord)
-		upgrade     func(sdk.Context, *keepers.AppKeepers, sdk.Coins, uint64, []types.TwapRecord)
+		pre_upgrade func(sdk.Context, *keepers.AppKeepers) (sdk.Coins, uint64, []types.TwapRecord, []types.TwapRecord)
+		upgrade     func(sdk.Context, *keepers.AppKeepers, sdk.Coins, uint64, []types.TwapRecord, []types.TwapRecord)
 	}{
 		{
 			"Test that the upgrade succeeds",
-			func(ctx sdk.Context, keepers *keepers.AppKeepers) (sdk.Coins, uint64, []types.TwapRecord) {
+			func(ctx sdk.Context, keepers *keepers.AppKeepers) (sdk.Coins, uint64, []types.TwapRecord, []types.TwapRecord) {
 				upgradeSetup()
 
 				var lastPoolID uint64 // To keep track of the last assigned pool ID
@@ -143,25 +142,27 @@ func (suite *UpgradeTestSuite) TestUpgrade() {
 					lastPoolID = poolID
 				}
 
-				// poolId = 1040
-				poolId := suite.PrepareConcentratedPoolWithCoins("ibc/1480B8FD20AD5FCAE81EA87584D269547DD4D436843C1D20F15E00EB64743EF4", "uosmo")
+				existingPool := suite.PrepareConcentratedPoolWithCoins("ibc/1480B8FD20AD5FCAE81EA87584D269547DD4D436843C1D20F15E00EB64743EF4", "uosmo")
 
 				// create few TWAP records for the pools
-				t1 := dummyTwapRecord(poolId.GetId(), time.Now().Add(-time.Hour*24), "ibc/1480B8FD20AD5FCAE81EA87584D269547DD4D436843C1D20F15E00EB64743EF4", "uosmo", sdk.NewDec(10),
+				t1 := dummyTwapRecord(existingPool.GetId(), time.Now().Add(-time.Hour*24), "ibc/1480B8FD20AD5FCAE81EA87584D269547DD4D436843C1D20F15E00EB64743EF4", "uosmo", sdk.NewDec(10),
 					sdk.OneDec().MulInt64(10*10),
 					sdk.OneDec().MulInt64(3),
 					sdk.ZeroDec())
 
 				suite.App.TwapKeeper.StoreNewRecord(suite.Ctx, t1)
 
-				clPoolTwapRecordPreUpgrade, err := keepers.TwapKeeper.GetAllMostRecentRecordsForPool(ctx, poolId.GetId())
+				clPoolTwapRecordPreUpgrade, err := keepers.TwapKeeper.GetAllMostRecentRecordsForPool(ctx, existingPool.GetId())
 				suite.Require().NoError(err)
 
-				return expectedCoinsUsedInUpgradeHandler, lastPoolID, clPoolTwapRecordPreUpgrade
+				clPoolTwapRecordHistoricalPoolIndex, err := keepers.TwapKeeper.GetAllHistoricalPoolIndexedTWAPsForPoolId(ctx, existingPool.GetId())
+				suite.Require().NoError(err)
+
+				return expectedCoinsUsedInUpgradeHandler, lastPoolID, clPoolTwapRecordPreUpgrade, clPoolTwapRecordHistoricalPoolIndex
 
 			},
-			func(ctx sdk.Context, keepers *keepers.AppKeepers, expectedCoinsUsedInUpgradeHandler sdk.Coins, lastPoolID uint64, twapRecord []types.TwapRecord) {
-				lastPoolID = twapRecord[0].PoolId
+			func(ctx sdk.Context, keepers *keepers.AppKeepers, expectedCoinsUsedInUpgradeHandler sdk.Coins, lastPoolID uint64, clPoolTwapRecordPreUpgrade []types.TwapRecord, clPoolTwapRecordHistoricalPoolIndex []types.TwapRecord) {
+				lastPoolID = clPoolTwapRecordPreUpgrade[0].GetPoolId()
 				stakingParams := suite.App.StakingKeeper.GetParams(suite.Ctx)
 				stakingParams.BondDenom = "uosmo"
 				suite.App.StakingKeeper.SetParams(suite.Ctx, stakingParams)
@@ -176,19 +177,6 @@ func (suite *UpgradeTestSuite) TestUpgrade() {
 					suite.App.BeginBlocker(suite.Ctx, abci.RequestBeginBlock{})
 				})
 
-				twapRecordHistoricalTimeIndexedPOST, err := keepers.TwapKeeper.GetAllHistoricalPoolIndexedTWAPs(ctx)
-				suite.Require().NoError(err)
-
-				fmt.Println("POST UPGRADE: ", twapRecordHistoricalTimeIndexedPOST)
-
-				clPoolTwapRecordPostUpgrade, err := keepers.TwapKeeper.GetAllMostRecentRecordsForPool(ctx, 1040)
-				suite.Require().NoError(err)
-
-				for idx := range clPoolTwapRecordPostUpgrade {
-					suite.Require().Equal(twapRecord[idx].Asset0Denom, clPoolTwapRecordPostUpgrade[idx].Asset1Denom)
-					suite.Require().Equal(twapRecord[idx].Asset1Denom, clPoolTwapRecordPostUpgrade[idx].Asset0Denom)
-				}
-
 				// Retrieve the community pool balance (and the feePool balance) after the upgrade
 				communityPoolBalancePost := suite.App.BankKeeper.GetAllBalances(suite.Ctx, communityPoolAddress)
 				feePoolCommunityPoolPost := suite.App.DistrKeeper.GetFeePool(suite.Ctx).CommunityPool
@@ -251,11 +239,11 @@ func (suite *UpgradeTestSuite) TestUpgrade() {
 		},
 		{
 			"Fails because CFMM pool is not found",
-			func(ctx sdk.Context, keepers *keepers.AppKeepers) (sdk.Coins, uint64, []types.TwapRecord) {
+			func(ctx sdk.Context, keepers *keepers.AppKeepers) (sdk.Coins, uint64, []types.TwapRecord, []types.TwapRecord) {
 				upgradeSetup()
-				return sdk.NewCoins(), 0, nil
+				return sdk.NewCoins(), 0, nil, nil
 			},
-			func(ctx sdk.Context, keepers *keepers.AppKeepers, expectedCoinsUsedInUpgradeHandler sdk.Coins, lastPoolID uint64, twapRecord []types.TwapRecord) {
+			func(ctx sdk.Context, keepers *keepers.AppKeepers, expectedCoinsUsedInUpgradeHandler sdk.Coins, lastPoolID uint64, clPoolTwapRecordPreUpgrade []types.TwapRecord, clPoolTwapRecordHistoricalPoolIndex []types.TwapRecord) {
 				dummyUpgrade(suite)
 				suite.Require().Panics(func() {
 					suite.App.BeginBlocker(suite.Ctx, abci.RequestBeginBlock{})
@@ -268,8 +256,8 @@ func (suite *UpgradeTestSuite) TestUpgrade() {
 		suite.Run(fmt.Sprintf("Case %s", tc.name), func() {
 			suite.SetupTest() // reset
 
-			expectedCoinsUsedInUpgradeHandler, lastPoolID, twapRecord := tc.pre_upgrade(suite.Ctx, &suite.App.AppKeepers)
-			tc.upgrade(suite.Ctx, &suite.App.AppKeepers, expectedCoinsUsedInUpgradeHandler, lastPoolID, twapRecord)
+			expectedCoinsUsedInUpgradeHandler, lastPoolID, clPoolTwapRecordPreUpgrade, clPoolTwapRecordHistoricalPoolIndex := tc.pre_upgrade(suite.Ctx, &suite.App.AppKeepers)
+			tc.upgrade(suite.Ctx, &suite.App.AppKeepers, expectedCoinsUsedInUpgradeHandler, lastPoolID, clPoolTwapRecordPreUpgrade, clPoolTwapRecordHistoricalPoolIndex)
 		})
 	}
 }

x/twap/store.go

@@ -162,6 +162,11 @@ func (k Keeper) GetAllHistoricalPoolIndexedTWAPs(ctx sdk.Context) ([]types.TwapR
 	return osmoutils.GatherValuesFromStorePrefix(ctx.KVStore(k.storeKey), []byte(types.HistoricalTWAPPoolIndexPrefix), types.ParseTwapFromBz)
 }
 
+// GetAllHistoricalPoolIndexedTWAPsForPoolId returns HistoricalTwapRecord for a pool give poolId.
+func (k Keeper) GetAllHistoricalPoolIndexedTWAPsForPoolId(ctx sdk.Context, poolId uint64) ([]types.TwapRecord, error) {
+	return osmoutils.GatherValuesFromStorePrefixWithKeyParser(ctx.KVStore(k.storeKey), types.FormatKeyPoolTwapRecords(poolId), types.ParseTwapHistoricalPoolIndexedRecordFromBz)
+}
+
 // StoreNewRecord stores a record, in both the most recent record store and historical stores.
 // NOTE: if it's a new Record make sure to update asset0 and asset1 as well.
 func (k Keeper) StoreNewRecord(ctx sdk.Context, twap types.TwapRecord) {

x/twap/store_test.go

@@ -574,3 +574,47 @@ func (s *TestSuite) TestAccumulatorOverflow() {
 		})
 	}
 }
+
+func (s *TestSuite) TestGetAllHistoricalPoolIndexedTWAPsForPooId() {
+	baseRecord := newEmptyPriceRecord(1, baseTime, denom0, denom1)
+	tPlusOneRecord := newEmptyPriceRecord(1, tPlusOne, denom0, denom1)
+	tests := map[string]struct {
+		recordsToSet    []types.TwapRecord
+		poolId          uint64
+		expectedRecords []types.TwapRecord
+	}{
+		"set single record": {
+			poolId:          1,
+			expectedRecords: []types.TwapRecord{baseRecord},
+		},
+		"query non-existent pool": {
+			poolId:          2,
+			expectedRecords: []types.TwapRecord{},
+		},
+		"set single record, different pool ID": {
+			poolId:          2,
+			expectedRecords: []types.TwapRecord{newEmptyPriceRecord(2, baseTime, denom0, denom1)},
+		},
+		"set two records": {
+			poolId:          1,
+			expectedRecords: []types.TwapRecord{baseRecord, tPlusOneRecord},
+		},
+	}
+
+	for name, test := range tests {
+		s.Run(name, func() {
+			s.SetupTest()
+			twapKeeper := s.twapkeeper
+			s.preSetRecords(test.expectedRecords)
+
+			// System under test.
+			actualRecords, err := twapKeeper.GetAllHistoricalPoolIndexedTWAPsForPoolId(s.Ctx, test.poolId)
+			s.NoError(err)
+
+			// Assertions.
+			s.Equal(test.expectedRecords, actualRecords)
+
+		})
+	}
+
+}

x/twap/types/keys.go

@@ -3,6 +3,8 @@ package types
 import (
 	"errors"
 	fmt "fmt"
+	"strconv"
+	"strings"
 	time "time"
 
 	sdk "github.com/cosmos/cosmos-sdk/types"
@@ -46,6 +48,10 @@ var (
 
 // TODO: make utility command to automatically interlace separators
 
+func FormatKeyPoolTwapRecords(poolId uint64) []byte {
+	return []byte(fmt.Sprintf("%s%d", HistoricalTWAPPoolIndexPrefix, poolId))
+}
+
 func FormatMostRecentTWAPKey(poolId uint64, denom1, denom2 string) []byte {
 	poolIdS := osmoutils.FormatFixedLengthU64(poolId)
 	return []byte(fmt.Sprintf("%s%s%s%s%s%s", mostRecentTWAPsPrefix, poolIdS, KeySeparator, denom1, KeySeparator, denom2))
@@ -92,3 +98,30 @@ func ParseTwapFromBz(bz []byte) (twap TwapRecord, err error) {
 	}
 	return twap, err
 }
+
+// ParseTwapHistoricalPoolIndexedRecordFromBz parses through an existing Twap Record by key and returns Twap Record.
+func ParseTwapHistoricalPoolIndexedRecordFromBz(key []byte, value []byte) (TwapRecord, error) {
+	if len(key) == 0 {
+		return TwapRecord{}, fmt.Errorf("Invalid twap record key")
+	}
+	if len(value) == 0 {
+		return TwapRecord{}, fmt.Errorf("Invalid twap record value")
+	}
+
+	keyStr := string(key)
+	twapRecordKeyComponent := strings.Split(keyStr, KeySeparator)
+
+	poolId, err := strconv.ParseUint(twapRecordKeyComponent[1], 10, 64)
+	if err != nil {
+		return TwapRecord{}, fmt.Errorf("cannot parse poolId key")
+	}
+
+	twapRecord, err := ParseTwapFromBz(value)
+	if err != nil {
+		return TwapRecord{}, fmt.Errorf("cannot parse TwapRecord value")
+	}
+
+	twapRecord.PoolId = poolId
+
+	return twapRecord, nil
+}