Convert EigenState models to use sqlite

Layr-Labs · Sep 6, 2024 · 3884e65 · 3884e65
1 parent b61be16
commit 3884e65
Show file tree

Hide file tree

Showing 10 changed files with 146 additions and 12 deletions.
diff --git a/internal/eigenState/avsOperators/avsOperators_test.go b/internal/eigenState/avsOperators/avsOperators_test.go
@@ -5,6 +5,7 @@ import (
 	"github.com/Layr-Labs/sidecar/internal/config"
 	"github.com/Layr-Labs/sidecar/internal/eigenState/stateManager"
 	"github.com/Layr-Labs/sidecar/internal/logger"
+	"github.com/Layr-Labs/sidecar/internal/sqlite/migrations"
 	"github.com/Layr-Labs/sidecar/internal/storage"
 	"github.com/Layr-Labs/sidecar/internal/tests"
 	"github.com/stretchr/testify/assert"
@@ -23,9 +24,16 @@ func setup() (
 	cfg := tests.GetConfig()
 	l, _ := logger.NewLogger(&logger.LoggerConfig{Debug: cfg.Debug})
 
-	_, grm, err := tests.GetDatabaseConnection(cfg)
+	db, err := tests.GetSqliteDatabaseConnection()
+	if err != nil {
+		panic(err)
+	}
+	sqliteMigrator := migrations.NewSqliteMigrator(db, l)
+	if err := sqliteMigrator.MigrateAll(); err != nil {
+		l.Sugar().Fatalw("Failed to migrate", "error", err)
+	}
 
-	return cfg, grm, l, err
+	return cfg, db, l, err
 }
 
 func teardown(model *AvsOperatorsModel) {

diff --git a/internal/eigenState/eigenstate_test.go b/internal/eigenState/eigenstate_test.go
@@ -7,6 +7,7 @@ import (
 	"github.com/Layr-Labs/sidecar/internal/eigenState/operatorShares"
 	"github.com/Layr-Labs/sidecar/internal/eigenState/stateManager"
 	"github.com/Layr-Labs/sidecar/internal/logger"
+	"github.com/Layr-Labs/sidecar/internal/sqlite/migrations"
 	"github.com/Layr-Labs/sidecar/internal/tests"
 	"github.com/stretchr/testify/assert"
 	"go.uber.org/zap"
@@ -23,9 +24,16 @@ func setup() (
 	cfg := tests.GetConfig()
 	l, _ := logger.NewLogger(&logger.LoggerConfig{Debug: cfg.Debug})
 
-	_, grm, err := tests.GetDatabaseConnection(cfg)
+	db, err := tests.GetSqliteDatabaseConnection()
+	if err != nil {
+		panic(err)
+	}
+	sqliteMigrator := migrations.NewSqliteMigrator(db, l)
+	if err := sqliteMigrator.MigrateAll(); err != nil {
+		l.Sugar().Fatalw("Failed to migrate", "error", err)
+	}
 
-	return cfg, grm, l, err
+	return cfg, db, l, err
 }
 
 func teardown(grm *gorm.DB) {

diff --git a/internal/eigenState/operatorShares/operatorShares_test.go b/internal/eigenState/operatorShares/operatorShares_test.go
@@ -6,6 +6,7 @@ import (
 	"github.com/Layr-Labs/sidecar/internal/config"
 	"github.com/Layr-Labs/sidecar/internal/eigenState/stateManager"
 	"github.com/Layr-Labs/sidecar/internal/logger"
+	"github.com/Layr-Labs/sidecar/internal/sqlite/migrations"
 	"github.com/Layr-Labs/sidecar/internal/storage"
 	"github.com/Layr-Labs/sidecar/internal/tests"
 	"github.com/stretchr/testify/assert"
@@ -25,9 +26,16 @@ func setup() (
 	cfg := tests.GetConfig()
 	l, _ := logger.NewLogger(&logger.LoggerConfig{Debug: cfg.Debug})
 
-	_, grm, err := tests.GetDatabaseConnection(cfg)
+	db, err := tests.GetSqliteDatabaseConnection()
+	if err != nil {
+		panic(err)
+	}
+	sqliteMigrator := migrations.NewSqliteMigrator(db, l)
+	if err := sqliteMigrator.MigrateAll(); err != nil {
+		l.Sugar().Fatalw("Failed to migrate", "error", err)
+	}
 
-	return cfg, grm, l, err
+	return cfg, db, l, err
 }
 
 func teardown(model *OperatorSharesModel) {

diff --git a/internal/eigenState/stakerDelegations/stakerDelegations_test.go b/internal/eigenState/stakerDelegations/stakerDelegations_test.go
@@ -5,6 +5,7 @@ import (
 	"github.com/Layr-Labs/sidecar/internal/config"
 	"github.com/Layr-Labs/sidecar/internal/eigenState/stateManager"
 	"github.com/Layr-Labs/sidecar/internal/logger"
+	"github.com/Layr-Labs/sidecar/internal/sqlite/migrations"
 	"github.com/Layr-Labs/sidecar/internal/storage"
 	"github.com/Layr-Labs/sidecar/internal/tests"
 	"github.com/stretchr/testify/assert"
@@ -23,9 +24,16 @@ func setup() (
 	cfg := tests.GetConfig()
 	l, _ := logger.NewLogger(&logger.LoggerConfig{Debug: cfg.Debug})
 
-	_, grm, err := tests.GetDatabaseConnection(cfg)
+	db, err := tests.GetSqliteDatabaseConnection()
+	if err != nil {
+		panic(err)
+	}
+	sqliteMigrator := migrations.NewSqliteMigrator(db, l)
+	if err := sqliteMigrator.MigrateAll(); err != nil {
+		l.Sugar().Fatalw("Failed to migrate", "error", err)
+	}
 
-	return cfg, grm, l, err
+	return cfg, db, l, err
 }
 
 func teardown(model *StakerDelegationsModel) {

diff --git a/internal/sqlite/migrations/202409061250_eigenlayerStateTables/up.go b/internal/sqlite/migrations/202409061250_eigenlayerStateTables/up.go
@@ -0,0 +1,34 @@
+package _202409061250_eigenlayerStateTables
+
+import (
+	"fmt"
+	"gorm.io/gorm"
+)
+
+type SqliteMigration struct {
+}
+
+func (m *SqliteMigration) Up(grm *gorm.DB) error {
+	queries := []string{
+		`create table if not exists registered_avs_operators (
+			operator TEXT NOT NULL,
+			avs TEXT NOT NULL,
+			block_number INTEGER NOT NULL,
+			created_at DATETIME default current_timestamp,
+			unique(operator, avs, block_number)
+		);
+		`,
+	}
+
+	for _, query := range queries {
+		if res := grm.Exec(query); res.Error != nil {
+			fmt.Printf("Failed to execute query: %s\n", query)
+			return res.Error
+		}
+	}
+	return nil
+}
+
+func (m *SqliteMigration) GetName() string {
+	return "202409061250_eigenlayerStateTables"
+}
diff --git a/internal/sqlite/migrations/202409061720_operatorShareChanges/up.go b/internal/sqlite/migrations/202409061720_operatorShareChanges/up.go
@@ -0,0 +1,31 @@
+package _202409061720_operatorShareChanges
+
+import (
+	"gorm.io/gorm"
+)
+
+type SqliteMigration struct {
+}
+
+func (m *SqliteMigration) Up(grm *gorm.DB) error {
+	queries := []string{
+		`create table if not exists operator_shares (
+			operator TEXT NOT NULL,
+			strategy TEXT NOT NULL,
+			shares NUMERIC NOT NULL,
+			block_number INTEGER NOT NULL,
+			created_at DATETIME default current_timestamp,
+			unique (operator, strategy, block_number)
+		)`,
+	}
+	for _, query := range queries {
+		if res := grm.Exec(query); res.Error != nil {
+			return res.Error
+		}
+	}
+	return nil
+}
+
+func (m *SqliteMigration) GetName() string {
+	return "202409061720_operatorShareChanges"
+}
diff --git a/internal/sqlite/migrations/202409062151_stakerDelegations/up.go b/internal/sqlite/migrations/202409062151_stakerDelegations/up.go
@@ -0,0 +1,30 @@
+package _202409062151_stakerDelegations
+
+import (
+	"gorm.io/gorm"
+)
+
+type SqliteMigration struct {
+}
+
+func (m *SqliteMigration) Up(grm *gorm.DB) error {
+	queries := []string{
+		`create table if not exists delegated_stakers (
+			staker TEXT NOT NULL,
+			operator TEXT NOT NULL,
+			block_number INTEGER NOT NULL,
+			created_at DATETIME default current_timestamp,
+			unique(staker, operator, block_number)
+		)`,
+	}
+	for _, query := range queries {
+		if res := grm.Exec(query); res.Error != nil {
+			return res.Error
+		}
+	}
+	return nil
+}
+
+func (m *SqliteMigration) GetName() string {
+	return "202409062151_stakerDelegations"
+}
diff --git a/...nal/storage/sqlite/migrations/migrator.go → internal/sqlite/migrations/migrator.go b/...nal/storage/sqlite/migrations/migrator.go → internal/sqlite/migrations/migrator.go
@@ -4,6 +4,9 @@ import (
 	"database/sql"
 	"fmt"
 	_202409061249_bootstrapDb "github.com/Layr-Labs/sidecar/internal/sqlite/migrations/202409061249_bootstrapDb"
+	_202409061250_eigenlayerStateTables "github.com/Layr-Labs/sidecar/internal/sqlite/migrations/202409061250_eigenlayerStateTables"
+	_202409061720_operatorShareChanges "github.com/Layr-Labs/sidecar/internal/sqlite/migrations/202409061720_operatorShareChanges"
+	_202409062151_stakerDelegations "github.com/Layr-Labs/sidecar/internal/sqlite/migrations/202409062151_stakerDelegations"
 	"go.uber.org/zap"
 	"gorm.io/gorm"
 	"time"
@@ -35,6 +38,9 @@ func (m *SqliteMigrator) MigrateAll() error {
 
 	migrations := []ISqliteMigration{
 		&_202409061249_bootstrapDb.SqliteMigration{},
+		&_202409061250_eigenlayerStateTables.SqliteMigration{},
+		&_202409061720_operatorShareChanges.SqliteMigration{},
+		&_202409062151_stakerDelegations.SqliteMigration{},
 	}
 
 	for _, migration := range migrations {

diff --git a/internal/sqlite/sqlite.go b/internal/sqlite/sqlite.go
@@ -3,6 +3,7 @@ package sqlite
 import (
 	"gorm.io/driver/sqlite"
 	"gorm.io/gorm"
+	"gorm.io/gorm/logger"
 )
 
 func NewSqlite(path string) gorm.Dialector {
@@ -11,7 +12,9 @@ func NewSqlite(path string) gorm.Dialector {
 }
 
 func NewGormSqliteFromSqlite(sqlite gorm.Dialector) (*gorm.DB, error) {
-	db, err := gorm.Open(sqlite, &gorm.Config{})
+	db, err := gorm.Open(sqlite, &gorm.Config{
+		Logger: logger.Default.LogMode(logger.Silent),
+	})
 	if err != nil {
 		return nil, err
 	}

diff --git a/internal/storage/sqlite/sqlite_test.go b/internal/storage/sqlite/sqlite_test.go
@@ -2,12 +2,11 @@ package sqlite
 
 import (
 	"encoding/json"
-	"fmt"
 	"github.com/Layr-Labs/sidecar/internal/config"
 	"github.com/Layr-Labs/sidecar/internal/logger"
 	"github.com/Layr-Labs/sidecar/internal/parser"
+	"github.com/Layr-Labs/sidecar/internal/sqlite/migrations"
 	"github.com/Layr-Labs/sidecar/internal/storage"
-	"github.com/Layr-Labs/sidecar/internal/storage/sqlite/migrations"
 	"github.com/Layr-Labs/sidecar/internal/tests"
 	"github.com/stretchr/testify/assert"
 	"go.uber.org/zap"
@@ -82,7 +81,6 @@ func Test_SqliteBlockstore(t *testing.T) {
 			_, err := sqliteStore.InsertBlockAtHeight(block.Number, block.Hash, uint64(block.BlockTime.Unix()))
 			assert.NotNil(t, err)
 			assert.Contains(t, err.Error(), "UNIQUE constraint failed")
-			fmt.Printf("Error: %v\n", err)
 		})
 	})
 	t.Run("Transactions", func(t *testing.T) {