1// Unless explicitly stated otherwise all files in this repository are licensed
2// under the Apache License Version 2.0.
3// This product includes software developed at Datadog (https://www.datadoghq.com/).
4// Copyright 2016 Datadog, Inc.
5
6package sql
7
8import (
9	"context"
10	"database/sql/driver"
11	"log"
12	"testing"
13
14	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/mocktracer"
15
16	"github.com/go-sql-driver/mysql"
17	"github.com/lib/pq"
18	"github.com/stretchr/testify/assert"
19)
20
21func TestWithSpanTags(t *testing.T) {
22	type sqlRegister struct {
23		name   string
24		dsn    string
25		driver driver.Driver
26		opts   []RegisterOption
27	}
28	type want struct {
29		opName  string
30		ctxTags map[string]string
31	}
32	testcases := []struct {
33		name        string
34		sqlRegister sqlRegister
35		want        want
36	}{
37		{
38			name: "mysql",
39			sqlRegister: sqlRegister{
40				name:   "mysql",
41				dsn:    "test:test@tcp(127.0.0.1:3306)/test",
42				driver: &mysql.MySQLDriver{},
43				opts:   []RegisterOption{},
44			},
45			want: want{
46				opName: "mysql.query",
47				ctxTags: map[string]string{
48					"mysql_tag1": "mysql_value1",
49					"mysql_tag2": "mysql_value2",
50					"mysql_tag3": "mysql_value3",
51				},
52			},
53		},
54		{
55			name: "postgres",
56			sqlRegister: sqlRegister{
57				name:   "postgres",
58				dsn:    "postgres://postgres:postgres@127.0.0.1:5432/postgres?sslmode=disable",
59				driver: &pq.Driver{},
60				opts: []RegisterOption{
61					WithServiceName("postgres-test"),
62					WithAnalyticsRate(0.2),
63				},
64			},
65			want: want{
66				opName: "postgres.query",
67				ctxTags: map[string]string{
68					"pg_tag1": "pg_value1",
69					"pg_tag2": "pg_value2",
70				},
71			},
72		},
73	}
74	mt := mocktracer.Start()
75	defer mt.Stop()
76	for _, tt := range testcases {
77		t.Run(tt.name, func(t *testing.T) {
78			Register(tt.sqlRegister.name, tt.sqlRegister.driver, tt.sqlRegister.opts...)
79			db, err := Open(tt.sqlRegister.name, tt.sqlRegister.dsn)
80			if err != nil {
81				log.Fatal(err)
82			}
83			defer db.Close()
84			mt.Reset()
85
86			ctx := WithSpanTags(context.Background(), tt.want.ctxTags)
87
88			rows, err := db.QueryContext(ctx, "SELECT 1")
89			assert.NoError(t, err)
90			rows.Close()
91
92			spans := mt.FinishedSpans()
93			assert.Len(t, spans, 1)
94
95			span := spans[0]
96			assert.Equal(t, tt.want.opName, span.OperationName())
97			for k, v := range tt.want.ctxTags {
98				assert.Equal(t, v, span.Tag(k), "Value mismatch on tag %s", k)
99			}
100		})
101	}
102}
103