Initial commit

2026-02-12 14:31:30 +08:00
commit e01cdc9889
25 changed files with 3227 additions and 0 deletions
--- a/database/database.go
+++ b/database/database.go
@@ -0,0 +1,120 @@
+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("数据库迁移成功")
+}