Files
Hightube/backend/internal/db/db.go
2026-04-09 00:14:57 +08:00

111 lines
2.7 KiB
Go

package db
import (
"log"
"os"
"time"
"gorm.io/driver/sqlite"
"gorm.io/gorm"
"gorm.io/gorm/logger"
"hightube/internal/model"
"hightube/internal/monitor"
"hightube/internal/utils"
)
var DB *gorm.DB
// InitDB initializes the SQLite database connection and auto-migrates models.
func InitDB() {
newLogger := logger.New(
log.New(os.Stdout, "\r\n", log.LstdFlags), // io writer
logger.Config{
SlowThreshold: time.Second, // Slow SQL threshold
LogLevel: logger.Warn, // Log level
IgnoreRecordNotFoundError: true, // Ignore ErrRecordNotFound error for logger
Colorful: true, // Disable color
},
)
var err error
// Use SQLite database stored in a local file named "hightube.db"
DB, err = gorm.Open(sqlite.Open("hightube.db"), &gorm.Config{
Logger: newLogger,
})
if err != nil {
log.Fatalf("Failed to connect database: %v", err)
}
// Auto-migrate the schema
err = DB.AutoMigrate(&model.User{}, &model.Room{})
if err != nil {
log.Fatalf("Failed to migrate database: %v", err)
}
// Phase 3.5 Fix: Reset all rooms to inactive on startup using explicit map to ensure false is updated
DB.Model(&model.Room{}).Where("1 = 1").Updates(map[string]interface{}{"is_active": false})
ensureAdminUser()
monitor.Infof("Database initialized successfully")
}
func ensureAdminUser() {
adminUsername := os.Getenv("HIGHTUBE_ADMIN_USER")
if adminUsername == "" {
adminUsername = "admin"
}
adminPassword := os.Getenv("HIGHTUBE_ADMIN_PASS")
if adminPassword == "" {
adminPassword = "admin123456"
}
var user model.User
err := DB.Where("username = ?", adminUsername).First(&user).Error
if err == nil {
updates := map[string]interface{}{}
if user.Role != "admin" {
updates["role"] = "admin"
}
if !user.Enabled {
updates["enabled"] = true
}
if len(updates) > 0 {
DB.Model(&user).Updates(updates)
monitor.Warnf("Admin account normalized for username=%s", adminUsername)
}
return
}
hash, hashErr := utils.HashPassword(adminPassword)
if hashErr != nil {
monitor.Errorf("Failed to hash default admin password: %v", hashErr)
return
}
newAdmin := model.User{
Username: adminUsername,
Password: hash,
Role: "admin",
Enabled: true,
}
if createErr := DB.Create(&newAdmin).Error; createErr != nil {
monitor.Errorf("Failed to create admin account: %v", createErr)
return
}
room := model.Room{
UserID: newAdmin.ID,
Title: newAdmin.Username + "'s Live Room",
StreamKey: utils.GenerateStreamKey(),
IsActive: false,
}
if roomErr := DB.Create(&room).Error; roomErr != nil {
monitor.Warnf("Failed to create default admin room: %v", roomErr)
}
monitor.Warnf("Default admin created: username=%s password=%s", adminUsername, adminPassword)
}