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	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer"
10
11	"github.com/globalsign/mgo"
12	"github.com/globalsign/mgo/bson"
13)
14
15// Collection provides a mgo.Collection along with
16// data used for APM Tracing.
17type Collection struct {
18	*mgo.Collection
19	cfg  *mongoConfig
20	tags map[string]string
21}
22
23// Create invokes and traces Collection.Create
24func (c *Collection) Create(info *mgo.CollectionInfo) error {
25	span := newChildSpanFromContext(c.cfg, c.tags)
26	err := c.Collection.Create(info)
27	span.Finish(tracer.WithError(err))
28	return err
29}
30
31// DropCollection invokes and traces Collection.DropCollection
32func (c *Collection) DropCollection() error {
33	span := newChildSpanFromContext(c.cfg, c.tags)
34	err := c.Collection.DropCollection()
35	span.Finish(tracer.WithError(err))
36	return err
37}
38
39// EnsureIndexKey invokes and traces Collection.EnsureIndexKey
40func (c *Collection) EnsureIndexKey(key ...string) error {
41	span := newChildSpanFromContext(c.cfg, c.tags)
42	err := c.Collection.EnsureIndexKey(key...)
43	span.Finish(tracer.WithError(err))
44	return err
45}
46
47// EnsureIndex invokes and traces Collection.EnsureIndex
48func (c *Collection) EnsureIndex(index mgo.Index) error {
49	span := newChildSpanFromContext(c.cfg, c.tags)
50	err := c.Collection.EnsureIndex(index)
51	span.Finish(tracer.WithError(err))
52	return err
53}
54
55// DropIndex invokes and traces Collection.DropIndex
56func (c *Collection) DropIndex(key ...string) error {
57	span := newChildSpanFromContext(c.cfg, c.tags)
58	err := c.Collection.DropIndex(key...)
59	span.Finish(tracer.WithError(err))
60	return err
61}
62
63// DropIndexName invokes and traces Collection.DropIndexName
64func (c *Collection) DropIndexName(name string) error {
65	span := newChildSpanFromContext(c.cfg, c.tags)
66	err := c.Collection.DropIndexName(name)
67	span.Finish(tracer.WithError(err))
68	return err
69}
70
71// Indexes invokes and traces Collection.Indexes
72func (c *Collection) Indexes() (indexes []mgo.Index, err error) {
73	span := newChildSpanFromContext(c.cfg, c.tags)
74	indexes, err = c.Collection.Indexes()
75	span.Finish(tracer.WithError(err))
76	return indexes, err
77}
78
79// Insert invokes and traces Collectin.Insert
80func (c *Collection) Insert(docs ...interface{}) error {
81	span := newChildSpanFromContext(c.cfg, c.tags)
82	err := c.Collection.Insert(docs...)
83	span.Finish(tracer.WithError(err))
84	return err
85}
86
87// Find invokes and traces Collection.Find
88func (c *Collection) Find(query interface{}) *Query {
89	return &Query{
90		Query: c.Collection.Find(query),
91		cfg:   c.cfg,
92		tags:  c.tags,
93	}
94}
95
96// FindId invokes and traces Collection.FindId
97func (c *Collection) FindId(id interface{}) *Query { // nolint
98	return &Query{
99		Query: c.Collection.FindId(id),
100		cfg:   c.cfg,
101		tags:  c.tags,
102	}
103}
104
105// Count invokes and traces Collection.Count
106func (c *Collection) Count() (n int, err error) {
107	span := newChildSpanFromContext(c.cfg, c.tags)
108	n, err = c.Collection.Count()
109	span.Finish(tracer.WithError(err))
110	return n, err
111}
112
113// Bulk creates a trace ready wrapper around Collection.Bulk
114func (c *Collection) Bulk() *Bulk {
115	return &Bulk{
116		Bulk: c.Collection.Bulk(),
117		cfg:  c.cfg,
118	}
119}
120
121// NewIter invokes and traces Collection.Iter
122func (c *Collection) NewIter(session *mgo.Session, firstBatch []bson.Raw, cursorId int64, err error) *Iter { // nolint
123	return &Iter{
124		Iter: c.Collection.NewIter(session, firstBatch, cursorId, err),
125		cfg:  c.cfg,
126	}
127}
128
129// Pipe invokes and traces Collection.Pipe
130func (c *Collection) Pipe(pipeline interface{}) *Pipe {
131	return &Pipe{
132		Pipe: c.Collection.Pipe(pipeline),
133		cfg:  c.cfg,
134		tags: c.tags,
135	}
136}
137
138// Update invokes and traces Collection.Update
139func (c *Collection) Update(selector interface{}, update interface{}) error {
140	span := newChildSpanFromContext(c.cfg, c.tags)
141	err := c.Collection.Update(selector, update)
142	span.Finish(tracer.WithError(err))
143	return err
144}
145
146// UpdateId invokes and traces Collection.UpdateId
147func (c *Collection) UpdateId(id interface{}, update interface{}) error { // nolint
148	span := newChildSpanFromContext(c.cfg, c.tags)
149	err := c.Collection.UpdateId(id, update)
150	span.Finish(tracer.WithError(err))
151	return err
152}
153
154// UpdateAll invokes and traces Collection.UpdateAll
155func (c *Collection) UpdateAll(selector interface{}, update interface{}) (info *mgo.ChangeInfo, err error) {
156	span := newChildSpanFromContext(c.cfg, c.tags)
157	info, err = c.Collection.UpdateAll(selector, update)
158	span.Finish(tracer.WithError(err))
159	return info, err
160}
161
162// Upsert invokes and traces Collection.Upsert
163func (c *Collection) Upsert(selector interface{}, update interface{}) (info *mgo.ChangeInfo, err error) {
164	span := newChildSpanFromContext(c.cfg, c.tags)
165	info, err = c.Collection.Upsert(selector, update)
166	span.Finish(tracer.WithError(err))
167	return info, err
168}
169
170// UpsertId invokes and traces Collection.UpsertId
171func (c *Collection) UpsertId(id interface{}, update interface{}) (info *mgo.ChangeInfo, err error) { // nolint
172	span := newChildSpanFromContext(c.cfg, c.tags)
173	info, err = c.Collection.UpsertId(id, update)
174	span.Finish(tracer.WithError(err))
175	return info, err
176}
177
178// Remove invokes and traces Collection.Remove
179func (c *Collection) Remove(selector interface{}) error {
180	span := newChildSpanFromContext(c.cfg, c.tags)
181	err := c.Collection.Remove(selector)
182	span.Finish(tracer.WithError(err))
183	return err
184}
185
186// RemoveId invokes and traces Collection.RemoveId
187func (c *Collection) RemoveId(id interface{}) error { // nolint
188	span := newChildSpanFromContext(c.cfg, c.tags)
189	err := c.Collection.RemoveId(id)
190	span.Finish(tracer.WithError(err))
191	return err
192}
193
194// RemoveAll invokes and traces Collection.RemoveAll
195func (c *Collection) RemoveAll(selector interface{}) (info *mgo.ChangeInfo, err error) {
196	span := newChildSpanFromContext(c.cfg, c.tags)
197	info, err = c.Collection.RemoveAll(selector)
198	span.Finish(tracer.WithError(err))
199	return info, err
200}
201
202// Repair invokes and traces Collection.Repair
203func (c *Collection) Repair() *Iter {
204	span := newChildSpanFromContext(c.cfg, c.tags)
205	iter := c.Collection.Repair()
206	span.Finish()
207	return &Iter{
208		Iter: iter,
209		cfg:  c.cfg,
210	}
211}
212