only pass through the URL params we need

slackhq · Mar 23, 2024 · 19d06e2 · 19d06e2
1 parent bb3a83c
commit 19d06e2
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 12 deletions.
diff --git 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
@@ -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())