-
Notifications
You must be signed in to change notification settings - Fork 20
/
tracer_test.go
133 lines (126 loc) · 3.27 KB
/
tracer_test.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
package otelpgx
import (
"strings"
"testing"
)
func TestTracer_sqlOperationName(t *testing.T) {
tests := []struct {
name string
tracer *Tracer
query string
expName string
}{
{
name: "Spaces only",
query: "SELECT * FROM users",
tracer: NewTracer(),
expName: "SELECT",
},
{
name: "Newline and tab",
query: "UPDATE\n\tfoo",
tracer: NewTracer(),
expName: "UPDATE",
},
{
name: "Additional whitespace",
query: " \n SELECT\n\t * FROM users ",
tracer: NewTracer(),
expName: "SELECT",
},
{
name: "Whitespace-only query",
query: " \n\t",
tracer: NewTracer(),
expName: sqlOperationUnknown,
},
{
name: "Empty query",
query: "",
tracer: NewTracer(),
expName: sqlOperationUnknown,
},
{
name: "Functional span name (-- comment style)",
query: "-- name: GetUsers :many\nSELECT * FROM users",
tracer: NewTracer(WithSpanNameFunc(defaultSpanNameFunc)),
expName: "GetUsers :many",
},
{
name: "Functional span name (/**/ comment style)",
query: "/* name: GetBooks :many */\nSELECT * FROM books",
tracer: NewTracer(WithSpanNameFunc(defaultSpanNameFunc)),
expName: "GetBooks :many",
},
{
name: "Functional span name (# comment style)",
query: "# name: GetRecords :many\nSELECT * FROM records",
tracer: NewTracer(WithSpanNameFunc(defaultSpanNameFunc)),
expName: "GetRecords :many",
},
{
name: "Functional span name (no annotation)",
query: "--\nSELECT * FROM user",
tracer: NewTracer(WithSpanNameFunc(defaultSpanNameFunc)),
expName: sqlOperationUnknown,
},
{
name: "Custom SQL name query (normal comment)",
query: "-- foo \nSELECT * FROM users",
tracer: NewTracer(WithSpanNameFunc(defaultSpanNameFunc)),
expName: sqlOperationUnknown,
},
{
name: "Custom SQL name query (invalid formatting)",
query: "foo \nSELECT * FROM users",
tracer: NewTracer(WithSpanNameFunc(defaultSpanNameFunc)),
expName: sqlOperationUnknown,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
tr := tt.tracer
if got := tr.sqlOperationName(tt.query); got != tt.expName {
t.Errorf("Tracer.sqlOperationName() = %v, want %v", got, tt.expName)
}
})
}
}
// defaultSpanNameFunc is an utility function for testing that attempts to get
// the first name of the query from a given SQL statement.
var defaultSpanNameFunc SpanNameFunc = func(query string) string {
for _, line := range strings.Split(query, "\n") {
var prefix string
switch {
case strings.HasPrefix(line, "--"):
prefix = "--"
case strings.HasPrefix(line, "/*"):
prefix = "/*"
case strings.HasPrefix(line, "#"):
prefix = "#"
default:
continue
}
rest := line[len(prefix):]
if !strings.HasPrefix(strings.TrimSpace(rest), "name") {
continue
}
if !strings.Contains(rest, ":") {
continue
}
if !strings.HasPrefix(rest, " name: ") {
return sqlOperationUnknown
}
part := strings.Split(strings.TrimSpace(line), " ")
if prefix == "/*" {
part = part[:len(part)-1] // removes the trailing "*/" element
}
if len(part) == 2 {
return sqlOperationUnknown
}
queryName := part[2]
queryType := strings.TrimSpace(part[3])
return queryName + " " + queryType
}
return sqlOperationUnknown
}