fix: add missing pruning logic, add missing block fk constraints

pkg/postgres/migrations/202506042023_addMissingBlockNumberForeignKeys/up.go

@@ -0,0 +1,31 @@
+package _202506042023_addMissingBlockNumberForeignKeys
+
+import (
+	"database/sql"
+	"fmt"
+	"github.com/Layr-Labs/sidecar/internal/config"
+	"gorm.io/gorm"
+)
+
+type Migration struct {
+}
+
+func (m *Migration) Up(db *sql.DB, grm *gorm.DB, cfg *config.Config) error {
+	tableNames := []string{
+		"staker_shares",
+		"operator_shares",
+	}
+
+	for _, tableName := range tableNames {
+		query := fmt.Sprintf(`alter table %s add constraint %s_block_number_fkey foreign key (block_number) references blocks (number) on delete cascade;`, tableName, tableName)
+		res := grm.Exec(query)
+		if res.Error != nil {
+			return fmt.Errorf("failed to add foreign key for table %s: %w", tableName, res.Error)
+		}
+	}
+	return nil
+}
+
+func (m *Migration) GetName() string {
+	return "202506042023_addMissingBlockNumberForeignKeys"
+}

pkg/postgres/migrations/migrator.go

@@ -71,6 +71,7 @@ import (
 	_202503311108_goldRewardHashIndex "github.com/Layr-Labs/sidecar/pkg/postgres/migrations/202503311108_goldRewardHashIndex"
 	_202504240743_fixQueuedSlashingWithdrawalsPk "github.com/Layr-Labs/sidecar/pkg/postgres/migrations/202504240743_fixQueuedSlashingWithdrawalsPk"
 	_202505092007_startupJobs "github.com/Layr-Labs/sidecar/pkg/postgres/migrations/202505092007_startupJobs"
+	_202506042023_addMissingBlockNumberForeignKeys "github.com/Layr-Labs/sidecar/pkg/postgres/migrations/202506042023_addMissingBlockNumberForeignKeys"
 	"go.uber.org/zap"
 	"gorm.io/gorm"
 	"time"
@@ -213,6 +214,7 @@ func (m *Migrator) MigrateAll() error {
 		&_202503171414_slashingWithdrawals.Migration{},
 		&_202504240743_fixQueuedSlashingWithdrawalsPk.Migration{},
 		&_202505092007_startupJobs.Migration{},
+		&_202506042023_addMissingBlockNumberForeignKeys.Migration{},
 	}
 
 	for _, migration := range migrations {

pkg/rewards/rewards.go

@@ -13,7 +13,6 @@ import (
 
 	"sync/atomic"
 
-	"slices"
 	"strings"
 
 	"strconv"
@@ -432,7 +431,42 @@ func (rc *RewardsCalculator) findRewardsTablesBySnapshotDate(snapshotDate string
 	return rewardsTables, nil
 }
 
+// FindRewardHashesToPrune finds all reward hashes that will be pruned when deleting corrupted state.
+func (rc *RewardsCalculator) FindRewardHashesToPrune(blockHeight uint64) ([]string, error) {
+	query := `
+		select
+			distinct(reward_hash) as reward_hash
+		from (
+			select reward_hash from combined_rewards where block_number >= @blockNumber
+			union all
+			select reward_hash from operator_directed_rewards where block_number >= @blockNumber
+			union all
+			select
+				odosrs.reward_hash
+			from operator_directed_operator_set_reward_submissions as odosrs
+			-- operator_directed_operator_set_reward_submissions lacks a block_time column, so we need to join blocks
+			join blocks as b on (b.number = odosrs.block_number)
+			where
+				b.number >= @blockNumber
+		) as t
+	`
+	var rewardHashes []string
+	res := rc.grm.Raw(query, sql.Named("blockNumber", blockHeight)).Scan(&rewardHashes)
+	if res.Error != nil {
+		rc.logger.Sugar().Errorw("Failed to find reward hashes to prune", "error", res.Error)
+		return nil, res.Error
+	}
+	return rewardHashes, nil
+}
+
+// DeleteCorruptedRewardsFromBlockHeight deletes rewards data that is >= the provided block height.
 func (rc *RewardsCalculator) DeleteCorruptedRewardsFromBlockHeight(blockHeight uint64) error {
+	rewardHashesToPrune, err := rc.FindRewardHashesToPrune(blockHeight)
+	if err != nil {
+		rc.logger.Sugar().Errorw("Failed to find reward hashes to prune", "error", err)
+		return err
+	}
+
 	generatedSnapshot, err := rc.findGeneratedRewardSnapshotByBlock(blockHeight)
 	if err != nil {
 		rc.logger.Sugar().Errorw("Failed to find generated snapshot", "error", err)
@@ -462,9 +496,12 @@ func (rc *RewardsCalculator) DeleteCorruptedRewardsFromBlockHeight(blockHeight u
 		lowerBoundSnapshot = nil
 	}
 
-	snapshotDates := make([]string, 0)
+	// Find all suffixed rewards tables that were generated as part of the rewards calculation process.
+	// e.g. gold_6_rfae_stakers_2024_12_01
+	// We since we're deleting the blocks where these rewards were applicable, we need to drop these tables
+	// so that they can be created again when we run the rewards calculation again.
 	for _, snapshot := range snapshotsToDelete {
-		snapshotDates = append(snapshotDates, snapshot.SnapshotDate)
+		// find and drop tables that were created for this snapshot (e.g. gold_6_rfae_stakers_2024_12_01)
 		tableNames, err := rc.findRewardsTablesBySnapshotDate(snapshot.SnapshotDate)
 		if err != nil {
 			rc.logger.Sugar().Errorw("Failed to find rewards tables", "error", err)
@@ -489,19 +526,49 @@ func (rc *RewardsCalculator) DeleteCorruptedRewardsFromBlockHeight(blockHeight u
 		}
 	}
 
-	// sort all snapshot dates in ascending order to purge from gold table
-	slices.SortFunc(snapshotDates, func(i, j string) int {
-		return strings.Compare(i, j)
-	})
+	// prune snapshot tables based on snapshot dates.
+	// - tables that have been migrated to the "new" single table format that is never dropped will end up getting backfilled.
+	// - tables that we still drop, dont need to be pruned since they'll get recreated.
+	tablesToPrune := []string{
+		"combined_rewards",
+		"default_operator_split_snapshots",
+		"operator_avs_registration_snapshots",
+		"operator_avs_split_snapshots",
+		"operator_avs_strategy_snapshots",
+		"operator_directed_rewards",
+		"operator_pi_split_snapshots",
+		"operator_share_snapshots",
+		"staker_delegation_snapshots",
+		"staker_share_snapshots",
+	}
+	for _, tableName := range tablesToPrune {
+		var query string
+
+		// if we have no lower bound, that means we're deleting everything
+		if lowerBoundSnapshot == nil {
+			query = fmt.Sprintf(`truncate table %s cascade`, tableName)
+		} else {
+			query = fmt.Sprintf(`delete from %s where snapshot >= '%s'`, tableName, lowerBoundSnapshot.SnapshotDate)
+		}
 
-	// purge from gold table
-	if lowerBoundSnapshot != nil {
-		rc.logger.Sugar().Infow("Purging rewards from gold table where snapshot >=", "snapshotDate", lowerBoundSnapshot.SnapshotDate)
-		res = rc.grm.Exec(`delete from gold_table where snapshot >= @snapshotDate`, sql.Named("snapshotDate", lowerBoundSnapshot.SnapshotDate))
-	} else {
-		// if the lower bound is nil, ther we're deleting everything
-		rc.logger.Sugar().Infow("Purging all rewards from gold table")
-		res = rc.grm.Exec(`delete from gold_table`)
+		res = rc.grm.Exec(query)
+		if res.Error != nil {
+			rc.logger.Sugar().Errorw("Failed to prune snapshot table", "tableName", tableName, "error", res.Error)
+			return res.Error
+		}
+	}
+
+	// Purge rewards from gold_table based on reward_hash.
+	// Since reward submissions are submitted at a block, but can be arbitrarily backwards looking,
+	// we cant just delete by snapshot date. We have to get all reward snapshots that were created by
+	// a specific reward hash that we're deleting.
+	goldPurgeQuery := `
+		delete from gold_table where reward_hash in @rewardHashes
+	`
+	res = rc.grm.Exec(goldPurgeQuery, sql.Named("rewardHashes", rewardHashesToPrune))
+	if res.Error != nil {
+		rc.logger.Sugar().Errorw("Failed to delete rewards from gold table by reward_hash", "error", res.Error)
+		return res.Error
 	}
 
 	if res.Error != nil {
@@ -510,7 +577,7 @@ func (rc *RewardsCalculator) DeleteCorruptedRewardsFromBlockHeight(blockHeight u
 	}
 	if lowerBoundSnapshot != nil {
 		rc.logger.Sugar().Infow("Deleted rewards from gold table",
-			zap.String("snapshotDate", lowerBoundSnapshot.SnapshotDate),
+			zap.Int("rewardHashCount", len(rewardHashesToPrune)),
 			zap.Int64("recordsDeleted", res.RowsAffected),
 		)
 	} else {