115 lines
2.5 KiB
Go
115 lines
2.5 KiB
Go
package store
|
|
|
|
import (
|
|
"fmt"
|
|
"log"
|
|
"opencatd-open/internal/model"
|
|
"opencatd-open/pkg/config"
|
|
"os"
|
|
"strings"
|
|
|
|
// "gocloud.dev/mysql"
|
|
// "gocloud.dev/postgres"
|
|
"github.com/glebarez/sqlite"
|
|
"gorm.io/driver/mysql"
|
|
"gorm.io/driver/postgres"
|
|
|
|
// "gorm.io/driver/sqlite"
|
|
|
|
"gorm.io/gorm"
|
|
)
|
|
|
|
var DB *gorm.DB
|
|
|
|
// var DBType consts.DBType
|
|
var IsPostgres bool
|
|
|
|
func GetDB() *gorm.DB {
|
|
return DB
|
|
}
|
|
|
|
// InitDB 初始化数据库连接
|
|
func InitDB(cfg *config.Config) (*gorm.DB, error) {
|
|
var db *gorm.DB
|
|
var err error
|
|
// 从环境变量获取DSN
|
|
dsn := cfg.DSN
|
|
|
|
if dsn == "" {
|
|
log.Println("No DSN provided, using SQLite as default")
|
|
db, err = initSQLite()
|
|
}
|
|
|
|
// 解析DSN来确定数据库类型
|
|
if strings.HasPrefix(dsn, "postgres://") {
|
|
IsPostgres = true
|
|
|
|
cfg.DB_Type = "postgres"
|
|
db, err = initPostgres(dsn)
|
|
} else if strings.HasPrefix(dsn, "mysql://") {
|
|
|
|
cfg.DB_Type = "mysql"
|
|
db, err = initMySQL(dsn)
|
|
} else {
|
|
if dsn != "" {
|
|
return nil, fmt.Errorf("unsupported database type in DSN: %s", dsn)
|
|
}
|
|
}
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
DB = db
|
|
|
|
if IsPostgres {
|
|
err = db.AutoMigrate(&model.User{}, &model.Token{}, &model.ApiKey_PG{}, &model.Usage{}, &model.DailyUsage{}, &model.Passkey{})
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
} else {
|
|
err = db.AutoMigrate(&model.User{}, &model.Token{}, &model.ApiKey{}, &model.Usage{}, &model.DailyUsage{}, &model.Passkey{})
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
}
|
|
|
|
return db, nil
|
|
}
|
|
|
|
// initSQLite 初始化 SQLite 数据库
|
|
func initSQLite() (*gorm.DB, error) {
|
|
if _, err := os.Stat("db"); os.IsNotExist(err) {
|
|
errDir := os.MkdirAll("db", 0755)
|
|
if errDir != nil {
|
|
log.Fatalln("Error creating directory:", err)
|
|
}
|
|
}
|
|
db, err := gorm.Open(sqlite.Open("./db/openteam.db"), &gorm.Config{})
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to connect to SQLite: %v", err)
|
|
}
|
|
// db, err := gorm.Open(sqlite.Open("file::memory:?cache=shared"), &gorm.Config{})
|
|
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
|
|
}
|