Skip to content

Commit

Permalink
contrib/database/sql: allow peer.service to be set explicitly and add…
Browse files Browse the repository at this point in the history 
… into propagated context (#2679)

Co-authored-by: Diana Shevchenko <40775148+dianashevchenko@users.noreply.github.com>
Co-authored-by: Dario Castañé <dario.castane@datadoghq.com>
  • Loading branch information
@tabgok @dianashevchenko @darccio
3 people committed May 22, 2024
1 parent 14d15d5 commit b0aa1b8
Show file tree
Hide file tree
Showing 4 changed files with 144 additions and 24 deletions.
26 changes: 25 additions & 1 deletion contrib/database/sql/conn.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -241,6 +241,29 @@ func WithSpanTags(ctx context.Context, tags map[string]string) context.Context {
return context.WithValue(ctx, spanTagsKey, tags)
}

// providedPeerService returns the peer service tag if provided manually by the user,
// derived from a possible sources (span tags, context, ...)
func (tc *TracedConn) providedPeerService(ctx context.Context) string {
// This occurs if the user sets peer.service explicitly while creating the connection
// through the use of sqltrace.WithSpanTags
if meta, ok := ctx.Value(spanTagsKey).(map[string]string); ok {
if peerServiceTag, ok := meta[ext.PeerService]; ok {
return peerServiceTag
}
}

// This occurs if the SQL Connection is opened or registered using
// WithCustomTags. This is lower precedence than above since it is
// less specific
if len(tc.cfg.tags) > 0 {
if v, ok := tc.cfg.tags[ext.PeerService].(string); ok {
return v
}
}

return ""
}

// 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.
Expand All @@ -253,7 +276,8 @@ func (tc *TracedConn) injectComments(ctx context.Context, query string, mode tra
if span, ok := tracer.SpanFromContext(ctx); ok {
spanCtx = span.Context()
}
carrier := tracer.SQLCommentCarrier{Query: query, Mode: mode, DBServiceName: tc.cfg.serviceName, PeerDBHostname: tc.meta[ext.TargetHost], PeerDBName: tc.meta[ext.DBName]}

carrier := tracer.SQLCommentCarrier{Query: query, Mode: mode, DBServiceName: tc.cfg.serviceName, PeerDBHostname: tc.meta[ext.TargetHost], PeerDBName: tc.meta[ext.DBName], PeerService: tc.providedPeerService(ctx)}
if err := carrier.Inject(spanCtx); err != nil {
// this should never happen
log.Warn("contrib/database/sql: failed to inject query comments: %v", err)
Expand Down
92 changes: 79 additions & 13 deletions contrib/database/sql/propagation_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ import (
"testing"

"gopkg.in/DataDog/dd-trace-go.v1/contrib/database/sql/internal"
"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/ext"
"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/mocktracer"
"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer"
"gopkg.in/DataDog/dd-trace-go.v1/internal/globalconfig"
Expand All @@ -33,12 +34,15 @@ func TestDBMPropagation(t *testing.T) {
defer globalconfig.SetServiceName(prevServiceName)

testCases := []struct {
name string
opts []RegisterOption
callDB func(ctx context.Context, db *sql.DB) error
prepared []string
dsn string
executed []*regexp.Regexp
name string
opts []RegisterOption
callDB func(ctx context.Context, db *sql.DB) error
prepared []string
dsn string
executed []*regexp.Regexp
peerServiceTag string
peerServiceCtx string
peerServiceCustomOpenTag string
}{
{
name: "prepare",
Expand Down Expand Up @@ -74,8 +78,9 @@ func TestDBMPropagation(t *testing.T) {
_, err := db.PrepareContext(ctx, "SELECT 1 from DUAL")
return err
},
dsn: "postgres://postgres:postgres@127.0.0.1:5432/fakepreparedb?sslmode=disable",
prepared: []string{"/*dddbs='test.db',dde='test-env',ddps='test-service',ddpv='1.0.0',ddh='127.0.0.1',dddb='fakepreparedb'*/ SELECT 1 from DUAL"},
dsn: "postgres://postgres:postgres@127.0.0.1:5432/fakepreparedb?sslmode=disable",
prepared: []string{"/*dddbs='test.db',dde='test-env',ddps='test-service',ddpv='1.0.0',ddh='127.0.0.1',dddb='fakepreparedb',ddprs='test-peer-service'*/ SELECT 1 from DUAL"},
peerServiceCtx: "test-peer-service",
},
{
name: "query",
Expand Down Expand Up @@ -111,8 +116,9 @@ func TestDBMPropagation(t *testing.T) {
_, err := db.QueryContext(ctx, "SELECT 1 from DUAL")
return err
},
dsn: "postgres://postgres:postgres@127.0.0.1:5432/fakequerydb?sslmode=disable",
executed: []*regexp.Regexp{regexp.MustCompile("/\\*dddbs='test.db',dde='test-env',ddps='test-service',ddpv='1.0.0',traceparent='00-00000000000000000000000000000001-[\\da-f]{16}-01',ddh='127.0.0.1',dddb='fakequerydb'\\*/ SELECT 1 from DUAL")},
dsn: "postgres://postgres:postgres@127.0.0.1:5432/fakequerydb?sslmode=disable",
executed: []*regexp.Regexp{regexp.MustCompile("/\\*dddbs='test.db',dde='test-env',ddps='test-service',ddpv='1.0.0',traceparent='00-00000000000000000000000000000001-[\\da-f]{16}-01',ddh='127.0.0.1',dddb='fakequerydb',ddprs='test-peer-service'\\*/ SELECT 1 from DUAL")},
peerServiceCtx: "test-peer-service",
},
{
name: "exec",
Expand Down Expand Up @@ -148,8 +154,55 @@ func TestDBMPropagation(t *testing.T) {
_, err := db.ExecContext(ctx, "SELECT 1 from DUAL")
return err
},
dsn: "postgres://postgres:postgres@127.0.0.1:5432/fakeexecdb?sslmode=disable",
executed: []*regexp.Regexp{regexp.MustCompile("/\\*dddbs='test.db',dde='test-env',ddps='test-service',ddpv='1.0.0',traceparent='00-00000000000000000000000000000001-[\\da-f]{16}-01',ddh='127.0.0.1',dddb='fakeexecdb'\\*/ SELECT 1 from DUAL")},
dsn: "postgres://postgres:postgres@127.0.0.1:5432/fakeexecdb?sslmode=disable",
executed: []*regexp.Regexp{regexp.MustCompile("/\\*dddbs='test.db',dde='test-env',ddps='test-service',ddpv='1.0.0',traceparent='00-00000000000000000000000000000001-[\\da-f]{16}-01',ddh='127.0.0.1',dddb='fakeexecdb',ddprs='test-peer-service'\\*/ SELECT 1 from DUAL")},
peerServiceCtx: "test-peer-service",
},
{
name: "exec-full-peer-service-tag",
opts: []RegisterOption{WithDBMPropagation(tracer.DBMPropagationModeFull)},
callDB: func(ctx context.Context, db *sql.DB) error {
_, err := db.ExecContext(ctx, "SELECT 1 from DUAL")
return err
},
dsn: "postgres://postgres:postgres@127.0.0.1:5432/fakeexecdb?sslmode=disable",
executed: []*regexp.Regexp{regexp.MustCompile("/\\*dddbs='test.db',dde='test-env',ddps='test-service',ddpv='1.0.0',traceparent='00-00000000000000000000000000000001-[\\da-f]{16}-01',ddh='127.0.0.1',dddb='fakeexecdb',ddprs='test-peer-service-tag'\\*/ SELECT 1 from DUAL")},
peerServiceTag: "test-peer-service-tag",
},
{
name: "exec-full-peer-service-custom-tag",
opts: []RegisterOption{WithDBMPropagation(tracer.DBMPropagationModeFull)},
callDB: func(ctx context.Context, db *sql.DB) error {
_, err := db.ExecContext(ctx, "SELECT 1 from DUAL")
return err
},
dsn: "postgres://postgres:postgres@127.0.0.1:5432/fakeexecdb?sslmode=disable",
executed: []*regexp.Regexp{regexp.MustCompile("/\\*dddbs='test.db',dde='test-env',ddps='test-service',ddpv='1.0.0',traceparent='00-00000000000000000000000000000001-[\\da-f]{16}-01',ddh='127.0.0.1',dddb='fakeexecdb',ddprs='test-peer-service-custom-tag'\\*/ SELECT 1 from DUAL")},
peerServiceCustomOpenTag: "test-peer-service-custom-tag",
},
{
name: "exec-full-peer-service-precedence-tag-over-conn-context",
opts: []RegisterOption{WithDBMPropagation(tracer.DBMPropagationModeFull)},
callDB: func(ctx context.Context, db *sql.DB) error {
_, err := db.ExecContext(ctx, "SELECT 1 from DUAL")
return err
},
dsn: "postgres://postgres:postgres@127.0.0.1:5432/fakeexecdb?sslmode=disable",
executed: []*regexp.Regexp{regexp.MustCompile("/\\*dddbs='test.db',dde='test-env',ddps='test-service',ddpv='1.0.0',traceparent='00-00000000000000000000000000000001-[\\da-f]{16}-01',ddh='127.0.0.1',dddb='fakeexecdb',ddprs='test-peer-service-tag'\\*/ SELECT 1 from DUAL")},
peerServiceCtx: "test-peer-service-ctx",
peerServiceTag: "test-peer-service-tag",
},
{
name: "exec-full-peer-service-precedence-conn-context-over-open-custom-tag",
opts: []RegisterOption{WithDBMPropagation(tracer.DBMPropagationModeFull)},
callDB: func(ctx context.Context, db *sql.DB) error {
_, err := db.ExecContext(ctx, "SELECT 1 from DUAL")
return err
},
dsn: "postgres://postgres:postgres@127.0.0.1:5432/fakeexecdb?sslmode=disable",
executed: []*regexp.Regexp{regexp.MustCompile("/\\*dddbs='test.db',dde='test-env',ddps='test-service',ddpv='1.0.0',traceparent='00-00000000000000000000000000000001-[\\da-f]{16}-01',ddh='127.0.0.1',dddb='fakeexecdb',ddprs='test-peer-service-ctx'\\*/ SELECT 1 from DUAL")},
peerServiceCtx: "test-peer-service-ctx",
peerServiceCustomOpenTag: "test-peer-service-custom-tag",
},
}

Expand All @@ -171,9 +224,22 @@ func TestDBMPropagation(t *testing.T) {
if tc.dsn != "" {
dsn = tc.dsn
}
db, err := Open("test", dsn)
var options = []Option{}
if tc.peerServiceCustomOpenTag != "" {
options = append(options, WithCustomTag(ext.PeerService, tc.peerServiceCustomOpenTag))
}
db, err := Open("test", dsn, options...)
require.NoError(t, err)
s, ctx := tracer.StartSpanFromContext(context.Background(), "test.call", tracer.WithSpanID(1))
if tc.peerServiceCtx != "" {
vars := map[string]string{
ext.PeerService: tc.peerServiceCtx,
}
ctx = WithSpanTags(ctx, vars)
}
if tc.peerServiceTag != "" {
s.SetTag(ext.PeerService, tc.peerServiceTag)
}
err = tc.callDB(ctx, db)
s.Finish()

Expand Down
10 changes: 9 additions & 1 deletion ddtrace/tracer/sqlcomment.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -61,6 +61,8 @@ const (
// "Peer" is the OpenTelemetry nomenclature for "thing I am talking to"
sqlCommentPeerHostname = "ddh"
sqlCommentPeerDBName = "dddb"
// This is for when peer.service is explicitly set as a tag
sqlCommentPeerService = "ddprs"
)

// Current trace context version (see https://www.w3.org/TR/trace-context/#version)
Expand All @@ -76,6 +78,7 @@ type SQLCommentCarrier struct {
SpanID uint64
PeerDBHostname string
PeerDBName string
PeerService string
}

// Inject injects a span context in the carrier's Query field as a comment.
Expand Down Expand Up @@ -117,6 +120,11 @@ func (c *SQLCommentCarrier) Inject(spanCtx ddtrace.SpanContext) error {
if c.PeerDBHostname != "" {
tags[sqlCommentPeerHostname] = c.PeerDBHostname
}
if v, ok := ctx.meta(ext.PeerService); ok && v != "" {
tags[sqlCommentPeerService] = v
} else if c.PeerService != "" {
tags[sqlCommentPeerService] = c.PeerService
}
}
if globalconfig.ServiceName() != "" {
tags[sqlCommentParentService] = globalconfig.ServiceName()
Expand Down Expand Up @@ -167,7 +175,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{sqlCommentDBService, sqlCommentEnv, sqlCommentParentService, sqlCommentParentVersion, sqlCommentTraceParent, sqlCommentPeerHostname, sqlCommentPeerDBName}
orderedKeys := []string{sqlCommentDBService, sqlCommentEnv, sqlCommentParentService, sqlCommentParentVersion, sqlCommentTraceParent, sqlCommentPeerHostname, sqlCommentPeerDBName, sqlCommentPeerService}
first := true
for _, k := range orderedKeys {
if v, ok := tags[k]; ok {
Expand Down
40 changes: 31 additions & 9 deletions ddtrace/tracer/sqlcomment_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,8 +26,9 @@ func TestSQLCommentCarrier(t *testing.T) {
mode DBMPropagationMode
injectSpan bool
samplingPriority int
peerDBHostname string
peerDBName string
peerDBHostname string
peerServiceName string
expectedQuery string
expectedSpanIDGen bool
expectedExtractErr error
Expand All @@ -37,8 +38,9 @@ func TestSQLCommentCarrier(t *testing.T) {
query: "SELECT * from FOO",
mode: DBMPropagationModeFull,
injectSpan: true,
peerDBHostname: "",
peerDBName: "",
peerDBHostname: "",
peerServiceName: "",
expectedQuery: "/*dddbs='whiskey-db',dde='test-env',ddps='whiskey-service%20%21%23%24%25%26%27%28%29%2A%2B%2C%2F%3A%3B%3D%3F%40%5B%5D',ddpv='1.0.0',traceparent='00-0000000000000000000000000000000a-<span_id>-00'*/ SELECT * from FOO",
expectedSpanIDGen: true,
expectedExtractErr: nil,
Expand All @@ -48,8 +50,9 @@ func TestSQLCommentCarrier(t *testing.T) {
query: "SELECT * from FOO",
mode: DBMPropagationModeService,
injectSpan: true,
peerDBHostname: "",
peerDBName: "",
peerDBHostname: "",
peerServiceName: "",
expectedQuery: "/*dddbs='whiskey-db',dde='test-env',ddps='whiskey-service%20%21%23%24%25%26%27%28%29%2A%2B%2C%2F%3A%3B%3D%3F%40%5B%5D',ddpv='1.0.0'*/ SELECT * from FOO",
expectedSpanIDGen: false,
expectedExtractErr: ErrSpanContextNotFound,
Expand All @@ -58,8 +61,9 @@ func TestSQLCommentCarrier(t *testing.T) {
name: "no-trace",
query: "SELECT * from FOO",
mode: DBMPropagationModeFull,
peerDBHostname: "",
peerDBName: "",
peerDBHostname: "",
peerServiceName: "",
expectedQuery: "/*dddbs='whiskey-db',ddps='whiskey-service%20%21%23%24%25%26%27%28%29%2A%2B%2C%2F%3A%3B%3D%3F%40%5B%5D',traceparent='00-0000000000000000<span_id>-<span_id>-00'*/ SELECT * from FOO",
expectedSpanIDGen: true,
expectedExtractErr: nil,
Expand All @@ -69,8 +73,9 @@ func TestSQLCommentCarrier(t *testing.T) {
query: "",
mode: DBMPropagationModeFull,
injectSpan: true,
peerDBHostname: "",
peerDBName: "",
peerDBHostname: "",
peerServiceName: "",
expectedQuery: "/*dddbs='whiskey-db',dde='test-env',ddps='whiskey-service%20%21%23%24%25%26%27%28%29%2A%2B%2C%2F%3A%3B%3D%3F%40%5B%5D',ddpv='1.0.0',traceparent='00-0000000000000000000000000000000a-<span_id>-00'*/",
expectedSpanIDGen: true,
expectedExtractErr: nil,
Expand Down Expand Up @@ -101,8 +106,9 @@ func TestSQLCommentCarrier(t *testing.T) {
mode: DBMPropagationModeFull,
injectSpan: true,
samplingPriority: 1,
peerDBHostname: "",
peerDBName: "",
peerDBHostname: "",
peerServiceName: "",
expectedQuery: "/*dddbs='whiskey-db',dde='test-env',ddps='whiskey-service%20%21%23%24%25%26%27%28%29%2A%2B%2C%2F%3A%3B%3D%3F%40%5B%5D',ddpv='1.0.0',traceparent='00-0000000000000000000000000000000a-<span_id>-01'*/ /* c */ SELECT * from FOO /**/",
expectedSpanIDGen: true,
expectedExtractErr: nil,
Expand All @@ -115,6 +121,7 @@ func TestSQLCommentCarrier(t *testing.T) {
samplingPriority: 1,
peerDBName: "fake-database",
peerDBHostname: "",
peerServiceName: "",
expectedQuery: "/*dddbs='whiskey-db',dde='test-env',ddps='whiskey-service%20%21%23%24%25%26%27%28%29%2A%2B%2C%2F%3A%3B%3D%3F%40%5B%5D',ddpv='1.0.0',traceparent='00-0000000000000000000000000000000a-<span_id>-01',dddb='fake-database'*/ /* c */ SELECT * from FOO /**/",
expectedSpanIDGen: true,
expectedExtractErr: nil,
Expand All @@ -125,8 +132,9 @@ func TestSQLCommentCarrier(t *testing.T) {
mode: DBMPropagationModeFull,
injectSpan: true,
samplingPriority: 1,
peerDBHostname: "fake-hostname",
peerDBName: "",
peerDBHostname: "fake-hostname",
peerServiceName: "",
expectedQuery: "/*dddbs='whiskey-db',dde='test-env',ddps='whiskey-service%20%21%23%24%25%26%27%28%29%2A%2B%2C%2F%3A%3B%3D%3F%40%5B%5D',ddpv='1.0.0',traceparent='00-0000000000000000000000000000000a-<span_id>-01',ddh='fake-hostname'*/ /* c */ SELECT * from FOO /**/",
expectedSpanIDGen: true,
expectedExtractErr: nil,
Expand All @@ -137,12 +145,26 @@ func TestSQLCommentCarrier(t *testing.T) {
mode: DBMPropagationModeFull,
injectSpan: true,
samplingPriority: 1,
peerDBHostname: "fake-hostname",
peerDBName: "fake-database",
peerDBHostname: "fake-hostname",
peerServiceName: "",
expectedQuery: "/*dddbs='whiskey-db',dde='test-env',ddps='whiskey-service%20%21%23%24%25%26%27%28%29%2A%2B%2C%2F%3A%3B%3D%3F%40%5B%5D',ddpv='1.0.0',traceparent='00-0000000000000000000000000000000a-<span_id>-01',ddh='fake-hostname',dddb='fake-database'*/ /* c */ SELECT * from FOO /**/",
expectedSpanIDGen: true,
expectedExtractErr: nil,
},
{
name: "peer_entity_tags_peer_service",
query: "/* c */ SELECT * from FOO /**/",
mode: DBMPropagationModeFull,
injectSpan: true,
samplingPriority: 1,
peerDBName: "",
peerDBHostname: "",
peerServiceName: "test-peer-service",
expectedQuery: "/*dddbs='whiskey-db',dde='test-env',ddps='whiskey-service%20%21%23%24%25%26%27%28%29%2A%2B%2C%2F%3A%3B%3D%3F%40%5B%5D',ddpv='1.0.0',traceparent='00-0000000000000000000000000000000a-<span_id>-01',ddprs='test-peer-service'*/ /* c */ SELECT * from FOO /**/",
expectedSpanIDGen: true,
expectedExtractErr: nil,
},
}

for _, tc := range testCases {
Expand All @@ -162,7 +184,7 @@ func TestSQLCommentCarrier(t *testing.T) {
spanCtx = root.Context()
}

carrier := SQLCommentCarrier{Query: tc.query, Mode: tc.mode, DBServiceName: "whiskey-db", PeerDBHostname: tc.peerDBHostname, PeerDBName: tc.peerDBName}
carrier := SQLCommentCarrier{Query: tc.query, Mode: tc.mode, DBServiceName: "whiskey-db", PeerDBHostname: tc.peerDBHostname, PeerDBName: tc.peerDBName, PeerService: tc.peerServiceName}
err := carrier.Inject(spanCtx)
require.NoError(t, err)
expected := strings.ReplaceAll(tc.expectedQuery, "<span_id>", fmt.Sprintf("%016s", strconv.FormatUint(carrier.SpanID, 16)))
Expand Down

0 comments on commit b0aa1b8

Please sign in to comment.