Online_whiteboard_backend/backend/global/mongo.go

package global

import (
	"context"
	"file_upload/backend/simple_zap"
	"fmt"
	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/gridfs"
	"go.mongodb.org/mongo-driver/mongo/options"
	"go.mongodb.org/mongo-driver/mongo/readpref"
	"io"
)

var MongoCon *mongo.Collection
var MongoGridFS *gridfs.Bucket

func SetupMongo() {
	clientOptions := options.Client().ApplyURI(MongoSetting.Url)
	connect, err := mongo.Connect(context.TODO(), clientOptions)
	if err != nil {
		simple_zap.WithCtx(context.Background()).Sugar().Warn(err, "mongo 出现异常")
		panic(err)
	}
	err = connect.Ping(context.TODO(), readpref.Primary())
	if err != nil {
		simple_zap.WithCtx(context.Background()).Sugar().Warn(err, "mongo 连接失败")
		panic(err)
	}
	database := connect.Database("bzd_file")
	MongoCon = database.Collection("file")
	MongoGridFS, err = gridfs.NewBucket(database)
	if err != nil {
		simple_zap.WithCtx(context.Background()).Sugar().Warn(err)
		panic(err)
	}
}
func UploadFileToGridFS(fileName string, fileData []byte) (string, error) {
	uploadStream, err := MongoGridFS.OpenUploadStream(fileName)
	if err != nil {
		return "", fmt.Errorf("无法打开上传流: %v", err)
	}
	defer uploadStream.Close()

	_, err = uploadStream.Write(fileData)
	if err != nil {
		return "", fmt.Errorf("写入文件数据时出错: %v", err)
	}

	return uploadStream.FileID.(string), nil
}

func DownloadFileFromGridFS(fileID string) ([]byte, error) {
	downloadStream, err := MongoGridFS.OpenDownloadStream(fileID)
	if err != nil {
		return nil, fmt.Errorf("无法打开下载流: %v", err)
	}
	defer downloadStream.Close()

	fileData, err := io.ReadAll(downloadStream)
	if err != nil {
		return nil, fmt.Errorf("读取文件数据时出错: %v", err)
	}

	return fileData, nil
}