llzai's picture
Upload 42 files
644c352 verified
package model
import (
"fmt"
"log"
"os"
"time"
"gorm.io/driver/postgres"
"gorm.io/gorm"
"gorm.io/gorm/logger"
)
var DB *gorm.DB
// InitDB 初始化数据库连接
func InitDB() error {
databaseURL := os.Getenv("DATABASE_URL")
if databaseURL == "" {
return fmt.Errorf("DATABASE_URL environment variable is not set")
}
var err error
DB, err = gorm.Open(postgres.Open(databaseURL), &gorm.Config{
Logger: logger.Default.LogMode(logger.Info),
NowFunc: func() time.Time {
return time.Now().UTC()
},
})
if err != nil {
return fmt.Errorf("failed to connect to database: %w", err)
}
// 配置连接池
sqlDB, err := DB.DB()
if err != nil {
return fmt.Errorf("failed to get database instance: %w", err)
}
sqlDB.SetMaxIdleConns(10)
sqlDB.SetMaxOpenConns(100)
sqlDB.SetConnMaxLifetime(time.Hour)
// 自动迁移
if err := autoMigrate(); err != nil {
return fmt.Errorf("failed to migrate database: %w", err)
}
log.Println("Database connected and migrated successfully")
return nil
}
// autoMigrate 自动迁移数据库表结构
func autoMigrate() error {
return DB.AutoMigrate(
&User{},
&MorphCookie{},
&UserSession{},
)
}
// CloseDB 关闭数据库连接
func CloseDB() error {
sqlDB, err := DB.DB()
if err != nil {
return err
}
return sqlDB.Close()
}