s_cherry_e/sr/nodes/map/internal/db/ver_helper.go

package db

import (
	"context"
	dbFacade "f1-game/internal/component/db/facade"

	clog "github.com/cherry-game/cherry/logger"
	"go.mongodb.org/mongo-driver/v2/bson"
	"go.mongodb.org/mongo-driver/v2/mongo"
	"go.mongodb.org/mongo-driver/v2/mongo/options"
)

// TODO 暂时放这里
func getAllTables[T dbFacade.Table](tableName string) []T {
	ctx := context.Background()
	collection := dbComponent.DB().Collection(tableName)
	filter := bson.M{}
	cur, err := collection.Find(ctx, filter)
	if err != nil {
		clog.Errorf("getAllTables: err = %v", err)
		return nil
	}

	list := []T{}
	err = cur.All(ctx, &list)
	if err != nil {
		clog.Errorf("getAllTables: err = %v", err)
		return nil
	}

	return list
}

func getAllTableMap[T dbFacade.Table](tableName string) map[any]T {
	list := getAllTables[T](tableName)

	m := map[any]T{}
	for _, table := range list {
		m[table.PrimaryValue()] = table
	}

	return m
}

func bulkWrite[T dbFacade.Table](collection *mongo.Collection, list []T) error {
	// // 构建批量更新对象
	models := make([]mongo.WriteModel, 0, len(list))
	for _, table := range list {
		tableBytes, err := bson.Marshal(table)
		if err != nil {
			return err
		}

		// set update
		update := bson.D{{Key: "$set", Value: bson.Raw(tableBytes)}}

		// set filter
		filter := bson.M{
			"_id": table.PrimaryValue(), // 利用自带的_id主键进行条件过滤
		}

		// build *mongo.UpdateOneModel
		model := mongo.NewUpdateOneModel().SetFilter(filter).SetUpdate(update).SetUpsert(true)
		models = append(models, model)
	}

	ops := options.BulkWrite().SetOrdered(false)
	result, err := collection.BulkWrite(context.TODO(), models, ops)
	if err != nil {
		return err
	}

	clog.Debugf("bulk write [table = %s, count = %d]", collection.Name(), result.ModifiedCount)
	return nil
}