From 77a8003475e0dc3cebb412db33ee5a1573831c97 Mon Sep 17 00:00:00 2001
From: Mitchell Hashimoto <mitchell.hashimoto@gmail.com>
Date: Wed, 6 Jun 2018 11:50:23 -0700
Subject: [PATCH] api: change Connect to a query option

---
 api/api.go            |  7 +++++++
 api/prepared_query.go | 23 -----------------------
 connect/resolver.go   |  5 ++++-
 3 files changed, 11 insertions(+), 24 deletions(-)

diff --git a/api/api.go b/api/api.go
index 6d64366377..6b359fef2b 100644
--- a/api/api.go
+++ b/api/api.go
@@ -112,6 +112,10 @@ type QueryOptions struct {
 	// a value from 0 to 5 (inclusive).
 	RelayFactor uint8
 
+	// Connect filters prepared query execution to only include Connect-capable
+	// services. This currently affects prepared query execution.
+	Connect bool
+
 	// ctx is an optional context pass through to the underlying HTTP
 	// request layer. Use Context() and WithContext() to manage this.
 	ctx context.Context
@@ -561,6 +565,9 @@ func (r *request) setQueryOptions(q *QueryOptions) {
 	if q.RelayFactor != 0 {
 		r.params.Set("relay-factor", strconv.Itoa(int(q.RelayFactor)))
 	}
+	if q.Connect {
+		r.params.Set("connect", "true")
+	}
 	r.ctx = q.ctx
 }
 
diff --git a/api/prepared_query.go b/api/prepared_query.go
index 4b4173d7fe..8bb1004eec 100644
--- a/api/prepared_query.go
+++ b/api/prepared_query.go
@@ -210,26 +210,3 @@ func (c *PreparedQuery) Execute(queryIDOrName string, q *QueryOptions) (*Prepare
 	}
 	return out, qm, nil
 }
-
-// ExecuteConnect is used to execute a specific prepared query and return
-// only Connect-capable nodes.
-func (c *PreparedQuery) ExecuteConnect(queryIDOrName string, q *QueryOptions) (*PreparedQueryExecuteResponse, *QueryMeta, error) {
-	r := c.c.newRequest("GET", "/v1/query/"+queryIDOrName+"/execute")
-	r.setQueryOptions(q)
-	r.params.Set("connect", "true")
-	rtt, resp, err := requireOK(c.c.doRequest(r))
-	if err != nil {
-		return nil, nil, err
-	}
-	defer resp.Body.Close()
-
-	qm := &QueryMeta{}
-	parseQueryMeta(resp, qm)
-	qm.RequestTime = rtt
-	var out PreparedQueryExecuteResponse
-	if err := decodeBody(resp, &out); err != nil {
-		return nil, nil, err
-	}
-
-	return &out, qm, nil
-}
diff --git a/connect/resolver.go b/connect/resolver.go
index 6506a9f77a..47b0063881 100644
--- a/connect/resolver.go
+++ b/connect/resolver.go
@@ -144,7 +144,7 @@ func (cr *ConsulResolver) resolveService(ctx context.Context) (string, connect.C
 }
 
 func (cr *ConsulResolver) resolveQuery(ctx context.Context) (string, connect.CertURI, error) {
-	resp, _, err := cr.Client.PreparedQuery().ExecuteConnect(cr.Name, cr.queryOptions(ctx))
+	resp, _, err := cr.Client.PreparedQuery().Execute(cr.Name, cr.queryOptions(ctx))
 	if err != nil {
 		return "", nil, err
 	}
@@ -193,6 +193,9 @@ func (cr *ConsulResolver) queryOptions(ctx context.Context) *api.QueryOptions {
 		// caching which is even more stale so...
 		AllowStale: true,
 		Datacenter: cr.Datacenter,
+
+		// For prepared queries
+		Connect: true,
 	}
 	return q.WithContext(ctx)
 }