From cd9c2eb0a877995072899760efa41188d48bf515 Mon Sep 17 00:00:00 2001 From: kunal Date: Thu, 26 Mar 2026 15:43:48 +0530 Subject: [PATCH 1/8] Cleanup gorm errors clutter --- core/database/database.go | 12 +++++++++++- core/database/hints.go | 3 +-- 2 files changed, 12 insertions(+), 3 deletions(-) diff --git a/core/database/database.go b/core/database/database.go index 3455ee29..54857b9f 100644 --- a/core/database/database.go +++ b/core/database/database.go @@ -3,6 +3,7 @@ package database import ( "crypto/rand" "fmt" + "gorm.io/gorm/logger" "os" "os/exec" "path/filepath" @@ -54,7 +55,16 @@ func LoadDbConfig() { func ConnectDatabase() error { LoadDbConfig() dsn := fmt.Sprintf("user=%s password=%s dbname=%s host=%s port=%s sslmode=%s", dbConfig.PsqlConf.User, dbConfig.PsqlConf.Password, dbConfig.PsqlConf.Dbname, dbConfig.PsqlConf.Host, dbConfig.PsqlConf.Port, dbConfig.PsqlConf.SslMode) - Db, dberr = gorm.Open(postgres.Open(dsn), &gorm.Config{}) + + Db, dberr = gorm.Open(postgres.Open(dsn), &gorm.Config{ + Logger: logger.New( + log.New(), + logger.Config{ + LogLevel: logger.Warn, + IgnoreRecordNotFoundError: true, + }, + )}) + if dberr != nil { log.Error("Error while initializing the database.", dberr) return dberr diff --git a/core/database/hints.go b/core/database/hints.go index 1b6561d0..a5e2a4d8 100644 --- a/core/database/hints.go +++ b/core/database/hints.go @@ -80,11 +80,10 @@ func UserHasTakenHint(userID, hintID uint) (bool, error) { var userHint UserHint if err := tx.Where("user_id = ? AND hint_id = ?", userID, hintID).First(&userHint).Error; err != nil { + tx.Rollback() if errors.Is(err, gorm.ErrRecordNotFound) { - tx.Rollback() return false, nil } - tx.Rollback() return false, fmt.Errorf("db_error") } From fc742d37521f1919b156e44388c0a5d5d6c1130d Mon Sep 17 00:00:00 2001 From: kunal Date: Thu, 26 Mar 2026 16:04:15 +0530 Subject: [PATCH 2/8] Add graceful database close --- cmd/beast/run.go | 8 ++++---- core/database/database.go | 22 +++++++++++++++++++++- 2 files changed, 25 insertions(+), 5 deletions(-) diff --git a/cmd/beast/run.go b/cmd/beast/run.go index d389ebc1..14035c48 100644 --- a/cmd/beast/run.go +++ b/cmd/beast/run.go @@ -76,13 +76,13 @@ func cleanupDatabaseConnections() { log.Infoln("Database backup completed successfully") } - log.Infoln("Terminating database connection...") + log.Infoln("Closing database connection...") - err = database.TerminateDatabaseConnections() + err = database.Close() if err != nil { - log.Errorln("Unable to terminate database connections:", err) + log.Errorln("Unable to close database connections:", err) } else { - log.Infoln("Database connections terminated successfully") + log.Infoln("Database connections close successfully") } } diff --git a/core/database/database.go b/core/database/database.go index 54857b9f..21c3967a 100644 --- a/core/database/database.go +++ b/core/database/database.go @@ -127,6 +127,27 @@ func Init() { } } +func Close() error { + if Db == nil { + log.Warnln(fmt.Sprintf("Trying to close database connection when no connection is established...")) + return nil + } + + sqlDb, err := Db.DB() + if err != nil { + log.Errorln(fmt.Sprintf("Error while closing database connection gracefully: %s, attempting to terminate forcefully", err.Error())) + return TerminateDatabaseConnections() + } + + err = sqlDb.Close() + if err != nil { + log.Errorln(fmt.Sprintf("Error while closing database connection gracefully: %s, attempting to terminate forcefully", err.Error())) + return TerminateDatabaseConnections() + } + + return nil +} + func BackupAndReset() { LoadDbConfig() @@ -177,7 +198,6 @@ func BackupDatabase() error { if dbConfig == (Config{}) { LoadDbConfig() } - backupPath := filepath.Join(core.BEAST_GLOBAL_DIR, "backup", "db") err := utils.CreateIfNotExistDir(backupPath) if err != nil { From a9b8b8304d9e08bb47ccc93e24974331461650c5 Mon Sep 17 00:00:00 2001 From: kunal Date: Thu, 26 Mar 2026 16:18:34 +0530 Subject: [PATCH 3/8] Clean terminate database connections --- core/constants.go | 4 +-- core/database/database.go | 51 +++++++++++++++++++++++---------------- 2 files changed, 32 insertions(+), 23 deletions(-) diff --git a/core/constants.go b/core/constants.go index 0fc10a06..7c8022ff 100644 --- a/core/constants.go +++ b/core/constants.go @@ -57,6 +57,7 @@ const ( //paths BEAST_SECRETS_DIR string = "secrets" BEAST_EXAMPLE_DIR string = "_examples" BEAST_CACHE_DIR string = "cache" + BEAST_BACKUP_DIR string = "backup" ) const ( //chall types @@ -203,11 +204,10 @@ var USER_STATUS = map[string]string{ const ( LEADERBOARD_SIZE = 25 LEADERBOARD_GRAPH_SIZE = 12 - SUBMISSIONS_PAGE_SIZE = 10 + SUBMISSIONS_PAGE_SIZE = 10 ) var NOTIFICATION_SERVICES = []string{ "slack", "discord", } - diff --git a/core/database/database.go b/core/database/database.go index 21c3967a..a3ccb1df 100644 --- a/core/database/database.go +++ b/core/database/database.go @@ -2,6 +2,7 @@ package database import ( "crypto/rand" + "database/sql" "fmt" "gorm.io/gorm/logger" "os" @@ -162,7 +163,7 @@ func BackupAndReset() { return } - backupPath := filepath.Join(core.BEAST_GLOBAL_DIR, "backup", core.BEAST_REMOTES_DIR) + backupPath := filepath.Join(core.BEAST_GLOBAL_DIR, core.BEAST_BACKUP_DIR, core.BEAST_REMOTES_DIR) err = utils.CreateIfNotExistDir(backupPath) if err != nil { log.Errorf("Error while creating backup directory: %s", err) @@ -177,7 +178,7 @@ func BackupAndReset() { return } - backupPath = filepath.Join(core.BEAST_GLOBAL_DIR, "backup", core.BEAST_STAGING_DIR) + backupPath = filepath.Join(core.BEAST_GLOBAL_DIR, core.BEAST_BACKUP_DIR, core.BEAST_STAGING_DIR) err = utils.CreateIfNotExistDir(backupPath) if err != nil { @@ -198,7 +199,7 @@ func BackupDatabase() error { if dbConfig == (Config{}) { LoadDbConfig() } - backupPath := filepath.Join(core.BEAST_GLOBAL_DIR, "backup", "db") + backupPath := filepath.Join(core.BEAST_GLOBAL_DIR, core.BEAST_BACKUP_DIR, "db") err := utils.CreateIfNotExistDir(backupPath) if err != nil { log.Errorf("Error while creating backup directory: %s", err) @@ -206,8 +207,20 @@ func BackupDatabase() error { } backupFile := fmt.Sprintf("%s_%s.bak", dbConfig.PsqlConf.Dbname, time.Now().Format("20060102150405")) - cmd := exec.Command("pg_dump", "-U", dbConfig.PsqlConf.User, "-h", dbConfig.PsqlConf.Host, "-p", dbConfig.PsqlConf.Port, "-F", "c", "-f", filepath.Join(backupPath, backupFile), dbConfig.PsqlConf.Dbname) + cmd := exec.Command( + "pg_dump", + "-U", + dbConfig.PsqlConf.User, + "-h", dbConfig.PsqlConf.Host, + "-p", + dbConfig.PsqlConf.Port, + "-F", + "c", + "-f", + filepath.Join(backupPath, backupFile), + dbConfig.PsqlConf.Dbname) cmd.Env = append(os.Environ(), fmt.Sprintf("PGPASSWORD=%s", dbConfig.PsqlConf.Password)) + output, err := cmd.CombinedOutput() if err != nil { log.Printf("Backup error: %s\n", string(output)) @@ -253,24 +266,21 @@ func TerminateDatabaseConnections() error { if dbConfig == (Config{}) { LoadDbConfig() } - terminateCmd := exec.Command( - "psql", - "-U", dbConfig.PsqlConf.User, - "-h", dbConfig.PsqlConf.Host, - "-p", dbConfig.PsqlConf.Port, - "-d", "postgres", - "-c", - fmt.Sprintf("SELECT pg_terminate_backend(pg_stat_activity.pid) FROM pg_stat_activity WHERE datname = '%s' AND pid <> pg_backend_pid();", dbConfig.PsqlConf.Dbname), - ) - terminateCmd.Env = append(os.Environ(), fmt.Sprintf("PGPASSWORD=%s", dbConfig.PsqlConf.Password)) - - output, err := terminateCmd.CombinedOutput() - outputStr := string(output) + + dsn := fmt.Sprintf("host=%s port=%s user=%s password=%s dbname=%s sslmode=%s", dbConfig.PsqlConf.Host, dbConfig.PsqlConf.Port, dbConfig.PsqlConf.User, dbConfig.PsqlConf.Password, "postgres", dbConfig.PsqlConf.SslMode) + db, err := sql.Open("pgx", dsn) + + if err != nil { + return err + } + defer db.Close() + + _, err = db.Exec("SELECT pg_terminate_backend(pg_stat_activity.pid) FROM pg_stat_activity WHERE datname = $1 AND pid <> pg_backend_pid();", dbConfig.PsqlConf.Dbname) if err != nil { - log.Errorf("Terminate connections error: %s\n", outputStr) + log.Errorf("Terminate connections error: %s\n", err.Error()) return err } - log.Debug(outputStr) + return nil } @@ -297,8 +307,7 @@ func RestoreDatabase(backupFile string) error { "--no-owner", "--clean", "--if-exists", - backupFile, - ) + backupFile) restoreCmd.Env = append(os.Environ(), fmt.Sprintf("PGPASSWORD=%s", dbConfig.PsqlConf.Password)) output, err := restoreCmd.CombinedOutput() From 5662c242d1c103ee7378b67f744eab4ac0c6d7e1 Mon Sep 17 00:00:00 2001 From: kunal Date: Sat, 28 Mar 2026 17:15:59 +0530 Subject: [PATCH 4/8] Fix postgres dns --- cmd/beast/init.go | 29 ++++++++++++++--------------- 1 file changed, 14 insertions(+), 15 deletions(-) diff --git a/cmd/beast/init.go b/cmd/beast/init.go index 90a0dd6f..020eb551 100644 --- a/cmd/beast/init.go +++ b/cmd/beast/init.go @@ -4,7 +4,6 @@ import ( "database/sql" "errors" "fmt" - "github.com/BurntSushi/toml" _ "github.com/jackc/pgx/v5/stdlib" "github.com/lib/pq" "github.com/sdslabs/beastv4/core" @@ -125,11 +124,11 @@ func dbUserCheck() (bool, error) { func initDb() error { log.Infoln("Initializing database...") - var configuration config.BeastConfig - _, err := toml.DecodeFile(BEAST_GLOBAL_CONFIG, &configuration) + err := config.ReloadBeastConfig() if err != nil { return err } + configuration := config.Cfg.PsqlConf isPostgres, err := dbUserCheck() if err != nil { @@ -140,7 +139,7 @@ func initDb() error { if isPostgres { log.Infoln("Attempting to connect to postgres as postgres super user...") - dsn := fmt.Sprintf("user=%s dbname=%s sslmode=%s", "postgres", "postgres", "disable") + dsn := fmt.Sprintf("user=%s dbname=%s host=%s port=%s sslmode=%s", "postgres", "postgres", configuration.Host, configuration.Port, "disable") db, err = sql.Open("pgx", dsn) if err != nil { @@ -152,7 +151,7 @@ func initDb() error { if utils.PromptBinary("Do you use password authentication for the postgres super user?") { password := utils.PromptSecret("Enter postgres super user password (leave blank if none):") - dsn := fmt.Sprintf("user=%s password=%s dbname=%s sslmode=%s", "postgres", password, "postgres", "disable") + dsn := fmt.Sprintf("user=%s password=%s dbname=%s host=%s port=%s sslmode=%s", "postgres", password, "postgres", configuration.Host, configuration.Port, "disable") db, err = sql.Open("pgx", dsn) if err != nil { @@ -167,41 +166,41 @@ func initDb() error { defer db.Close() var exists int - err = db.QueryRow("SELECT 1 FROM pg_roles WHERE rolname = $1", configuration.PsqlConf.User).Scan(&exists) + err = db.QueryRow("SELECT 1 FROM pg_roles WHERE rolname = $1", configuration.User).Scan(&exists) if errors.Is(err, sql.ErrNoRows) { - if err = createBeastDbUser(db, &configuration.PsqlConf); err != nil { + if err = createBeastDbUser(db, &configuration); err != nil { return err } } else if err != nil { return err } else { - log.Infoln(fmt.Sprintf("User %s already exists", configuration.PsqlConf.User)) + log.Infoln(fmt.Sprintf("User %s already exists", configuration.User)) } - log.Infoln(fmt.Sprintf("Changing password for user %s", configuration.PsqlConf.User)) - query := fmt.Sprintf("ALTER USER %s WITH PASSWORD %s", pq.QuoteIdentifier(configuration.PsqlConf.User), utils.QuoteLiteral(configuration.PsqlConf.Password)) + log.Infoln(fmt.Sprintf("Changing password for user %s", configuration.User)) + query := fmt.Sprintf("ALTER USER %s WITH PASSWORD %s", pq.QuoteIdentifier(configuration.User), utils.QuoteLiteral(configuration.Password)) _, err = db.Exec(query) if err != nil { return err } - err = db.QueryRow("SELECT 1 FROM pg_database WHERE datname = $1", configuration.PsqlConf.Dbname).Scan(&exists) + err = db.QueryRow("SELECT 1 FROM pg_database WHERE datname = $1", configuration.Dbname).Scan(&exists) if errors.Is(err, sql.ErrNoRows) { - if err = createBeastDatabase(db, &configuration.PsqlConf); err != nil { + if err = createBeastDatabase(db, &configuration); err != nil { return err } } else if err != nil { return err } else { - log.Infoln(fmt.Sprintf("Database %s already exists", configuration.PsqlConf.Dbname)) + log.Infoln(fmt.Sprintf("Database %s already exists", configuration.Dbname)) } - _, err = db.Exec(fmt.Sprintf("ALTER DATABASE %s OWNER TO %s", pq.QuoteIdentifier(configuration.PsqlConf.Dbname), pq.QuoteIdentifier(configuration.PsqlConf.User))) + _, err = db.Exec(fmt.Sprintf("ALTER DATABASE %s OWNER TO %s", pq.QuoteIdentifier(configuration.Dbname), pq.QuoteIdentifier(configuration.User))) if err != nil { return err } - log.Infoln(fmt.Sprintf("%s set as owner of database %s", configuration.PsqlConf.User, configuration.PsqlConf.Dbname)) + log.Infoln(fmt.Sprintf("%s set as owner of database %s", configuration.User, configuration.Dbname)) return nil } From 6693449f2736ece106830770ac12f58eed02d4e5 Mon Sep 17 00:00:00 2001 From: kunal Date: Sat, 28 Mar 2026 20:09:30 +0530 Subject: [PATCH 5/8] Remove unused var --- core/database/database.go | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/core/database/database.go b/core/database/database.go index a3ccb1df..41755b70 100644 --- a/core/database/database.go +++ b/core/database/database.go @@ -27,8 +27,7 @@ var ( ) var ( - BEAST_GLOBAL_DIR string = filepath.Join(os.Getenv("HOME"), ".beast") - dbConfig Config + dbConfig Config ) type Config struct { From 2e584c20f24d140bd93560435a5a3833b64e095b Mon Sep 17 00:00:00 2001 From: kunal Date: Wed, 1 Apr 2026 00:00:51 +0530 Subject: [PATCH 6/8] Cleanup Db var --- core/database/database.go | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/core/database/database.go b/core/database/database.go index 41755b70..f462ad1f 100644 --- a/core/database/database.go +++ b/core/database/database.go @@ -133,6 +133,9 @@ func Close() error { return nil } + DBMux.Lock() + defer DBMux.Unlock() + sqlDb, err := Db.DB() if err != nil { log.Errorln(fmt.Sprintf("Error while closing database connection gracefully: %s, attempting to terminate forcefully", err.Error())) @@ -145,6 +148,7 @@ func Close() error { return TerminateDatabaseConnections() } + Db = nil return nil } From 1fb2fb03addf6b259813c0383f8390bfa0eee8c9 Mon Sep 17 00:00:00 2001 From: kunal Date: Thu, 2 Apr 2026 15:12:50 +0530 Subject: [PATCH 7/8] Fix formatting --- core/database/database.go | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/core/database/database.go b/core/database/database.go index f462ad1f..8a85920b 100644 --- a/core/database/database.go +++ b/core/database/database.go @@ -221,7 +221,8 @@ func BackupDatabase() error { "c", "-f", filepath.Join(backupPath, backupFile), - dbConfig.PsqlConf.Dbname) + dbConfig.PsqlConf.Dbname, + ) cmd.Env = append(os.Environ(), fmt.Sprintf("PGPASSWORD=%s", dbConfig.PsqlConf.Password)) output, err := cmd.CombinedOutput() @@ -310,7 +311,8 @@ func RestoreDatabase(backupFile string) error { "--no-owner", "--clean", "--if-exists", - backupFile) + backupFile, + ) restoreCmd.Env = append(os.Environ(), fmt.Sprintf("PGPASSWORD=%s", dbConfig.PsqlConf.Password)) output, err := restoreCmd.CombinedOutput() From 4fa83f88e057273649aa48f416bf7a032e74be64 Mon Sep 17 00:00:00 2001 From: kunal Date: Thu, 2 Apr 2026 15:34:00 +0530 Subject: [PATCH 8/8] Exec remove from reset database --- core/database/database.go | 20 +++++++++++++++----- 1 file changed, 15 insertions(+), 5 deletions(-) diff --git a/core/database/database.go b/core/database/database.go index 8a85920b..30c09ac7 100644 --- a/core/database/database.go +++ b/core/database/database.go @@ -4,6 +4,7 @@ import ( "crypto/rand" "database/sql" "fmt" + "github.com/lib/pq" "gorm.io/gorm/logger" "os" "os/exec" @@ -253,14 +254,23 @@ func ResetDatabase() error { return err } - createCmd := exec.Command("psql", "-U", dbConfig.PsqlConf.User, "-h", dbConfig.PsqlConf.Host, "-p", dbConfig.PsqlConf.Port, "-d", "postgres", "-c", "CREATE DATABASE "+dbConfig.PsqlConf.Dbname+";") - createCmd.Env = append(os.Environ(), fmt.Sprintf("PGPASSWORD=%s", dbConfig.PsqlConf.Password)) + dsn := fmt.Sprintf("user=%s password=%s dbname=%s host=%s port=%s sslmode=%s", dbConfig.PsqlConf.User, dbConfig.PsqlConf.Password, "postgres", dbConfig.PsqlConf.Host, dbConfig.PsqlConf.Port, "disable") + db, err := sql.Open("pgx", dsn) + if err != nil { + return fmt.Errorf("unable to connect to database: %s", err) + } + defer db.Close() - output, err = createCmd.CombinedOutput() + _, err = db.Exec(fmt.Sprintf("CREATE DATABASE %s", pq.QuoteIdentifier(dbConfig.PsqlConf.Dbname))) if err != nil { - log.Printf("Create DB error: %s\n", string(output)) - return err + return fmt.Errorf("unable to create database: %s", err) } + + _, err = db.Exec(fmt.Sprintf("ALTER DATABASE %s OWNER TO %s", pq.QuoteIdentifier(dbConfig.PsqlConf.Dbname), pq.QuoteIdentifier(dbConfig.PsqlConf.User))) + if err != nil { + return fmt.Errorf("unable to alter database owner: %s", err) + } + log.Debug("Reset successful.") return nil }