Send connection attributes (#1389)

Co-authored-by: Inada Naoki <[email protected]>

Author: Daemonxiao

.github/workflows/test.yml

@@ -79,6 +79,7 @@ jobs:
  ; TestConcurrent fails if max_connections is too large
  max_connections=50
  local_infile=1
+ performance_schema=on
  - name: setup database
 run: |
  mysql --user 'root' --host '127.0.0.1' -e 'create database gotest;'

README.md

@@ -393,6 +393,15 @@ Default: 0
 
 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"*.
 
+##### `connectionAttributes`
+
+```
+Type: comma-delimited string of user-defined "key:value" pairs
+Valid Values: (<name1>:<value1>,<name2>:<value2>,...)
+Default: none
+```
+
+[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.
 
 ##### System Variables
 

connection.go

@@ -27,6 +27,7 @@ type mysqlConn struct{
 affectedRowsuint64
 insertIduint64
 cfg*Config
+connector*connector
 maxAllowedPacketint
 maxWriteSizeint
 writeTimeout time.Duration

connector.go

@@ -11,11 +11,54 @@ package mysql
 import (
 "context"
 "database/sql/driver"
+"fmt"
 "net"
+"os"
+"strconv"
+"strings"
 )
 
 typeconnectorstruct{
-cfg*Config// immutable private copy.
+cfg*Config// immutable private copy.
+encodedAttributesstring// Encoded connection attributes.
+}
+
+funcencodeConnectionAttributes(textAttributesstring) string{
+connAttrsBuf:=make([]byte, 0, 251)
+
+// default connection attributes
+connAttrsBuf=appendLengthEncodedString(connAttrsBuf, connAttrClientName)
+connAttrsBuf=appendLengthEncodedString(connAttrsBuf, connAttrClientNameValue)
+connAttrsBuf=appendLengthEncodedString(connAttrsBuf, connAttrOS)
+connAttrsBuf=appendLengthEncodedString(connAttrsBuf, connAttrOSValue)
+connAttrsBuf=appendLengthEncodedString(connAttrsBuf, connAttrPlatform)
+connAttrsBuf=appendLengthEncodedString(connAttrsBuf, connAttrPlatformValue)
+connAttrsBuf=appendLengthEncodedString(connAttrsBuf, connAttrPid)
+connAttrsBuf=appendLengthEncodedString(connAttrsBuf, strconv.Itoa(os.Getpid()))
+
+// user-defined connection attributes
+for_, connAttr:=rangestrings.Split(textAttributes, ","){
+attr:=strings.SplitN(connAttr, ":", 2)
+iflen(attr) !=2{
+continue
+ }
+for_, v:=rangeattr{
+connAttrsBuf=appendLengthEncodedString(connAttrsBuf, v)
+ }
+ }
+
+returnstring(connAttrsBuf)
+}
+
+funcnewConnector(cfg*Config) (*connector, error){
+encodedAttributes:=encodeConnectionAttributes(cfg.ConnectionAttributes)
+iflen(encodedAttributes) >250{
+returnnil, fmt.Errorf("connection attributes are longer than 250 bytes: %dbytes (%q)", len(encodedAttributes), cfg.ConnectionAttributes)
+ }
+return&connector{
+cfg: cfg,
+encodedAttributes: encodedAttributes,
+ }, nil
 }
 
 // Connect implements driver.Connector interface.
@@ -29,6 +72,7 @@ func (c *connector) Connect(ctx context.Context) (driver.Conn, error){
 maxWriteSize: maxPacketSize-1,
 closech: make(chanstruct{}),
 cfg: c.cfg,
+connector: c,
  }
 mc.parseTime=mc.cfg.ParseTime
 

connector_test.go

@@ -8,13 +8,16 @@ import (
 )
 
 funcTestConnectorReturnsTimeout(t*testing.T){
-connector:=&connector{&Config{
+connector, err:=newConnector(&Config{
 Net: "tcp",
 Addr: "1.1.1.1:1234",
 Timeout: 10*time.Millisecond,
- }}
+ })
+iferr!=nil{
+t.Fatal(err)
+ }
 
-_, err:=connector.Connect(context.Background())
+_, err=connector.Connect(context.Background())
 iferr==nil{
 t.Fatal("error expected")
  }

const.go

@@ -8,12 +8,24 @@
 
 package mysql
 
+import"runtime"
+
 const (
 defaultAuthPlugin="mysql_native_password"
 defaultMaxAllowedPacket=64<<20// 64 MiB. See https://github.com/go-sql-driver/mysql/issues/1355
 minProtocolVersion=10
 maxPacketSize=1<<24-1
 timeFormat="2006-01-02 15:04:05.999999"
+
+// Connection attributes
+// See https://dev.mysql.com/doc/refman/8.0/en/performance-schema-connection-attribute-tables.html#performance-schema-connection-attributes-available
+connAttrClientName="_client_name"
+connAttrClientNameValue="Go-MySQL-Driver"
+connAttrOS="_os"
+connAttrOSValue=runtime.GOOS
+connAttrPlatform="_platform"
+connAttrPlatformValue=runtime.GOARCH
+connAttrPid="_pid"
 )
 
 // MySQL constants documentation:

driver.go

@@ -85,8 +85,9 @@ func (d MySQLDriver) Open(dsn string) (driver.Conn, error){
 iferr!=nil{
 returnnil, err
  }
-c:=&connector{
-cfg: cfg,
+c, err:=newConnector(cfg)
+iferr!=nil{
+returnnil, err
  }
 returnc.Connect(context.Background())
 }
@@ -103,7 +104,7 @@ func NewConnector(cfg *Config) (driver.Connector, error){
 iferr:=cfg.normalize(); err!=nil{
 returnnil, err
  }
-return&connector{cfg: cfg}, nil
+returnnewConnector(cfg)
 }
 
 // OpenConnector implements driver.DriverContext.
@@ -112,7 +113,5 @@ func (d MySQLDriver) OpenConnector(dsn string) (driver.Connector, error){
 iferr!=nil{
 returnnil, err
  }
-return&connector{
-cfg: cfg,
- }, nil
+returnnewConnector(cfg)
 }

driver_test.go

@@ -3214,3 +3214,50 @@ func TestConnectorTimeoutsWatchCancel(t *testing.T){
 t.Errorf("connection not closed")
  }
 }
+
+funcTestConnectionAttributes(t*testing.T){
+if!available{
+t.Skipf("MySQL server not running on %s", netAddr)
+ }
+
+attr1:="attr1"
+value1:="value1"
+attr2:="foo"
+value2:="boo"
+dsn+=fmt.Sprintf("&connectionAttributes=%s:%s,%s:%s", attr1, value1, attr2, value2)
+
+vardb*sql.DB
+if_, err:=ParseDSN(dsn); err!=errInvalidDSNUnsafeCollation{
+db, err=sql.Open("mysql", dsn)
+iferr!=nil{
+t.Fatalf("error connecting: %s", err.Error())
+ }
+deferdb.Close()
+ }
+
+dbt:=&DBTest{t, db}
+
+varattrValuestring
+queryString:="SELECT ATTR_VALUE FROM performance_schema.session_account_connect_attrs WHERE PROCESSLIST_ID = CONNECTION_ID() and ATTR_NAME = ?"
+rows:=dbt.mustQuery(queryString, connAttrClientName)
+ifrows.Next(){
+rows.Scan(&attrValue)
+ifattrValue!=connAttrClientNameValue{
+dbt.Errorf("expected %q, got %q", connAttrClientNameValue, attrValue)
+ }
+ } else{
+dbt.Errorf("no data")
+ }
+rows.Close()
+
+rows=dbt.mustQuery(queryString, attr2)
+ifrows.Next(){
+rows.Scan(&attrValue)
+ifattrValue!=value2{
+dbt.Errorf("expected %q, got %q", value2, attrValue)
+ }
+ } else{
+dbt.Errorf("no data")
+ }
+rows.Close()
+}

dsn.go

@@ -34,23 +34,24 @@ var (
 // If a new Config is created instead of being parsed from a DSN string,
 // the NewConfig function should be used, which sets default values.
 typeConfigstruct{
-Userstring// Username
-Passwdstring// Password (requires User)
-Netstring// Network type
-Addrstring// Network address (requires Net)
-DBNamestring// Database name
-Paramsmap[string]string// Connection parameters
-Collationstring// Connection collation
-Loc*time.Location// Location for time.Time values
-MaxAllowedPacketint// Max packet size allowed
-ServerPubKeystring// Server public key name
-pubKey*rsa.PublicKey// Server public key
-TLSConfigstring// TLS configuration name
-TLS*tls.Config// TLS configuration, its priority is higher than TLSConfig
-Timeout time.Duration// Dial timeout
-ReadTimeout time.Duration// I/O read timeout
-WriteTimeout time.Duration// I/O write timeout
-LoggerLogger// Logger
+Userstring// Username
+Passwdstring// Password (requires User)
+Netstring// Network type
+Addrstring// Network address (requires Net)
+DBNamestring// Database name
+Paramsmap[string]string// Connection parameters
+ConnectionAttributesstring// Connection Attributes, comma-delimited string of user-defined "key:value" pairs
+Collationstring// Connection collation
+Loc*time.Location// Location for time.Time values
+MaxAllowedPacketint// Max packet size allowed
+ServerPubKeystring// Server public key name
+pubKey*rsa.PublicKey// Server public key
+TLSConfigstring// TLS configuration name
+TLS*tls.Config// TLS configuration, its priority is higher than TLSConfig
+Timeout time.Duration// Dial timeout
+ReadTimeout time.Duration// I/O read timeout
+WriteTimeout time.Duration// I/O write timeout
+LoggerLogger// Logger
 
 AllowAllFilesbool// Allow all files to be used with LOAD DATA LOCAL INFILE
 AllowCleartextPasswordsbool// Allows the cleartext client side plugin
@@ -560,6 +561,11 @@ func parseDSNParams(cfg *Config, params string) (err error){
 iferr!=nil{
 return
  }
+
+// Connection attributes
+case"connectionAttributes":
+cfg.ConnectionAttributes=value
+
 default:
 // lazy init
 ifcfg.Params==nil{

packets.go

@@ -285,6 +285,7 @@ func (mc *mysqlConn) writeHandshakeResponsePacket(authResp []byte, plugin string
 clientLocalFiles|
 clientPluginAuth|
 clientMultiResults|
+clientConnectAttrs|
 mc.flags&clientLongFlag
 
 ifmc.cfg.ClientFoundRows{
@@ -318,6 +319,13 @@ func (mc *mysqlConn) writeHandshakeResponsePacket(authResp []byte, plugin string
 pktLen+=n+1
  }
 
+// 1 byte to store length of all key-values
+// NOTE: Actually, this is length encoded integer.
+// But we support only len(connAttrBuf) < 251 for now because takeSmallBuffer
+// doesn't support buffer size more than 4096 bytes.
+// TODO(methane): Rewrite buffer management.
+pktLen+=1+len(mc.connector.encodedAttributes)
+
 // Calculate packet length and get buffer with that size
 data, err:=mc.buf.takeSmallBuffer(pktLen+4)
 iferr!=nil{
@@ -394,6 +402,11 @@ func (mc *mysqlConn) writeHandshakeResponsePacket(authResp []byte, plugin string
 data[pos] =0x00
 pos++
 
+// Connection Attributes
+data[pos] =byte(len(mc.connector.encodedAttributes))
+pos++
+pos+=copy(data[pos:], []byte(mc.connector.encodedAttributes))
+
 // Send Auth packet
 returnmc.writePacket(data[:pos])
 }