From 295243e63a380393480a6c289af460bd330a0cf8 Mon Sep 17 00:00:00 2001
From: Ronald Ekambi <ronekambi@gmail.com>
Date: Tue, 21 Nov 2023 12:56:35 -0500
Subject: [PATCH] kv keys endpoint

---
 dependency/kv_keys.go      | 25 +++++++++++-----
 dependency/kv_keys_test.go | 61 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 79 insertions(+), 7 deletions(-)

diff --git a/dependency/kv_keys.go b/dependency/kv_keys.go
index 32ae98100..cd20a6ef6 100644
--- a/dependency/kv_keys.go
+++ b/dependency/kv_keys.go
@@ -18,15 +18,17 @@ var (
 	_ Dependency = (*KVKeysQuery)(nil)
 
 	// KVKeysQueryRe is the regular expression to use.
-	KVKeysQueryRe = regexp.MustCompile(`\A` + prefixRe + dcRe + `\z`)
+	KVKeysQueryRe = regexp.MustCompile(`\A` + prefixRe + queryRe + dcRe + `\z`)
 )
 
 // KVKeysQuery queries the KV store for a single key.
 type KVKeysQuery struct {
 	stopCh chan struct{}
 
-	dc     string
-	prefix string
+	dc        string
+	prefix    string
+	namespace string
+	partition string
 }
 
 // NewKVKeysQuery parses a string into a dependency.
@@ -36,10 +38,17 @@ func NewKVKeysQuery(s string) (*KVKeysQuery, error) {
 	}
 
 	m := regexpMatch(KVKeysQueryRe, s)
+	queryParams, err := GetConsulQueryOpts(m, "kv.keys")
+	if err != nil {
+		return nil, err
+	}
+
 	return &KVKeysQuery{
-		stopCh: make(chan struct{}, 1),
-		dc:     m["dc"],
-		prefix: m["prefix"],
+		stopCh:    make(chan struct{}, 1),
+		dc:        m["dc"],
+		prefix:    m["prefix"],
+		namespace: queryParams.Get(QueryNamespace),
+		partition: queryParams.Get(QueryPartition),
 	}, nil
 }
 
@@ -52,7 +61,9 @@ func (d *KVKeysQuery) Fetch(clients *ClientSet, opts *QueryOptions) (interface{}
 	}
 
 	opts = opts.Merge(&QueryOptions{
-		Datacenter: d.dc,
+		Datacenter:      d.dc,
+		ConsulPartition: d.partition,
+		ConsulNamespace: d.namespace,
 	})
 
 	log.Printf("[TRACE] %s: GET %s", d, &url.URL{
diff --git a/dependency/kv_keys_test.go b/dependency/kv_keys_test.go
index 4c02dd669..6e018665f 100644
--- a/dependency/kv_keys_test.go
+++ b/dependency/kv_keys_test.go
@@ -30,6 +30,18 @@ func TestNewKVKeysQuery(t *testing.T) {
 			nil,
 			true,
 		},
+		{
+			"query_only",
+			"?ns=foo",
+			nil,
+			true,
+		},
+		{
+			"invalid query param (unsupported key)",
+			"prefix?unsupported=foo",
+			nil,
+			true,
+		},
 		{
 			"prefix",
 			"prefix",
@@ -47,6 +59,55 @@ func TestNewKVKeysQuery(t *testing.T) {
 			},
 			false,
 		},
+		{
+			"partition",
+			"prefix?partition=foo",
+			&KVKeysQuery{
+				prefix:    "prefix",
+				partition: "foo",
+			},
+			false,
+		},
+		{
+			"namespace",
+			"prefix?ns=foo",
+			&KVKeysQuery{
+				prefix:    "prefix",
+				namespace: "foo",
+			},
+			false,
+		},
+		{
+			"namespace_and_partition",
+			"prefix?ns=foo&partition=bar",
+			&KVKeysQuery{
+				prefix:    "prefix",
+				namespace: "foo",
+				partition: "bar",
+			},
+			false,
+		},
+		{
+			"namespace_and_partition_and_dc",
+			"prefix?ns=foo&partition=bar@dc1",
+			&KVKeysQuery{
+				prefix:    "prefix",
+				namespace: "foo",
+				partition: "bar",
+				dc:        "dc1",
+			},
+			false,
+		},
+		{
+			"empty_query",
+			"prefix?ns=&partition=",
+			&KVKeysQuery{
+				prefix:    "prefix",
+				namespace: "",
+				partition: "",
+			},
+			false,
+		},
 		{
 			"dots",
 			"prefix.with.dots",