copter-train/src/db/repo/repo.go

package repo

import (
	"context"
	"copter-train/db"
	"copter-train/log"
	"fmt"

	"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"
)

type RepoSession struct {
	Ctx    context.Context
	Client *db.MongoDB
}

const (
	CollectionCategory = "category"
	CollectionUser     = "users"
)

type Map map[string]interface{}

type PageResult struct {
	List  []map[string]interface{} `json:"list"`
	Total int64                    `json:"total"`
	Page  int64                    `json:"page"`
	Size  int64                    `json:"size"`
}

type PageSearchOptions struct {
	CollectName string
	Page        int64
	Size        int64
	Query       map[string]interface{}
	Project     []string
	Sort        interface{}
}

type DocSearchOptions struct {
	Db          string
	CollectName string
	Query       Map
	Project     []string
	Sort        bson.M
}

func NewDocSearchOptions(filter Map, project []string) *DocSearchOptions {
	return &DocSearchOptions{
		Query:   filter,
		Project: project,
	}
}

func RepoAddDoc(ctx *RepoSession, collectName string, doc interface{}) (string, error) {
	users := ctx.Client.GetCollection(collectName)
	result, err := users.InsertOne(ctx.Ctx, doc)
	if err != nil {
		return "", err
	}
	return result.InsertedID.(primitive.ObjectID).Hex(), nil
}

func RepoDeleteDoc(ctx *RepoSession, collectName string, id string) (interface{}, error) {

	uid, _ := primitive.ObjectIDFromHex(id)

	colls := ctx.Client.GetCollection(collectName)

	return colls.DeleteOne(ctx.Ctx, &bson.M{"_id": uid})
}

func RepoDeleteDocs(ctx *RepoSession, collectName string, query interface{}) (interface{}, error) {
	colls := ctx.Client.GetCollection(collectName)

	return colls.DeleteMany(ctx.Ctx, query)
}

func RepoUpdateSetDocs(ctx *RepoSession, collectName string, filter interface{}, model interface{}) (*mongo.UpdateResult, error) {

	colls := ctx.Client.GetCollection(collectName)

	update := bson.M{"$set": model}

	return colls.UpdateOne(ctx.Ctx, filter, update)
}

func RepoUpdateSetDoc(ctx *RepoSession, collectName string, idstr string, model interface{}) (*mongo.UpdateResult, error) {

	colls := ctx.Client.GetCollection(collectName)

	update := bson.M{"$set": model}

	uid, _ := primitive.ObjectIDFromHex(idstr)

	return colls.UpdateByID(ctx.Ctx, uid, update)
}

func RepoUpsertSetDoc(ctx *RepoSession, collectName string, filter interface{}, model interface{}) (*mongo.UpdateResult, error) {

	coll := ctx.Client.GetCollection(collectName)

	update := bson.M{"$set": model}

	// filter := bson.D{{"type", "Oolong"}}
	// update := bson.D{{"$set", bson.D{}}}
	opts := options.Update().SetUpsert(true)

	return coll.UpdateOne(ctx.Ctx, filter, update, opts)
}

func RepoUpdateSetDocProps(ctx *RepoSession, collectName string, idstr string, update interface{}) (*mongo.UpdateResult, error) {

	colls := ctx.Client.GetCollection(collectName)
	// update := bson.M{"$set": model}
	uid, _ := primitive.ObjectIDFromHex(idstr)
	return colls.UpdateByID(ctx.Ctx, uid, update)
}

func RepoSeachDoc(ctx *RepoSession, param *DocSearchOptions, v interface{}) (bool, error) {

	colls := ctx.Client.GetDbCollection(param.Db, param.CollectName)

	opt := &options.FindOneOptions{}
	if param.Sort != nil {
		opt.SetSort(param.Sort)
	}
	if len(param.Project) > 0 {
		prj := bson.M{}
		for _, v := range param.Project {
			prj[v] = 1
		}
		opt.SetProjection(prj)
	}
	filter := bson.M{}
	if len(param.Query) > 0 {
		for k, v := range param.Query {
			if k == "_id" {
				if uid, ok := v.(string); ok {
					docId, _ := primitive.ObjectIDFromHex(uid)
					filter["_id"] = docId
					continue
				}
			}
			filter[k] = v
		}
	}
	err := colls.FindOne(ctx.Ctx, filter, opt).Decode(v)
	if err == mongo.ErrNoDocuments {
		return false, nil
	}
	if err != nil {
		return false, err
	}
	return true, nil
}

func RepoSeachDocs(ctx *RepoSession, param *DocSearchOptions, out interface{}) error {

	colls := ctx.Client.GetCollection(param.CollectName)
	opt := &options.FindOptions{}
	if len(param.Project) > 0 {
		prj := bson.M{}
		for _, v := range param.Project {
			prj[v] = 1
		}
		opt.SetProjection(prj)
	}
	filter := bson.M{}
	if len(param.Query) > 0 {
		for k, v := range param.Query {
			if k == "_id" {
				if uid, ok := v.(string); ok {
					docId, _ := primitive.ObjectIDFromHex(uid)
					filter["_id"] = docId
					continue
				}
			}
			filter[k] = v
		}
	}
	cur, err := colls.Find(ctx.Ctx, filter, opt)
	if err != nil {
		return err
	}

	defer cur.Close(ctx.Ctx)

	err = cur.All(ctx.Ctx, out)

	return err
}

func RepoSeachDocMap(ctx *RepoSession, param *DocSearchOptions) (bool, map[string]interface{}) {

	ret := map[string]interface{}{}
	ok := true

	colls := ctx.Client.GetCollection(param.CollectName)
	opt := &options.FindOneOptions{}
	if len(param.Project) > 0 {
		prj := bson.M{}
		for _, v := range param.Project {
			prj[v] = 1
		}
		opt.SetProjection(prj)
	}
	filter := bson.M{}
	if len(param.Query) > 0 {
		for k, v := range param.Query {
			if k == "_id" {
				if uid, ok := v.(string); ok {
					docId, _ := primitive.ObjectIDFromHex(uid)
					filter["_id"] = docId
					continue
				}
			}
			filter[k] = v
		}
	}
	ok = true
	err := colls.FindOne(ctx.Ctx, filter, opt).Decode(ret)
	if err == mongo.ErrNoDocuments {
		ok = false
	}
	if err != nil {
		ok = false
	}

	return ok, ret
}

// PageSearch 单表分页查询
func RepoPageSearch(ctx *RepoSession, para *PageSearchOptions) (out *PageResult, err error) {

	colls := ctx.Client.GetCollection(para.CollectName)
	findoptions := &options.FindOptions{}

	if para.Size > 0 {
		findoptions.SetLimit(para.Size)
		findoptions.SetSkip(para.Size * (para.Page - 1))
	}
	if para.Sort != nil {
		findoptions.SetSort(para.Sort)
	}

	if len(para.Project) > 0 {
		prj := bson.M{}
		for _, v := range para.Project {
			prj[v] = 1
		}
		findoptions.SetProjection(prj)
	}

	filter := bson.M{}
	if len(para.Query) > 0 {
		for k, v := range para.Query {

			if value, ok := v.(string); ok {
				if len(value) > 0 {
					filter[k] = v
					continue
				}
			} else if v != nil {
				filter[k] = v
			}
		}
	}

	cur, err := colls.Find(ctx.Ctx, filter, findoptions)

	out = &PageResult{
		List:  []map[string]interface{}{},
		Total: 0,
		Page:  para.Page,
		Size:  para.Size,
	}

	if err != nil {
		return out, err
	}

	defer cur.Close(ctx.Ctx)

	err = cur.All(ctx.Ctx, &out.List)

	out.Total, _ = colls.CountDocuments(ctx.Ctx, filter)
	return
}

// PageSearch 单表分页查询
func RepoCountDoc(ctx *RepoSession, collectionName string, Query Map) (int64, error) {

	colls := ctx.Client.GetCollection(collectionName)
	filter := bson.M{}
	if len(Query) > 0 {
		for k, v := range Query {
			if value, ok := v.(string); ok {
				if len(value) > 0 {
					filter[k] = v
					continue
				}
			} else {
				filter[k] = v
			}
		}
	}
	return colls.CountDocuments(ctx.Ctx, filter)
}

// PageSearch 单表分页查询
func RepoDocsSearch(ctx *RepoSession, para *PageSearchOptions, out interface{}) (err error) {

	colls := ctx.Client.GetCollection(para.CollectName)
	findoptions := &options.FindOptions{}

	if para.Size > 0 {
		findoptions.SetLimit(para.Size)
		findoptions.SetSkip(para.Size * (para.Page - 1))
	}
	if para.Sort != nil {
		findoptions.SetSort(para.Sort)
	}

	if len(para.Project) > 0 {
		prj := bson.M{}
		for _, v := range para.Project {
			prj[v] = 1
		}
		findoptions.SetProjection(prj)
	}

	filter := bson.M{}
	if len(para.Query) > 0 {
		for k, v := range para.Query {

			if value, ok := v.(string); ok {
				if len(value) > 0 {
					filter[k] = v
					continue
				}
			} else {
				filter[k] = v
			}
		}
	}

	cur, err := colls.Find(ctx.Ctx, filter, findoptions)
	if err != nil {
		return err
	}
	defer cur.Close(ctx.Ctx)

	err = cur.All(ctx.Ctx, out)

	return
}

func RepoDocArrayAppend(ctx *RepoSession, collectName string, idstr string, fieldpath string, arrayItem interface{}) (*mongo.UpdateResult, error) {

	colls := ctx.Client.GetCollection(collectName)
	arrayOp := bson.M{}
	arrayOp[fieldpath] = arrayItem

	update := bson.M{"$push": arrayOp}

	uid, _ := primitive.ObjectIDFromHex(idstr)

	return colls.UpdateByID(ctx.Ctx, uid, update)
}

// { _id: 4, "grades.grade": 85 },
// { $set: { "grades.$.std" : 6 } }

type ArrayOneUpdateOption struct {
	Query       Map
	Set         Map
	CollectName string
	Id          string
}

func RepoDocArrayOneUpdate(ctx *RepoSession, options *ArrayOneUpdateOption) (*mongo.UpdateResult, error) {

	colls := ctx.Client.GetCollection(options.CollectName)
	docId, _ := primitive.ObjectIDFromHex(options.Id)

	query := bson.M{"_id": docId}
	if len(options.Query) > 0 {
		for k, v := range options.Query {
			query[k] = v
		}
	}

	setOp := bson.M{}
	for k, v := range options.Set {
		setOp[k] = v
	}
	update := bson.M{"$set": setOp}

	return colls.UpdateOne(ctx.Ctx, query, update)
}

type ArrayOneRemoveOption struct {
	ArrayQuery  Map
	CollectName string
	Id          string
}

// { $pull: { "items" : { id: 23 } } }
func RepoDocArrayOneRemove(ctx *RepoSession, options *ArrayOneRemoveOption) (*mongo.UpdateResult, error) {

	colls := ctx.Client.GetCollection(options.CollectName)
	docId, _ := primitive.ObjectIDFromHex(options.Id)

	query := bson.M{"_id": docId}

	arrayQuery := bson.M{}
	if len(options.ArrayQuery) > 0 {
		for k, v := range options.ArrayQuery {
			arrayQuery[k] = v
		}
	}
	update := bson.M{"$pull": arrayQuery}

	return colls.UpdateOne(ctx.Ctx, query, update)
}

type ArrayOneSearchOption struct {
	ArrayQuery  Map
	CollectName string
	Id          string
	Field       string
}

func RepoDocArraySearch(ctx *RepoSession, options *ArrayOneSearchOption) (ok bool, ret map[string]interface{}) {

	colls := ctx.Client.GetCollection(options.CollectName)
	docId, _ := primitive.ObjectIDFromHex(options.Id)

	match := []bson.E{}
	match = append(match, bson.E{Key: "_id", Value: docId})
	if len(options.ArrayQuery) > 0 {
		for k, v := range options.ArrayQuery {
			match = append(match, bson.E{Key: k, Value: v})
		}
	}
	matchStage := bson.D{
		{Key: "$match", Value: match},
	}

	unwindStage := bson.D{
		{Key: "$unwind", Value: fmt.Sprintf("%s%s", "$", options.Field)},
	}

	curr, err := colls.Aggregate(ctx.Ctx, mongo.Pipeline{matchStage, unwindStage})

	if err != nil {
		ok = false
		return
	}
	defer curr.Close(ctx.Ctx)

	var list []map[string]interface{}

	err = curr.All(ctx.Ctx, list)
	if err != nil {
		ok = false
		return
	}
	ok = true
	ret = list[0]
	return
}

type DocsSearchOptions struct {
	CollectName string
	Query       map[string]interface{}
	Project     []string
	Sort        interface{} //bson.D{ bson.E{"update_time", -1},	bson.E{"goods_id", -1},}
}

func RepoSeachDocsMap(ctx *RepoSession, param *DocsSearchOptions) (ok bool, list []map[string]interface{}) {

	colls := ctx.Client.GetCollection(param.CollectName)
	findoptions := &options.FindOptions{}

	if len(param.Project) > 0 {
		prj := bson.M{}
		for _, v := range param.Project {
			prj[v] = 1
		}
		findoptions.SetProjection(prj)
	}

	if param.Sort != nil {
		findoptions.SetSort(param.Sort)
	}

	filter := bson.M{}
	if len(param.Query) > 0 {
		for k, v := range param.Query {

			if value, ok := v.(string); ok {
				if len(value) > 0 {
					filter[k] = v
					continue
				}
			} else {
				filter[k] = v
			}
		}
	}
	cur, err := colls.Find(ctx.Ctx, filter, findoptions)

	if err != nil {
		ok = false
		return
	}

	defer cur.Close(ctx.Ctx)

	listRes := []map[string]interface{}{}

	err = cur.All(ctx.Ctx, &listRes)

	if err != nil {
		log.Error(err)
		ok = false
		return
	}

	list = listRes
	ok = true
	return
}