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 6// Package mgo provides functions and types which allow tracing of the MGO MongoDB client (https://github.com/globalsign/mgo) 7package mgo // import "gopkg.in/DataDog/dd-trace-go.v1/contrib/globalsign/mgo" 8 9import ( 10 "math" 11 "strings" 12 13 "gopkg.in/DataDog/dd-trace-go.v1/ddtrace" 14 "gopkg.in/DataDog/dd-trace-go.v1/ddtrace/ext" 15 "gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer" 16 "gopkg.in/DataDog/dd-trace-go.v1/internal/log" 17 18 "github.com/globalsign/mgo" 19) 20 21// Dial opens a connection to a MongoDB server and configures it 22// for tracing. 23func Dial(url string, opts ...DialOption) (*Session, error) { 24 session, err := mgo.Dial(url) 25 info, _ := session.BuildInfo() 26 s := &Session{ 27 Session: session, 28 cfg: newConfig(), 29 tags: map[string]string{ 30 "hosts": strings.Join(session.LiveServers(), ", "), 31 "mgo_version": info.Version, 32 }, 33 } 34 for _, fn := range opts { 35 fn(s.cfg) 36 } 37 log.Debug("contrib/globalsign/mgo: Dialing: %s, %#v", url, s.cfg) 38 return s, err 39} 40 41// Session is an mgo.Session instance that will be traced. 42type Session struct { 43 *mgo.Session 44 cfg *mongoConfig 45 tags map[string]string 46} 47 48func newChildSpanFromContext(cfg *mongoConfig, tags map[string]string) ddtrace.Span { 49 opts := []ddtrace.StartSpanOption{ 50 tracer.SpanType(ext.SpanTypeMongoDB), 51 tracer.ServiceName(cfg.serviceName), 52 tracer.ResourceName("mongodb.query"), 53 } 54 if !math.IsNaN(cfg.analyticsRate) { 55 opts = append(opts, tracer.Tag(ext.EventSampleRate, cfg.analyticsRate)) 56 } 57 span, _ := tracer.StartSpanFromContext(cfg.ctx, "mongodb.query", opts...) 58 for key, value := range tags { 59 span.SetTag(key, value) 60 } 61 return span 62} 63 64// Run invokes and traces Session.Run 65func (s *Session) Run(cmd interface{}, result interface{}) (err error) { 66 span := newChildSpanFromContext(s.cfg, s.tags) 67 err = s.Session.Run(cmd, result) 68 span.Finish(tracer.WithError(err)) 69 return 70} 71 72// Database is an mgo.Database along with the data necessary for tracing. 73type Database struct { 74 *mgo.Database 75 cfg *mongoConfig 76 tags map[string]string 77} 78 79// DB returns a new database for this Session. 80func (s *Session) DB(name string) *Database { 81 tags := make(map[string]string, len(s.tags)+1) 82 for k, v := range s.tags { 83 tags[k] = v 84 } 85 tags["name"] = name 86 return &Database{ 87 Database: s.Session.DB(name), 88 cfg: s.cfg, 89 tags: tags, 90 } 91} 92 93// C returns a new Collection from this Database. 94func (db *Database) C(name string) *Collection { 95 return &Collection{ 96 Collection: db.Database.C(name), 97 cfg: db.cfg, 98 tags: db.tags, 99 } 100} 101 102// Iter is an mgo.Iter instance that will be traced. 103type Iter struct { 104 *mgo.Iter 105 cfg *mongoConfig 106 tags map[string]string 107} 108 109// Next invokes and traces Iter.Next 110func (iter *Iter) Next(result interface{}) bool { 111 span := newChildSpanFromContext(iter.cfg, iter.tags) 112 r := iter.Iter.Next(result) 113 span.Finish() 114 return r 115} 116 117// For invokes and traces Iter.For 118func (iter *Iter) For(result interface{}, f func() error) (err error) { 119 span := newChildSpanFromContext(iter.cfg, iter.tags) 120 err = iter.Iter.For(result, f) 121 span.Finish(tracer.WithError(err)) 122 return err 123} 124 125// All invokes and traces Iter.All 126func (iter *Iter) All(result interface{}) (err error) { 127 span := newChildSpanFromContext(iter.cfg, iter.tags) 128 err = iter.Iter.All(result) 129 span.Finish(tracer.WithError(err)) 130 return err 131} 132 133// Close invokes and traces Iter.Close 134func (iter *Iter) Close() (err error) { 135 span := newChildSpanFromContext(iter.cfg, iter.tags) 136 err = iter.Iter.Close() 137 span.Finish(tracer.WithError(err)) 138 return err 139} 140 141// Bulk is an mgo.Bulk instance that will be traced. 142type Bulk struct { 143 *mgo.Bulk 144 tags map[string]string 145 cfg *mongoConfig 146} 147 148// Run invokes and traces Bulk.Run 149func (b *Bulk) Run() (result *mgo.BulkResult, err error) { 150 span := newChildSpanFromContext(b.cfg, b.tags) 151 result, err = b.Bulk.Run() 152 span.Finish(tracer.WithError(err)) 153 154 return result, err 155} 156