-
Notifications
You must be signed in to change notification settings - Fork 415
/
sqlcomment.go
271 lines (234 loc) · 8.97 KB
/
sqlcomment.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
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
// Unless explicitly stated otherwise all files in this repository are licensed
// under the Apache License Version 2.0.
// This product includes software developed at Datadog (https://www.datadoghq.com/).
// Copyright 2016 Datadog, Inc.
package tracer
import (
"fmt"
"net/url"
"sort"
"strconv"
"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/internal/globalconfig"
)
// SQLCommentInjectionMode represents the mode of sql comment injection.
type SQLCommentInjectionMode int
const (
// CommentInjectionDisabled represents the comment injection mode where all injection is disabled.
CommentInjectionDisabled SQLCommentInjectionMode = 0
// StaticTagsSQLCommentInjection represents the comment injection mode where only static tags are injected. Static tags include values that are set once during the lifetime of an application: service name, env, version.
StaticTagsSQLCommentInjection SQLCommentInjectionMode = 1
// FullSQLCommentInjection represents the comment injection mode where both static and dynamic tags are injected. Dynamic tags include values like span id, trace id and sampling priority.
FullSQLCommentInjection SQLCommentInjectionMode = 2
)
// Key names for SQL comment tags.
const (
SamplingPrioritySQLCommentKey = "ddsp"
TraceIDSQLCommentKey = "ddtid"
SpanIDSQLCommentKey = "ddsid"
ServiceNameSQLCommentKey = "ddsn"
ServiceVersionSQLCommentKey = "ddsv"
ServiceEnvironmentSQLCommentKey = "dde"
)
// QueryCommentCarrier is a more specific interface implemented by carriers that implement the TextMapWriter
// as well as CommentQuery, AddSpanID and SetDynamicTag methods. It is compatible with Datadog's SQLCommentPropagator.
// Note that Datadog's TextMapPropagator is compatible with QueryCommentCarriers but only in the fact that it
// ignores it without returning errors.
type QueryCommentCarrier interface {
TextMapWriter
// SetDynamicTag sets the given dynamic key/value pair. This method is a variation on TextMapWriter's Set method
// that exists for dynamic tag values.
SetDynamicTag(key, val string)
// CommentQuery returns the given query with any injected tags prepended in the form of a sqlcommenter-formatted
// SQL comment. It also returns a non-zero span ID if a new span ID was generated as part of the inject operation.
// See https://google.github.io/sqlcommenter/spec for the full sqlcommenter spec.
CommentQuery(query string) (commented string, spanID uint64)
// AddSpanID adds a new span ID to the carrier. This happens if no existing trace is found when the
// injection happens.
AddSpanID(spanID uint64)
}
// SQLCommentPropagator implements the Propagator interface to inject tags
// in sql comments. It is only compatible with QueryCommentCarrier implementations
// but allows other carriers to flow through without returning errors.
type SQLCommentPropagator struct {
mode SQLCommentInjectionMode
}
// SQLCommentWithDynamicTagsDiscarded enables control discarding dynamic tags on a SQLCommentCarrier.
// Its main purpose is to allow discarding dynamic tags per SQL operation when they aren't relevant
// (i.e. Prepared statements).
func SQLCommentWithDynamicTagsDiscarded(discard bool) SQLCommentCarrierOption {
return func(c *SQLCommentCarrierConfig) {
c.discardDynamicTags = discard
}
}
// NewSQLCommentPropagator returns a new SQLCommentPropagator with the given injection mode
func NewSQLCommentPropagator(mode SQLCommentInjectionMode) *SQLCommentPropagator {
return &SQLCommentPropagator{mode: mode}
}
// Inject injects the span context in the given carrier. Note that it is only compatible
// with QueryCommentCarriers and no-ops if the carrier is of any other type.
func (p *SQLCommentPropagator) Inject(spanCtx ddtrace.SpanContext, carrier interface{}) error {
switch c := carrier.(type) {
case QueryCommentCarrier:
return p.injectWithCommentCarrier(spanCtx, c)
default:
// SQLCommentPropagator only handles QueryCommentCarrier carriers but lets any other carrier
// flow through without returning errors
return nil
}
}
func (p *SQLCommentPropagator) injectWithCommentCarrier(spanCtx ddtrace.SpanContext, carrier QueryCommentCarrier) error {
if p.mode == CommentInjectionDisabled {
return nil
}
spanID := random.Uint64()
carrier.AddSpanID(spanID)
if p.mode == StaticTagsSQLCommentInjection || p.mode == FullSQLCommentInjection {
ctx, ok := spanCtx.(*spanContext)
var env, version string
if ok {
if e, ok := ctx.meta(ext.Environment); ok {
env = e
}
if v, ok := ctx.meta(ext.Version); ok {
version = v
}
}
if globalconfig.ServiceName() != "" {
carrier.Set(ServiceNameSQLCommentKey, globalconfig.ServiceName())
}
if env != "" {
carrier.Set(ServiceEnvironmentSQLCommentKey, env)
}
if version != "" {
carrier.Set(ServiceVersionSQLCommentKey, version)
}
}
if p.mode == FullSQLCommentInjection {
samplingPriority := 0
var traceID uint64
ctx, ok := spanCtx.(*spanContext)
if ok {
if sp, ok := ctx.samplingPriority(); ok {
samplingPriority = sp
}
if ctx.TraceID() > 0 {
traceID = ctx.TraceID()
}
}
if traceID == 0 {
traceID = spanID
}
carrier.SetDynamicTag(TraceIDSQLCommentKey, strconv.FormatUint(traceID, 10))
carrier.SetDynamicTag(SpanIDSQLCommentKey, strconv.FormatUint(spanID, 10))
carrier.SetDynamicTag(SamplingPrioritySQLCommentKey, strconv.Itoa(samplingPriority))
}
return nil
}
// Extract is a no-op for the SQLCommentPropagator.
func (p *SQLCommentPropagator) Extract(carrier interface{}) (ddtrace.SpanContext, error) {
return nil, nil
}
// SQLCommentCarrierConfig holds configuration for a SQLCommentCarrier
type SQLCommentCarrierConfig struct {
discardDynamicTags bool
}
// SQLCommentCarrierOption represents a function that can be provided as a parameter to NewSQLCommentCarrier.
type SQLCommentCarrierOption func(c *SQLCommentCarrierConfig)
// SQLCommentCarrier implements QueryCommentCarrier by holding tags, configuration
// and a potential new span id to generate a SQL comment injected in queries.
type SQLCommentCarrier struct {
tags map[string]string
cfg SQLCommentCarrierConfig
spanID uint64
}
// NewSQLCommentCarrier returns a new SQLCommentCarrier.
func NewSQLCommentCarrier(opts ...SQLCommentCarrierOption) (s *SQLCommentCarrier) {
s = new(SQLCommentCarrier)
for _, apply := range opts {
apply(&s.cfg)
}
return s
}
// Set implements TextMapWriter. In the context of SQL comment injection, this method
// is used for static tags only. See SetDynamicTag for dynamic tags.
func (c *SQLCommentCarrier) Set(key, val string) {
if c.tags == nil {
c.tags = make(map[string]string)
}
c.tags[key] = val
}
// ForeachKey implements TextMapReader.
func (c SQLCommentCarrier) ForeachKey(handler func(key, val string) error) error {
for k, v := range c.tags {
if err := handler(k, v); err != nil {
return err
}
}
return nil
}
// SetDynamicTag implements QueryCommentCarrier. This method is used to inject dynamic tags only
// (i.e. span id, trace id, sampling priority). See Set for static tags.
func (c *SQLCommentCarrier) SetDynamicTag(key, val string) {
if c.cfg.discardDynamicTags {
return
}
if c.tags == nil {
c.tags = make(map[string]string)
}
c.tags[key] = val
}
// AddSpanID implements QueryCommentCarrier and is used to save a span id generated during SQL comment injection
// in the case where there is no active trace.
func (c *SQLCommentCarrier) AddSpanID(spanID uint64) {
c.spanID = spanID
}
// CommentQuery returns the given query with the tags from the SQLCommentCarrier applied to it as a
// prepended SQL comment. The format of the comment follows the sqlcommenter spec.
// See https://google.github.io/sqlcommenter/spec/ for more details.
func (c *SQLCommentCarrier) CommentQuery(query string) (commented string, spanID uint64) {
comment := commentWithTags(c.tags)
if comment == "" {
return query, c.spanID
}
if query == "" {
return comment, c.spanID
}
return fmt.Sprintf("%s %s", comment, query), c.spanID
}
func commentWithTags(tags map[string]string) (comment string) {
if len(tags) == 0 {
return ""
}
serializedTags := make([]string, 0, len(tags))
for k, v := range tags {
serializedTags = append(serializedTags, serializeTag(k, v))
}
sort.Strings(serializedTags)
comment = strings.Join(serializedTags, ",")
return fmt.Sprintf("/*%s*/", comment)
}
func serializeTag(key string, value string) (serialized string) {
sKey := serializeKey(key)
sValue := serializeValue(value)
return fmt.Sprintf("%s=%s", sKey, sValue)
}
func serializeKey(key string) (encoded string) {
urlEncoded := url.PathEscape(key)
escapedMeta := escapeMetaChars(urlEncoded)
return escapedMeta
}
func serializeValue(value string) (encoded string) {
urlEncoded := url.PathEscape(value)
escapedMeta := escapeMetaChars(urlEncoded)
escaped := escapeSQL(escapedMeta)
return escaped
}
func escapeSQL(value string) (escaped string) {
return fmt.Sprintf("'%s'", value)
}
func escapeMetaChars(value string) (escaped string) {
return strings.ReplaceAll(value, "'", "\\'")
}