go/mysql: server.go: Add a callback on Handler, ConnectionAuthenticated, which is called immediately after the connection is authenticated.

This allows a server implementation to know the authenticated user without waiting for the first command interactions, such as ComQuery or ComInitDB.

Author: reltuk

go/mysql/server.go

@@ -93,10 +93,17 @@ type Handler interface {
 	// ConnectionClosed is called when a connection is closed.
 	ConnectionClosed(c *Conn)
 
+	// ConnectionAuthenticated is called when a connection is authenticated.
+	// Always called after NewConnection and before ConnectionClosed.
+	ConnectionAuthenticated(*Conn) error
+
 	// ConnectionAborted is called when a new connection cannot be fully established. For
 	// example, if a client connects to the server, but fails authentication, or can't
 	// negotiate an authentication handshake, this method will be called to let integrators
 	// know about the failed connection attempt.
+	//
+	// ConnectionClosed will still be called for the connection after ConnectionAborted is
+	// called.
 	ConnectionAborted(c *Conn, reason string) error
 
 	// ComInitDB is called once at the beginning to set db name,
@@ -560,6 +567,13 @@ func (l *Listener) handle(ctx context.Context, conn net.Conn, connectionID uint3
 		defer connCountPerUser.Add(c.User, -1)
 	}
 
+	if err = l.handler.ConnectionAuthenticated(c); err != nil {
+		log.Errorf("failed to register the connection as authenticated %s: %v", c, err)
+
+		c.writeErrorPacketFromError(err)
+		return
+	}
+
 	// Set initial db name.
 	if c.schemaName != "" {
 		if err = l.handler.ComInitDB(c, c.schemaName); err != nil {

go/mysql/server_test.go

@@ -119,6 +119,10 @@ func (th *testHandler) ConnectionAborted(c *Conn, reason string) error {
 	return nil
 }
 
+func (th *testHandler) ConnectionAuthenticated(c *Conn) error {
+	return nil
+}
+
 func (th *testHandler) ParserOptionsForConnection(c *Conn) (sqlparser.ParserOptions, error) {
 	return sqlparser.ParserOptions{}, nil
 }

go/sqltypes/arithmetic_test.go

@@ -21,6 +21,7 @@ import (
 	"fmt"
 	"math"
 	"reflect"
+	"runtime"
 	"strconv"
 	"testing"
 
@@ -486,6 +487,7 @@ func TestNullsafeAdd(t *testing.T) {
 		v1, v2 Value
 		out    Value
 		err    error
+		skip   bool
 	}{{
 		// All nulls.
 		v1:  NULL,
@@ -527,7 +529,20 @@ func TestNullsafeAdd(t *testing.T) {
 		v2:  NewFloat64(2),
 		out: NewInt64(3),
 	}}
+	if runtime.GOARCH != "amd64" {
+		// This test relies on float64 -> int64 conversion, where the float64 is larger than the maximum int64.
+		// According to the Golang spec:
+		//
+		// > In all non-constant conversions involving floating-point or complex values, if the result
+		// type cannot represent the value the conversion succeeds but the result value is implementation-dependent.
+		//
+		// And indeed, the test fails on arm64 but passes on amd64 because the computed conversion is different.
+		tcases[6].skip = true
+	}
 	for _, tcase := range tcases {
+		if tcase.skip {
+			t.SkipNow()
+		}
 		got := NullsafeAdd(tcase.v1, tcase.v2, Int64)
 
 		if !reflect.DeepEqual(got, tcase.out) {