package database import ( "os" "path/filepath" "gorm.io/driver/postgres" "gorm.io/driver/sqlite" "gorm.io/gorm" "git.beifan.cn/trace-system/backend-go/config" "git.beifan.cn/trace-system/backend-go/logger" "git.beifan.cn/trace-system/backend-go/models" _ "github.com/mattn/go-sqlite3" ) var DB *gorm.DB // InitDB 初始化数据库连接 func InitDB() { var err error cfg := config.GetAppConfig() dbConfig := cfg.Database logger.Info("数据库配置信息", logger.String("驱动", dbConfig.Driver), ) switch dbConfig.Driver { case "sqlite": logger.Info("使用 SQLite 数据库", logger.String("路径", dbConfig.SQLite.Path)) DB, err = initSQLite(dbConfig.SQLite.Path) case "postgres": logger.Info("使用 PostgreSQL 数据库", logger.String("主机", dbConfig.Postgres.Host), logger.String("端口", dbConfig.Postgres.Port), logger.String("数据库", dbConfig.Postgres.DBName), ) DB, err = initPostgres(dbConfig.Postgres) default: logger.Fatal("不支持的数据库驱动: " + dbConfig.Driver) } if err != nil { logger.Fatal("数据库连接失败", logger.Err(err)) } logger.Info("数据库连接成功") } // initSQLite 初始化 SQLite 连接 func initSQLite(path string) (*gorm.DB, error) { // 确保数据目录存在 dataDir := filepath.Dir(path) if err := os.MkdirAll(dataDir, 0755); err != nil { return nil, err } db, err := gorm.Open(sqlite.Open(path+"?_pragma=foreign_keys(1)"), &gorm.Config{}) if err != nil { logger.Warn("SQLite 连接失败,尝试使用内存数据库", logger.Err(err)) // 尝试使用内存数据库作为备选方案 db, err = gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) } if err != nil { return nil, err } // 配置连接池 sqlDB, _ := db.DB() sqlDB.SetMaxOpenConns(25) sqlDB.SetMaxIdleConns(5) sqlDB.SetConnMaxLifetime(5 * 60 * 1000000000) return db, nil } // initPostgres 初始化 PostgreSQL 连接 func initPostgres(pgConfig struct { Host string `mapstructure:"host"` Port string `mapstructure:"port"` User string `mapstructure:"user"` Password string `mapstructure:"password"` DBName string `mapstructure:"dbname"` SSLMode string `mapstructure:"sslmode"` }) (*gorm.DB, error) { dsn := "host=" + pgConfig.Host + " port=" + pgConfig.Port + " user=" + pgConfig.User + " password=" + pgConfig.Password + " dbname=" + pgConfig.DBName + " sslmode=" + pgConfig.SSLMode db, err := gorm.Open(postgres.Open(dsn), &gorm.Config{}) if err != nil { logger.Fatal("PostgreSQL 连接失败", logger.Err(err)) } // 配置连接池 sqlDB, _ := db.DB() sqlDB.SetMaxOpenConns(100) sqlDB.SetMaxIdleConns(10) sqlDB.SetConnMaxLifetime(30 * 60 * 1000000000) return db, nil } // AutoMigrate 自动迁移数据库 func AutoMigrate() { if err := DB.AutoMigrate( &models.User{}, &models.Company{}, &models.Serial{}, ); err != nil { logger.Fatal("数据库迁移失败", logger.Err(err)) } logger.Info("数据库迁移成功") }