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)
	conn, err := gorm.Open(mysql.Open(dsn), &gorm.Config{})
	if err != nil {
		return err
	}
	Db = conn
	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,
	})
	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!")
}