Skip to content

Commit 325704c

Browse files
SpencerTorresSpencerTorres
authored
Merge pull request #1736 from ClickHouse/client_info_query_option
Add context option to append more ClientInfo to the system.query_log
2 parents be5a492 + 55bfc3e commit 325704c

File tree

9 files changed

+247
-25
lines changed

9 files changed

+247
-25
lines changed

‎clickhouse_std.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -203,7 +203,7 @@ func (std *stdDriver) Open(dsn string) (_ driver.Conn, err error) {
203203
if o.Debug {
204204
debugf = log.New(os.Stdout, "[clickhouse-std][opener] ", 0).Printf
205205
}
206-
o.ClientInfo.comment = []string{"database/sql"}
206+
o.ClientInfo.Comment = []string{"database/sql"}
207207
return (&stdConnOpener{opt: o, debugf: debugf}).Connect(context.Background())
208208
}
209209

‎client_info.go‎

Lines changed: 33 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,10 +2,11 @@ package clickhouse
22

33
import (
44
"fmt"
5-
"github.com/ClickHouse/clickhouse-go/v2/lib/proto"
65
"runtime"
76
"sort"
87
"strings"
8+
9+
"github.com/ClickHouse/clickhouse-go/v2/lib/proto"
910
)
1011

1112
const ClientName = "clickhouse-go"
@@ -23,7 +24,34 @@ type ClientInfo struct {
2324
Version string
2425
}
2526

26-
comment []string
27+
Comment []string
28+
}
29+
30+
// Append returns a new copy of the combined ClientInfo structs
31+
func (a ClientInfo) Append(b ClientInfo) ClientInfo {
32+
c := ClientInfo{
33+
Products: make([]struct {
34+
Name string
35+
Version string
36+
}, 0, len(a.Products)+len(b.Products)),
37+
Comment: make([]string, 0, len(a.Comment)+len(b.Comment)),
38+
}
39+
40+
for _, p := range a.Products {
41+
c.Products = append(c.Products, p)
42+
}
43+
for _, p := range b.Products {
44+
c.Products = append(c.Products, p)
45+
}
46+
47+
for _, cm := range a.Comment {
48+
c.Comment = append(c.Comment, cm)
49+
}
50+
for _, cm := range b.Comment {
51+
c.Comment = append(c.Comment, cm)
52+
}
53+
54+
return c
2755
}
2856

2957
func (o ClientInfo) String() string {
@@ -45,7 +73,9 @@ func (o ClientInfo) String() string {
4573
lvMeta := "lv:go/" + runtime.Version()[2:]
4674
osMeta := "os:" + runtime.GOOS
4775

48-
chunks := append(info.comment, lvMeta, osMeta) // nolint:gocritic
76+
chunks := make([]string, 0, len(info.Comment) + 2)
77+
chunks = append(chunks, info.Comment...)
78+
chunks = append(chunks, lvMeta, osMeta)
4979

5080
s.WriteByte(' ')
5181
s.WriteByte('(')

‎client_info_test.go‎

Lines changed: 61 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,11 +2,69 @@ package clickhouse
22

33
import (
44
"fmt"
5-
"github.com/stretchr/testify/assert"
65
"runtime"
76
"testing"
7+
8+
"github.com/stretchr/testify/assert"
9+
"github.com/stretchr/testify/require"
810
)
911

12+
func TestClientInfoAppend(t *testing.T) {
13+
a := ClientInfo{
14+
Products: []struct {
15+
Name string
16+
Version string
17+
}{
18+
{
19+
Name: "product",
20+
Version: "1.0.0",
21+
},
22+
},
23+
Comment: []string{"comment_a"},
24+
}
25+
26+
b := ClientInfo{
27+
Products: []struct {
28+
Name string
29+
Version string
30+
}{
31+
{
32+
Name: "product2",
33+
Version: "2.0.0",
34+
},
35+
},
36+
Comment: []string{"comment_b"},
37+
}
38+
39+
c := a.Append(b)
40+
41+
// Check first ClientInfo unchanged
42+
require.Len(t, a.Products, 1)
43+
require.Equal(t, "product", a.Products[0].Name)
44+
require.Equal(t, "1.0.0", a.Products[0].Version)
45+
require.Len(t, a.Comment, 1)
46+
require.Equal(t, "comment_a", a.Comment[0])
47+
48+
// Check second ClientInfo unchanged
49+
require.Len(t, b.Products, 1)
50+
require.Equal(t, "product2", b.Products[0].Name)
51+
require.Equal(t, "2.0.0", b.Products[0].Version)
52+
require.Len(t, b.Comment, 1)
53+
require.Equal(t, "comment_b", b.Comment[0])
54+
55+
// Verify third ClientInfo is merged correctly
56+
require.Len(t, c.Products, 2)
57+
require.Equal(t, "product", c.Products[0].Name)
58+
require.Equal(t, "1.0.0", c.Products[0].Version)
59+
require.Equal(t, "product2", c.Products[1].Name)
60+
require.Equal(t, "2.0.0", c.Products[1].Version)
61+
62+
require.Len(t, c.Comment, 2)
63+
require.Equal(t, "comment_a", c.Comment[0])
64+
require.Equal(t, "comment_b", c.Comment[1])
65+
66+
}
67+
1068
func TestClientInfoString(t *testing.T) {
1169
// e.g. clickhouse-go/2.5.1
1270
expectedClientProduct := fmt.Sprintf("%s/%d.%d.%d", ClientName, ClientVersionMajor, ClientVersionMinor, ClientVersionPatch)
@@ -25,7 +83,7 @@ func TestClientInfoString(t *testing.T) {
2583
},
2684
"client with comment": {
2785
ClientInfo{
28-
comment: []string{"database/sql"},
86+
Comment: []string{"database/sql"},
2987
},
3088
// e.g. clickhouse-go/2.5.1 (database/sql; lv:go/1.19.5; os:darwin)
3189
fmt.Sprintf("%s (database/sql; %s)", expectedClientProduct, expectedDefaultMeta),
@@ -51,7 +109,7 @@ func TestClientInfoString(t *testing.T) {
51109
{Name: "grafana", Version: "6.1"},
52110
{Name: "grafana-datasource", Version: "0.1.1"},
53111
},
54-
comment: []string{"database/sql"},
112+
Comment: []string{"database/sql"},
55113
},
56114
// e.g. grafana/6.1 grafana-datasource/0.1.1 clickhouse-go/2.5.1 (database/sql; lv:go/1.19.5; os:darwin)
57115
fmt.Sprintf("grafana/6.1 grafana-datasource/0.1.1 %s (database/sql; %s)", expectedClientProduct, expectedDefaultMeta),

‎conn_http.go‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -103,7 +103,9 @@ func (rw *HTTPReaderWriter) reset(pw *io.PipeWriter) io.WriteCloser {
103103

104104
// applyOptionsToRequest applies the client Options (such as auth, headers, client info) to the given http.Request
105105
func applyOptionsToRequest(ctx context.Context, req *http.Request, opt *Options) error {
106-
jwt := queryOptionsJWT(ctx)
106+
queryOpt := queryOptions(ctx)
107+
108+
jwt := queryOpt.jwt
107109
useJWT := jwt != "" || useJWTAuth(opt)
108110

109111
if opt.TLS != nil && useJWT {
@@ -133,7 +135,7 @@ func applyOptionsToRequest(ctx context.Context, req *http.Request, opt *Options)
133135
}
134136
}
135137

136-
req.Header.Set("User-Agent", opt.ClientInfo.String())
138+
req.Header.Set("User-Agent", opt.ClientInfo.Append(queryOpt.clientInfo).String())
137139

138140
for k, v := range opt.HttpHeaders {
139141
req.Header.Set(k, v)

‎conn_send_query.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@ func (c *connect) sendQuery(body string, o *QueryOptions) error {
1111
c.buffer.PutByte(proto.ClientQuery)
1212
q := proto.Query{
1313
ClientTCPProtocolVersion: ClientTCPProtocolVersion,
14-
ClientName: c.opt.ClientInfo.String(),
14+
ClientName: c.opt.ClientInfo.Append(o.clientInfo).String(),
1515
ClientVersion: proto.Version{ClientVersionMajor, ClientVersionMinor, ClientVersionPatch}, //nolint:govet
1616
ID: o.queryID,
1717
Body: body,

‎context.go‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,7 @@ type (
5656
blockBufferSize uint8
5757
userLocation *time.Location
5858
columnNamesAndTypes []ColumnNameAndType
59+
clientInfo ClientInfo
5960
}
6061
)
6162

@@ -107,6 +108,16 @@ func WithColumnNamesAndTypes(columnNamesAndTypes []ColumnNameAndType) QueryOptio
107108
}
108109
}
109110

111+
// WithClientInfo appends client info data to the query, visible in the system.query_log table.
112+
// This does not replace the client info provided in the connection options, it appends to it.
113+
// Can be called multiple times to append more info.
114+
func WithClientInfo(ci ClientInfo) QueryOption {
115+
return func(o *QueryOptions) error {
116+
o.clientInfo = o.clientInfo.Append(ci)
117+
return nil
118+
}
119+
}
120+
110121
func WithSettings(settings Settings) QueryOption {
111122
return func(o *QueryOptions) error {
112123
o.settings = settings
@@ -318,5 +329,9 @@ func (q *QueryOptions) clone() QueryOptions {
318329
c.columnNamesAndTypes = slices.Clone(q.columnNamesAndTypes)
319330
}
320331

332+
if q.clientInfo.Products != nil || q.clientInfo.Comment != nil {
333+
c.clientInfo = q.clientInfo.Append(ClientInfo{})
334+
}
335+
321336
return c
322337
}

‎context_test.go‎

Lines changed: 58 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,9 +2,10 @@ package clickhouse
22

33
import (
44
"context"
5-
"github.com/stretchr/testify/require"
65
"testing"
76
"time"
7+
8+
"github.com/stretchr/testify/require"
89
)
910

1011
func TestContext(t *testing.T) {
@@ -124,4 +125,60 @@ func TestContext(t *testing.T) {
124125
require.Nil(t, loc)
125126
},
126127
)
128+
129+
t.Run("correctly appends client info on multiple calls",
130+
func(t *testing.T) {
131+
// First context
132+
firstContext := Context(context.Background(), WithClientInfo(ClientInfo{
133+
Products: []struct {
134+
Name string
135+
Version string
136+
}{
137+
{
138+
Name: "product",
139+
Version: "1.0.0",
140+
},
141+
},
142+
Comment: []string{"comment_a"},
143+
}))
144+
145+
firstOpts := queryOptions(firstContext)
146+
require.Len(t, firstOpts.clientInfo.Products, 1)
147+
require.Equal(t, "product", firstOpts.clientInfo.Products[0].Name)
148+
require.Equal(t, "1.0.0", firstOpts.clientInfo.Products[0].Version)
149+
require.Len(t, firstOpts.clientInfo.Comment, 1)
150+
require.Equal(t, "comment_a", firstOpts.clientInfo.Comment[0])
151+
152+
// Second context
153+
secondContext := Context(firstContext, WithClientInfo(ClientInfo{
154+
Products: []struct {
155+
Name string
156+
Version string
157+
}{
158+
{
159+
Name: "product2",
160+
Version: "2.0.0",
161+
},
162+
},
163+
Comment: []string{"comment_b"},
164+
}))
165+
166+
// Product and comment values should be merged from the first+second contexts
167+
secondOpts := queryOptions(secondContext)
168+
169+
// Check first context values still present
170+
require.Len(t, secondOpts.clientInfo.Products, 2)
171+
require.Equal(t, "product", secondOpts.clientInfo.Products[0].Name)
172+
require.Equal(t, "1.0.0", secondOpts.clientInfo.Products[0].Version)
173+
require.Len(t, secondOpts.clientInfo.Comment, 2)
174+
require.Equal(t, "comment_a", secondOpts.clientInfo.Comment[0])
175+
176+
// Check second context values present
177+
require.Len(t, secondOpts.clientInfo.Products, 2)
178+
require.Equal(t, "product2", secondOpts.clientInfo.Products[1].Name)
179+
require.Equal(t, "2.0.0", secondOpts.clientInfo.Products[1].Version)
180+
require.Len(t, secondOpts.clientInfo.Comment, 2)
181+
require.Equal(t, "comment_b", secondOpts.clientInfo.Comment[1])
182+
},
183+
)
127184
}

‎tests/client_info_test.go‎

Lines changed: 34 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -25,16 +25,19 @@ func TestClientInfo(t *testing.T) {
2525

2626
testCases := map[string]struct {
2727
expectedClientInfo string
28+
ctx context.Context
2829
clientInfo clickhouse.ClientInfo
2930
}{
3031
"no additional products": {
3132
// e.g. clickhouse-go/2.5.1 (database/sql; lv:go/1.19.3; os:darwin)
3233
expectedClientProduct,
34+
context.Background(),
3335
clickhouse.ClientInfo{},
3436
},
3537
"one additional product": {
3638
// e.g. tests/dev clickhouse-go/2.5.1 (database/sql; lv:go/1.19.3; os:darwin)
3739
fmt.Sprintf("tests/dev %s", expectedClientProduct),
40+
context.Background(),
3841
clickhouse.ClientInfo{
3942
Products: []struct {
4043
Name string
@@ -50,6 +53,7 @@ func TestClientInfo(t *testing.T) {
5053
"two additional products": {
5154
// e.g. product/version tests/dev clickhouse-go/2.5.1 (database/sql; lv:go/1.19.3; os:darwin)
5255
fmt.Sprintf("product/version tests/dev %s", expectedClientProduct),
56+
context.Background(),
5357
clickhouse.ClientInfo{
5458
Products: []struct {
5559
Name string
@@ -66,6 +70,31 @@ func TestClientInfo(t *testing.T) {
6670
},
6771
},
6872
},
73+
"additional product from context": {
74+
// e.g. ctxProduct/1.2.3 clickhouse-go/2.41.0 (ctxComment; lv:go/1.25.5; os:linux)
75+
fmt.Sprintf(
76+
"ctxProduct/1.2.3 %s/%d.%d.%d (ctxComment; lv:go/%s; os:%s)",
77+
clickhouse.ClientName,
78+
clickhouse.ClientVersionMajor,
79+
clickhouse.ClientVersionMinor,
80+
clickhouse.ClientVersionPatch,
81+
runtime.Version()[2:],
82+
runtime.GOOS,
83+
),
84+
clickhouse.Context(context.Background(), clickhouse.WithClientInfo(clickhouse.ClientInfo{
85+
Products: []struct {
86+
Name string
87+
Version string
88+
}{
89+
{
90+
Name: "ctxProduct",
91+
Version: "1.2.3",
92+
},
93+
},
94+
Comment: []string{"ctxComment"},
95+
})),
96+
clickhouse.ClientInfo{},
97+
},
6998
}
7099

71100
env, err := GetTestEnvironment(testSet)
@@ -79,23 +108,23 @@ func TestClientInfo(t *testing.T) {
79108
conn, err := clickhouse.Open(&opts)
80109
require.NoError(t, err)
81110

82-
actualClientInfo := getConnectedClientInfo(t, conn)
111+
actualClientInfo := getConnectedClientInfo(t, conn, testCase.ctx)
83112
assert.Equal(t, testCase.expectedClientInfo, actualClientInfo)
84113
})
85114
}
86115
}
87116

88-
func getConnectedClientInfo(t *testing.T, conn driver.Conn) string {
117+
func getConnectedClientInfo(t *testing.T, conn driver.Conn, ctx context.Context) string {
89118
var queryID string
90-
row := conn.QueryRow(context.TODO(), "SELECT queryID()")
119+
row := conn.QueryRow(ctx, "SELECT queryID()")
91120
require.NoError(t, row.Err())
92121
require.NoError(t, row.Scan(&queryID))
93122

94-
err := conn.Exec(context.TODO(), "SYSTEM FLUSH LOGS")
123+
err := conn.Exec(ctx, "SYSTEM FLUSH LOGS")
95124
require.NoError(t, err)
96125

97126
var clientName string
98-
row = conn.QueryRow(context.TODO(), fmt.Sprintf("SELECT IF(interface = 2, http_user_agent, client_name) as client_name FROM system.query_log WHERE query_id = '%s'", queryID))
127+
row = conn.QueryRow(ctx, fmt.Sprintf("SELECT IF(interface = 2, http_user_agent, client_name) as client_name FROM system.query_log WHERE query_id = '%s'", queryID))
99128
require.NoError(t, row.Err())
100129
require.NoError(t, row.Scan(&clientName))
101130

0 commit comments

Comments
 (0)