You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
208 lines
5.2 KiB
Go
208 lines
5.2 KiB
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)
|
|
}
|