mgocrud/crud.go

package mgocrud

import (
	"fmt"
	"reflect"
	"strings"
	"time"

	"github.com/davecgh/go-spew/spew"
	mgo "gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
)

// CreateDocument creates a document from specified model
func CreateDocument(db *mgo.Database, m ModelInterface) error {
	m.PrepareInsert()

	c := db.C(GetCollectionName(m))
	err := c.Insert(m)

	return err
}

// ReadDocument gets one document via its id
func ReadDocument(db *mgo.Database, m ModelInterface, selector bson.M) error {
	c := db.C(GetCollectionName(m))

	q := c.FindId(m.GetID())
	if selector != nil {
		q = q.Select(selector)
	}
	err := q.One(m)

	return err
}

// PipelineModifierFunction is a function to modify mongodb query
type PipelineModifierFunction func(pipeline []bson.M) []bson.M

// ReadCollection gets the filtered collection of the model
func ReadCollection(db *mgo.Database, results interface{}, filter bson.M, selector bson.M, offset int, limit int, sort []string, pipelineModifier PipelineModifierFunction) (err error) {
	defer func() {
		if r := recover(); r != nil {
			err = fmt.Errorf("%v", r)
		}
	}()

	// get pointer to model (element of slice in results) to get collection name
	m := reflect.New(reflect.TypeOf(
		reflect.Indirect(reflect.ValueOf(results)).Interface(), // get indirection of slice pointer
	).Elem()).Interface().(ModelInterface) // it must be a ModelInterface here

	c := db.C(GetCollectionName(m))

	if pipelineModifier != nil {
		// search via pipeline

		pipeline := []bson.M{}

		if filter != nil {
			pipeline = append(pipeline, bson.M{
				"$match": filter,
			})
		}

		if len(sort) > 0 {
			sortM := bson.M{}
			for _, s := range sort {
				if strings.HasPrefix(s, "-") {
					s = s[1:]
					sortM[s] = -1
				} else {
					sortM[s] = 1
				}
			}
			spew.Dump(sortM)
			pipeline = append(pipeline, bson.M{
				"$sort": sortM,
			})
		}

		if offset > 0 {
			pipeline = append(pipeline, bson.M{
				"$skip": offset,
			})
		}
		if limit > 0 {
			pipeline = append(pipeline, bson.M{
				"$limit": limit,
			})
		}
		if selector != nil {
			pipeline = append(pipeline, bson.M{
				"$project": selector,
			})
		}

		if pipelineModifier != nil {
			pipeline = pipelineModifier(pipeline)
		}

		q := c.Pipe(pipeline).AllowDiskUse().Iter()
		err = q.All(results)
	} else {
		// search without pipe is faster

		q := c.Find(filter)
		if selector != nil {
			q = q.Select(selector)
		}
		if len(sort) > 0 {
			q = q.Sort(sort...)
		}
		if offset > 0 {
			q = q.Skip(offset)
		}
		if limit > 0 {
			q = q.Limit(limit)
		}

		err = q.All(results)
	}

	return err
}

// ReadCollectionCount gets the count of elements in filtered collection
func ReadCollectionCount(db *mgo.Database, m ModelInterface, filter bson.M) (count int, err error) {
	defer func() {
		if r := recover(); r != nil {
			err = fmt.Errorf("%v", r)
		}
	}()

	c := db.C(GetCollectionName(m))
	return c.Find(filter).Count()
}

// UpdateDocument updates a document from specified model
func UpdateDocument(db *mgo.Database, m ModelInterface, changes bson.M) error {
	m.PrepareUpdate()
	changes["updateTime"] = time.Now()

	c := db.C(GetCollectionName(m))
	err := c.UpdateId(m.GetID(), bson.M{"$set": changes})

	return err
}

// UpsertDocument updates a document from specified model or inserts it, of not found
func UpsertDocument(db *mgo.Database, m ModelInterface, changes bson.M) error {
	m.PrepareUpdate()
	changes["updateTime"] = time.Now()

	c := db.C(GetCollectionName(m))
	_, err := c.Upsert(m, bson.M{"$set": changes})

	return err
}

// DeleteDocument deletes one document via its id
func DeleteDocument(db *mgo.Database, m ModelInterface) error {
	c := db.C(GetCollectionName(m))

	err := c.RemoveId(m.GetID())

	return err
}