EMS_Platform/backend/models/init.go

package models

import (
	"fmt"
	"log"
	"os"
	"time"

	"gorm.io/driver/postgres"
	"gorm.io/gorm"
)

var DB *gorm.DB

func InitDB() {
	dsn := os.Getenv("DB_DSN")
	if dsn == "" {
		// Fallback for local development if not in env
		dsn = "host=localhost user=ems password=ems_pass dbname=ems port=5433 sslmode=disable TimeZone=Asia/Shanghai"
	}

	var err error
	maxRetries := 30
	for i := 0; i < maxRetries; i++ {
		DB, err = gorm.Open(postgres.Open(dsn), &gorm.Config{})
		if err == nil {
			break
		}
		log.Printf("Failed to connect to database (attempt %d/%d): %v. Retrying in 2 seconds...", i+1, maxRetries, err)
		time.Sleep(2 * time.Second)
	}

	if err != nil {
		log.Fatal("Failed to connect to database after multiple attempts:", err)
	}

	// Set connection pool settings
	sqlDB, err := DB.DB()
	if err != nil {
		log.Fatal("Failed to get generic database object:", err)
	}

	// SetMaxIdleConns sets the maximum number of connections in the idle connection pool.
	sqlDB.SetMaxIdleConns(10)

	// SetMaxOpenConns sets the maximum number of open connections to the database.
	sqlDB.SetMaxOpenConns(100)

	// SetConnMaxLifetime sets the maximum amount of time a connection may be reused.
	sqlDB.SetConnMaxLifetime(time.Hour)

	fmt.Println("Database connected successfully")

	// Auto Migrate
	err = DB.AutoMigrate(
		&IntegrationSource{},
		&Device{},
		&SysLocation{},
		&AlarmLog{},
		&SysOperLog{},
		&InspectionLedger{},
		&User{},
		&SysRole{},
		&SysUserRole{},
		&SysConfig{},
		&SysMenu{},
		&SysRoleMenu{},
		&EquipmentCleaningFormulaTemplate{},
		&BackupLog{},
		&AlarmRule{},
		&AlarmRuleBinding{},
	)
	if err != nil {
		log.Fatal("Failed to migrate database:", err)
	}
	fmt.Println("Database migration completed")

	// Initialize default user and role
	InitDefaultUser(DB)

	// Initialize basic menu data
	InitSysMenuData(DB)
}

func InitDefaultUser(db *gorm.DB) {
	// Find or create admin role
	var adminRole SysRole
	if err := db.Where("role_key = ?", "admin").First(&adminRole).Error; err != nil {
		fmt.Println("Initializing Admin Role...")
		adminRole = SysRole{
			Name:      "超级管理员",
			RoleKey:   "admin",
			DataScope: "1",
			Status:    "0",
		}
		if err := db.Create(&adminRole).Error; err != nil {
			log.Printf("Error creating admin role: %v", err)
			return
		}
	}

	// Find or create admin user
	var adminUser User
	if err := db.Where("username = ?", "admin").First(&adminUser).Error; err != nil {
		fmt.Println("Initializing Admin User...")
		adminUser = User{
			Username: "admin",
			Password: "admin123", // Default password
			Name:     "系统管理员",
			Role:     "ADMIN",
			Status:   "0",
		}
		if err := db.Create(&adminUser).Error; err != nil {
			log.Printf("Error creating admin user: %v", err)
			return
		}

		// Assign admin role to admin user
		db.Create(&SysUserRole{
			UserID: adminUser.ID,
			RoleID: adminRole.ID,
		})
	}
}