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"), ) }