package store import ( "fmt" "git.linuxforward.com/byom/byom-core-gateway/logger" "github.com/sirupsen/logrus" "gorm.io/gorm" ) type DataStore struct { logger *logrus.Entry db *gorm.DB } func NewDataStore(db *gorm.DB) *DataStore { err := InitTables(db) if err != nil { logger.NewLogger("data-store").Fatal("Failed to initialize tables: ", err) } return &DataStore{ db: db, logger: logger.NewLogger("data-store"), } } func InitTables(db *gorm.DB) error { err := db.AutoMigrate(&UserHost{}) if err != nil { return err } return nil } //Tx methods func (s *DataStore) BeginTx() *gorm.DB { return s.db.Begin() } func (s *DataStore) CommitTx(tx *gorm.DB) error { return tx.Commit().Error } func (s *DataStore) Close() error { s.logger.Info("Closing database connections") sqlDB, err := s.db.DB() if err != nil { return fmt.Errorf("get sql.DB instance: %w", err) } if err := sqlDB.Close(); err != nil { return fmt.Errorf("close database connections: %w", err) } s.logger.Info("Database connections closed successfully") return nil }