From caae034f3b4783a8d0be60e50f64dfedd92fea92 Mon Sep 17 00:00:00 2001
From: Mitchell Hashimoto <mitchell.hashimoto@gmail.com>
Date: Tue, 5 Jun 2018 20:27:30 -0700
Subject: [PATCH] command/connect/proxy: can specify prepared query upstream
 types

---
 command/connect/proxy/flag_upstreams.go      | 19 ++++++++++++++++++-
 command/connect/proxy/flag_upstreams_test.go | 20 ++++++++++++++++++++
 2 files changed, 38 insertions(+), 1 deletion(-)

diff --git a/command/connect/proxy/flag_upstreams.go b/command/connect/proxy/flag_upstreams.go
index 02ce3a3bb1..b726a20667 100644
--- a/command/connect/proxy/flag_upstreams.go
+++ b/command/connect/proxy/flag_upstreams.go
@@ -34,6 +34,23 @@ func (f *FlagUpstreams) Set(value string) error {
 		portRaw = portRaw[idx+1:]
 	}
 
+	destinationType := "service"
+	if idx := strings.Index(name, "."); idx != -1 {
+		typ := name[idx+1:]
+		name = name[:idx]
+		switch typ {
+		case "", "service":
+			destinationType = "service"
+
+		case "query":
+			destinationType = "prepared_query"
+
+		default:
+			return fmt.Errorf(
+				"Upstream type must be blank, 'service', or 'query'. Got: %q", typ)
+		}
+	}
+
 	port, err := strconv.ParseInt(portRaw, 0, 0)
 	if err != nil {
 		return err
@@ -47,7 +64,7 @@ func (f *FlagUpstreams) Set(value string) error {
 		LocalBindAddress: addr,
 		LocalBindPort:    int(port),
 		DestinationName:  name,
-		DestinationType:  "service",
+		DestinationType:  destinationType,
 	}
 
 	return nil
diff --git a/command/connect/proxy/flag_upstreams_test.go b/command/connect/proxy/flag_upstreams_test.go
index d43c49d036..16eab61dcb 100644
--- a/command/connect/proxy/flag_upstreams_test.go
+++ b/command/connect/proxy/flag_upstreams_test.go
@@ -53,6 +53,26 @@ func TestFlagUpstreams(t *testing.T) {
 			"",
 		},
 
+		{
+			"single value prepared query",
+			[]string{"db.query:8181"},
+			map[string]proxy.UpstreamConfig{
+				"db": proxy.UpstreamConfig{
+					LocalBindPort:   8181,
+					DestinationName: "db",
+					DestinationType: "prepared_query",
+				},
+			},
+			"",
+		},
+
+		{
+			"invalid type",
+			[]string{"db.bad:8181"},
+			nil,
+			"Upstream type",
+		},
+
 		{
 			"address specified",
 			[]string{"db:127.0.0.55:8181"},