Skip to content

Commit

Permalink
contrib/database/sql: trace connection time (#1154)
Browse files Browse the repository at this point in the history 
Creating a connection to a sql database will now be included in traces

Fixes #760
  • Loading branch information
@ajgajg1134
ajgajg1134 committed Feb 4, 2022
1 parent f985744 commit d1b2aba
Show file tree
Hide file tree
Showing 5 changed files with 112 additions and 21 deletions.
3 changes: 2 additions & 1 deletion contrib/database/sql/conn.go
View file
Expand Up @@ -22,7 +22,8 @@ var _ driver.Conn = (*tracedConn)(nil)
type queryType string

const (
queryTypeQuery queryType = "Query"
queryTypeConnect queryType = "Connect"
queryTypeQuery = "Query"
queryTypePing = "Ping"
queryTypePrepare = "Prepare"
queryTypeExec = "Exec"
Expand Down
11 changes: 9 additions & 2 deletions contrib/database/sql/conn_test.go
View file
Expand Up @@ -92,9 +92,16 @@ func TestWithSpanTags(t *testing.T) {
rows.Close()

spans := mt.FinishedSpans()
assert.Len(t, spans, 1)
assert.Len(t, spans, 2)

span := spans[0]
connectSpan := spans[0]
assert.Equal(t, tt.want.opName, connectSpan.OperationName())
assert.Equal(t, "Connect", connectSpan.Tag("sql.query_type"))
for k, v := range tt.want.ctxTags {
assert.Equal(t, v, connectSpan.Tag(k), "Value mismatch on tag %s", k)
}

span := spans[1]
assert.Equal(t, tt.want.opName, span.OperationName())
for k, v := range tt.want.ctxTags {
assert.Equal(t, v, span.Tag(k), "Value mismatch on tag %s", k)
Expand Down
17 changes: 10 additions & 7 deletions contrib/database/sql/sql.go
View file
Expand Up @@ -23,6 +23,7 @@ import (
"errors"
"math"
"reflect"
"time"

"gopkg.in/DataDog/dd-trace-go.v1/contrib/database/sql/internal"
"gopkg.in/DataDog/dd-trace-go.v1/internal/log"
Expand Down Expand Up @@ -128,11 +129,7 @@ type tracedConnector struct {
cfg *config
}

func (t *tracedConnector) Connect(c context.Context) (driver.Conn, error) {
conn, err := t.connector.Connect(c)
if err != nil {
return nil, err
}
func (t *tracedConnector) Connect(ctx context.Context) (driver.Conn, error) {
tp := &traceParams{
driverName: t.driverName,
cfg: t.cfg,
Expand All @@ -142,6 +139,12 @@ func (t *tracedConnector) Connect(c context.Context) (driver.Conn, error) {
} else if t.cfg.dsn != "" {
tp.meta, _ = internal.ParseDSN(t.driverName, t.cfg.dsn)
}
start := time.Now()
conn, err := t.connector.Connect(ctx)
tp.tryTrace(ctx, queryTypeConnect, "", start, err)
if err != nil {
return nil, err
}
return &tracedConn{conn, tp}, err
}

Expand All @@ -163,7 +166,7 @@ func (t dsnConnector) Driver() driver.Driver {
return t.driver
}

// OpenDB returns connection to a DB using a the traced version of the given driver. In order for OpenDB
// OpenDB returns connection to a DB using the traced version of the given driver. In order for OpenDB
// to work, the driver must first be registered using Register. If this did not occur, OpenDB will panic.
func OpenDB(c driver.Connector, opts ...Option) *sql.DB {
name, ok := registeredDrivers.name(c.Driver())
Expand Down Expand Up @@ -192,7 +195,7 @@ func OpenDB(c driver.Connector, opts ...Option) *sql.DB {
return sql.OpenDB(tc)
}

// Open returns connection to a DB using a the traced version of the given driver. In order for Open
// Open returns connection to a DB using the traced version of the given driver. In order for Open
// to work, the driver must first be registered using Register. If this did not occur, Open will
// return an error.
func Open(driverName, dataSourceName string, opts ...Option) (*sql.DB, error) {
Expand Down
44 changes: 44 additions & 0 deletions contrib/database/sql/sql_test.go
View file
Expand Up @@ -6,14 +6,19 @@
package sql

import (
"context"
"database/sql/driver"
"errors"
"fmt"
"log"
"math"
"os"
"testing"
"time"

"gopkg.in/DataDog/dd-trace-go.v1/contrib/internal/sqltest"
"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/ext"
"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/mocktracer"

"github.com/go-sql-driver/mysql"
"github.com/lib/pq"
Expand Down Expand Up @@ -190,3 +195,42 @@ func TestMySQLUint64(t *testing.T) {
assert.NoError(rows.Err())
assert.NoError(rows.Close())
}

// hangingConnector hangs on Connect until ctx is cancelled.
type hangingConnector struct{}

func (h *hangingConnector) Connect(ctx context.Context) (driver.Conn, error) {
select {
case <-ctx.Done():
return nil, errors.New("context cancelled")
}
}

func (h *hangingConnector) Driver() driver.Driver {
panic("hangingConnector: Driver() not implemented")
}

func TestConnectCancelledCtx(t *testing.T) {
mockTracer := mocktracer.Start()
defer mockTracer.Stop()
assert := assert.New(t)
tc := tracedConnector{
connector: &hangingConnector{},
driverName: "hangingConnector",
cfg: new(config),
}
ctx, cancelFunc := context.WithCancel(context.Background())

go func() {
tc.Connect(ctx)
}()
time.Sleep(time.Millisecond * 100)
cancelFunc()
time.Sleep(time.Millisecond * 100)

spans := mockTracer.FinishedSpans()
assert.Len(spans, 1)
s := spans[0]
assert.Equal("hangingConnector.query", s.OperationName())
assert.Equal("Connect", s.Tag("sql.query_type"))
}
58 changes: 47 additions & 11 deletions contrib/internal/sqltest/sqltest.go
View file
Expand Up @@ -48,8 +48,10 @@ func Prepare(tableName string) func() {
func RunAll(t *testing.T, cfg *Config) {
cfg.mockTracer = mocktracer.Start()
defer cfg.mockTracer.Stop()
cfg.DB.SetMaxIdleConns(0)

for name, test := range map[string]func(*Config) func(*testing.T){
"Connect": testConnect,
"Ping": testPing,
"Query": testQuery,
"Statement": testStatement,
Expand All @@ -60,17 +62,37 @@ func RunAll(t *testing.T, cfg *Config) {
}
}

func testPing(cfg *Config) func(*testing.T) {
func testConnect(cfg *Config) func(*testing.T) {
return func(t *testing.T) {
cfg.mockTracer.Reset()
assert := assert.New(t)
err := cfg.DB.Ping()
assert.Nil(err)
spans := cfg.mockTracer.FinishedSpans()
assert.Len(spans, 1)
assert.Len(spans, 2)

span := spans[0]
assert.Equal(cfg.ExpectName, span.OperationName())
cfg.ExpectTags["sql.query_type"] = "Connect"
for k, v := range cfg.ExpectTags {
assert.Equal(v, span.Tag(k), "Value mismatch on tag %s", k)
}
}
}

func testPing(cfg *Config) func(*testing.T) {
return func(t *testing.T) {
cfg.mockTracer.Reset()
assert := assert.New(t)
err := cfg.DB.Ping()
assert.Nil(err)
spans := cfg.mockTracer.FinishedSpans()
assert.Len(spans, 2)

verifyConnectSpan(spans[0], assert, cfg)

span := spans[1]
assert.Equal(cfg.ExpectName, span.OperationName())
cfg.ExpectTags["sql.query_type"] = "Ping"
for k, v := range cfg.ExpectTags {
assert.Equal(v, span.Tag(k), "Value mismatch on tag %s", k)
Expand All @@ -88,9 +110,11 @@ func testQuery(cfg *Config) func(*testing.T) {
assert.Nil(err)

spans := cfg.mockTracer.FinishedSpans()
assert.Len(spans, 1)
assert.Len(spans, 2)

span := spans[0]
verifyConnectSpan(spans[0], assert, cfg)

span := spans[1]
cfg.ExpectTags["sql.query_type"] = "Query"
assert.Equal(cfg.ExpectName, span.OperationName())
for k, v := range cfg.ExpectTags {
Expand All @@ -114,9 +138,11 @@ func testStatement(cfg *Config) func(*testing.T) {
assert.Equal(nil, err)

spans := cfg.mockTracer.FinishedSpans()
assert.Len(spans, 1)
assert.Len(spans, 3)

span := spans[0]
verifyConnectSpan(spans[0], assert, cfg)

span := spans[1]
assert.Equal(cfg.ExpectName, span.OperationName())
cfg.ExpectTags["sql.query_type"] = "Prepare"
for k, v := range cfg.ExpectTags {
Expand All @@ -128,8 +154,8 @@ func testStatement(cfg *Config) func(*testing.T) {
assert.Equal(nil, err2)

spans = cfg.mockTracer.FinishedSpans()
assert.Len(spans, 1)
span = spans[0]
assert.Len(spans, 4)
span = spans[2]
assert.Equal(cfg.ExpectName, span.OperationName())
cfg.ExpectTags["sql.query_type"] = "Exec"
for k, v := range cfg.ExpectTags {
Expand All @@ -147,9 +173,11 @@ func testBeginRollback(cfg *Config) func(*testing.T) {
assert.Equal(nil, err)

spans := cfg.mockTracer.FinishedSpans()
assert.Len(spans, 1)
assert.Len(spans, 2)

span := spans[0]
verifyConnectSpan(spans[0], assert, cfg)

span := spans[1]
assert.Equal(cfg.ExpectName, span.OperationName())
cfg.ExpectTags["sql.query_type"] = "Begin"
for k, v := range cfg.ExpectTags {
Expand Down Expand Up @@ -192,7 +220,7 @@ func testExec(cfg *Config) func(*testing.T) {
parent.Finish() // flush children

spans := cfg.mockTracer.FinishedSpans()
assert.Len(spans, 4)
assert.Len(spans, 5)

var span mocktracer.Span
for _, s := range spans {
Expand All @@ -218,6 +246,14 @@ func testExec(cfg *Config) func(*testing.T) {
}
}

func verifyConnectSpan(span mocktracer.Span, assert *assert.Assertions, cfg *Config) {
assert.Equal(cfg.ExpectName, span.OperationName())
cfg.ExpectTags["sql.query_type"] = "Connect"
for k, v := range cfg.ExpectTags {
assert.Equal(v, span.Tag(k), "Value mismatch on tag %s", k)
}
}

// Config holds the test configuration.
type Config struct {
*sql.DB
Expand Down

0 comments on commit d1b2aba

Please sign in to comment.