go-policy-service/utils/init_service.go

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)
		}
	}
}