116 lines
2.8 KiB
Go
116 lines
2.8 KiB
Go
package database
|
||
|
||
import (
|
||
"fmt"
|
||
"log"
|
||
"os"
|
||
"strings"
|
||
|
||
"gorm.io/driver/mysql"
|
||
"gorm.io/gorm"
|
||
"gorm.io/gorm/logger"
|
||
)
|
||
|
||
// loadDotEnv 在当前工作目录查找 .env 文件并加载到进程环境变量中（已存在的变量不覆盖）。
|
||
// 仅用于本地开发场景：生产部署通过 docker-compose / 容器环境变量直接注入，没有 .env 文件时静默跳过。
|
||
func loadDotEnv() {
|
||
data, err := os.ReadFile(".env")
|
||
if err != nil {
|
||
return
|
||
}
|
||
for _, line := range strings.Split(string(data), "\n") {
|
||
line = strings.TrimSpace(line)
|
||
if line == "" || strings.HasPrefix(line, "#") {
|
||
continue
|
||
}
|
||
key, value, ok := strings.Cut(line, "=")
|
||
if !ok {
|
||
continue
|
||
}
|
||
key = strings.TrimSpace(key)
|
||
if key == "" {
|
||
continue
|
||
}
|
||
if _, exists := os.LookupEnv(key); exists {
|
||
continue
|
||
}
|
||
os.Setenv(key, strings.TrimSpace(value))
|
||
}
|
||
}
|
||
|
||
// DBConfig holds MySQL connection parameters.
|
||
type DBConfig struct {
|
||
Host string
|
||
Port string
|
||
User string
|
||
Password string
|
||
DBName string
|
||
}
|
||
|
||
func (c DBConfig) DSN() string {
|
||
return fmt.Sprintf("%s:%s@tcp(%s:%s)/%s?charset=utf8mb4&parseTime=True&loc=Local",
|
||
c.User, c.Password, c.Host, c.Port, c.DBName)
|
||
}
|
||
|
||
// envConfig 从环境变量读取连接参数。Host/User/Password 没有内置默认值，必须由
|
||
// 部署环境（本地 shell、.env、docker-compose 等）显式提供；DBName 默认统一为
|
||
// "sproutgate"，开发和生产环境共用同一个库名。
|
||
func envConfig() (DBConfig, error) {
|
||
host := os.Getenv("DB_HOST")
|
||
user := os.Getenv("DB_USER")
|
||
password := os.Getenv("DB_PASSWORD")
|
||
if host == "" || user == "" || password == "" {
|
||
return DBConfig{}, fmt.Errorf("缺少数据库环境变量，请设置 DB_HOST/DB_USER/DB_PASSWORD（或直接设置 DB_DSN）")
|
||
}
|
||
|
||
port := os.Getenv("DB_PORT")
|
||
if port == "" {
|
||
port = "3306"
|
||
}
|
||
dbName := os.Getenv("DB_NAME")
|
||
if dbName == "" {
|
||
dbName = "sproutgate"
|
||
}
|
||
|
||
return DBConfig{Host: host, Port: port, User: user, Password: password, DBName: dbName}, nil
|
||
}
|
||
|
||
// Open 根据环境变量打开数据库连接。
|
||
// 优先级：DB_DSN（完整自定义连接串）> DB_HOST/DB_PORT/DB_USER/DB_PASSWORD/DB_NAME
|
||
func Open() (*gorm.DB, error) {
|
||
loadDotEnv()
|
||
|
||
dsn := os.Getenv("DB_DSN")
|
||
if dsn == "" {
|
||
cfg, err := envConfig()
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
dsn = cfg.DSN()
|
||
log.Printf("[database] 使用数据库 %s:%s/%s", cfg.Host, cfg.Port, cfg.DBName)
|
||
} else {
|
||
log.Println("[database] 使用自定义 DB_DSN")
|
||
}
|
||
|
||
logLevel := logger.Info
|
||
if os.Getenv("GIN_MODE") == "release" {
|
||
logLevel = logger.Warn
|
||
}
|
||
|
||
db, err := gorm.Open(mysql.Open(dsn), &gorm.Config{
|
||
Logger: logger.Default.LogMode(logLevel),
|
||
})
|
||
if err != nil {
|
||
return nil, fmt.Errorf("连接数据库失败: %w", err)
|
||
}
|
||
|
||
sqlDB, err := db.DB()
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
sqlDB.SetMaxOpenConns(25)
|
||
sqlDB.SetMaxIdleConns(5)
|
||
|
||
return db, nil
|
||
}
|