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