package db import ( "fmt" "gpt-load/internal/models" "log" "os" "time" "gorm.io/driver/mysql" "gorm.io/gorm" "gorm.io/gorm/logger" ) var DB *gorm.DB func InitDB() (*gorm.DB, error) { // TODO: 从配置中心读取DSN dsn := "root:1236@tcp(127.0.0.1:3306)/gpt_load?charset=utf8mb4&parseTime=True&loc=Local" newLogger := logger.New( log.New(os.Stdout, "\r\n", log.LstdFlags), // io writer logger.Config{ SlowThreshold: time.Second, // Slow SQL threshold LogLevel: logger.Info, // Log level IgnoreRecordNotFoundError: true, // Ignore ErrRecordNotFound error for logger Colorful: true, // Disable color }, ) var err error DB, err = gorm.Open(mysql.Open(dsn), &gorm.Config{ Logger: newLogger, }) 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 sql.DB: %w", err) } // Set connection pool parameters sqlDB.SetMaxIdleConns(10) sqlDB.SetMaxOpenConns(100) sqlDB.SetConnMaxLifetime(time.Hour) // Auto-migrate models err = DB.AutoMigrate( &models.SystemSetting{}, &models.Group{}, &models.APIKey{}, &models.RequestLog{}, ) if err != nil { return nil, fmt.Errorf("failed to auto-migrate database: %w", err) } fmt.Println("Database connection initialized and models migrated.") return DB, nil }