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