From 19d06e2e02e2696de2f847573eba0cd14486e12d Mon Sep 17 00:00:00 2001
From: Michael Demmer <mdemmer@slack-corp.com>
Date: Fri, 22 Mar 2024 17:00:14 -0700
Subject: [PATCH] only pass through the URL params we need

---
 go/vt/vtgateproxy/discovery.go   |  3 ---
 go/vt/vtgateproxy/vtgateproxy.go | 27 ++++++++++++++++++---------
 2 files changed, 18 insertions(+), 12 deletions(-)

diff --git a/go/vt/vtgateproxy/discovery.go b/go/vt/vtgateproxy/discovery.go
index be34ca3ebd4..ec708e3a46a 100644
--- a/go/vt/vtgateproxy/discovery.go
+++ b/go/vt/vtgateproxy/discovery.go
@@ -113,9 +113,6 @@ func RegisterJsonDiscovery() {
 	log.Infof("Registered JSON discovery scheme %v to watch: %v\n", jsonDiscovery.Scheme(), *jsonDiscoveryConfig)
 }
 
-type hostFilters = map[string]string
-
-// exampleResolver is a
 // Resolver(https://godoc.org/google.golang.org/grpc/resolver#Resolver).
 type JSONGateConfigResolver struct {
 	target   resolver.Target
diff --git a/go/vt/vtgateproxy/vtgateproxy.go b/go/vt/vtgateproxy/vtgateproxy.go
index c91000a90df..d1c28ebce9f 100644
--- a/go/vt/vtgateproxy/vtgateproxy.go
+++ b/go/vt/vtgateproxy/vtgateproxy.go
@@ -91,22 +91,31 @@ func (proxy *VTGateProxy) getConnection(ctx context.Context, target string) (*vt
 
 func (proxy *VTGateProxy) NewSession(ctx context.Context, options *querypb.ExecuteOptions, connectionAttributes map[string]string) (*vtgateconn.VTGateSession, error) {
 
-	if *poolTypeAttr != "" {
-		_, ok := connectionAttributes[*poolTypeAttr]
-		if !ok {
-			return nil, vterrors.Errorf(vtrpcpb.Code_UNAVAILABLE, "pool type attribute %s not supplied by client", *poolTypeAttr)
-		}
-	}
-
 	targetUrl := url.URL{
 		Scheme: "vtgate",
 		Host:   "pool",
 	}
 
 	values := url.Values{}
-	for k, v := range connectionAttributes {
-		values.Set(k, v)
+
+	if *poolTypeAttr != "" {
+		poolType, ok := connectionAttributes[*poolTypeAttr]
+		if ok {
+			values.Set(*poolTypeAttr, poolType)
+		} else {
+			return nil, vterrors.Errorf(vtrpcpb.Code_UNAVAILABLE, "pool type attribute %s not supplied by client", *poolTypeAttr)
+		}
+	}
+
+	if *affinityAttr != "" {
+		affinity, ok := connectionAttributes[*affinityAttr]
+		if ok {
+			values.Set(*affinityAttr, affinity)
+		} else {
+			return nil, vterrors.Errorf(vtrpcpb.Code_UNAVAILABLE, "pool type attribute %s not supplied by client", *affinityAttr)
+		}
 	}
+
 	targetUrl.RawQuery = values.Encode()
 
 	conn, err := proxy.getConnection(ctx, targetUrl.String())