85 lines
1.6 KiB
Go
85 lines
1.6 KiB
Go
package database
|
|
|
|
import (
|
|
"fmt"
|
|
"os"
|
|
"path/filepath"
|
|
"time"
|
|
|
|
"github.com/drama-generator/backend/domain/models"
|
|
"github.com/drama-generator/backend/pkg/config"
|
|
"gorm.io/driver/mysql"
|
|
"gorm.io/driver/sqlite"
|
|
"gorm.io/gorm"
|
|
)
|
|
|
|
func NewDatabase(cfg config.DatabaseConfig) (*gorm.DB, error) {
|
|
dsn := cfg.DSN()
|
|
|
|
if cfg.Type == "sqlite" {
|
|
dbDir := filepath.Dir(dsn)
|
|
if err := os.MkdirAll(dbDir, 0755); err != nil {
|
|
return nil, fmt.Errorf("failed to create database directory: %w", err)
|
|
}
|
|
}
|
|
|
|
gormConfig := &gorm.Config{
|
|
Logger: NewCustomLogger(),
|
|
}
|
|
|
|
var db *gorm.DB
|
|
var err error
|
|
|
|
if cfg.Type == "sqlite" {
|
|
db, err = gorm.Open(sqlite.Open(dsn), gormConfig)
|
|
} else {
|
|
db, err = gorm.Open(mysql.Open(dsn), gormConfig)
|
|
}
|
|
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to connect to database: %w", err)
|
|
}
|
|
|
|
sqlDB, err := db.DB()
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to get database instance: %w", err)
|
|
}
|
|
|
|
sqlDB.SetMaxIdleConns(cfg.MaxIdle)
|
|
sqlDB.SetMaxOpenConns(cfg.MaxOpen)
|
|
sqlDB.SetConnMaxLifetime(time.Hour)
|
|
|
|
if err := sqlDB.Ping(); err != nil {
|
|
return nil, fmt.Errorf("failed to ping database: %w", err)
|
|
}
|
|
|
|
return db, nil
|
|
}
|
|
|
|
func AutoMigrate(db *gorm.DB) error {
|
|
return db.AutoMigrate(
|
|
// 核心模型
|
|
&models.Drama{},
|
|
&models.Episode{},
|
|
&models.Character{},
|
|
&models.Scene{},
|
|
&models.Storyboard{},
|
|
|
|
// 生成相关
|
|
&models.ImageGeneration{},
|
|
&models.VideoGeneration{},
|
|
&models.VideoMerge{},
|
|
|
|
// AI配置
|
|
&models.AIServiceConfig{},
|
|
&models.AIServiceProvider{},
|
|
|
|
// 资源管理
|
|
&models.Asset{},
|
|
&models.CharacterLibrary{},
|
|
|
|
// 任务管理
|
|
&models.AsyncTask{},
|
|
)
|
|
}
|