1// Unless explicitly stated otherwise all files in this repository are licensed
2// under the Apache License Version 2.0.
3// This product includes software developed at Datadog (https://www.datadoghq.com/).
4// Copyright 2016 Datadog, Inc.
5
6package mgo
7
8import (
9	"time"
10
11	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer"
12
13	"github.com/globalsign/mgo"
14)
15
16// Query is an mgo.Query instance along with the data necessary for tracing.
17type Query struct {
18	*mgo.Query
19	cfg  *mongoConfig
20	tags map[string]string
21}
22
23// Iter invokes and traces Query.Iter
24func (q *Query) Iter() *Iter {
25	span := newChildSpanFromContext(q.cfg, q.tags)
26	iter := q.Query.Iter()
27	span.Finish()
28	return &Iter{
29		Iter: iter,
30		cfg:  q.cfg,
31	}
32}
33
34// All invokes and traces Query.All
35func (q *Query) All(result interface{}) error {
36	span := newChildSpanFromContext(q.cfg, q.tags)
37	err := q.Query.All(result)
38	span.Finish(tracer.WithError(err))
39	return err
40}
41
42// Apply invokes and traces Query.Apply
43func (q *Query) Apply(change mgo.Change, result interface{}) (info *mgo.ChangeInfo, err error) {
44	span := newChildSpanFromContext(q.cfg, q.tags)
45	info, err = q.Query.Apply(change, result)
46	span.Finish(tracer.WithError(err))
47	return info, err
48}
49
50// Count invokes and traces Query.Count
51func (q *Query) Count() (n int, err error) {
52	span := newChildSpanFromContext(q.cfg, q.tags)
53	n, err = q.Query.Count()
54	span.Finish(tracer.WithError(err))
55	return n, err
56}
57
58// Distinct invokes and traces Query.Distinct
59func (q *Query) Distinct(key string, result interface{}) error {
60	span := newChildSpanFromContext(q.cfg, q.tags)
61	err := q.Query.Distinct(key, result)
62	span.Finish(tracer.WithError(err))
63	return err
64}
65
66// Explain invokes and traces Query.Explain
67func (q *Query) Explain(result interface{}) error {
68	span := newChildSpanFromContext(q.cfg, q.tags)
69	err := q.Query.Explain(result)
70	span.Finish(tracer.WithError(err))
71	return err
72}
73
74// For invokes and traces Query.For
75func (q *Query) For(result interface{}, f func() error) error {
76	span := newChildSpanFromContext(q.cfg, q.tags)
77	err := q.Query.For(result, f)
78	span.Finish(tracer.WithError(err))
79	return err
80}
81
82// MapReduce invokes and traces Query.MapReduce
83func (q *Query) MapReduce(job *mgo.MapReduce, result interface{}) (info *mgo.MapReduceInfo, err error) {
84	span := newChildSpanFromContext(q.cfg, q.tags)
85	info, err = q.Query.MapReduce(job, result)
86	span.Finish(tracer.WithError(err))
87	return info, err
88}
89
90// One invokes and traces Query.One
91func (q *Query) One(result interface{}) error {
92	span := newChildSpanFromContext(q.cfg, q.tags)
93	err := q.Query.One(result)
94	span.Finish(tracer.WithError(err))
95	return err
96}
97
98// Tail invokes and traces Query.Tail
99func (q *Query) Tail(timeout time.Duration) *Iter {
100	span := newChildSpanFromContext(q.cfg, q.tags)
101	iter := q.Query.Tail(timeout)
102	span.Finish()
103	return &Iter{
104		Iter: iter,
105		cfg:  q.cfg,
106	}
107}
108