如何使用go-mongo-driver调用MongoDB JS函数实现自增计数器

我来帮你把Mongo Shell里的自增计数器逻辑转换成Go语言版本，用官方的go-mongo-driver实现完全相同的功能。核心思路是用驱动的FindOneAndUpdate方法（对应Shell里的findAndModify）来保证操作的原子性，避免并发场景下的重复值问题。

步骤1：定义计数器结构体 #

首先我们定义一个对应counters集合文档的结构体，方便后续解码MongoDB的返回结果：

type Counter struct {
    ID           primitive.ObjectID `bson:"_id,omitempty"`
    SequenceName string             `bson:"sequenceName"`
    Sequence     int                `bson:"sequence"`
}

步骤2：实现获取下一个序列值的函数 #

这个函数完全复刻你原来JS函数的逻辑：检查当前计数器值，若超过9998则重置为1000，否则自增1，最后返回新的序列值：

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

func GetNextSequenceValue(ctx context.Context, countersColl *mongo.Collection, fieldName string) (int, error) {
    // 尝试自增：仅当当前sequence <= 9998时执行
    updateInc := bson.M{"$inc": bson.M{"sequence": 1}}
    opts := options.FindOneAndUpdate().SetReturnDocument(options.After)
    
    var updatedCounter Counter
    err := countersColl.FindOneAndUpdate(
        ctx,
        bson.M{
            "sequenceName": fieldName,
            "sequence":     bson.M{"$lte": 9998},
        },
        updateInc,
        opts,
    ).Decode(&updatedCounter)

    if err == nil {
        // 自增成功，直接返回新值
        return updatedCounter.Sequence, nil
    }

    // 若自增失败是因为没有匹配文档（说明sequence > 9998），执行重置操作
    if err != mongo.ErrNoDocuments {
        return 0, err
    }

    // 重置sequence为1000
    updateSet := bson.M{"$set": bson.M{"sequence": 1000}}
    err = countersColl.FindOneAndUpdate(
        ctx,
        bson.M{"sequenceName": fieldName},
        updateSet,
        opts,
    ).Decode(&updatedCounter)

    if err == mongo.ErrNoDocuments {
        // 处理计数器文档不存在的情况：插入初始值后递归调用
        initialCounter := Counter{
            SequenceName: fieldName,
            Sequence:     0, // 和你初始数据一致
        }
        _, err = countersColl.InsertOne(ctx, initialCounter)
        if err != nil {
            return 0, err
        }
        // 插入后重新获取序列值
        return GetNextSequenceValue(ctx, countersColl, fieldName)
    }

    if err != nil {
        return 0, err
    }

    return updatedCounter.Sequence, nil
}

步骤3：用序列值插入数据到STORE集合 #

现在可以调用上面的函数获取自增ID，然后插入数据到STORE集合，和你Shell里的操作逻辑一致：

func InsertStoreItem(ctx context.Context, storeColl *mongo.Collection, countersColl *mongo.Collection) error {
    // 获取下一个shortID的值
    shortID, err := GetNextSequenceValue(ctx, countersColl, "shortID")
    if err != nil {
        return err
    }

    // 构造待插入的文档
    item := bson.M{
        "shortID":         shortID,
        "item_short_name": "DEF",
        "specification":   "book",
        "category":        "fictional",
        "seller":          "best_buy",
        "network":         "XYZ",
        "plan":            "regular",
    }

    // 执行插入
    _, err = storeColl.InsertOne(ctx, item)
    return err
}

关键注意事项 #

• 原子性保证：FindOneAndUpdate是原子操作，和Shell的findAndModify一样，能避免多个请求同时修改计数器导致的重复值问题。
• 初始数据：确保counters集合中已存在初始文档{sequenceName: "shortID", sequence: 0}，如果不存在，上面的函数会自动插入初始值。
• Context使用：所有MongoDB操作都传入了context，你可以根据需求设置超时、取消信号等，比如ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)，记得调用cancel()。

内容的提问来源于stack exchange，提问作者Eu Ga