encapsulate migrations in a managed struct object

cmd/database.go

@@ -88,7 +88,10 @@ var runDatabaseCmd = &cobra.Command{
 			log.Fatalln(err)
 		}
 
-		smig := stateMigrator.NewStateMigrator(grm, cfg, l)
+		smig, err := stateMigrator.NewStateMigrator(grm, cfg, l)
+		if err != nil {
+			l.Sugar().Fatalw("Failed to create state migrator", zap.Error(err))
+		}
 
 		sm := stateManager.NewEigenStateManager(smig, l, grm)
 		msm := metaStateManager.NewMetaStateManager(grm, l, cfg)

cmd/debugger/main.go

@@ -84,7 +84,10 @@ func main() {
 		log.Fatalln(err)
 	}
 
-	smig := stateMigrator.NewStateMigrator(grm, cfg, l)
+	smig, err := stateMigrator.NewStateMigrator(grm, cfg, l)
+	if err != nil {
+		l.Sugar().Fatalw("Failed to create state migrator", zap.Error(err))
+	}
 
 	sm := stateManager.NewEigenStateManager(smig, l, grm)
 	msm := metaStateManager.NewMetaStateManager(grm, l, cfg)

cmd/operatorRestakedStrategies.go

@@ -76,7 +76,10 @@ var runOperatorRestakedStrategiesCmd = &cobra.Command{
 			log.Fatalln(err)
 		}
 
-		smig := stateMigrator.NewStateMigrator(grm, cfg, l)
+		smig, err := stateMigrator.NewStateMigrator(grm, cfg, l)
+		if err != nil {
+			l.Sugar().Fatalw("Failed to create state migrator", zap.Error(err))
+		}
 
 		sm := stateManager.NewEigenStateManager(smig, l, grm)
 

cmd/rpc.go

@@ -83,7 +83,10 @@ var rpcCmd = &cobra.Command{
 			log.Fatalln(err)
 		}
 
-		smig := stateMigrator.NewStateMigrator(grm, cfg, l)
+		smig, err := stateMigrator.NewStateMigrator(grm, cfg, l)
+		if err != nil {
+			l.Sugar().Fatalw("Failed to create state migrator", zap.Error(err))
+		}
 
 		sm := stateManager.NewEigenStateManager(smig, l, grm)
 

cmd/run.go

@@ -103,7 +103,10 @@ var runCmd = &cobra.Command{
 			log.Fatalln(err)
 		}
 
-		smig := stateMigrator.NewStateMigrator(grm, cfg, l)
+		smig, err := stateMigrator.NewStateMigrator(grm, cfg, l)
+		if err != nil {
+			l.Sugar().Fatalw("Failed to create state migrator", zap.Error(err))
+		}
 
 		sm := stateManager.NewEigenStateManager(smig, l, grm)
 		msm := metaStateManager.NewMetaStateManager(grm, l, cfg)

pkg/eigenState/stateMigrator/stateMigrator.go

@@ -11,50 +11,78 @@ import (
 	"gorm.io/gorm"
 )
 
+type Migrations struct {
+	Migrations map[uint64][]types.IStateMigration
+}
+
+func NewMigrations() *Migrations {
+	return &Migrations{
+		Migrations: make(map[uint64][]types.IStateMigration),
+	}
+}
+
+func (m *Migrations) RegisterMigrations(blockNumber uint64, migrations []types.IStateMigration) {
+	if _, ok := m.Migrations[blockNumber]; !ok {
+		m.Migrations[blockNumber] = []types.IStateMigration{}
+	}
+	m.Migrations[blockNumber] = append(m.Migrations[blockNumber], migrations...)
+}
+
+func (m *Migrations) GetMigrationsForBlock(blockNumber uint64) []types.IStateMigration {
+	if migrations, ok := m.Migrations[blockNumber]; ok {
+		return migrations
+	}
+	return nil
+}
+
 type StateMigrator struct {
 	logger       *zap.Logger
 	globalConfig *config.Config
 	db           *gorm.DB
+	migrations   *Migrations
 }
 
-func NewStateMigrator(grm *gorm.DB, cfg *config.Config, l *zap.Logger) *StateMigrator {
-	return &StateMigrator{
+func NewStateMigrator(grm *gorm.DB, cfg *config.Config, l *zap.Logger) (*StateMigrator, error) {
+	sm := &StateMigrator{
 		globalConfig: cfg,
 		logger:       l,
 		db:           grm,
 	}
+
+	err := sm.initialize()
+	if err != nil {
+		return nil, err
+	}
+	return sm, nil
 }
 
 // getMigrations returns a map of block numbers to migrations that should be run at that block number
 // taking the chain into account
-func (sm *StateMigrator) getMigrations() (map[uint64][]types.IStateMigration, error) {
+func (sm *StateMigrator) initialize() error {
 	forks, err := sm.globalConfig.GetRewardsSqlForkDates()
 	if err != nil {
-		return nil, err
+		return err
 	}
-
-	return map[uint64][]types.IStateMigration{
+	migrations := NewMigrations()
+
+	if mississippiFork := forks[config.RewardsFork_Mississippi]; mississippiFork.BlockNumber != 0 {
 		// Mississippi hard fork. Applies to preprod and holesky to materialize EigenState models for operator sets
-		forks[config.RewardsFork_Mississippi].BlockNumber: {
+		migrations.RegisterMigrations(forks[config.RewardsFork_Mississippi].BlockNumber, []types.IStateMigration{
 			_02502031222_operatorSets.NewStateMigration(sm.db, sm.logger, sm.globalConfig),
-		},
-	}, nil
+		})
+	}
+	sm.migrations = migrations
+	return nil
 }
 
 // GetMigrationsForBlock returns a list of migrations that should be run at the given block number
 func (sm *StateMigrator) GetMigrationsForBlock(blockNumber uint64) ([]types.IStateMigration, error) {
-	allMigrations, err := sm.getMigrations()
-	if err != nil {
-		sm.logger.Sugar().Error(err)
-		return nil, err
-	}
-
-	migrations, ok := allMigrations[blockNumber]
-	if !ok {
+	blockMigrations := sm.migrations.GetMigrationsForBlock(blockNumber)
+	if len(blockMigrations) == 0 {
 		sm.logger.Sugar().Debugw("No migrations found for block", blockNumber)
 		return nil, nil
 	}
-	return migrations, nil
+	return blockMigrations, nil
 }
 
 // RunMigrationsForBlock runs all migrations for the given block number and returns