go_blog/models/db.go

package models

import (
	"log"
	"os"
	"path/filepath"

	"go_blog/config"

	"gorm.io/driver/mysql"
	"gorm.io/gorm"
	"gorm.io/gorm/logger"

	"github.com/glebarez/sqlite"
)

// DB is the global database connection, initialized by InitDB.
var DB *gorm.DB

// InitDB opens the database connection, runs migrations, and seeds the admin user.
func InitDB(cfg *config.Config) *gorm.DB {
	// Ensure the storage path exists.
	if err := os.MkdirAll(cfg.Path, 0755); err != nil {
		log.Fatalf("Failed to create storage directory %s: %v", cfg.Path, err)
	}

	var dialector gorm.Dialector

	switch cfg.Database.Type {
	case "mysql":
		dsn := cfg.Database.DSN
		if dsn == "" {
			dsn = "root:password@tcp(127.0.0.1:3306)/blog_go?charset=utf8mb4&parseTime=True&loc=Local"
		}
		dialector = mysql.Open(dsn)
	default:
		dbPath := filepath.Join(cfg.Path, "blog.db")
		dialector = sqlite.Open(dbPath)
	}

	db, err := gorm.Open(dialector, &gorm.Config{
		Logger: logger.Default.LogMode(logger.Warn),
	})
	if err != nil {
		log.Fatalf("Failed to connect to database: %v", err)
	}

	// Auto-migrate the User table (idempotent).
	if err := db.AutoMigrate(&User{}); err != nil {
		log.Fatalf("Failed to auto-migrate database: %v", err)
	}

	// First-run seed: create admin user if no users exist.
	var count int64
	db.Model(&User{}).Count(&count)
	if count == 0 {
		admin := &User{
			Username:    "admin",
			DisplayName: "Administrator",
			Gender:      "other",
			Status:      StatusNormal,
			Role:        RoleAdmin,
		}
		if err := admin.SetPassword("admin"); err != nil {
			log.Fatalf("Failed to hash admin password: %v", err)
		}
		if err := db.Create(admin).Error; err != nil {
			log.Fatalf("Failed to create admin user: %v", err)
		}
		log.Println("==============================================")
		log.Println("  First run: created default admin user.")
		log.Println("  Username: admin")
		log.Println("  Password: admin")
		log.Println("  Please change this password immediately!")
		log.Println("==============================================")
	}

	// Migration fix: always set admin role on the admin user.
	result := db.Model(&User{}).Where("username = ?", "admin").Update("role", RoleAdmin)
	if result.RowsAffected > 0 {
		log.Printf("Migration: set admin role for existing admin user (rows affected: %d)", result.RowsAffected)
	}

	DB = db
	return db
}