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

package mon

import (
	"context"
	"log"
	"strings"

	"github.com/zeromicro/go-zero/core/breaker"
	"github.com/zeromicro/go-zero/core/timex"
	"go.mongodb.org/mongo-driver/mongo"
	mopt "go.mongodb.org/mongo-driver/mongo/options"
)

type (
	// Model is a mongodb store model that represents a collection.
	Model struct {
		Collection
		name string
		cli  *mongo.Client
		brk  breaker.Breaker
		opts []Option
	}

	wrapSession struct {
		mongo.Session
		brk breaker.Breaker
	}
)

// MustNewModel returns a Model, exits on errors.
func MustNewModel(uri, db, collection string, opts ...Option) *Model {
	model, err := NewModel(uri, db, collection, opts...)
	if err != nil {
		log.Fatal(err)
	}

	return model
}

// NewModel returns a Model.
func NewModel(uri, db, collection string, opts ...Option) (*Model, error) {
	cli, err := getClient(uri)
	if err != nil {
		return nil, err
	}

	name := strings.Join([]string{uri, collection}, "/")
	brk := breaker.GetBreaker(uri)
	coll := newCollection(cli.Database(db).Collection(collection), brk)
	return newModel(name, cli, coll, brk, opts...), nil
}

func newModel(name string, cli *mongo.Client, coll Collection, brk breaker.Breaker,
	opts ...Option) *Model {
	return &Model{
		name:       name,
		Collection: coll,
		cli:        cli,
		brk:        brk,
		opts:       opts,
	}
}

// StartSession starts a new session.
func (m *Model) StartSession(opts ...*mopt.SessionOptions) (sess mongo.Session, err error) {
	err = m.brk.DoWithAcceptable(func() error {
		starTime := timex.Now()
		defer func() {
			logDuration(m.name, "StartSession", starTime, err)
		}()

		session, sessionErr := m.cli.StartSession(opts...)
		if sessionErr != nil {
			return sessionErr
		}

		sess = &wrapSession{Session: session, brk: m.brk}

		return nil
	}, acceptable)
	return
}

// Aggregate executes an aggregation pipeline.
func (m *Model) Aggregate(ctx context.Context, v, pipeline interface{}, opts ...*mopt.AggregateOptions) error {
	cur, err := m.Collection.Aggregate(ctx, pipeline, opts...)
	if err != nil {
		return err
	}
	defer cur.Close(ctx)

	return cur.All(ctx, v)
}

// DeleteMany deletes documents that match the filter.
func (m *Model) DeleteMany(ctx context.Context, filter interface{}, opts ...*mopt.DeleteOptions) (int64, error) {
	res, err := m.Collection.DeleteMany(ctx, filter, opts...)
	if err != nil {
		return 0, err
	}

	return res.DeletedCount, nil
}

// DeleteOne deletes the first document that matches the filter.
func (m *Model) DeleteOne(ctx context.Context, filter interface{}, opts ...*mopt.DeleteOptions) (int64, error) {
	res, err := m.Collection.DeleteOne(ctx, filter, opts...)
	if err != nil {
		return 0, err
	}

	return res.DeletedCount, nil
}

// Find finds documents that match the filter.
func (m *Model) Find(ctx context.Context, v, filter interface{}, opts ...*mopt.FindOptions) error {
	cur, err := m.Collection.Find(ctx, filter, opts...)
	if err != nil {
		return err
	}
	defer cur.Close(ctx)

	return cur.All(ctx, v)
}

// FindOne finds the first document that matches the filter.
func (m *Model) FindOne(ctx context.Context, v, filter interface{}, opts ...*mopt.FindOneOptions) error {
	res, err := m.Collection.FindOne(ctx, filter, opts...)
	if err != nil {
		return err
	}

	return res.Decode(v)
}

// FindOneAndDelete finds a single document and deletes it.
func (m *Model) FindOneAndDelete(ctx context.Context, v, filter interface{},
	opts ...*mopt.FindOneAndDeleteOptions) error {
	res, err := m.Collection.FindOneAndDelete(ctx, filter, opts...)
	if err != nil {
		return err
	}

	return res.Decode(v)
}

// FindOneAndReplace finds a single document and replaces it.
func (m *Model) FindOneAndReplace(ctx context.Context, v, filter interface{}, replacement interface{},
	opts ...*mopt.FindOneAndReplaceOptions) error {
	res, err := m.Collection.FindOneAndReplace(ctx, filter, replacement, opts...)
	if err != nil {
		return err
	}

	return res.Decode(v)
}

// FindOneAndUpdate finds a single document and updates it.
func (m *Model) FindOneAndUpdate(ctx context.Context, v, filter interface{}, update interface{},
	opts ...*mopt.FindOneAndUpdateOptions) error {
	res, err := m.Collection.FindOneAndUpdate(ctx, filter, update, opts...)
	if err != nil {
		return err
	}

	return res.Decode(v)
}

func (w *wrapSession) AbortTransaction(ctx context.Context) error {
	ctx, span := startSpan(ctx)
	defer span.End()

	return w.brk.DoWithAcceptable(func() error {
		return w.Session.AbortTransaction(ctx)
	}, acceptable)
}

func (w *wrapSession) CommitTransaction(ctx context.Context) error {
	ctx, span := startSpan(ctx)
	defer span.End()

	return w.brk.DoWithAcceptable(func() error {
		return w.Session.CommitTransaction(ctx)
	}, acceptable)
}

func (w *wrapSession) WithTransaction(ctx context.Context, fn func(sessCtx mongo.SessionContext) (interface{}, error), opts ...*mopt.TransactionOptions) (res interface{}, err error) {
	ctx, span := startSpan(ctx)
	defer span.End()

	err = w.brk.DoWithAcceptable(func() error {
		res, err = w.Session.WithTransaction(ctx, fn, opts...)
		return err
	}, acceptable)

	return
}

func (w *wrapSession) EndSession(ctx context.Context) {
	ctx, span := startSpan(ctx)
	defer span.End()

	_ = w.brk.DoWithAcceptable(func() error {
		w.Session.EndSession(ctx)
		return nil
	}, acceptable)
}