184 lines
4.1 KiB
Go
184 lines
4.1 KiB
Go
|
package mgo
|
||
|
|
||
|
import (
|
||
|
"time"
|
||
|
|
||
|
"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer"
|
||
|
|
||
|
"github.com/globalsign/mgo"
|
||
|
)
|
||
|
|
||
|
// Query is an mgo.Query instance along with the data necessary for tracing.
|
||
|
type Query struct {
|
||
|
*mgo.Query
|
||
|
cfg mongoConfig
|
||
|
}
|
||
|
|
||
|
// Iter invokes and traces Query.Iter
|
||
|
func (q *Query) Iter() *Iter {
|
||
|
span := newChildSpanFromContext(q.cfg)
|
||
|
iter := q.Query.Iter()
|
||
|
span.Finish()
|
||
|
return &Iter{
|
||
|
Iter: iter,
|
||
|
cfg: q.cfg,
|
||
|
}
|
||
|
}
|
||
|
|
||
|
// All invokes and traces Query.All
|
||
|
func (q *Query) All(result interface{}) error {
|
||
|
span := newChildSpanFromContext(q.cfg)
|
||
|
err := q.All(result)
|
||
|
span.Finish(tracer.WithError(err))
|
||
|
return err
|
||
|
}
|
||
|
|
||
|
// Apply invokes and traces Query.Apply
|
||
|
func (q *Query) Apply(change mgo.Change, result interface{}) (info *mgo.ChangeInfo, err error) {
|
||
|
span := newChildSpanFromContext(q.cfg)
|
||
|
info, err = q.Apply(change, result)
|
||
|
span.Finish(tracer.WithError(err))
|
||
|
return info, err
|
||
|
}
|
||
|
|
||
|
// Batch invokes and traces Query.Batch
|
||
|
func (q *Query) Batch(n int) *Query {
|
||
|
return &Query{
|
||
|
Query: q.Query.Batch(n),
|
||
|
cfg: q.cfg,
|
||
|
}
|
||
|
}
|
||
|
|
||
|
// Collation invokes and traces Query.Collation
|
||
|
func (q *Query) Collation(collation *mgo.Collation) *Query {
|
||
|
return &Query{
|
||
|
Query: q.Query.Collation(collation),
|
||
|
cfg: q.cfg,
|
||
|
}
|
||
|
}
|
||
|
|
||
|
// Comment invokes and traces Query.Comment
|
||
|
func (q *Query) Comment(comment string) *Query {
|
||
|
return &Query{
|
||
|
Query: q.Query.Comment(comment),
|
||
|
cfg: q.cfg,
|
||
|
}
|
||
|
}
|
||
|
|
||
|
// Count invokes and traces Query.Count
|
||
|
func (q *Query) Count() (n int, err error) {
|
||
|
span := newChildSpanFromContext(q.cfg)
|
||
|
n, err = q.Count()
|
||
|
span.Finish(tracer.WithError(err))
|
||
|
return n, err
|
||
|
}
|
||
|
|
||
|
// Distinct invokes and traces Query.Distinct
|
||
|
func (q *Query) Distinct(key string, result interface{}) error {
|
||
|
span := newChildSpanFromContext(q.cfg)
|
||
|
err := q.Distinct(key, result)
|
||
|
span.Finish(tracer.WithError(err))
|
||
|
return err
|
||
|
}
|
||
|
|
||
|
// Explain invokes and traces Query.Explain
|
||
|
func (q *Query) Explain(result interface{}) error {
|
||
|
span := newChildSpanFromContext(q.cfg)
|
||
|
err := q.Explain(result)
|
||
|
span.Finish(tracer.WithError(err))
|
||
|
return err
|
||
|
}
|
||
|
|
||
|
// For invokes and traces Query.For
|
||
|
func (q *Query) For(result interface{}, f func() error) error {
|
||
|
span := newChildSpanFromContext(q.cfg)
|
||
|
err := q.For(result, f)
|
||
|
span.Finish(tracer.WithError(err))
|
||
|
return err
|
||
|
}
|
||
|
|
||
|
// MapReduce invokes and traces Query.MapReduce
|
||
|
func (q *Query) MapReduce(job *mgo.MapReduce, result interface{}) (info *mgo.MapReduceInfo, err error) {
|
||
|
span := newChildSpanFromContext(q.cfg)
|
||
|
info, err = q.MapReduce(job, result)
|
||
|
span.Finish(tracer.WithError(err))
|
||
|
return info, err
|
||
|
}
|
||
|
|
||
|
// One invokes and traces Query.One
|
||
|
func (q *Query) One(result interface{}) error {
|
||
|
span := newChildSpanFromContext(q.cfg)
|
||
|
err := q.One(result)
|
||
|
span.Finish(tracer.WithError(err))
|
||
|
return err
|
||
|
}
|
||
|
|
||
|
// Prefetch invokes Query.Prefetch and configures the
|
||
|
// returned *Query for tracing.
|
||
|
func (q *Query) Prefetch(p float64) *Query {
|
||
|
return &Query{
|
||
|
Query: q.Query.Prefetch(p),
|
||
|
cfg: q.cfg,
|
||
|
}
|
||
|
}
|
||
|
|
||
|
// Select invokes Query.Select and configures the
|
||
|
// returned *Query for tracing.
|
||
|
func (q *Query) Select(selector interface{}) *Query {
|
||
|
return &Query{
|
||
|
Query: q.Query.Select(selector),
|
||
|
cfg: q.cfg,
|
||
|
}
|
||
|
}
|
||
|
|
||
|
// SetMaxScan invokes and traces Query.SetMaxScan
|
||
|
func (q *Query) SetMaxScan(n int) *Query {
|
||
|
return &Query{
|
||
|
Query: q.Query.SetMaxScan(n),
|
||
|
cfg: q.cfg,
|
||
|
}
|
||
|
}
|
||
|
|
||
|
// SetMaxTime invokes and traces Query.SetMaxTime
|
||
|
func (q *Query) SetMaxTime(d time.Duration) *Query {
|
||
|
return &Query{
|
||
|
Query: q.Query.SetMaxTime(d),
|
||
|
cfg: q.cfg,
|
||
|
}
|
||
|
}
|
||
|
|
||
|
// Skip invokes and traces Query.Skip
|
||
|
func (q *Query) Skip(n int) *Query {
|
||
|
return &Query{
|
||
|
Query: q.Query.Skip(n),
|
||
|
cfg: q.cfg,
|
||
|
}
|
||
|
}
|
||
|
|
||
|
// Snapshot invokes and traces Query.Snapshot
|
||
|
func (q *Query) Snapshot() *Query {
|
||
|
return &Query{
|
||
|
Query: q.Query.Snapshot(),
|
||
|
cfg: q.cfg,
|
||
|
}
|
||
|
}
|
||
|
|
||
|
// Sort invokes and traces Query.Sort
|
||
|
func (q *Query) Sort(fields ...string) *Query {
|
||
|
return &Query{
|
||
|
Query: q.Query.Sort(fields...),
|
||
|
cfg: q.cfg,
|
||
|
}
|
||
|
}
|
||
|
|
||
|
// Tail invokes and traces Query.Tail
|
||
|
func (q *Query) Tail(timeout time.Duration) *Iter {
|
||
|
span := newChildSpanFromContext(q.cfg)
|
||
|
iter := q.Query.Tail(timeout)
|
||
|
span.Finish()
|
||
|
return &Iter{
|
||
|
Iter: iter,
|
||
|
cfg: q.cfg,
|
||
|
}
|
||
|
}
|