add operator_set_id

Author: serichoi65

pkg/postgres/migrations/202511141700_withdrawalQueueAndAllocationRounding/up.go

@@ -13,37 +13,23 @@ type Migration struct {
 func (m *Migration) Up(db *sql.DB, grm *gorm.DB, cfg *config.Config) error {
 	queries := []string{
 		// =============================================================================
-		// PART 1: Enhance queued_slashing_withdrawals to support withdrawal queue
+		// PART 1: Withdrawal queue - no schema changes needed
 		// =============================================================================
-
-		// Add completion tracking (timestamps can be derived from blocks table via FK)
-		`alter table queued_slashing_withdrawals add column if not exists completed boolean default false`,
-		`alter table queued_slashing_withdrawals add column if not exists completion_block_number bigint`,
-
-		// Add FK constraint for completion block
-		`alter table queued_slashing_withdrawals add constraint fk_completion_block foreign key (completion_block_number) references blocks(number) on delete set null`,
-
-		// =============================================================================
-		// PART 2: Create indexes for efficient withdrawal queue queries
-		// =============================================================================
-
-		// Index for finding active (non-completed) withdrawals
-		`create index if not exists idx_queued_withdrawals_active on queued_slashing_withdrawals(staker, operator, completed) where completed = false`,
-
-		// Index for completion queries
-		`create index if not exists idx_queued_withdrawals_completed on queued_slashing_withdrawals(completion_block_number) where completed = true`,
+		// Note: Withdrawal queue logic uses withdrawable_date to determine when
+		// shares should stop earning rewards. No completion tracking needed for
+		// rewards calculation as withdrawable_date is derived from queued_date + 14 days.
 
 		// =============================================================================
-		// PART 3: Update operator_allocations table for rounding logic
+		// PART 2: Update operator_allocations table for rounding logic
 		// =============================================================================
 
 		// Add effective_date column for allocation/deallocation rounding
 		// This is a computed value based on magnitude changes (round UP for increases, DOWN for decreases)
 		// block_timestamp can be derived from block_number FK to blocks table
 		`alter table operator_allocations add column if not exists effective_date date`,
 
-		// Create index for effective_date queries
-		`create index if not exists idx_operator_allocations_effective_date on operator_allocations(operator, avs, strategy, effective_date)`,
+		// Create index for effective_date queries (includes operator_set_id for proper partitioning)
+		`create index if not exists idx_operator_allocations_effective_date on operator_allocations(operator, avs, strategy, operator_set_id, effective_date)`,
 	}
 
 	for _, query := range queries {

pkg/rewards/deallocationQueueShareSnapshots.go

@@ -25,18 +25,29 @@ import (
 // NOTE: This handles DECREASES (deallocations). Increases (allocations) are simpler
 // because they round UP, so they naturally start earning on effective_date.
 const deallocationQueueShareSnapshotsQuery = `
+	insert into deallocation_queue_snapshots(
+		operator,
+		avs,
+		operator_set_id,
+		strategy,
+		magnitude_decrease,
+		block_date,
+		effective_date,
+		snapshot
+	)
 	with deallocation_adjustments as (
 		select
 			oa.operator,
 			oa.avs,
+			oa.operator_set_id,
 			oa.strategy,
 			oa.magnitude as new_magnitude,
 			oa.effective_date,
 			date(b.block_time) as block_date,
 			oa.block_number,
 			-- Get previous allocation to calculate the difference
 			lag(oa.magnitude) over (
-				partition by oa.operator, oa.avs, oa.strategy
+				partition by oa.operator, oa.avs, oa.operator_set_id, oa.strategy
 				order by oa.block_number, oa.log_index
 			) as prev_magnitude
 		from operator_allocations oa
@@ -52,6 +63,7 @@ const deallocationQueueShareSnapshotsQuery = `
 		select
 			operator,
 			avs,
+			operator_set_id,
 			strategy,
 			new_magnitude,
 			prev_magnitude,
@@ -64,18 +76,10 @@ const deallocationQueueShareSnapshotsQuery = `
 			prev_magnitude is not null
 			and new_magnitude::numeric < prev_magnitude::numeric
 	)
-	insert into deallocation_queue_snapshots(
-		operator,
-		avs,
-		strategy,
-		magnitude_decrease,
-		block_date,
-		effective_date,
-		snapshot
-	)
 	select
 		operator,
 		avs,
+		operator_set_id,
 		strategy,
 		magnitude_decrease,
 		block_date,
@@ -88,6 +92,7 @@ const deallocationQueueShareSnapshotsQuery = `
 type DeallocationQueueSnapshot struct {
 	Operator          string `gorm:"column:operator;primaryKey"`
 	Avs               string `gorm:"column:avs;primaryKey"`
+	OperatorSetId     string `gorm:"column:operator_set_id;primaryKey"`
 	Strategy          string `gorm:"column:strategy;primaryKey"`
 	MagnitudeDecrease string `gorm:"column:magnitude_decrease"`
 	BlockDate         string `gorm:"column:block_date"`