Skip to content

Commit 924f833

Browse files
Daemonxiaomethane
Daemonxiao
and
methane
authored
Send connection attributes (#1389)
Co-authored-by: Inada Naoki <[email protected]>
1 parent 72e78ee commit 924f833

12 files changed

+173
-28
lines changed

Diff for: ‎.github/workflows/test.yml

+1
Original file line numberDiff line numberDiff line change
@@ -79,6 +79,7 @@ jobs:
7979
; TestConcurrent fails if max_connections is too large
8080
max_connections=50
8181
local_infile=1
82+
performance_schema=on
8283
- name: setup database
8384
run: |
8485
mysql --user 'root' --host '127.0.0.1' -e 'create database gotest;'

Diff for: ‎README.md

+9
Original file line numberDiff line numberDiff line change
@@ -393,6 +393,15 @@ Default: 0
393393

394394
I/O write timeout. The value must be a decimal number with a unit suffix (*"ms"*, *"s"*, *"m"*, *"h"*), such as *"30s"*, *"0.5m"* or *"1m30s"*.
395395

396+
##### `connectionAttributes`
397+
398+
```
399+
Type: comma-delimited string of user-defined "key:value" pairs
400+
Valid Values: (<name1>:<value1>,<name2>:<value2>,...)
401+
Default: none
402+
```
403+
404+
[Connection attributes](https://dev.mysql.com/doc/refman/8.0/en/performance-schema-connection-attribute-tables.html) are key-value pairs that application programs can pass to the server at connect time.
396405

397406
##### System Variables
398407

Diff for: ‎connection.go

+1
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ type mysqlConn struct {
2727
affectedRows uint64
2828
insertId uint64
2929
cfg *Config
30+
connector *connector
3031
maxAllowedPacket int
3132
maxWriteSize int
3233
writeTimeout time.Duration

Diff for: ‎connector.go

+45-1
Original file line numberDiff line numberDiff line change
@@ -11,11 +11,54 @@ package mysql
1111
import (
1212
"context"
1313
"database/sql/driver"
14+
"fmt"
1415
"net"
16+
"os"
17+
"strconv"
18+
"strings"
1519
)
1620

1721
type connector struct {
18-
cfg *Config // immutable private copy.
22+
cfg *Config // immutable private copy.
23+
encodedAttributes string // Encoded connection attributes.
24+
}
25+
26+
func encodeConnectionAttributes(textAttributes string) string {
27+
connAttrsBuf := make([]byte, 0, 251)
28+
29+
// default connection attributes
30+
connAttrsBuf = appendLengthEncodedString(connAttrsBuf, connAttrClientName)
31+
connAttrsBuf = appendLengthEncodedString(connAttrsBuf, connAttrClientNameValue)
32+
connAttrsBuf = appendLengthEncodedString(connAttrsBuf, connAttrOS)
33+
connAttrsBuf = appendLengthEncodedString(connAttrsBuf, connAttrOSValue)
34+
connAttrsBuf = appendLengthEncodedString(connAttrsBuf, connAttrPlatform)
35+
connAttrsBuf = appendLengthEncodedString(connAttrsBuf, connAttrPlatformValue)
36+
connAttrsBuf = appendLengthEncodedString(connAttrsBuf, connAttrPid)
37+
connAttrsBuf = appendLengthEncodedString(connAttrsBuf, strconv.Itoa(os.Getpid()))
38+
39+
// user-defined connection attributes
40+
for _, connAttr := range strings.Split(textAttributes, ",") {
41+
attr := strings.SplitN(connAttr, ":", 2)
42+
if len(attr) != 2 {
43+
continue
44+
}
45+
for _, v := range attr {
46+
connAttrsBuf = appendLengthEncodedString(connAttrsBuf, v)
47+
}
48+
}
49+
50+
return string(connAttrsBuf)
51+
}
52+
53+
func newConnector(cfg *Config) (*connector, error) {
54+
encodedAttributes := encodeConnectionAttributes(cfg.ConnectionAttributes)
55+
if len(encodedAttributes) > 250 {
56+
return nil, fmt.Errorf("connection attributes are longer than 250 bytes: %dbytes (%q)", len(encodedAttributes), cfg.ConnectionAttributes)
57+
}
58+
return &connector{
59+
cfg: cfg,
60+
encodedAttributes: encodedAttributes,
61+
}, nil
1962
}
2063

2164
// Connect implements driver.Connector interface.
@@ -29,6 +72,7 @@ func (c *connector) Connect(ctx context.Context) (driver.Conn, error) {
2972
maxWriteSize: maxPacketSize - 1,
3073
closech: make(chan struct{}),
3174
cfg: c.cfg,
75+
connector: c,
3276
}
3377
mc.parseTime = mc.cfg.ParseTime
3478

Diff for: ‎connector_test.go

+6-3
Original file line numberDiff line numberDiff line change
@@ -8,13 +8,16 @@ import (
88
)
99

1010
func TestConnectorReturnsTimeout(t *testing.T) {
11-
connector := &connector{&Config{
11+
connector, err := newConnector(&Config{
1212
Net: "tcp",
1313
Addr: "1.1.1.1:1234",
1414
Timeout: 10 * time.Millisecond,
15-
}}
15+
})
16+
if err != nil {
17+
t.Fatal(err)
18+
}
1619

17-
_, err := connector.Connect(context.Background())
20+
_, err = connector.Connect(context.Background())
1821
if err == nil {
1922
t.Fatal("error expected")
2023
}

Diff for: ‎const.go

+12
Original file line numberDiff line numberDiff line change
@@ -8,12 +8,24 @@
88

99
package mysql
1010

11+
import "runtime"
12+
1113
const (
1214
defaultAuthPlugin = "mysql_native_password"
1315
defaultMaxAllowedPacket = 64 << 20 // 64 MiB. See https://github.com/go-sql-driver/mysql/issues/1355
1416
minProtocolVersion = 10
1517
maxPacketSize = 1<<24 - 1
1618
timeFormat = "2006-01-02 15:04:05.999999"
19+
20+
// Connection attributes
21+
// See https://dev.mysql.com/doc/refman/8.0/en/performance-schema-connection-attribute-tables.html#performance-schema-connection-attributes-available
22+
connAttrClientName = "_client_name"
23+
connAttrClientNameValue = "Go-MySQL-Driver"
24+
connAttrOS = "_os"
25+
connAttrOSValue = runtime.GOOS
26+
connAttrPlatform = "_platform"
27+
connAttrPlatformValue = runtime.GOARCH
28+
connAttrPid = "_pid"
1729
)
1830

1931
// MySQL constants documentation:

Diff for: ‎driver.go

+5-6
Original file line numberDiff line numberDiff line change
@@ -85,8 +85,9 @@ func (d MySQLDriver) Open(dsn string) (driver.Conn, error) {
8585
if err != nil {
8686
return nil, err
8787
}
88-
c := &connector{
89-
cfg: cfg,
88+
c, err := newConnector(cfg)
89+
if err != nil {
90+
return nil, err
9091
}
9192
return c.Connect(context.Background())
9293
}
@@ -103,7 +104,7 @@ func NewConnector(cfg *Config) (driver.Connector, error) {
103104
if err := cfg.normalize(); err != nil {
104105
return nil, err
105106
}
106-
return &connector{cfg: cfg}, nil
107+
return newConnector(cfg)
107108
}
108109

109110
// OpenConnector implements driver.DriverContext.
@@ -112,7 +113,5 @@ func (d MySQLDriver) OpenConnector(dsn string) (driver.Connector, error) {
112113
if err != nil {
113114
return nil, err
114115
}
115-
return &connector{
116-
cfg: cfg,
117-
}, nil
116+
return newConnector(cfg)
118117
}

Diff for: ‎driver_test.go

+47
Original file line numberDiff line numberDiff line change
@@ -3214,3 +3214,50 @@ func TestConnectorTimeoutsWatchCancel(t *testing.T) {
32143214
t.Errorf("connection not closed")
32153215
}
32163216
}
3217+
3218+
func TestConnectionAttributes(t *testing.T) {
3219+
if !available {
3220+
t.Skipf("MySQL server not running on %s", netAddr)
3221+
}
3222+
3223+
attr1 := "attr1"
3224+
value1 := "value1"
3225+
attr2 := "foo"
3226+
value2 := "boo"
3227+
dsn += fmt.Sprintf("&connectionAttributes=%s:%s,%s:%s", attr1, value1, attr2, value2)
3228+
3229+
var db *sql.DB
3230+
if _, err := ParseDSN(dsn); err != errInvalidDSNUnsafeCollation {
3231+
db, err = sql.Open("mysql", dsn)
3232+
if err != nil {
3233+
t.Fatalf("error connecting: %s", err.Error())
3234+
}
3235+
defer db.Close()
3236+
}
3237+
3238+
dbt := &DBTest{t, db}
3239+
3240+
var attrValue string
3241+
queryString := "SELECT ATTR_VALUE FROM performance_schema.session_account_connect_attrs WHERE PROCESSLIST_ID = CONNECTION_ID() and ATTR_NAME = ?"
3242+
rows := dbt.mustQuery(queryString, connAttrClientName)
3243+
if rows.Next() {
3244+
rows.Scan(&attrValue)
3245+
if attrValue != connAttrClientNameValue {
3246+
dbt.Errorf("expected %q, got %q", connAttrClientNameValue, attrValue)
3247+
}
3248+
} else {
3249+
dbt.Errorf("no data")
3250+
}
3251+
rows.Close()
3252+
3253+
rows = dbt.mustQuery(queryString, attr2)
3254+
if rows.Next() {
3255+
rows.Scan(&attrValue)
3256+
if attrValue != value2 {
3257+
dbt.Errorf("expected %q, got %q", value2, attrValue)
3258+
}
3259+
} else {
3260+
dbt.Errorf("no data")
3261+
}
3262+
rows.Close()
3263+
}

Diff for: ‎dsn.go

+23-17
Original file line numberDiff line numberDiff line change
@@ -34,23 +34,24 @@ var (
3434
// If a new Config is created instead of being parsed from a DSN string,
3535
// the NewConfig function should be used, which sets default values.
3636
type Config struct {
37-
User string // Username
38-
Passwd string // Password (requires User)
39-
Net string // Network type
40-
Addr string // Network address (requires Net)
41-
DBName string // Database name
42-
Params map[string]string // Connection parameters
43-
Collation string // Connection collation
44-
Loc *time.Location // Location for time.Time values
45-
MaxAllowedPacket int // Max packet size allowed
46-
ServerPubKey string // Server public key name
47-
pubKey *rsa.PublicKey // Server public key
48-
TLSConfig string // TLS configuration name
49-
TLS *tls.Config // TLS configuration, its priority is higher than TLSConfig
50-
Timeout time.Duration // Dial timeout
51-
ReadTimeout time.Duration // I/O read timeout
52-
WriteTimeout time.Duration // I/O write timeout
53-
Logger Logger // Logger
37+
User string // Username
38+
Passwd string // Password (requires User)
39+
Net string // Network type
40+
Addr string // Network address (requires Net)
41+
DBName string // Database name
42+
Params map[string]string // Connection parameters
43+
ConnectionAttributes string // Connection Attributes, comma-delimited string of user-defined "key:value" pairs
44+
Collation string // Connection collation
45+
Loc *time.Location // Location for time.Time values
46+
MaxAllowedPacket int // Max packet size allowed
47+
ServerPubKey string // Server public key name
48+
pubKey *rsa.PublicKey // Server public key
49+
TLSConfig string // TLS configuration name
50+
TLS *tls.Config // TLS configuration, its priority is higher than TLSConfig
51+
Timeout time.Duration // Dial timeout
52+
ReadTimeout time.Duration // I/O read timeout
53+
WriteTimeout time.Duration // I/O write timeout
54+
Logger Logger // Logger
5455

5556
AllowAllFiles bool // Allow all files to be used with LOAD DATA LOCAL INFILE
5657
AllowCleartextPasswords bool // Allows the cleartext client side plugin
@@ -560,6 +561,11 @@ func parseDSNParams(cfg *Config, params string) (err error) {
560561
if err != nil {
561562
return
562563
}
564+
565+
// Connection attributes
566+
case "connectionAttributes":
567+
cfg.ConnectionAttributes = value
568+
563569
default:
564570
// lazy init
565571
if cfg.Params == nil {

Diff for: ‎packets.go

+13
Original file line numberDiff line numberDiff line change
@@ -285,6 +285,7 @@ func (mc *mysqlConn) writeHandshakeResponsePacket(authResp []byte, plugin string
285285
clientLocalFiles |
286286
clientPluginAuth |
287287
clientMultiResults |
288+
clientConnectAttrs |
288289
mc.flags&clientLongFlag
289290

290291
if mc.cfg.ClientFoundRows {
@@ -318,6 +319,13 @@ func (mc *mysqlConn) writeHandshakeResponsePacket(authResp []byte, plugin string
318319
pktLen += n + 1
319320
}
320321

322+
// 1 byte to store length of all key-values
323+
// NOTE: Actually, this is length encoded integer.
324+
// But we support only len(connAttrBuf) < 251 for now because takeSmallBuffer
325+
// doesn't support buffer size more than 4096 bytes.
326+
// TODO(methane): Rewrite buffer management.
327+
pktLen += 1 + len(mc.connector.encodedAttributes)
328+
321329
// Calculate packet length and get buffer with that size
322330
data, err := mc.buf.takeSmallBuffer(pktLen + 4)
323331
if err != nil {
@@ -394,6 +402,11 @@ func (mc *mysqlConn) writeHandshakeResponsePacket(authResp []byte, plugin string
394402
data[pos] = 0x00
395403
pos++
396404

405+
// Connection Attributes
406+
data[pos] = byte(len(mc.connector.encodedAttributes))
407+
pos++
408+
pos += copy(data[pos:], []byte(mc.connector.encodedAttributes))
409+
397410
// Send Auth packet
398411
return mc.writePacket(data[:pos])
399412
}

Diff for: ‎packets_test.go

+6-1
Original file line numberDiff line numberDiff line change
@@ -96,9 +96,14 @@ var _ net.Conn = new(mockConn)
9696

9797
func newRWMockConn(sequence uint8) (*mockConn, *mysqlConn) {
9898
conn := new(mockConn)
99+
connector, err := newConnector(NewConfig())
100+
if err != nil {
101+
panic(err)
102+
}
99103
mc := &mysqlConn{
100104
buf: newBuffer(conn),
101-
cfg: NewConfig(),
105+
cfg: connector.cfg,
106+
connector: connector,
102107
netConn: conn,
103108
closech: make(chan struct{}),
104109
maxAllowedPacket: defaultMaxAllowedPacket,

Diff for: ‎utils.go

+5
Original file line numberDiff line numberDiff line change
@@ -616,6 +616,11 @@ func appendLengthEncodedInteger(b []byte, n uint64) []byte {
616616
byte(n>>32), byte(n>>40), byte(n>>48), byte(n>>56))
617617
}
618618

619+
func appendLengthEncodedString(b []byte, s string) []byte {
620+
b = appendLengthEncodedInteger(b, uint64(len(s)))
621+
return append(b, s...)
622+
}
623+
619624
// reserveBuffer checks cap(buf) and expand buffer to len(buf) + appendSize.
620625
// If cap(buf) is not enough, reallocate new buffer.
621626
func reserveBuffer(buf []byte, appendSize int) []byte {

0 commit comments

Comments
 (0)