gd_crontab/common.in/mongo/mgo.go

package mongo

import (
	"encoding/json"
	"errors"
	"fmt"
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
)

var Session *mgo.Session

func MgoInit(addr string, username string, password string) error {
	var err error
	var session *mgo.Session
	if addr == "" {
		return errors.New("mongo init need addr")
	}
	if username == "" && password == "" {
		fmt.Printf("1:%v,%v,%v\n", addr, username, password)
		session, err = mgo.Dial(addr)
	} else {
		fmt.Printf("2:%v,%v,%v\n", addr, username, password)
		url := fmt.Sprintf("mongodb://%s:%s@%s", username, password, addr)
		session, err = mgo.Dial(url)
	}

	if err != nil {
		return err
	}
	Session = session
	return nil
}

func MgoSelectField(collection *mgo.Collection, filter interface{}, skip int, limit int, result interface{}, field ...string) error {
	var fields = bson.M{}
	for _, v := range field {
		fields[v] = 1
	}
	if limit == 0 {
		limit = 1000
	}
	query := collection.Find(filter).Select(fields)
	err := query.Skip(skip).Limit(limit).All(result)
	return err
}

func MgoSelectAllField(collection *mgo.Collection, filter interface{}, result interface{}, skip int, limit int, sort string) error {
	if limit == 0 {
		limit = 1000
	}
	query := collection.Find(filter)
	if sort == "" {
		return query.Skip(skip).Limit(limit).Sort("_id").All(result)
	}
	return query.Skip(skip).Limit(limit).Sort("_id").Sort(sort).All(result)
}

func MgoInsertMulti(collection *mgo.Collection, autoid bool, values ...interface{}) error {
	if autoid == false {
		return collection.Insert(values...)
	}
	insertValues := []interface{}{}
	for _, v := range values {
		bytes, _ := json.Marshal(v)
		tmp := map[string]interface{}{}
		json.Unmarshal(bytes, &tmp)
		delete(tmp, "_id")
		insertValues = append(insertValues, tmp)
	}
	return collection.Insert(insertValues...)
}

func MgoUpdateByIds(collection *mgo.Collection, ids []bson.ObjectId, set bson.M) error {
	if len(ids) == 0 {
		return nil
	}
	_, err := collection.UpdateAll(bson.M{"_id": bson.M{"$in": ids}}, bson.M{"$set": set})
	return err
}

func MgoUpdateAll(collection *mgo.Collection, selector interface{}, set interface{}) error {
	if count, _ := collection.Find(selector).Count(); count == 0 {
		return MgoInsertMulti(collection, true, set)
	}
	_, err := collection.UpdateAll(selector, bson.M{"$set": set})
	return err
}