From d634d2df39943cbd629c339c1a3bc1e435fbc42a Mon Sep 17 00:00:00 2001
From: Abhi Shah <abhis@splunk.com>
Date: Tue, 21 Sep 2021 16:30:16 +0530
Subject: [PATCH] Updated KVStore Methods

Added dictionary support for KVStore "insert" and "update" methods. Change with reference to issue #355 . KVStore example also updated
---
 examples/kvstore.py | 17 +++++++++++++++--
 splunklib/client.py |  4 ++++
 2 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/examples/kvstore.py b/examples/kvstore.py
index 291858701..c6733d255 100644
--- a/examples/kvstore.py
+++ b/examples/kvstore.py
@@ -51,9 +51,10 @@ def main():
     # Let's make sure it doesn't have any data
     print("Should be empty: %s" % json.dumps(collection.data.query()))
     
-    # Let's add some data
+    # Let's add some json data
     collection.data.insert(json.dumps({"_key": "item1", "somekey": 1, "otherkey": "foo"}))
-    collection.data.insert(json.dumps({"_key": "item2", "somekey": 2, "otherkey": "foo"}))
+    #Let's add data as a dictionary object
+    collection.data.insert({"_key": "item2", "somekey": 2, "otherkey": "foo"})
     collection.data.insert(json.dumps({"somekey": 3, "otherkey": "bar"}))
     
     # Let's make sure it has the data we just entered
@@ -61,6 +62,18 @@ def main():
     
     # Let's run some queries
     print("Should return item1: %s" % json.dumps(collection.data.query_by_id("item1"), indent=1))
+
+    #Let's update some data
+    data = collection.data.query_by_id("item2")
+    data['otherkey'] = "bar"
+    #Passing data using 'json.dumps'
+    collection.data.update("item2", json.dumps(data))
+    print("Should return item2 with updated data: %s" % json.dumps(collection.data.query_by_id("item2"), indent=1))
+    data['otherkey'] = "foo"
+    # Passing data as a dictionary instance
+    collection.data.update("item2", data)
+    print("Should return item2 with updated data: %s" % json.dumps(collection.data.query_by_id("item2"), indent=1))
+
     
     query = json.dumps({"otherkey": "foo"})
     print("Should return item1 and item2: %s" % json.dumps(collection.data.query(query=query), indent=1))
diff --git a/splunklib/client.py b/splunklib/client.py
index b44d90bcf..283a2fb5d 100644
--- a/splunklib/client.py
+++ b/splunklib/client.py
@@ -3664,6 +3664,8 @@ def insert(self, data):
         :return: _id of inserted object
         :rtype: ``dict``
         """
+        if isinstance(data, dict):
+            data = json.dumps(data)
         return json.loads(self._post('', headers=KVStoreCollectionData.JSON_HEADER, body=data).body.read().decode('utf-8'))
 
     def delete(self, query=None):
@@ -3700,6 +3702,8 @@ def update(self, id, data):
         :return: id of replaced document
         :rtype: ``dict``
         """
+        if isinstance(data, dict):
+            data = json.dumps(data)
         return json.loads(self._post(UrlEncoded(str(id)), headers=KVStoreCollectionData.JSON_HEADER, body=data).body.read().decode('utf-8'))
 
     def batch_find(self, *dbqueries):