81 lines
1.9 KiB
Go
81 lines
1.9 KiB
Go
package store
|
|
|
|
import (
|
|
"fmt"
|
|
"log"
|
|
"opencatd-open/team/model"
|
|
"os"
|
|
"strings"
|
|
|
|
// "gocloud.dev/mysql"
|
|
// "gocloud.dev/postgres"
|
|
"github.com/glebarez/sqlite"
|
|
"gorm.io/driver/mysql"
|
|
"gorm.io/driver/postgres"
|
|
"gorm.io/gorm"
|
|
)
|
|
|
|
var IsPostgres bool
|
|
|
|
// InitDB 初始化数据库连接
|
|
func InitDB() (*gorm.DB, error) {
|
|
var db *gorm.DB
|
|
var err error
|
|
// 从环境变量获取DSN
|
|
dsn := os.Getenv("DSN")
|
|
|
|
if dsn == "" {
|
|
log.Println("No DSN provided, using SQLite as default")
|
|
db, err = initSQLite()
|
|
}
|
|
|
|
// 解析DSN来确定数据库类型
|
|
if strings.HasPrefix(dsn, "postgres://") {
|
|
IsPostgres = true
|
|
db, err = initPostgres(dsn)
|
|
} else if strings.HasPrefix(dsn, "mysql://") {
|
|
db, err = initMySQL(dsn)
|
|
}
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if IsPostgres {
|
|
err = db.AutoMigrate(&model.User{}, &model.ApiKey_PG{}, &model.Token{}, &model.Session{}, &model.Usage{}, &model.DailyUsage{})
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
}
|
|
|
|
return nil, fmt.Errorf("unsupported database type in DSN: %s", dsn)
|
|
}
|
|
|
|
// initSQLite 初始化 SQLite 数据库
|
|
func initSQLite() (*gorm.DB, error) {
|
|
db, err := gorm.Open(sqlite.Open("openteam.db"), &gorm.Config{})
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to connect to SQLite: %v", err)
|
|
}
|
|
return db, nil
|
|
}
|
|
|
|
// initPostgres 初始化 PostgreSQL 数据库
|
|
func initPostgres(dsn string) (*gorm.DB, error) {
|
|
db, err := gorm.Open(postgres.Open(dsn), &gorm.Config{})
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to connect to PostgreSQL: %v", err)
|
|
}
|
|
return db, nil
|
|
}
|
|
|
|
// initMySQL 初始化 MySQL 数据库
|
|
func initMySQL(dsn string) (*gorm.DB, error) {
|
|
// 移除 "mysql://" 前缀，因为 MySQL 驱动不需要这个前缀
|
|
dsn = strings.TrimPrefix(dsn, "mysql://")
|
|
|
|
db, err := gorm.Open(mysql.Open(dsn), &gorm.Config{})
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to connect to MySQL: %v", err)
|
|
}
|
|
return db, nil
|
|
}
|