96 lines
2.5 KiB
Go
96 lines
2.5 KiB
Go
package config
|
||
|
||
import (
|
||
"fmt"
|
||
"log"
|
||
"os"
|
||
"path/filepath"
|
||
|
||
"github.com/joho/godotenv"
|
||
"gorm.io/driver/mysql"
|
||
"gorm.io/gorm"
|
||
"gorm.io/gorm/logger"
|
||
)
|
||
|
||
var DB *gorm.DB
|
||
|
||
// JWTSecret is the secret key used for signing JWT tokens
|
||
var JWTSecret = "your-secret-key" // Default value, should be set via JWT_SECRET environment variable in production
|
||
|
||
// LoadEnvFile 从可执行文件同级目录加载 .env 文件
|
||
func LoadEnvFile() {
|
||
// 获取可执行文件路径
|
||
execPath, err := os.Executable()
|
||
if err != nil {
|
||
log.Printf("Warning: Failed to get executable path: %v. Using current working directory.", err)
|
||
// 如果获取可执行文件路径失败，使用当前工作目录
|
||
if err := godotenv.Load(".env"); err != nil {
|
||
log.Printf("Info: .env file not found in current directory: %v", err)
|
||
} else {
|
||
log.Println("Loaded .env file from current directory")
|
||
}
|
||
return
|
||
}
|
||
|
||
// 获取可执行文件所在目录
|
||
execDir := filepath.Dir(execPath)
|
||
envPath := filepath.Join(execDir, ".env")
|
||
|
||
// 检查 .env 文件是否存在
|
||
if _, err := os.Stat(envPath); os.IsNotExist(err) {
|
||
log.Printf("Info: .env file not found at %s, using environment variables or defaults", envPath)
|
||
return
|
||
}
|
||
|
||
// 加载 .env 文件
|
||
if err := godotenv.Load(envPath); err != nil {
|
||
log.Printf("Warning: Failed to load .env file from %s: %v", envPath, err)
|
||
} else {
|
||
log.Printf("Loaded .env file from %s", envPath)
|
||
}
|
||
}
|
||
|
||
func InitDB() {
|
||
var err error
|
||
|
||
// Try to get DSN from environment variable, otherwise use default
|
||
dsn := os.Getenv("DB_DSN")
|
||
if dsn == "" {
|
||
dsn = "root:root@tcp(127.0.0.1:3306)/nl_blog?charset=utf8mb4&parseTime=True&loc=Local"
|
||
//dsn = "root:mysql_PKC65h@tcp(127.0.0.1:3306)/nl_blog?charset=utf8mb4&parseTime=True&loc=Local"
|
||
}
|
||
|
||
// Try to get JWT secret from environment variable
|
||
if secret := os.Getenv("JWT_SECRET"); secret != "" {
|
||
JWTSecret = secret
|
||
}
|
||
|
||
DB, err = gorm.Open(mysql.Open(dsn), &gorm.Config{
|
||
Logger: logger.Default.LogMode(logger.Silent), // 禁用 GORM 日志，使用自定义日志
|
||
// 不使用 NowFunc，因为我们使用模型钩子手动处理 int64 时间戳
|
||
})
|
||
if err != nil {
|
||
log.Fatal("Failed to connect to database:", err)
|
||
}
|
||
|
||
sqlDB, err := DB.DB()
|
||
if err != nil {
|
||
log.Fatal("Failed to get underlying sql.DB:", err)
|
||
}
|
||
|
||
if err = sqlDB.Ping(); err != nil {
|
||
log.Fatal("Failed to ping database:", err)
|
||
}
|
||
|
||
fmt.Println("Database connected successfully")
|
||
}
|
||
|
||
func CloseDB() {
|
||
if DB != nil {
|
||
sqlDB, err := DB.DB()
|
||
if err == nil {
|
||
sqlDB.Close()
|
||
}
|
||
}
|
||
}
|