/
sqlcomment.go
102 lines (93 loc) · 4.19 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
// 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 (
v2 "github.com/DataDog/dd-trace-go/v2/ddtrace/tracer"
"gopkg.in/DataDog/dd-trace-go.v1/ddtrace"
"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/internal"
)
// SQLCommentInjectionMode represents the mode of SQL comment injection.
//
// Deprecated: Use DBMPropagationMode instead.
type SQLCommentInjectionMode DBMPropagationMode
const (
// SQLInjectionUndefined represents the comment injection mode is not set. This is the same as SQLInjectionDisabled.
SQLInjectionUndefined SQLCommentInjectionMode = SQLCommentInjectionMode(DBMPropagationModeUndefined)
// SQLInjectionDisabled represents the comment injection mode where all injection is disabled.
SQLInjectionDisabled SQLCommentInjectionMode = SQLCommentInjectionMode(DBMPropagationModeDisabled)
// SQLInjectionModeService represents the comment injection mode where only service tags (name, env, version) are injected.
SQLInjectionModeService SQLCommentInjectionMode = SQLCommentInjectionMode(DBMPropagationModeService)
// SQLInjectionModeFull represents the comment injection mode where both service tags and tracing tags. Tracing tags include span id, trace id and sampling priority.
SQLInjectionModeFull SQLCommentInjectionMode = SQLCommentInjectionMode(DBMPropagationModeFull)
)
// DBMPropagationMode represents the mode of dbm propagation.
//
// Note that enabling sql comment propagation results in potentially confidential data (service names)
// being stored in the databases which can then be accessed by other 3rd parties that have been granted
// access to the database.
type DBMPropagationMode string
const (
// DBMPropagationModeUndefined represents the dbm propagation mode not being set. This is the same as DBMPropagationModeDisabled.
DBMPropagationModeUndefined DBMPropagationMode = ""
// DBMPropagationModeDisabled represents the dbm propagation mode where all propagation is disabled.
DBMPropagationModeDisabled DBMPropagationMode = "disabled"
// DBMPropagationModeService represents the dbm propagation mode where only service tags (name, env, version) are propagated to dbm.
DBMPropagationModeService DBMPropagationMode = "service"
// DBMPropagationModeFull represents the dbm propagation mode where both service tags and tracing tags are propagated. Tracing tags include span id, trace id and the sampled flag.
DBMPropagationModeFull DBMPropagationMode = "full"
)
// SQLCommentCarrier is a carrier implementation that injects a span context in a SQL query in the form
// of a sqlcommenter formatted comment prepended to the original query text.
// See https://google.github.io/sqlcommenter/spec/ for more details.
type SQLCommentCarrier struct {
Query string
Mode DBMPropagationMode
DBServiceName string
SpanID uint64
v2carrier *v2.SQLCommentCarrier
}
// Inject injects a span context in the carrier's Query field as a comment.
func (c *SQLCommentCarrier) Inject(spanCtx ddtrace.SpanContext) error {
if c.v2carrier == nil {
c.v2carrier = &v2.SQLCommentCarrier{}
}
if c.v2carrier.Query != c.Query {
c.v2carrier.Query = c.Query
}
mode := v2.DBMPropagationMode(c.Mode)
if c.v2carrier.Mode != mode {
c.v2carrier.Mode = mode
}
if c.v2carrier.DBServiceName != c.DBServiceName {
c.v2carrier.DBServiceName = c.DBServiceName
}
if c.v2carrier.SpanID != c.SpanID {
c.v2carrier.SpanID = c.SpanID
}
var ctx *v2.SpanContext
if spanCtx != nil {
ctx = spanCtx.(internal.SpanContextV2Adapter).Ctx
}
if err := c.v2carrier.Inject(ctx); err != nil {
return err
}
c.Query = c.v2carrier.Query
c.SpanID = c.v2carrier.SpanID
return nil
}
// Extract parses for key value attributes in a sql query injected with trace information in order to build a span context
func (c *SQLCommentCarrier) Extract() (ddtrace.SpanContext, error) {
if c.v2carrier == nil {
c.v2carrier = &v2.SQLCommentCarrier{}
}
if c.v2carrier.Query != c.Query {
c.v2carrier.Query = c.Query
}
ctx, err := c.v2carrier.Extract()
if err != nil {
return nil, err
}
return internal.SpanContextV2Adapter{Ctx: ctx}, nil
}