box-cost/db/db.go

package db

import (
	"context"
	"log"
	"time"

	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"
)

const (
	MONGODB_URI = "mongodb://root:boxcost@124.71.139.24:37030/box-cost?authSource=admin"
	DB_NAME     = "box-cost"
	COLL_NAME   = "product-plan"
)

var (
	client     *mongo.Client
	collection *mongo.Collection
)

// 初始化MongoDB连接
func InitMongoConn() {
	// 创建一个新的上下文
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	// 连接到MongoDB
	var err error
	client, err = mongo.Connect(ctx, options.Client().ApplyURI(MONGODB_URI))
	if err != nil {
		log.Fatal(err)
	}

	// 选择数据库和集合
	collection = client.Database(DB_NAME).Collection(COLL_NAME)
}

// 获取已初始化的集合
func GetCollection() *mongo.Collection {
	if collection == nil {
		log.Fatal("MongoDB connection not initialized")
	}
	return collection
}

// 关闭MongoDB连接
func CloseMongoConn() {
	if client != nil {
		// 创建一个新的上下文
		ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
		defer cancel()

		// 关闭连接
		err := client.Disconnect(ctx)
		if err != nil {
			log.Fatal(err)
		}
	}
}