go-zero/core/stores/mongo/model.go

package mongo

import (
	"log"
	"time"

	"github.com/globalsign/mgo"
)

type (
	options struct {
		timeout time.Duration
	}

	Option func(opts *options)

	Model struct {
		session    *concurrentSession
		db         *mgo.Database
		collection string
		opts       []Option
	}
)

func MustNewModel(url, database, collection string, opts ...Option) *Model {
	model, err := NewModel(url, database, collection, opts...)
	if err != nil {
		log.Fatal(err)
	}

	return model
}

func NewModel(url, database, collection string, opts ...Option) (*Model, error) {
	session, err := getConcurrentSession(url)
	if err != nil {
		return nil, err
	}

	return &Model{
		session:    session,
		db:         session.DB(database),
		collection: collection,
		opts:       opts,
	}, nil
}

func (mm *Model) Find(query interface{}) (Query, error) {
	return mm.query(func(c Collection) Query {
		return c.Find(query)
	})
}

func (mm *Model) FindId(id interface{}) (Query, error) {
	return mm.query(func(c Collection) Query {
		return c.FindId(id)
	})
}

func (mm *Model) GetCollection(session *mgo.Session) Collection {
	return newCollection(mm.db.C(mm.collection).With(session))
}

func (mm *Model) Insert(docs ...interface{}) error {
	return mm.execute(func(c Collection) error {
		return c.Insert(docs...)
	})
}

func (mm *Model) Pipe(pipeline interface{}) (Pipe, error) {
	return mm.pipe(func(c Collection) Pipe {
		return c.Pipe(pipeline)
	})
}

func (mm *Model) PutSession(session *mgo.Session) {
	mm.session.putSession(session)
}

func (mm *Model) Remove(selector interface{}) error {
	return mm.execute(func(c Collection) error {
		return c.Remove(selector)
	})
}

func (mm *Model) RemoveAll(selector interface{}) (*mgo.ChangeInfo, error) {
	return mm.change(func(c Collection) (*mgo.ChangeInfo, error) {
		return c.RemoveAll(selector)
	})
}

func (mm *Model) RemoveId(id interface{}) error {
	return mm.execute(func(c Collection) error {
		return c.RemoveId(id)
	})
}

func (mm *Model) TakeSession() (*mgo.Session, error) {
	return mm.session.takeSession(mm.opts...)
}

func (mm *Model) Update(selector, update interface{}) error {
	return mm.execute(func(c Collection) error {
		return c.Update(selector, update)
	})
}

func (mm *Model) UpdateId(id, update interface{}) error {
	return mm.execute(func(c Collection) error {
		return c.UpdateId(id, update)
	})
}

func (mm *Model) Upsert(selector, update interface{}) (*mgo.ChangeInfo, error) {
	return mm.change(func(c Collection) (*mgo.ChangeInfo, error) {
		return c.Upsert(selector, update)
	})
}

func (mm *Model) change(fn func(c Collection) (*mgo.ChangeInfo, error)) (*mgo.ChangeInfo, error) {
	session, err := mm.TakeSession()
	if err != nil {
		return nil, err
	}
	defer mm.PutSession(session)

	return fn(mm.GetCollection(session))
}

func (mm *Model) execute(fn func(c Collection) error) error {
	session, err := mm.TakeSession()
	if err != nil {
		return err
	}
	defer mm.PutSession(session)

	return fn(mm.GetCollection(session))
}

func (mm *Model) pipe(fn func(c Collection) Pipe) (Pipe, error) {
	session, err := mm.TakeSession()
	if err != nil {
		return nil, err
	}
	defer mm.PutSession(session)

	return fn(mm.GetCollection(session)), nil
}

func (mm *Model) query(fn func(c Collection) Query) (Query, error) {
	session, err := mm.TakeSession()
	if err != nil {
		return nil, err
	}
	defer mm.PutSession(session)

	return fn(mm.GetCollection(session)), nil
}

func WithTimeout(timeout time.Duration) Option {
	return func(opts *options) {
		opts.timeout = timeout
	}
}