This repository has been archived by the owner on May 27, 2021. It is now read-only.
forked from DataDog/dd-trace-go
-
Notifications
You must be signed in to change notification settings - Fork 0
/
mgo.go
144 lines (118 loc) · 3.39 KB
/
mgo.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
// Package mgo provides functions and types which allow tracing of the MGO MongoDB client (https://github.com/globalsign/mgo)
package mgo // import "gopkg.in/DataDog/dd-trace-go.v1/contrib/globalsign/mgo"
import (
"strings"
"gopkg.in/DataDog/dd-trace-go.v1/ddtrace"
"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/ext"
"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer"
"github.com/globalsign/mgo"
)
// Dial opens a connection to a MongoDB server and configures it
// for tracing.
func Dial(url string, opts ...DialOption) (*Session, error) {
session, err := mgo.Dial(url)
s := &Session{Session: session}
defaults(&s.cfg)
for _, fn := range opts {
fn(&s.cfg)
}
// Record metadata so that it can be added to recorded traces
s.cfg.tags["hosts"] = strings.Join(session.LiveServers(), ", ")
info, _ := session.BuildInfo()
s.cfg.tags["mgo_version"] = info.Version
return s, err
}
// Session is an mgo.Session instance that will be traced.
type Session struct {
*mgo.Session
cfg mongoConfig
}
func newChildSpanFromContext(config mongoConfig) ddtrace.Span {
span, _ := tracer.StartSpanFromContext(
config.ctx,
"mongodb.query",
tracer.SpanType(ext.SpanTypeMongoDB),
tracer.ServiceName(config.serviceName),
tracer.ResourceName("mongodb.query"))
for key, value := range config.tags {
span.SetTag(key, value)
}
return span
}
// Run invokes and traces Session.Run
func (s *Session) Run(cmd interface{}, result interface{}) (err error) {
span := newChildSpanFromContext(s.cfg)
err = s.Session.Run(cmd, result)
span.Finish(tracer.WithError(err))
return
}
// Database is an mgo.Database along with the data necessary for tracing.
type Database struct {
*mgo.Database
cfg mongoConfig
}
// DB returns a new database for this Session.
func (s *Session) DB(name string) *Database {
dbCfg := mongoConfig{
ctx: s.cfg.ctx,
serviceName: s.cfg.serviceName,
tags: s.cfg.tags,
}
dbCfg.tags["database"] = name
return &Database{
Database: s.Session.DB(name),
cfg: dbCfg,
}
}
// C returns a new Collection from this Database.
func (db *Database) C(name string) *Collection {
return &Collection{
Collection: db.Database.C(name),
cfg: db.cfg,
}
}
// Iter is an mgo.Iter instance that will be traced.
type Iter struct {
*mgo.Iter
cfg mongoConfig
}
// Next invokes and traces Iter.Next
func (iter *Iter) Next(result interface{}) bool {
span := newChildSpanFromContext(iter.cfg)
r := iter.Iter.Next(result)
span.Finish()
return r
}
// For invokes and traces Iter.For
func (iter *Iter) For(result interface{}, f func() error) (err error) {
span := newChildSpanFromContext(iter.cfg)
err = iter.Iter.For(result, f)
span.Finish(tracer.WithError(err))
return err
}
// All invokes and traces Iter.All
func (iter *Iter) All(result interface{}) (err error) {
span := newChildSpanFromContext(iter.cfg)
err = iter.Iter.All(result)
span.Finish(tracer.WithError(err))
return err
}
// Close invokes and traces Iter.Close
func (iter *Iter) Close() (err error) {
span := newChildSpanFromContext(iter.cfg)
err = iter.Iter.Close()
span.Finish(tracer.WithError(err))
return err
}
// Bulk is an mgo.Bulk instance that will be traced.
type Bulk struct {
*mgo.Bulk
cfg mongoConfig
}
// Run invokes and traces Bulk.Run
func (b *Bulk) Run() (result *mgo.BulkResult, err error) {
span := newChildSpanFromContext(b.cfg)
result, err = b.Bulk.Run()
span.Finish(tracer.WithError(err))
return result, err
}