Skip to content

Commit

Permalink
contrib/database/sql: add db service name tag and _dd.dbm_trace_injec…
Browse files Browse the repository at this point in the history 
…ted span tag (#1448)
  • Loading branch information
@alexandre-normand
alexandre-normand committed Sep 2, 2022
1 parent 1e0594a commit e611cb3
Show file tree
Hide file tree
Showing 4 changed files with 170 additions and 29 deletions.
38 changes: 25 additions & 13 deletions contrib/database/sql/conn.go
View file
Expand Up @@ -34,6 +34,10 @@ const (
queryTypeRollback = "Rollback"
)

const (
keyDBMTraceInjected = "_dd.dbm_trace_injected"
)

type tracedConn struct {
driver.Conn
*traceParams
Expand Down Expand Up @@ -64,17 +68,17 @@ func (tc *tracedConn) PrepareContext(ctx context.Context, query string) (stmt dr
// no context other than service in prepared statements
mode = tracer.SQLInjectionModeService
}
cquery, spanID := injectComments(ctx, query, mode)
cquery, spanID := tc.injectComments(ctx, query, mode)
if connPrepareCtx, ok := tc.Conn.(driver.ConnPrepareContext); ok {
stmt, err := connPrepareCtx.PrepareContext(ctx, cquery)
tc.tryTrace(ctx, queryTypePrepare, query, start, err, tracer.WithSpanID(spanID))
tc.tryTrace(ctx, queryTypePrepare, query, start, err, append(withDBMTraceInjectedTag(mode), tracer.WithSpanID(spanID))...)
if err != nil {
return nil, err
}
return &tracedStmt{Stmt: stmt, traceParams: tc.traceParams, ctx: ctx, query: query}, nil
}
stmt, err = tc.Prepare(cquery)
tc.tryTrace(ctx, queryTypePrepare, query, start, err, tracer.WithSpanID(spanID))
tc.tryTrace(ctx, queryTypePrepare, query, start, err, append(withDBMTraceInjectedTag(mode), tracer.WithSpanID(spanID))...)
if err != nil {
return nil, err
}
Expand All @@ -84,9 +88,9 @@ func (tc *tracedConn) PrepareContext(ctx context.Context, query string) (stmt dr
func (tc *tracedConn) ExecContext(ctx context.Context, query string, args []driver.NamedValue) (r driver.Result, err error) {
start := time.Now()
if execContext, ok := tc.Conn.(driver.ExecerContext); ok {
cquery, spanID := injectComments(ctx, query, tc.cfg.commentInjectionMode)
cquery, spanID := tc.injectComments(ctx, query, tc.cfg.commentInjectionMode)
r, err := execContext.ExecContext(ctx, cquery, args)
tc.tryTrace(ctx, queryTypeExec, query, start, err, tracer.WithSpanID(spanID))
tc.tryTrace(ctx, queryTypeExec, query, start, err, append(withDBMTraceInjectedTag(tc.cfg.commentInjectionMode), tracer.WithSpanID(spanID))...)
return r, err
}
if execer, ok := tc.Conn.(driver.Execer); ok {
Expand All @@ -99,9 +103,9 @@ func (tc *tracedConn) ExecContext(ctx context.Context, query string, args []driv
return nil, ctx.Err()
default:
}
cquery, spanID := injectComments(ctx, query, tc.cfg.commentInjectionMode)
cquery, spanID := tc.injectComments(ctx, query, tc.cfg.commentInjectionMode)
r, err = execer.Exec(cquery, dargs)
tc.tryTrace(ctx, queryTypeExec, query, start, err, tracer.WithSpanID(spanID))
tc.tryTrace(ctx, queryTypeExec, query, start, err, append(withDBMTraceInjectedTag(tc.cfg.commentInjectionMode), tracer.WithSpanID(spanID))...)
return r, err
}
return nil, driver.ErrSkip
Expand All @@ -120,9 +124,9 @@ func (tc *tracedConn) Ping(ctx context.Context) (err error) {
func (tc *tracedConn) QueryContext(ctx context.Context, query string, args []driver.NamedValue) (rows driver.Rows, err error) {
start := time.Now()
if queryerContext, ok := tc.Conn.(driver.QueryerContext); ok {
cquery, spanID := injectComments(ctx, query, tc.cfg.commentInjectionMode)
cquery, spanID := tc.injectComments(ctx, query, tc.cfg.commentInjectionMode)
rows, err := queryerContext.QueryContext(ctx, cquery, args)
tc.tryTrace(ctx, queryTypeQuery, query, start, err, tracer.WithSpanID(spanID))
tc.tryTrace(ctx, queryTypeQuery, query, start, err, append(withDBMTraceInjectedTag(tc.cfg.commentInjectionMode), tracer.WithSpanID(spanID))...)
return rows, err
}
if queryer, ok := tc.Conn.(driver.Queryer); ok {
Expand All @@ -135,9 +139,9 @@ func (tc *tracedConn) QueryContext(ctx context.Context, query string, args []dri
return nil, ctx.Err()
default:
}
cquery, spanID := injectComments(ctx, query, tc.cfg.commentInjectionMode)
cquery, spanID := tc.injectComments(ctx, query, tc.cfg.commentInjectionMode)
rows, err = queryer.Query(cquery, dargs)
tc.tryTrace(ctx, queryTypeQuery, query, start, err, tracer.WithSpanID(spanID))
tc.tryTrace(ctx, queryTypeQuery, query, start, err, append(withDBMTraceInjectedTag(tc.cfg.commentInjectionMode), tracer.WithSpanID(spanID))...)
return rows, err
}
return nil, driver.ErrSkip
Expand Down Expand Up @@ -181,7 +185,7 @@ func WithSpanTags(ctx context.Context, tags map[string]string) context.Context {
// injectComments returns the query with SQL comments injected according to the comment injection mode along
// with a span ID injected into SQL comments. The returned span ID should be used when the SQL span is created
// following the traced database call.
func injectComments(ctx context.Context, query string, mode tracer.SQLCommentInjectionMode) (cquery string, spanID uint64) {
func (tc *tracedConn) injectComments(ctx context.Context, query string, mode tracer.SQLCommentInjectionMode) (cquery string, spanID uint64) {
// The sql span only gets created after the call to the database because we need to be able to skip spans
// when a driver returns driver.ErrSkip. In order to work with those constraints, a new span id is generated and
// used during SQL comment injection and returned for the sql span to be used later when/if the span
Expand All @@ -190,14 +194,22 @@ func injectComments(ctx context.Context, query string, mode tracer.SQLCommentInj
if span, ok := tracer.SpanFromContext(ctx); ok {
spanCtx = span.Context()
}
carrier := tracer.SQLCommentCarrier{Query: query, Mode: mode}
carrier := tracer.SQLCommentCarrier{Query: query, Mode: mode, DBServiceName: tc.cfg.serviceName}
if err := carrier.Inject(spanCtx); err != nil {
// this should never happen
log.Warn("contrib/database/sql: failed to inject query comments: %v", err)
}
return carrier.Query, carrier.SpanID
}

func withDBMTraceInjectedTag(mode tracer.SQLCommentInjectionMode) []tracer.StartSpanOption {
if mode == tracer.SQLInjectionModeFull {
return []tracer.StartSpanOption{tracer.Tag(keyDBMTraceInjected, true)}
}

return nil
}

// tryTrace will create a span using the given arguments, but will act as a no-op when err is driver.ErrSkip.
func (tp *traceParams) tryTrace(ctx context.Context, qtype queryType, query string, startTime time.Time, err error, spanOpts ...ddtrace.StartSpanOption) {
if err == driver.ErrSkip {
Expand Down
136 changes: 130 additions & 6 deletions contrib/database/sql/injection_test.go
View file
Expand Up @@ -15,6 +15,7 @@ import (
"github.com/stretchr/testify/require"

"gopkg.in/DataDog/dd-trace-go.v1/contrib/database/sql/internal"
"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/mocktracer"
"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer"
)

Expand Down Expand Up @@ -51,7 +52,7 @@ func TestCommentInjection(t *testing.T) {
_, err := db.PrepareContext(ctx, "SELECT 1 from DUAL")
return err
},
prepared: []string{"/*dde='test-env',ddsn='test-service',ddsv='1.0.0'*/ SELECT 1 from DUAL"},
prepared: []string{"/*dddbs='test.db',dde='test-env',ddsn='test-service',ddsv='1.0.0'*/ SELECT 1 from DUAL"},
},
{
name: "prepare-full",
Expand All @@ -60,7 +61,7 @@ func TestCommentInjection(t *testing.T) {
_, err := db.PrepareContext(ctx, "SELECT 1 from DUAL")
return err
},
prepared: []string{"/*dde='test-env',ddsn='test-service',ddsv='1.0.0'*/ SELECT 1 from DUAL"},
prepared: []string{"/*dddbs='test.db',dde='test-env',ddsn='test-service',ddsv='1.0.0'*/ SELECT 1 from DUAL"},
},
{
name: "query",
Expand All @@ -87,7 +88,7 @@ func TestCommentInjection(t *testing.T) {
_, err := db.QueryContext(ctx, "SELECT 1 from DUAL")
return err
},
executed: []*regexp.Regexp{regexp.MustCompile("/\\*dde='test-env',ddsn='test-service',ddsv='1.0.0'\\*/ SELECT 1 from DUAL")},
executed: []*regexp.Regexp{regexp.MustCompile("/\\*dddbs='test.db',dde='test-env',ddsn='test-service',ddsv='1.0.0'\\*/ SELECT 1 from DUAL")},
},
{
name: "query-full",
Expand All @@ -96,7 +97,7 @@ func TestCommentInjection(t *testing.T) {
_, err := db.QueryContext(ctx, "SELECT 1 from DUAL")
return err
},
executed: []*regexp.Regexp{regexp.MustCompile("/\\*dde='test-env',ddsid='\\d+',ddsn='test-service',ddsp='1',ddsv='1.0.0',ddtid='1'\\*/ SELECT 1 from DUAL")},
executed: []*regexp.Regexp{regexp.MustCompile("/\\*dddbs='test.db',dde='test-env',ddsid='\\d+',ddsn='test-service',ddsp='1',ddsv='1.0.0',ddtid='1'\\*/ SELECT 1 from DUAL")},
},
{
name: "exec",
Expand All @@ -123,7 +124,7 @@ func TestCommentInjection(t *testing.T) {
_, err := db.ExecContext(ctx, "SELECT 1 from DUAL")
return err
},
executed: []*regexp.Regexp{regexp.MustCompile("/\\*dde='test-env',ddsn='test-service',ddsv='1.0.0'\\*/ SELECT 1 from DUAL")},
executed: []*regexp.Regexp{regexp.MustCompile("/\\*dddbs='test.db',dde='test-env',ddsn='test-service',ddsv='1.0.0'\\*/ SELECT 1 from DUAL")},
},
{
name: "exec-full",
Expand All @@ -132,7 +133,7 @@ func TestCommentInjection(t *testing.T) {
_, err := db.ExecContext(ctx, "SELECT 1 from DUAL")
return err
},
executed: []*regexp.Regexp{regexp.MustCompile("/\\*dde='test-env',ddsid='\\d+',ddsn='test-service',ddsp='1',ddsv='1.0.0',ddtid='1'\\*/ SELECT 1 from DUAL")},
executed: []*regexp.Regexp{regexp.MustCompile("/\\*dddbs='test.db',dde='test-env',ddsid='\\d+',ddsn='test-service',ddsp='1',ddsv='1.0.0',ddtid='1'\\*/ SELECT 1 from DUAL")},
},
}

Expand Down Expand Up @@ -167,3 +168,126 @@ func TestCommentInjection(t *testing.T) {
})
}
}

func TestDBMTraceContextTagging(t *testing.T) {
testCases := []struct {
name string
opts []RegisterOption
callDB func(ctx context.Context, db *sql.DB) error
spanType string
traceContextInjectedTag bool
}{
{
name: "prepare",
opts: []RegisterOption{WithSQLCommentInjection(tracer.SQLInjectionModeFull)},
callDB: func(ctx context.Context, db *sql.DB) error {
_, err := db.PrepareContext(ctx, "SELECT 1 from DUAL")
return err
},
spanType: queryTypePrepare,
traceContextInjectedTag: false,
},
{
name: "query-disabled",
opts: []RegisterOption{WithSQLCommentInjection(tracer.SQLInjectionDisabled)},
callDB: func(ctx context.Context, db *sql.DB) error {
_, err := db.QueryContext(ctx, "SELECT 1 from DUAL")
return err
},
spanType: queryTypeQuery,
traceContextInjectedTag: false,
},
{
name: "query-service",
opts: []RegisterOption{WithSQLCommentInjection(tracer.SQLInjectionModeService)},
callDB: func(ctx context.Context, db *sql.DB) error {
_, err := db.QueryContext(ctx, "SELECT 1 from DUAL")
return err
},
spanType: queryTypeQuery,
traceContextInjectedTag: false,
},
{
name: "query-full",
opts: []RegisterOption{WithSQLCommentInjection(tracer.SQLInjectionModeFull)},
callDB: func(ctx context.Context, db *sql.DB) error {
_, err := db.QueryContext(ctx, "SELECT 1 from DUAL")
return err
},
spanType: queryTypeQuery,
traceContextInjectedTag: true,
},
{
name: "exec-disabled",
opts: []RegisterOption{WithSQLCommentInjection(tracer.SQLInjectionDisabled)},
callDB: func(ctx context.Context, db *sql.DB) error {
_, err := db.ExecContext(ctx, "SELECT 1 from DUAL")
return err
},
spanType: queryTypeExec,
traceContextInjectedTag: false,
},
{
name: "exec-service",
opts: []RegisterOption{WithSQLCommentInjection(tracer.SQLInjectionModeService)},
callDB: func(ctx context.Context, db *sql.DB) error {
_, err := db.ExecContext(ctx, "SELECT 1 from DUAL")
return err
},
spanType: queryTypeExec,
traceContextInjectedTag: false,
},
{
name: "exec-full",
opts: []RegisterOption{WithSQLCommentInjection(tracer.SQLInjectionModeFull)},
callDB: func(ctx context.Context, db *sql.DB) error {
_, err := db.ExecContext(ctx, "SELECT 1 from DUAL")
return err
},
spanType: queryTypeExec,
traceContextInjectedTag: true,
},
}

for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
tr := mocktracer.Start()
defer tr.Stop()

d := &internal.MockDriver{}
Register("test", d, tc.opts...)
defer unregister("test")

db, err := Open("test", "dn")
require.NoError(t, err)

s, ctx := tracer.StartSpanFromContext(context.Background(), "test.call", tracer.WithSpanID(1))
err = tc.callDB(ctx, db)
s.Finish()

require.NoError(t, err)
spans := tr.FinishedSpans()

sps := spansOfType(spans, tc.spanType)
for _, s := range sps {
tags := s.Tags()
if tc.traceContextInjectedTag {
assert.Equal(t, tags[keyDBMTraceInjected], true)
} else {
_, ok := tags[keyDBMTraceInjected]
assert.False(t, ok)
}
}
})
}
}

func spansOfType(spans []mocktracer.Span, spanType string) (filtered []mocktracer.Span) {
filtered = make([]mocktracer.Span, 0)
for _, s := range spans {
if s.Tag("sql.query_type") == spanType {
filtered = append(filtered, s)
}
}
return filtered
}
13 changes: 9 additions & 4 deletions ddtrace/tracer/sqlcomment.go
View file
Expand Up @@ -12,6 +12,7 @@ import (
"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"
"gopkg.in/DataDog/dd-trace-go.v1/internal/log"
)

// SQLCommentInjectionMode represents the mode of SQL comment injection.
Expand All @@ -34,6 +35,7 @@ const (
sqlCommentTraceID = "ddtid"
sqlCommentSpanID = "ddsid"
sqlCommentService = "ddsn"
sqlCommentDBService = "dddbs"
sqlCommentVersion = "ddsv"
sqlCommentEnv = "dde"
)
Expand All @@ -42,9 +44,10 @@ const (
// 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 SQLCommentInjectionMode
SpanID uint64
Query string
Mode SQLCommentInjectionMode
DBServiceName string
SpanID uint64
}

// Inject injects a span context in the carrier's Query field as a comment.
Expand Down Expand Up @@ -93,6 +96,7 @@ func (c *SQLCommentCarrier) Inject(spanCtx ddtrace.SpanContext) error {
if version != "" {
tags[sqlCommentVersion] = version
}
tags[sqlCommentDBService] = c.DBServiceName
}
c.Query = commentQuery(c.Query, tags)
return nil
Expand All @@ -113,7 +117,7 @@ func commentQuery(query string, tags map[string]string) string {
var b strings.Builder
// the sqlcommenter specification dictates that tags should be sorted. Since we know all injected keys,
// we skip a sorting operation by specifying the order of keys statically
orderedKeys := []string{sqlCommentEnv, sqlCommentSpanID, sqlCommentService, sqlCommentKeySamplingPriority, sqlCommentVersion, sqlCommentTraceID}
orderedKeys := []string{sqlCommentDBService, sqlCommentEnv, sqlCommentSpanID, sqlCommentService, sqlCommentKeySamplingPriority, sqlCommentVersion, sqlCommentTraceID}
first := true
for _, k := range orderedKeys {
if v, ok := tags[k]; ok {
Expand Down Expand Up @@ -141,6 +145,7 @@ func commentQuery(query string, tags map[string]string) string {
if query == "" {
return b.String()
}
log.Debug("Injected sql comment: %s", b.String())
b.WriteRune(' ')
b.WriteString(query)
return b.String()
Expand Down

0 comments on commit e611cb3

Please sign in to comment.