From 3d421889a5a9ba261887ae115a60df923b495195 Mon Sep 17 00:00:00 2001
From: Praveen Kukreja <90465691+praveen-elastic@users.noreply.github.com>
Date: Thu, 7 Nov 2024 16:22:04 +0530
Subject: [PATCH] add query params 'ns' and 'partition' to kv store

---
 consul/api/kv.py | 51 +++++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 50 insertions(+), 1 deletion(-)

diff --git a/consul/api/kv.py b/consul/api/kv.py
index b162d3f..be59695 100644
--- a/consul/api/kv.py
+++ b/consul/api/kv.py
@@ -22,6 +22,8 @@ def get(
         keys=False,
         separator=None,
         dc=None,
+        ns=None,
+        partition=None,
         connections_timeout=None,
     ):
         """
@@ -43,6 +45,14 @@ def get(
         *dc* is the optional datacenter that you wish to communicate with.
         If None is provided, defaults to the agent's datacenter.
 
+        *ns* is the optional namespace to use for the request.
+        If None is provided, the namespace is inherited from the namespace
+        of request's ACL token, or defaults to the default namespace.
+
+        *partition* is the optional Admin Partition to use for the request.
+        If None is provided, the partition is inferred from the
+        request's ACL token, or defaults to the default partition.
+
         The *value* returned is for the specified key, or if *recurse* is
         True a list of *values* for all keys with the given prefix is
         returned.
@@ -81,6 +91,10 @@ def get(
         consistency = consistency or self.agent.consistency
         if consistency in ("consistent", "stale"):
             params.append((consistency, "1"))
+        if ns:
+            params.append(("ns", ns))
+        if partition:
+            params.append(("partition", partition))
 
         one = False
         decode = False
@@ -108,6 +122,8 @@ def put(
         release=None,
         token=None,
         dc=None,
+        ns=None,
+        partition=None,
         connections_timeout=None,
     ):
         """
@@ -140,6 +156,14 @@ def put(
         *dc* is the optional datacenter that you wish to communicate with.
         If None is provided, defaults to the agent's datacenter.
 
+        *ns* is the optional namespace to use for the request.
+        If None is provided, the namespace is inherited from the namespace
+        of request's ACL token, or defaults to the default namespace.
+
+        *partition* is the optional Admin Partition to use for the request.
+        If None is provided, the partition is inferred from the
+        request's ACL token, or defaults to the default partition.
+
         The return value is simply either True or False. If False is
         returned, then the update has not taken place.
         """
@@ -158,6 +182,10 @@ def put(
         dc = dc or self.agent.dc
         if dc:
             params.append(("dc", dc))
+        if ns:
+            params.append(("ns", ns))
+        if partition:
+            params.append(("partition", partition))
         http_kwargs = {}
         if connections_timeout:
             http_kwargs["connections_timeout"] = connections_timeout
@@ -166,7 +194,16 @@ def put(
             CB.json(), f"/v1/kv/{key}", params=params, headers=headers, data=value, **http_kwargs
         )
 
-    def delete(self, key, recurse=None, cas=None, token=None, dc=None, connections_timeout=None):
+    def delete(
+        self,
+        key,
+        recurse=None,
+        cas=None,
+        token=None,
+        dc=None,
+        ns=None,
+        partition=None,
+        connections_timeout=None):
         """
         Deletes a single key or if *recurse* is True, all keys sharing a
         prefix.
@@ -184,6 +221,14 @@ def delete(self, key, recurse=None, cas=None, token=None, dc=None, connections_t
 
         *dc* is the optional datacenter that you wish to communicate with.
         If None is provided, defaults to the agent's datacenter.
+
+        *ns* is the optional namespace to use for the request.
+        If None is provided, the namespace is inherited from the namespace
+        of request's ACL token, or defaults to the default namespace.
+
+        *partition* is the optional Admin Partition to use for the request.
+        If None is provided, the partition is inferred from the
+        request's ACL token, or defaults to the default partition.
         """
         assert not key.startswith("/"), "keys should not start with a forward slash"
 
@@ -195,6 +240,10 @@ def delete(self, key, recurse=None, cas=None, token=None, dc=None, connections_t
         dc = dc or self.agent.dc
         if dc:
             params.append(("dc", dc))
+        if ns:
+            params.append(("ns", ns))
+        if partition:
+            params.append(("partition", partition))
         http_kwargs = {}
         if connections_timeout:
             http_kwargs["connections_timeout"] = connections_timeout