go 预处理 mongodb 的insert

1.如何预处理

如果我们想预处理mongo的数据创建方法，如添加 公共字段 public1，public2

可以重写collection来实现

具体如下：

package mongo

import (
	"context"
	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"
)

type Collection struct {
	*mongo.Collection
}

func (col *Collection) InsertOne(ctx context.Context, document interface{},
	opts ...*options.InsertOneOptions) (*mongo.InsertOneResult, error) {
	b, err := bson.Marshal(document)
	if err != nil {
		return nil, err
	}
	newDoc := bson.M{}
	err = bson.Unmarshal(b, &newDoc)
	if err != nil {
		return nil, err
	}
	newDoc["public1"] = "public1"
	newDoc["public2"] = "public2"
	return col.Collection.InsertOne(ctx, newDoc, opts...)
}

func (col *Collection) InsertMany(ctx context.Context, documents []interface{},
	opts ...*options.InsertManyOptions) (*mongo.InsertManyResult, error) {

	for i, v := range documents {
		b, err := bson.Marshal(v)
		if err != nil {
			return nil, err
		}
		newDoc := bson.M{}
		err = bson.Unmarshal(b, &newDoc)
		if err != nil {
			return nil, err
		}
		newDoc["public1"] = "public1"
		newDoc["public2"] = "public2"
		documents[i] = newDoc
	}
	return col.Collection.InsertMany(ctx, documents, opts...)
}

2.如何使用？

下面写了一个go test 来举例如何使用

package mongo

import (
	"context"
	"fmt"
	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"
	"go.mongodb.org/mongo-driver/mongo/writeconcern"
	"log"
	"testing"
)

func cli() *mongo.Client {
	// 设置客户端选项
	clientOptions := options.Client().ApplyURI("mongodb://localhost:27017")
	// 连接 MongoDB
	client, err := mongo.Connect(context.TODO(), clientOptions)
	if err != nil {
		log.Fatal(err)
	}
	// 检查连接
	err = client.Ping(context.TODO(), nil)
	if err != nil {
		log.Fatal(err)
	}
	return client
}

func TestCreateMongoCollection(t *testing.T) {
	c := cli()
	collectionStudent := &Collection{c.Database("mongodb_study").Collection("student")}
	collectionStudent2 := &Collection{c.Database("mongodb_study").Collection("student2")}
	ctx := context.TODO()
	collectionStudent.DeleteMany(ctx, bson.M{})
	collectionStudent2.DeleteMany(ctx, bson.M{})

	type Student struct {
		Name string
		Age  int
	}
	//1.插入一条数据
	stu1 := Student{"李四", 133}
	_, err := collectionStudent.InsertOne(ctx, stu1)
	if err != nil {
		log.Println(err)
	}
	//2.插入多条文档
	students := []interface{}{Student{"张三", 12}, Student{"王五", 16}}
	_, err = collectionStudent.InsertMany(ctx, students)
	if err != nil {
		log.Println(err)
	}
	//3.更新单个文档，如果有多个满足条件的，只更新第一条
	filter := bson.D{{"name", "张三"}}
	// 修改name 为hhhh
	update := bson.D{
		{"$set", bson.D{{"name", "张三三"}}},
	}
	collectionStudent.UpdateOne(ctx, filter, update)

	//4.开启事务
	session, err := c.StartSession()
	if err != nil {
		panic(err)
	}
	defer session.EndSession(context.TODO())
	wc := writeconcern.New(writeconcern.WMajority())
	txnOptions := options.Transaction().SetWriteConcern(wc)
	session.WithTransaction(context.TODO(), func(ctx mongo.SessionContext) (interface{}, error) {
		students := []interface{}{Student{"张三", 15}, Student{"王五", 17}}
		_, err := collectionStudent.InsertMany(ctx, students)
		if err != nil {
			fmt.Println("InsertMany err:", err)
		}
		_, err = collectionStudent2.InsertMany(ctx, students)
		if err != nil {
			fmt.Println("InsertMany2 err:", err)
		}
		update := bson.D{
			{"$set", bson.M{"name": "张三三", "age": 18}},
		}
		_, err = collectionStudent.UpdateOne(ctx, filter, update)
		if err != nil {
			fmt.Println("UpdateOne err:", err)
		}

		return nil, nil
	}, txnOptions)

}