Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

contrib/database/sql: add db service name tag and _dd.dbm_trace_injected span tag #1448

Merged
merged 7 commits into from Sep 2, 2022
Merged

contrib/database/sql: add db service name tag and _dd.dbm_trace_injected span tag #1448

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
12cbdbd
contrib/database/sql: add db service name tag and _dd.dbm_trace_injec…
alexandre-normand Aug 26, 2022
23ec6e0
Add debug log
alexandre-normand Aug 30, 2022
29cd203
Address review comments
alexandre-normand Aug 31, 2022
75a1e93
Merge branch 'main' into alex.normand/add-db-service-name-to-injected…
alexandre-normand Aug 31, 2022
d1791d5
Remove obsolete line
alexandre-normand Aug 31, 2022
1dc064a
Merge branch 'alex.normand/add-db-service-name-to-injected-sql-tags' …
alexandre-normand Aug 31, 2022
9cd4841
Merge branch 'main' into alex.normand/add-db-service-name-to-injected…
ajgajg1134 Sep 2, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
38 changes: 25 additions & 13 deletions contrib/database/sql/conn.go
View file
Open in desktop
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
Open in desktop
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
}
18 changes: 14 additions & 4 deletions ddtrace/tracer/sqlcomment.go
View file
Open in desktop
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 @@ -83,6 +86,9 @@ func (c *SQLCommentCarrier) Inject(spanCtx ddtrace.SpanContext) error {
if v, ok := ctx.meta(ext.Version); ok {
version = v
}
if s := ctx.span; s != nil {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Any reason to create s here? Perhaps just if ctx.span != nil ?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good catch. Updated ✅

tags[sqlCommentDBService] = ctx.span.Service
}
}
if globalconfig.ServiceName() != "" {
tags[sqlCommentService] = globalconfig.ServiceName()
Expand All @@ -93,6 +99,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 +120,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 +148,9 @@ func commentQuery(query string, tags map[string]string) string {
if query == "" {
return b.String()
}
if log.DebugEnabled() {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't think you need to check if DebugEnabled, you can just call log.Debug and it will not log if debug is disabled

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, I was wondering about that. The comment on that DebugEnabled said it was for things in "hot paths" to avoid the interface{} allocations. I wasn't sure if this would be considered a hot path or not. I'm fine to just switch to log.Debug.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Updated ✅

log.Debug("Injected sql comment %s", b.String())
}
b.WriteRune(' ')
b.WriteString(query)
return b.String()
Expand Down