Merge remote-tracking branch 'origin/vtgateproxy' into jas_loadbalancing

slackhq · Mar 8, 2024 · a59dff4 · a59dff4
2 parents b5c48da + 82258ab
commit a59dff4
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 26 deletions.
diff --git a/go/mysql/server.go b/go/mysql/server.go
@@ -19,7 +19,6 @@ package mysql
 import (
 	"context"
 	"crypto/tls"
-	"fmt"
 	"io"
 	"net"
 	"strings"
@@ -317,11 +316,6 @@ func (l *Listener) handle(conn net.Conn, connectionID uint32, acceptTime time.Ti
 	c := newServerConn(conn, l)
 	c.ConnectionID = connectionID
 
-	t := time.Now()
-	defer func() {
-		fmt.Printf("CLOSE CONNECTION %d [%v] \n", connectionID, time.Since(t))
-	}()
-
 	// Catch panics, and close the connection in any case.
 	defer func() {
 		if x := recover(); x != nil {
@@ -472,11 +466,9 @@ func (l *Listener) handle(conn net.Conn, connectionID uint32, acceptTime time.Ti
 
 	// Set initial db name.
 	if c.schemaName != "" {
-		fmt.Printf("handshake use %s\n", c.schemaName)
 		err = l.handler.ComQuery(c, "use "+sqlescape.EscapeID(c.schemaName), func(result *sqltypes.Result) error {
 			return nil
 		})
-		fmt.Printf("handshake use %s %v [%s]\n", c.schemaName, err, time.Now().Sub(t))
 		if err != nil {
 			c.writeErrorPacketFromError(err)
 			return
@@ -499,8 +491,6 @@ func (l *Listener) handle(conn net.Conn, connectionID uint32, acceptTime time.Ti
 		log.Warningf("Slow connection from %s: %v", c, connectTime)
 	}
 
-	fmt.Printf("Connection %d connectTime %s\n", connectionID, connectTime)
-
 	// Tell our handler that we're finished handshake and are ready to
 	// process commands.
 	l.handler.ConnectionReady(c)

diff --git a/go/vt/vtgateproxy/mysql_server.go b/go/vt/vtgateproxy/mysql_server.go
@@ -179,15 +179,6 @@ func (ph *proxyHandler) ComQuery(c *mysql.Conn, query string, callback func(*sql
 		defer cancel()
 	}
 
-	t := time.Now()
-	defer func() {
-		logSql := query
-		if len(logSql) > 40 {
-			logSql = logSql[:40]
-		}
-		fmt.Printf("ComQuery conn %d %s [%s]\n", c.ConnectionID, logSql, time.Since(t))
-	}()
-
 	span, ctx, err := startSpan(ctx, query, "proxyHandler.ComQuery")
 	if err != nil {
 		return vterrors.Wrap(err, "failed to extract span")

diff --git a/go/vt/vtgateproxy/vtgateproxy.go b/go/vt/vtgateproxy/vtgateproxy.go
@@ -24,10 +24,12 @@ import (
 	"fmt"
 	"io"
 	"net/url"
+	"strings"
 	"sync"
 	"time"
 
 	"google.golang.org/grpc"
+	"vitess.io/vitess/go/sqlescape"
 	"vitess.io/vitess/go/sqltypes"
 	"vitess.io/vitess/go/vt/grpcclient"
 	querypb "vitess.io/vitess/go/vt/proto/query"
@@ -132,15 +134,15 @@ func (proxy *VTGateProxy) Prepare(ctx context.Context, session *vtgateconn.VTGat
 }
 
 func (proxy *VTGateProxy) Execute(ctx context.Context, session *vtgateconn.VTGateSession, sql string, bindVariables map[string]*querypb.BindVariable) (qr *sqltypes.Result, err error) {
-	t := time.Now()
-	qr, err = session.Execute(WithSlackAZAffinityContext(ctx, proxy.azID, proxy.gateType), sql, bindVariables)
-	logSql := sql
-	if len(logSql) > 40 {
-		logSql = logSql[:40]
+
+	// Intercept "use" statements since they just have to update the local session
+	if strings.HasPrefix(sql, "use ") {
+		targetString := sqlescape.UnescapeID(sql[4:])
+		session.SessionPb().TargetString = targetString
+		return &sqltypes.Result{}, nil
 	}
-	fmt.Printf("Execute %s [%s]\n", logSql, time.Since(t))
-	return qr, err
 
+	return session.Execute(WithSlackAZAffinityContext(ctx, proxy.azID, proxy.gateType), sql, bindVariables)
 }
 
 func (proxy *VTGateProxy) StreamExecute(ctx context.Context, session *vtgateconn.VTGateSession, sql string, bindVariables map[string]*querypb.BindVariable, callback func(*sqltypes.Result) error) error {