backend-go/database/database.go

package database

import (
	"os"
	"path/filepath"

	"golang.org/x/crypto/bcrypt"
	"gorm.io/driver/postgres"
	"gorm.io/driver/sqlite"
	"gorm.io/gorm"

	"git.beifan.cn/trace-system/backend-go/config"
	"git.beifan.cn/trace-system/backend-go/logger"
	"git.beifan.cn/trace-system/backend-go/models"
	_ "github.com/mattn/go-sqlite3"
)

var DB *gorm.DB

// InitDB 初始化数据库连接
func InitDB() {
	var err error
	cfg := config.GetAppConfig()
	dbConfig := cfg.Database

	logger.Info("数据库配置信息",
		logger.String("驱动", dbConfig.Driver),
	)

	switch dbConfig.Driver {
	case "sqlite":
		logger.Info("使用 SQLite 数据库", logger.String("路径", dbConfig.SQLite.Path))
		DB, err = initSQLite(dbConfig.SQLite.Path)
	case "postgres":
		logger.Info("使用 PostgreSQL 数据库",
			logger.String("主机", dbConfig.Postgres.Host),
			logger.String("端口", dbConfig.Postgres.Port),
			logger.String("数据库", dbConfig.Postgres.DBName),
		)
		DB, err = initPostgres(dbConfig.Postgres)
	default:
		logger.Fatal("不支持的数据库驱动: " + dbConfig.Driver)
	}

	if err != nil {
		logger.Fatal("数据库连接失败", logger.Err(err))
	}

	logger.Info("数据库连接成功")
}

// initSQLite 初始化 SQLite 连接
func initSQLite(path string) (*gorm.DB, error) {
	// 确保数据目录存在
	dataDir := filepath.Dir(path)
	if err := os.MkdirAll(dataDir, 0755); err != nil {
		return nil, err
	}

	db, err := gorm.Open(sqlite.Open(path+"?_pragma=foreign_keys(1)"), &gorm.Config{})
	if err != nil {
		logger.Warn("SQLite 连接失败，尝试使用内存数据库", logger.Err(err))
		// 尝试使用内存数据库作为备选方案
		db, err = gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
	}

	if err != nil {
		return nil, err
	}

	// 配置连接池
	sqlDB, _ := db.DB()
	sqlDB.SetMaxOpenConns(25)
	sqlDB.SetMaxIdleConns(5)
	sqlDB.SetConnMaxLifetime(5 * 60 * 1000000000)

	return db, nil
}

// initPostgres 初始化 PostgreSQL 连接
func initPostgres(pgConfig struct {
	Host     string `mapstructure:"host"`
	Port     string `mapstructure:"port"`
	User     string `mapstructure:"user"`
	Password string `mapstructure:"password"`
	DBName   string `mapstructure:"dbname"`
	SSLMode  string `mapstructure:"sslmode"`
}) (*gorm.DB, error) {
	dsn := "host=" + pgConfig.Host +
		" port=" + pgConfig.Port +
		" user=" + pgConfig.User +
		" password=" + pgConfig.Password +
		" dbname=" + pgConfig.DBName +
		" sslmode=" + pgConfig.SSLMode

	db, err := gorm.Open(postgres.Open(dsn), &gorm.Config{})
	if err != nil {
		logger.Fatal("PostgreSQL 连接失败", logger.Err(err))
	}

	// 配置连接池
	sqlDB, _ := db.DB()
	sqlDB.SetMaxOpenConns(100)
	sqlDB.SetMaxIdleConns(10)
	sqlDB.SetConnMaxLifetime(30 * 60 * 1000000000)

	return db, nil
}

// AutoMigrate 自动迁移数据库
func AutoMigrate() {
	if err := DB.AutoMigrate(
		&models.User{},
		&models.Company{},
		&models.Serial{},
		&models.EmployeeSerial{},
	); err != nil {
		logger.Fatal("数据库迁移失败", logger.Err(err))
	}

	logger.Info("数据库迁移成功")

	// 创建默认管理员用户
	seedAdminUser()
}

// seedAdminUser 创建默认管理员用户
func seedAdminUser() {
	var count int64
	DB.Model(&models.User{}).Count(&count)

	if count > 0 {
		return
	}

	// 默认管理员账号
	adminUsername := "admin"
	adminPassword := "Beifan@2026"
	adminName := "管理员"
	adminEmail := "admin@example.com"

	// 加密密码
	hashedPassword, err := bcrypt.GenerateFromPassword([]byte(adminPassword), bcrypt.DefaultCost)
	if err != nil {
		logger.Warn("管理员密码加密失败", logger.Err(err))
		return
	}

	admin := models.User{
		Username: adminUsername,
		Password: string(hashedPassword),
		Name:     adminName,
		Email:    adminEmail,
		Role:     "admin",
	}

	if err := DB.Create(&admin).Error; err != nil {
		logger.Warn("创建默认管理员用户失败", logger.Err(err))
		return
	}

	logger.Info("已创建默认管理员用户",
		logger.String("username", adminUsername),
		logger.String("password", adminPassword),
		logger.String("role", "admin"),
	)
}