package utils import ( "context" "fmt" "go-policy-service/config" "log" "os" "time" "github.com/go-redis/redis/v8" "go.mongodb.org/mongo-driver/mongo" "go.mongodb.org/mongo-driver/mongo/options" "gorm.io/driver/mysql" "gorm.io/gorm" ) var Db *gorm.DB var Rdb *redis.Client var Mongo *mongo.Client func InitDB(cfg *config.Config) error { dsn := getDSN(cfg) db, err := gorm.Open(mysql.Open(dsn), &gorm.Config{}) if err != nil { return err } sqlDB, _ := db.DB() sqlDB.SetMaxIdleConns(cfg.DBMaxIdle) sqlDB.SetMaxOpenConns(cfg.DBMaxOpen) sqlDB.SetConnMaxLifetime(time.Duration(cfg.DBMaxLifetime) * time.Minute) Db = db return nil } func getDSN(cfg *config.Config) string { return fmt.Sprintf("%s:%s@tcp(%s:%d)/%s?charset=utf8mb4", cfg.DBUser, cfg.DBPassword, cfg.DBHost, cfg.DBPort, cfg.DBName) } // 初始化 Redis func InitRedis(cfg *config.Config) error { Rdb = redis.NewClient(&redis.Options{ Addr: fmt.Sprintf("%s:%d", cfg.RedisHost, cfg.RedisPort), Password: cfg.RedisPassword, DB: cfg.RedisDB, DialTimeout: 5 * time.Second, }) ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() if _, err := Rdb.Ping(ctx).Result(); err != nil { return fmt.Errorf("redis connection failed: %w", err) } return nil } func InitMongoDB(cfg *config.Config) { // 从环境变量中获取 MongoDB 连接信息 dsn := os.Getenv("MONGO_DSN") host := os.Getenv("MONGO_HOST") port := os.Getenv("MONGO_PORT") // 如果存在 MONGO_DSN 环境变量,则使用它来建立连接 var client *mongo.Client var err error if dsn != "" { client, err = mongo.Connect(context.TODO(), options.Client().ApplyURI(dsn)) } else { // 否则,使用 host 和 port 构建连接 URI uri := fmt.Sprintf("mongodb://%s:%s", host, port) client, err = mongo.Connect(context.TODO(), options.Client().ApplyURI(uri)) } if err != nil { log.Fatal(err) } // 检查连接 err = client.Ping(context.TODO(), nil) if err != nil { log.Fatal(err) } Mongo = client fmt.Println("Connected to MongoDB!") } func CloseResources() { if Db != nil { sqlDB, err := Db.DB() if err != nil { Logger.Error("Failed to get SQL DB: ", err) return } if err := sqlDB.Close(); err != nil { Logger.Error("Failed to close MySQL connection: ", err) } } if Rdb != nil { if err := Rdb.Close(); err != nil { Logger.Error("Failed to close Redis connection: ", err) } } if Mongo != nil { if err := Mongo.Disconnect(context.TODO()); err != nil { Logger.Error("Failed to close MongoDB connection: ", err) } } }