mdbc/bulk_write_scope_test.go

package mdbc

import (
	"encoding/json"
	"fmt"
	"testing"
	"time"

	"github.com/google/uuid"
	"github.com/sirupsen/logrus"
	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/mongo"
)

func TestBulkWriteScope(t *testing.T) {
	cfg := &Config{
		URI:         "mongodb://admin:admin@10.0.0.135:27017/admin",
		MinPoolSize: 32,
		ConnTimeout: 10,
	}
	cfg.RegistryBuilder = RegisterTimestampCodec(nil)
	client, err := ConnInit(cfg)

	if err != nil {
		logrus.Fatalf("get err: %+v", err)
	}
	InitDB(client.Database("heywoods_golang_jingliao_crm_dev"))

	var m = NewModel(&ModelRobotFriend{})

	var record ModelRobotFriend
	if err := m.FindOne().SetFilter(bson.M{
		ModelRobotFriendField_Id.DbFieldName: "498b1c85be41266efb29b6a79560ec7f",
	}).Get(&record); err != nil {
		panic(err)
	}

	record.AddAt = 12345
	record.UpdateTime = time.Now().Unix()

	var updateData = bson.M{
		"$set": &record,
		"$setOnInsert": bson.M{
			"create_time": 1000000,
		},
	}
	b, _ := bson.MarshalExtJSON(updateData, true, true)
	logrus.Infof("updateData: %+v", string(b))
	SetMapOmitInsertField(&updateData)
	b, _ = json.Marshal(updateData)
	logrus.Infof("updateData: %+v", string(b))

	var opData []mongo.WriteModel
	opData = append(opData, mongo.NewUpdateOneModel().SetFilter(bson.M{
		ModelRobotFriendField_Id.DbFieldName: "498b1c85be41266efb29b6a79560ec7f",
	}).SetUpsert(true).SetUpdate(updateData))

	res, err := m.SetDebug(true).BulkWrite().
		SetWriteModel(opData).Do()

	if err != nil {
		panic(err)
	}
	fmt.Printf("res %+v\n", res)
	logrus.Infof("res %+v\n", res)
}

func TestSliceChunk(t *testing.T) {
	a := []int{1, 2, 3}
	chunkSize := 2

	var b [][]int
	for i := 0; i < len(a); i += chunkSize {
		end := i + chunkSize

		if end > len(a) {
			end = len(a)
		}

		b = append(b, a[i:end])
	}

	fmt.Println(b)
}

func TestBulkWriteScope_SetChunkSize(t *testing.T) {
	cfg := &Config{
		URI:         "mongodb://admin:admin@10.0.0.135:27017/admin",
		MinPoolSize: 32,
		ConnTimeout: 10,
	}
	cfg.RegistryBuilder = RegisterTimestampCodec(nil)
	client, err := ConnInit(cfg)

	if err != nil {
		logrus.Fatalf("get err: %+v", err)
	}
	InitDB(client.Database("heywoods_golang_jingliao_crm_dev"))

	var m = NewModel(&ModelRobotFriend{})
	var opData []mongo.WriteModel

	for i := 0; i < 10; i++ {
		opData = append(opData, mongo.NewInsertOneModel().SetDocument(&ModelRobotFriend{Id: uuid.NewString()}))
	}

	res, err := m.SetDebug(true).BulkWrite().SetChunkSize(3).SetWriteModel(opData).Do()

	if err != nil {
		panic(err)
	}

	fmt.Println("res", res)
}