From 62d4aaa33e1a954fc2d856f0c2c20738ed79389b Mon Sep 17 00:00:00 2001
From: Mitchell Hashimoto <mitchell.hashimoto@gmail.com>
Date: Tue, 5 Jun 2018 16:04:50 -0700
Subject: [PATCH] agent: accept connect param for execute

---
 agent/prepared_query_endpoint.go      | 10 ++++++++++
 agent/prepared_query_endpoint_test.go | 23 +++++++++++++++++++++++
 2 files changed, 33 insertions(+)

diff --git a/agent/prepared_query_endpoint.go b/agent/prepared_query_endpoint.go
index 38c9b4f39a..68b1403e29 100644
--- a/agent/prepared_query_endpoint.go
+++ b/agent/prepared_query_endpoint.go
@@ -107,6 +107,16 @@ func (s *HTTPServer) preparedQueryExecute(id string, resp http.ResponseWriter, r
 		return nil, fmt.Errorf("Bad limit: %s", err)
 	}
 
+	params := req.URL.Query()
+	if raw := params.Get("connect"); raw != "" {
+		val, err := strconv.ParseBool(raw)
+		if err != nil {
+			return nil, fmt.Errorf("Error parsing 'connect' value: %s", err)
+		}
+
+		args.Connect = val
+	}
+
 	var reply structs.PreparedQueryExecuteResponse
 	defer setMeta(resp, &reply.QueryMeta)
 RETRY_ONCE:
diff --git a/agent/prepared_query_endpoint_test.go b/agent/prepared_query_endpoint_test.go
index e58e48fbbf..ff86d17a9a 100644
--- a/agent/prepared_query_endpoint_test.go
+++ b/agent/prepared_query_endpoint_test.go
@@ -11,6 +11,7 @@ import (
 
 	"github.com/hashicorp/consul/agent/structs"
 	"github.com/hashicorp/consul/types"
+	"github.com/stretchr/testify/require"
 )
 
 // MockPreparedQuery is a fake endpoint that we inject into the Consul server
@@ -678,6 +679,28 @@ func TestPreparedQuery_Explain(t *testing.T) {
 			t.Fatalf("bad code: %d", resp.Code)
 		}
 	})
+
+	// Ensure that Connect is passed through
+	t.Run("", func(t *testing.T) {
+		a := NewTestAgent(t.Name(), "")
+		defer a.Shutdown()
+		require := require.New(t)
+
+		m := MockPreparedQuery{
+			executeFn: func(args *structs.PreparedQueryExecuteRequest, reply *structs.PreparedQueryExecuteResponse) error {
+				require.True(args.Connect)
+				return nil
+			},
+		}
+		require.NoError(a.registerEndpoint("PreparedQuery", &m))
+
+		body := bytes.NewBuffer(nil)
+		req, _ := http.NewRequest("GET", "/v1/query/my-id/execute?connect=true", body)
+		resp := httptest.NewRecorder()
+		_, err := a.srv.PreparedQuerySpecific(resp, req)
+		require.NoError(err)
+		require.Equal(200, resp.Code)
+	})
 }
 
 func TestPreparedQuery_Get(t *testing.T) {