From 82258abc4fe955d878e8e8bdc04d519a94f93cf6 Mon Sep 17 00:00:00 2001
From: Michael Demmer <mdemmer@slack-corp.com>
Date: Fri, 3 Nov 2023 11:46:04 -0700
Subject: [PATCH] intercept use statements

---
 go/vt/vtgateproxy/vtgateproxy.go | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/go/vt/vtgateproxy/vtgateproxy.go b/go/vt/vtgateproxy/vtgateproxy.go
index c103e9341bb..be5d91d430b 100644
--- a/go/vt/vtgateproxy/vtgateproxy.go
+++ b/go/vt/vtgateproxy/vtgateproxy.go
@@ -22,10 +22,12 @@ import (
 	"context"
 	"flag"
 	"io"
+	"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"
@@ -117,6 +119,14 @@ 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) {
+
+	// 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
+	}
+
 	return session.Execute(ctx, sql, bindVariables)
 }