MongoDB Golang 示例代码

package main
 
import (
    "context"
    "flag"
    "fmt"
    "go.mongodb.org/mongo-driver/bson"
    "log"
    "math/rand"
    "sync"
    "time"
 
    "go.mongodb.org/mongo-driver/mongo"
    "go.mongodb.org/mongo-driver/mongo/options"
    //"go.mongodb.org/mongo-driver/bson/primitive"
)
 
type testStruct struct {
    ID uint64
    Name string
}
 
var ID uint64
var l sync.Mutex
func genID() uint64 {
 
    l.Lock()
    defer  l.Unlock()
    ID += 1
    return ID
}
 
func main() {
 
    var rcot int    // insert 协程数
    var frcot int    // find 协程数
 
    flag.IntVar(&rcot, "c", 1, "insert go routine count")
    flag.IntVar(&frcot, "f", 1, "find go routine count")
    flag.Parse()
 
    clientOptions := options.Client().ApplyURI("mongodb://localhost:27017")
    client, err := mongo.Connect(context.TODO(),clientOptions)
    if err != nil {
        return
    }
 
 
    db := client.Database("test")
    col := db.Collection("colA")
 
    // 初始化最大ID, 使用两种方式获取最大值
    //initID()    // 使用聚合，效率较低
    initIDEx(col)    // 有索引，使用sort效率较高
 
    // 压力测试插入
    testInsert( col, rcot)
 
    // 各种情况查询
    find2(col)
 
    // 更新
    update(col)
 
    // 测试删除
    count(col)
    delete(col)
    count(col)
 
    // 测试事务， 需要开启replica set 集群
    transaction(client)
 
    // 设置索引
    testIndex(col)
 
    //  压力测试 查询
    find(col, frcot)
 
}
 
// 更新
func update(col *mongo.Collection) {
    rand.Seed(int64(time.Now().UnixNano()))
    id := rand.Uint64()%ID
    filter := bson.D{{"id", id}}
    update := bson.D{{"$set", bson.D{{"name", "xxxxxxx"}}}}
 
    result,err := col.UpdateOne(context.TODO(), filter, update )
    if err != nil {
        log.Println(err)
        return
    }
    fmt.Println("update count:", result.ModifiedCount)
 
}
 
// 删除
func delete(col *mongo.Collection) {
    rand.Seed(int64(time.Now().Second()) )
    id := rand.Uint64()%ID
    filter := bson.D{{"id", id}}
    res, err := col.DeleteOne(context.TODO(), filter)
    if err != nil {
        fmt.Println(err)
        return
    }
    fmt.Println( "delete count", res.DeletedCount)
}
 
// Count 数量统计
func count(col *mongo.Collection) {
    cot,_ := col.CountDocuments(context.TODO(), bson.D{})
    fmt.Println("count:", cot)
}
 
// 事务
func transaction(client *mongo.Client) {
 
    fmt.Println("start test transaction")
 
    db := client.Database("test")
    col := db.Collection("colA")
 
    count(col)
 
    ctx := context.Background()
 
    err := client.UseSession(ctx, func(sessioinContext mongo.SessionContext) error {
 
        if err := sessioinContext.StartTransaction(); err !=nil {
            return er
        }
 
        res, err := col.InsertOne(sessioinContext, bson.M{"id":genID(), "name":GenRandStr(12)})
        if err != nil {
            fmt.Println("insert err in transaction:", err)
            return er
        }
        fmt.Println("insert result:", res.InsertedID)
 
        err = sessioinContext.CommitTransaction(sessioinContext)
        if err != nil {
            fmt.Println("commit transaction err:", err)
        }
        return nil
    })
    if err != nil {
        fmt.Println("trans err:", err)
        return
    }
 
    count(col)
    fmt.Println("end test transaction")
 
}
 
 
// 查询
func find2(col *mongo.Collection){
 
    // 查询ID小于等于3的项
    fo := options.FindOptions{}
    fo.SetLimit(2)    // 不设置limit,会查处3个值
    //fo.SetSort(bson.D{{"id", -1}})
    cur, err:= col.Find(context.TODO(), bson.D{{"id",bson.D{{"$lte",3}}}}, &fo)
    if err != nil {
        panic(err)
    }
    defer cur.Close(context.TODO())
 
    // 批量decode
    var a []testStruct
    cur.All(context.TODO(), &a)
    fmt.Println("find all result:", a)
 
    // 逐个迭代decode (这里并不会进去，除非注释上面的All()，具体原因看All接口的注释说明)
    for cur.Next(context.TODO()){
        temp := testStruct{}
        cur.Decode(&temp)
        println(temp.ID, temp.Name)
    }
 
    // 使用findOne 进行多条件查询
    res := col.FindOne(context.TODO(), bson.D{{"id", a[0].ID},{"name", a[0].Name}})
 
    t := testStruct{}
    res.Decode(&t)
    fmt.Println( "multiple conditions query:", t.ID, t.Name)
}
 
// 生成随机字符串
func GenRandStr(len int) string {
 
    bs := make([]byte, 0, len)
    for i:=0; i<len;i++{
        r := rand.Int()%26
        bs = append(bs, 'a'+byte(r))
    }
    return string(bs)
}
 
 
// max 聚合
func initID(col *mongo.Collection) {
    pip := mongo.Pipeline{bson.D{
        {"$group", bson.D{
            {"_id", ""},
            {"maxid", bson.D{{"$max", "$id"}}},
        },
        }},
    }
    cur, err := col.Aggregate(context.TODO(), pip)
    if err != nil {
        log.Println(err)
        return
    }
    defer cur.Close(context.TODO())
 
    var results []bson.M
    cur.All(context.TODO(), &results)
    for _, result := range results {
        fmt.Println(result)
        if v, ok := result["maxid"]; ok {
            println("max id is ", v.(int64))
            ID = uint64(v.(int64))
        }
    }
}
 
// max findOne sort
func initIDEx(col *mongo.Collection) {
    findOp := options.FindOne()
    findOp.SetSort(bson.D{{"id", -1}})
    rst := col.FindOne(context.TODO(), bson.D{}, findOp)
 
    var a testStruct
    if err := rst.Decode(&a); err == nil {
        ID = a.ID
        println("max id:", ID)
    }
 
}
 
 
// 查询
func testFind(col *mongo.Collection, wg *sync.WaitGroup){
    t0:=time.Now()
    for k:=0;k<10000;k++{
        id := rand.Uint64()%ID
        cur, err:= col.Find(context.TODO(), bson.D{{"id",id}})
        if err != nil {
            panic(err)
        }
 
        for cur.Next(context.TODO()){
            a := testStruct{}
            cur.Decode(&a)
            if a.ID != id {
                println(a.ID, id, ID)
 
            }
        }
 
        cur.Close(context.TODO())
    }
    fmt.Println("find use time:", time.Since(t0))
    wg.Done()
}
 
// 查询压力测试
func find(col *mongo.Collection, cot int) {
    wg := sync.WaitGroup{}
    wg.Add(cot)
    t0:=time.Now()
    for k:=0;k<cot;k++{
        go testFind(col, &wg)
    }
    wg.Wait()
    fmt.Println(cot, "go routine find total use time ", time.Since(t0))
 
}
 
// 索引
func testIndex(col *mongo.Collection) {
    indexName, err := col.Indexes().CreateOne(
        context.Background(),
        mongo.IndexModel{
            Keys: bson.M{"id": 1},
            Options: options.Index().SetUnique(true),
        },
    )
    if err != nil {
        log.Fatal(err)
    }
    fmt.Println("index name:", indexName)
 
}
 
func testInsert( col *mongo.Collection, rcot int) {
    t0:=time.Now()
    wg := sync.WaitGroup{}
    for k:=0;k<rcot;k++{
        wg.Add(1)
        go insert(col, &wg)
    }
    wg.Wait()
    fmt.Println(rcot, "go routines insert total use time:", time.Since(t0))
}
 
// 插入
func insert(col *mongo.Collection, wg *sync.WaitGroup) {
    defer wg.Done()
 
    t0:=time.Now()
    for i:=0;i<10000;i++{
        //tdata := testStruct{
        //    ID: genID(),
        //    Name: GenRandStr(12),
        //}
        //_,err := col.InsertOne(context.TODO(), &tdata)
 
        // 上面的方式也可以
        _, err := col.InsertOne(context.TODO(), bson.M{"id":genID(), "name":GenRandStr(12)})
        if err != nil {
            fmt.Println(err)
            return
        }
    }
 
    //fmt.Println("Last insert id:", res.InsertedID)
    fmt.Println("insert 10000 data use time:", time.Since(t0))
}