mongo.go 6.17 KB
package db

import (
	"context"
	"fmt"
	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"
	"go.mongodb.org/mongo-driver/mongo/readpref"
	"go.mongodb.org/mongo-driver/x/bsonx"
	"pro2d/src/utils"
	"sort"
	"strconv"
	"strings"
	"time"
)

var (
	MongoClient *mongo.Client
	MongoDatabase *mongo.Database
)

//初始化
func Connect(user, password, host string,port int, MaxNum int, timeOut int) error {
	var uri string
	if user!= "" {
		//uri = fmt.Sprintf("mongodb://%s:%s@%s:%d/%s?w=majority", conf.User, conf.Password, conf.Host, conf.Port, conf.DBName)
		uri = fmt.Sprintf("mongodb://%s:%s@%s:%d/?w=majority", user, password, host, port)
	}else {
		//uri = fmt.Sprintf("mongodb://%s:%d/%s?w=majority", conf.Host, conf.Port, conf.DBName)
		uri = fmt.Sprintf("mongodb://%s:%d/?w=majority", host, port)
	}
	// 设置连接超时时间
	ctx, cancel := context.WithTimeout(context.Background(), time.Duration(timeOut))
	defer cancel()
	// 通过传进来的uri连接相关的配置
	o := options.Client().ApplyURI(uri)
	// 设置最大连接数 - 默认是100 ,不设置就是最大 max 64
	o.SetMaxPoolSize(uint64(MaxNum))
	// 发起链接
	var err error
	MongoClient, err = mongo.Connect(ctx, o)
	if err != nil {
		return err
	}
	// 判断服务是不是可用
	if err = MongoClient.Ping(context.Background(), readpref.Primary()); err != nil {
		return err
	}

	//MongoDatabase = MongoClient.Database(dbname)
	return nil
}

func CreateCollection(collection string) error {
	colls, _ := MongoDatabase.ListCollectionNames(context.TODO(), bson.D{})
	pos := sort.SearchStrings(colls, collection)
	if pos != len(colls) {
		if collection == colls[pos] {
			return MongoDatabase.CreateCollection(context.TODO(), collection)
		}
	}
	return MongoDatabase.CreateCollection(context.TODO(), collection)
}

func SetUnique(collection string, key string) (string, error)  {
	return MongoDatabase.Collection(collection).Indexes().CreateOne(
		context.TODO(),
		mongo.IndexModel{
			Keys   : bsonx.Doc{{key, bsonx.Int32(1)}},
			Options: options.Index().SetUnique(true),
		},
	)
}

type MgoColl struct {
	collection *mongo.Collection

	schema *Schema
}

func GetBsonD(key string, value interface{}) interface{} {
	return bson.D{ {key, value}}
}
func GetBsonM(key string, value interface{}) interface{}  {
	return bson.M{key: value}
}

func NewMongoColl(schema *Schema) *MgoColl {
	return &MgoColl{
		collection: MongoDatabase.Collection(schema.GetCollName()),
		schema:     schema,
	}
}

func FindOne(pri interface{}, schema interface{}) error {
	r := MongoDatabase.Collection(utils.GetCollName(schema)).FindOne(context.TODO(), pri)
	return r.Decode(schema)
}

// 查询单个
func (m *MgoColl) FindOneKV(key string, value interface{}) *mongo.SingleResult {
	//collection.
	filter := bson.D{ {key, value}}
	singleResult := m.collection.FindOne(context.TODO(), filter)
	return singleResult
}

//查询集合里有多少数据
func (m *MgoColl) CollectionCount() (string, int64) {
	size, _ := m.collection.EstimatedDocumentCount(context.TODO())
	return m.collection.Name(), size
}

//按选项查询集合 Skip 跳过 Limit 读取数量 sort 1 ,-1 . 1 为最初时间读取 , -1 为最新时间读取
func (m *MgoColl) CollectionDocuments(Skip, Limit int64, sort int) *mongo.Cursor {
	SORT := bson.D{
		{"_id", sort}} //filter := bson.D{ {key,value}}
	filter := bson.D{
	{}}
	findOptions := options.Find().SetSort(SORT).SetLimit(Limit).SetSkip(Skip)
	//findOptions.SetLimit(i)
	temp, _ := m.collection.Find(context.Background(), filter, findOptions)
	return temp
}

//获取集合创建时间和编号
func (m *MgoColl) ParsingId(result string) (time.Time, uint64) {
	temp1 := result[:8]
	timestamp, _ := strconv.ParseInt(temp1, 16, 64)
	dateTime := time.Unix(timestamp, 0) //这是截获情报时间 时间格式 2019-04-24 09:23:39 +0800 CST
	temp2 := result[18:]
	count, _ := strconv.ParseUint(temp2, 16, 64) //截获情报的编号
	return dateTime, count
}

//删除文章和查询文章
func (m *MgoColl) DeleteAndFind(key string, value interface{}) (int64, *mongo.SingleResult) {
	filter := bson.D{
		{key, value}}
	singleResult := m.collection.FindOne(context.TODO(), filter)
	DeleteResult, err := m.collection.DeleteOne(context.TODO(), filter, nil)
	if err != nil {
		fmt.Println("删除时出现错误,你删不掉的~")
	}
	return DeleteResult.DeletedCount, singleResult
}

//删除文章
func (m *MgoColl) Delete(key string, value interface{}) int64 {
	filter := bson.D{ {key, value}}
	count, err := m.collection.DeleteOne(context.TODO(), filter, nil)
	if err != nil {
		fmt.Println(err)
	}
	return count.DeletedCount

}

//删除多个
func (m *MgoColl) DeleteMany(key string, value interface{}) int64 {
	filter := bson.D{ {key, value}}

	count, err := m.collection.DeleteMany(context.TODO(), filter)
	if err != nil {
		fmt.Println(err)
	}
	return count.DeletedCount
}

//索引
func (m *MgoColl) SetUnique(key string){
	m.collection.Indexes().CreateOne(
		context.Background(),
		mongo.IndexModel{
			Keys   : bsonx.Doc{{key, bsonx.Int32(1)}},
			Options: options.Index().SetUnique(true),
		},
	)
}

//更新&保存
func (m *MgoColl) FindOneAndUpdate(filter interface{}, update interface{})*mongo.SingleResult {
	//filter := bson.M{"name": "x", "array.name": "b"}
	//update := bson.M{"array.$[item].detail": "test"}

	res := m.collection.FindOneAndUpdate(context.Background(),
		filter,
		bson.M{"$set": update})
	if res.Err() != nil {
		return nil
	}
	return res
}

func (m *MgoColl) UpdateOne(filter interface{}, update interface{})*mongo.UpdateResult {
	res, err := m.collection.UpdateOne(context.TODO(), filter, bson.D{{"$set", update}})
	if err != nil {
		return nil
	}

	return res
}

func (m *MgoColl) Load() error{
	r := m.collection.FindOne(context.TODO(), m.schema.GetPri())
	err := r.Decode(m.schema.GetSchema())
	if err != nil {
		return err
	}
	return nil
}

func (m *MgoColl) Create() (*mongo.InsertOneResult, error){
	return m.collection.InsertOne(context.TODO(), m.schema.GetSchema())
}

func (m *MgoColl) UpdateProperty(key string, val interface{}) {
	m.UpdateOne(m.schema.GetPri(), bson.M{strings.ToLower(key): val})
}

func (m *MgoColl) UpdateProperties(properties map[string]interface{}) {
	m.UpdateOne(m.schema.GetPri(), properties)
}