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 buntdb
7
8import (
9	"context"
10	"math"
11	"time"
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/tidwall/buntdb"
19)
20
21// A DB wraps a buntdb.DB, automatically tracing any transactions.
22type DB struct {
23	*buntdb.DB
24	opts []Option
25	ctx  context.Context
26}
27
28// Open calls buntdb.Open and wraps the result.
29func Open(path string, opts ...Option) (*DB, error) {
30	db, err := buntdb.Open(path)
31	if err != nil {
32		return nil, err
33	}
34	return WrapDB(db, opts...), nil
35}
36
37// WrapDB wraps a buntdb.DB so it can be traced.
38func WrapDB(db *buntdb.DB, opts ...Option) *DB {
39	return &DB{
40		DB:   db,
41		opts: opts,
42		ctx:  context.Background(),
43	}
44}
45
46// Begin calls the underlying DB.Begin and traces the transaction.
47func (db *DB) Begin(writable bool) (*Tx, error) {
48	tx, err := db.DB.Begin(writable)
49	if err != nil {
50		return nil, err
51	}
52	return WrapTx(tx, db.opts...), nil
53}
54
55// Update calls the underlying DB.Update and traces the transaction.
56func (db *DB) Update(fn func(tx *Tx) error) error {
57	return db.DB.Update(func(tx *buntdb.Tx) error {
58		return fn(WrapTx(tx, db.opts...))
59	})
60}
61
62// View calls the underlying DB.View and traces the transaction.
63func (db *DB) View(fn func(tx *Tx) error) error {
64	return db.DB.View(func(tx *buntdb.Tx) error {
65		return fn(WrapTx(tx, db.opts...))
66	})
67}
68
69// WithContext sets the context for the DB.
70func (db *DB) WithContext(ctx context.Context) *DB {
71	newdb := *db
72	newdb.opts = append(newdb.opts[:len(newdb.opts):len(newdb.opts)], WithContext(ctx))
73	return &newdb
74}
75
76// A Tx wraps a buntdb.Tx, automatically tracing any queries.
77type Tx struct {
78	*buntdb.Tx
79	cfg *config
80}
81
82// WrapTx wraps a buntdb.Tx so it can be traced.
83func WrapTx(tx *buntdb.Tx, opts ...Option) *Tx {
84	cfg := new(config)
85	defaults(cfg)
86	for _, opt := range opts {
87		opt(cfg)
88	}
89	log.Debug("contrib/tidwall/buntdb: Wrapping Transaction: %#v", cfg)
90	return &Tx{
91		Tx:  tx,
92		cfg: cfg,
93	}
94}
95
96func (tx *Tx) startSpan(name string) ddtrace.Span {
97	opts := []ddtrace.StartSpanOption{
98		tracer.SpanType(ext.AppTypeDB),
99		tracer.ServiceName(tx.cfg.serviceName),
100		tracer.ResourceName(name),
101	}
102	if !math.IsNaN(tx.cfg.analyticsRate) {
103		opts = append(opts, tracer.Tag(ext.EventSampleRate, tx.cfg.analyticsRate))
104	}
105	span, _ := tracer.StartSpanFromContext(tx.cfg.ctx, "buntdb.query", opts...)
106	return span
107}
108
109// WithContext sets the context for the Tx.
110func (tx *Tx) WithContext(ctx context.Context) *Tx {
111	newcfg := *tx.cfg
112	newcfg.ctx = ctx
113	return &Tx{
114		Tx:  tx.Tx,
115		cfg: &newcfg,
116	}
117}
118
119// Ascend calls the underlying Tx.Ascend and traces the query.
120func (tx *Tx) Ascend(index string, iterator func(key, value string) bool) error {
121	span := tx.startSpan("Ascend")
122	err := tx.Tx.Ascend(index, iterator)
123	span.Finish(tracer.WithError(err))
124	return err
125}
126
127// AscendEqual calls the underlying Tx.AscendEqual and traces the query.
128func (tx *Tx) AscendEqual(index, pivot string, iterator func(key, value string) bool) error {
129	span := tx.startSpan("AscendEqual")
130	err := tx.Tx.AscendEqual(index, pivot, iterator)
131	span.Finish(tracer.WithError(err))
132	return err
133}
134
135// AscendGreaterOrEqual calls the underlying Tx.AscendGreaterOrEqual and traces the query.
136func (tx *Tx) AscendGreaterOrEqual(index, pivot string, iterator func(key, value string) bool) error {
137	span := tx.startSpan("AscendGreaterOrEqual")
138	err := tx.Tx.AscendGreaterOrEqual(index, pivot, iterator)
139	span.Finish(tracer.WithError(err))
140	return err
141}
142
143// AscendKeys calls the underlying Tx.AscendKeys and traces the query.
144func (tx *Tx) AscendKeys(pattern string, iterator func(key, value string) bool) error {
145	span := tx.startSpan("AscendKeys")
146	err := tx.Tx.AscendKeys(pattern, iterator)
147	span.Finish(tracer.WithError(err))
148	return err
149}
150
151// AscendLessThan calls the underlying Tx.AscendLessThan and traces the query.
152func (tx *Tx) AscendLessThan(index, pivot string, iterator func(key, value string) bool) error {
153	span := tx.startSpan("AscendLessThan")
154	err := tx.Tx.AscendLessThan(index, pivot, iterator)
155	span.Finish(tracer.WithError(err))
156	return err
157}
158
159// AscendRange calls the underlying Tx.AscendRange and traces the query.
160func (tx *Tx) AscendRange(index, greaterOrEqual, lessThan string, iterator func(key, value string) bool) error {
161	span := tx.startSpan("AscendRange")
162	err := tx.Tx.AscendRange(index, greaterOrEqual, lessThan, iterator)
163	span.Finish(tracer.WithError(err))
164	return err
165}
166
167// CreateIndex calls the underlying Tx.CreateIndex and traces the query.
168func (tx *Tx) CreateIndex(name, pattern string, less ...func(a, b string) bool) error {
169	span := tx.startSpan("CreateIndex")
170	err := tx.Tx.CreateIndex(name, pattern, less...)
171	span.Finish(tracer.WithError(err))
172	return err
173}
174
175// CreateIndexOptions calls the underlying Tx.CreateIndexOptions and traces the query.
176func (tx *Tx) CreateIndexOptions(name, pattern string, opts *buntdb.IndexOptions, less ...func(a, b string) bool) error {
177	span := tx.startSpan("CreateIndexOptions")
178	err := tx.Tx.CreateIndexOptions(name, pattern, opts, less...)
179	span.Finish(tracer.WithError(err))
180	return err
181}
182
183// CreateSpatialIndex calls the underlying Tx.CreateSpatialIndex and traces the query.
184func (tx *Tx) CreateSpatialIndex(name, pattern string, rect func(item string) (min, max []float64)) error {
185	span := tx.startSpan("CreateSpatialIndex")
186	err := tx.Tx.CreateSpatialIndex(name, pattern, rect)
187	span.Finish(tracer.WithError(err))
188	return err
189}
190
191// CreateSpatialIndexOptions calls the underlying Tx.CreateSpatialIndexOptions and traces the query.
192func (tx *Tx) CreateSpatialIndexOptions(name, pattern string, opts *buntdb.IndexOptions, rect func(item string) (min, max []float64)) error {
193	span := tx.startSpan("CreateSpatialIndexOptions")
194	err := tx.Tx.CreateSpatialIndexOptions(name, pattern, opts, rect)
195	span.Finish(tracer.WithError(err))
196	return err
197}
198
199// Delete calls the underlying Tx.Delete and traces the query.
200func (tx *Tx) Delete(key string) (val string, err error) {
201	span := tx.startSpan("Delete")
202	val, err = tx.Tx.Delete(key)
203	span.Finish(tracer.WithError(err))
204	return val, err
205}
206
207// DeleteAll calls the underlying Tx.DeleteAll and traces the query.
208func (tx *Tx) DeleteAll() error {
209	span := tx.startSpan("DeleteAll")
210	err := tx.Tx.DeleteAll()
211	span.Finish(tracer.WithError(err))
212	return err
213}
214
215// Descend calls the underlying Tx.Descend and traces the query.
216func (tx *Tx) Descend(index string, iterator func(key, value string) bool) error {
217	span := tx.startSpan("Descend")
218	err := tx.Tx.Descend(index, iterator)
219	span.Finish(tracer.WithError(err))
220	return err
221}
222
223// DescendEqual calls the underlying Tx.DescendEqual and traces the query.
224func (tx *Tx) DescendEqual(index, pivot string, iterator func(key, value string) bool) error {
225	span := tx.startSpan("DescendEqual")
226	err := tx.Tx.DescendEqual(index, pivot, iterator)
227	span.Finish(tracer.WithError(err))
228	return err
229}
230
231// DescendGreaterThan calls the underlying Tx.DescendGreaterThan and traces the query.
232func (tx *Tx) DescendGreaterThan(index, pivot string, iterator func(key, value string) bool) error {
233	span := tx.startSpan("DescendGreaterThan")
234	err := tx.Tx.DescendGreaterThan(index, pivot, iterator)
235	span.Finish(tracer.WithError(err))
236	return err
237}
238
239// DescendKeys calls the underlying Tx.DescendKeys and traces the query.
240func (tx *Tx) DescendKeys(pattern string, iterator func(key, value string) bool) error {
241	span := tx.startSpan("DescendKeys")
242	err := tx.Tx.DescendKeys(pattern, iterator)
243	span.Finish(tracer.WithError(err))
244	return err
245}
246
247// DescendLessOrEqual calls the underlying Tx.DescendLessOrEqual and traces the query.
248func (tx *Tx) DescendLessOrEqual(index, pivot string, iterator func(key, value string) bool) error {
249	span := tx.startSpan("DescendLessOrEqual")
250	err := tx.Tx.DescendLessOrEqual(index, pivot, iterator)
251	span.Finish(tracer.WithError(err))
252	return err
253}
254
255// DescendRange calls the underlying Tx.DescendRange and traces the query.
256func (tx *Tx) DescendRange(index, lessOrEqual, greaterThan string, iterator func(key, value string) bool) error {
257	span := tx.startSpan("DescendRange")
258	err := tx.Tx.DescendRange(index, lessOrEqual, greaterThan, iterator)
259	span.Finish(tracer.WithError(err))
260	return err
261}
262
263// DropIndex calls the underlying Tx.DropIndex and traces the query.
264func (tx *Tx) DropIndex(name string) error {
265	span := tx.startSpan("DropIndex")
266	err := tx.Tx.DropIndex(name)
267	span.Finish(tracer.WithError(err))
268	return err
269}
270
271// Get calls the underlying Tx.Get and traces the query.
272func (tx *Tx) Get(key string, ignoreExpired ...bool) (val string, err error) {
273	span := tx.startSpan("Get")
274	val, err = tx.Tx.Get(key, ignoreExpired...)
275	span.Finish(tracer.WithError(err))
276	return val, err
277}
278
279// Indexes calls the underlying Tx.Indexes and traces the query.
280func (tx *Tx) Indexes() ([]string, error) {
281	span := tx.startSpan("Indexes")
282	indexes, err := tx.Tx.Indexes()
283	span.Finish(tracer.WithError(err))
284	return indexes, err
285}
286
287// Intersects calls the underlying Tx.Intersects and traces the query.
288func (tx *Tx) Intersects(index, bounds string, iterator func(key, value string) bool) error {
289	span := tx.startSpan("Intersects")
290	err := tx.Tx.Intersects(index, bounds, iterator)
291	span.Finish(tracer.WithError(err))
292	return err
293}
294
295// Len calls the underlying Tx.Len and traces the query.
296func (tx *Tx) Len() (int, error) {
297	span := tx.startSpan("Len")
298	n, err := tx.Tx.Len()
299	span.Finish(tracer.WithError(err))
300	return n, err
301}
302
303// Nearby calls the underlying Tx.Nearby and traces the query.
304func (tx *Tx) Nearby(index, bounds string, iterator func(key, value string, dist float64) bool) error {
305	span := tx.startSpan("Nearby")
306	err := tx.Tx.Nearby(index, bounds, iterator)
307	span.Finish(tracer.WithError(err))
308	return err
309}
310
311// Set calls the underlying Tx.Set and traces the query.
312func (tx *Tx) Set(key, value string, opts *buntdb.SetOptions) (previousValue string, replaced bool, err error) {
313	span := tx.startSpan("Set")
314	previousValue, replaced, err = tx.Tx.Set(key, value, opts)
315	span.Finish(tracer.WithError(err))
316	return previousValue, replaced, err
317}
318
319// TTL calls the underlying Tx.TTL and traces the query.
320func (tx *Tx) TTL(key string) (time.Duration, error) {
321	span := tx.startSpan("TTL")
322	duration, err := tx.Tx.TTL(key)
323	span.Finish(tracer.WithError(err))
324	return duration, err
325}
326