Merge pull request splunk#126 from splunk/feature/kvstore

itay · itay · commit bd8d1425872f · 2015-11-20T18:40:30.000-08:00
Add KV Store support to the Python SDK
diff --git a/docs/client.rst b/docs/client.rst
@@ -68,6 +68,18 @@ splunklib.client
     :members: create, export, itemmeta, oneshot
     :inherited-members:
 
+.. autoclass:: KVStoreCollection
+    :members: data, update_index, update_field
+    :inherited-members:
+
+.. autoclass:: KVStoreCollectionData
+    :members: query, query_by_id, insert, delete, delete_by_id, update, batch_save
+    :inherited-members:
+
+.. autoclass:: KVStoreCollections
+    :members: create
+    :inherited-members:
+
 .. autoclass:: Loggers
     :members: itemmeta
     :inherited-members:
@@ -110,7 +122,7 @@ splunklib.client
     :inherited-members:
 
 .. autoclass:: Service
-    :members: apps, confs, capabilities, event_types, fired_alerts, indexes, info, inputs, job, jobs, loggers, messages, modular_input_kinds, parse, restart, restart_required, roles, search, saved_searches, settings, splunk_version, storage_passwords, users
+    :members: apps, confs, capabilities, event_types, fired_alerts, indexes, info, inputs, job, jobs, kvstore, loggers, messages, modular_input_kinds, parse, restart, restart_required, roles, search, saved_searches, settings, splunk_version, storage_passwords, users
     :inherited-members:
 
 .. autoclass:: Settings
diff --git a/docs/index.rst b/docs/index.rst
@@ -69,6 +69,12 @@ For more information, see the `Splunk Developer Portal <http://dev.splunk.com/vi
 
     :class:`~splunklib.client.Jobs` class
 
+    :class:`~splunklib.client.KVStoreCollection` class
+
+    :class:`~splunklib.client.KVStoreCollectionData` class
+
+    :class:`~splunklib.client.KVStoreCollections` class
+
     :class:`~splunklib.client.Loggers` class
 
     :class:`~splunklib.client.Message` class
diff --git a/examples/kvstore.py b/examples/kvstore.py
@@ -0,0 +1,75 @@
+#!/usr/bin/env python
+#
+# Copyright 2011-2015 Splunk, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"): you may
+# not use this file except in compliance with the License. You may obtain
+# a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+# License for the specific language governing permissions and limitations
+# under the License.
+
+"""A command line utility for interacting with Splunk KV Store Collections."""
+
+import sys, os, json
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
+
+from splunklib.client import connect
+
+try:
+    from utils import parse
+except ImportError:
+    raise Exception("Add the SDK repository to your PYTHONPATH to run the examples "
+                    "(e.g., export PYTHONPATH=~/splunk-sdk-python.")
+
+def main():
+    opts = parse(sys.argv[1:], {}, ".splunkrc")
+    opts.kwargs["owner"] = "nobody"
+    opts.kwargs["app"] = "search"
+    service = connect(**opts.kwargs)
+
+    print "KV Store Collections:"
+    for collection in service.kvstore:
+        print "  %s" % collection.name
+    
+    # Let's delete a collection if it already exists, and then create it
+    collection_name = "example_collection"
+    if collection_name in service.kvstore:
+        service.kvstore.delete(collection_name)
+    
+    # Let's create it and then make sure it exists    
+    service.kvstore.create(collection_name)
+    collection = service.kvstore[collection_name]
+    
+    # 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
+    collection.data.insert(json.dumps({"_key": "item1", "somekey": 1, "otherkey": "foo"}))
+    collection.data.insert(json.dumps({"_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
+    print "Should have our data: %s" % json.dumps(collection.data.query(), indent=1)
+    
+    # Let's run some queries
+    print "Should return item1: %s" % json.dumps(collection.data.query_by_id("item1"), indent=1)
+    
+    query = json.dumps({"otherkey": "foo"})
+    print "Should return item1 and item2: %s" % json.dumps(collection.data.query(query=query), indent=1)
+    
+    query = json.dumps({"otherkey": "bar"})
+    print "Should return third item with auto-generated _key: %s" % json.dumps(collection.data.query(query=query), indent=1)
+    
+    # Let's delete the collection
+    collection.delete()
+
+if __name__ == "__main__":
+    main()
+
+
diff --git a/splunklib/binding.py b/splunklib/binding.py
@@ -1168,10 +1168,16 @@ def post(self, url, headers=None, **kwargs):
         :rtype: ``dict``
         """
         if headers is None: headers = []
-        headers.append(("Content-Type", "application/x-www-form-urlencoded")),
+
         # We handle GET-style arguments and an unstructured body. This is here
         # to support the receivers/stream endpoint.
         if 'body' in kwargs:
+            # We only use application/x-www-form-urlencoded if there is no other
+            # Content-Type header present. This can happen in cases where we 
+            # send requests as application/json, e.g. for KV Store.
+            if len(filter(lambda x: x[0].lower() == "content-type", headers)) == 0:
+                headers.append(("Content-Type", "application/x-www-form-urlencoded"))
+
             body = kwargs.pop('body')
             if len(kwargs) > 0:
                 url = url + UrlEncoded('?' + _encode(**kwargs), skip_encode=True)
diff --git a/splunklib/client.py b/splunklib/client.py
@@ -653,6 +653,14 @@ def splunk_version(self):
             self._splunk_version = tuple([int(p) for p in self.info['version'].split('.')])
         return self._splunk_version
 
+    @property
+    def kvstore(self):
+        """Returns the collection of KV Store collections.
+
+        :return: A :class:`KVStoreCollections` collection of :class:`KVStoreCollection` entities.
+        """
+        return KVStoreCollections(self)
+
     @property
     def users(self):
         """Returns the collection of users.
@@ -3518,3 +3526,194 @@ def package(self):
     def updateInfo(self):
         """Returns any update information that is available for the app."""
         return self._run_action("update")
+
+class KVStoreCollections(Collection):
+    def __init__(self, service):
+        Collection.__init__(self, service, 'storage/collections/config', item=KVStoreCollection)
+
+    def create(self, name, indexes = {}, fields = {}, **kwargs):
+        """Creates a KV Store Collection.
+
+        :param name: name of collection to create
+        :type name: ``string``
+        :param indexes: dictionary of index definitions
+        :type indexes: ``dict``
+        :param fields: dictionary of field definitions
+        :type fields: ``dict``
+        :param kwargs: a dictionary of additional parameters specifying indexes and field definitions
+        :type kwargs: ``dict``
+
+        :return: Result of POST request
+        """
+        for k, v in indexes.iteritems():
+            if isinstance(v, dict):
+                v = json.dumps(v)
+            kwargs['index.' + k] = v
+        for k, v in fields.iteritems():
+            kwargs['field.' + k] = v
+        return self.post(name=name, **kwargs)
+
+class KVStoreCollection(Entity):
+    @property
+    def data(self):
+        """Returns data object for this Collection.
+
+        :rtype: :class:`KVStoreData`
+        """
+        return KVStoreCollectionData(self)
+
+    def update_index(self, name, value):
+        """Changes the definition of a KV Store index.
+
+        :param name: name of index to change
+        :type name: ``string``
+        :param value: new index definition
+        :type value: ``dict`` or ``string``
+
+        :return: Result of POST request
+        """
+        kwargs = {}
+        kwargs['index.' + name] = value if isinstance(value, basestring) else json.dumps(value)
+        return self.post(**kwargs)
+
+    def update_field(self, name, value):
+        """Changes the definition of a KV Store field.
+
+        :param name: name of field to change
+        :type name: ``string``
+        :param value: new field definition
+        :type value: ``string``
+
+        :return: Result of POST request
+        """
+        kwargs = {}
+        kwargs['field.' + name] = value
+        return self.post(**kwargs)
+
+class KVStoreCollectionData(object):
+    """This class represents the data endpoint for a KVStoreCollection.
+
+    Retrieve using :meth:`KVStoreCollection.data`
+    """
+    JSON_HEADER = [('Content-Type', 'application/json')]
+
+    def __init__(self, collection):
+        self.service = collection.service
+        self.collection = collection
+        self.owner, self.app, self.sharing = collection._proper_namespace()
+        self.path = 'storage/collections/data/' + UrlEncoded(self.collection.name) + '/'
+
+    def _get(self, url, **kwargs):
+        return self.service.get(self.path + url, owner=self.owner, app=self.app, sharing=self.sharing, **kwargs)
+
+    def _post(self, url, **kwargs):
+        return self.service.post(self.path + url, owner=self.owner, app=self.app, sharing=self.sharing, **kwargs)
+
+    def _delete(self, url, **kwargs):
+        return self.service.delete(self.path + url, owner=self.owner, app=self.app, sharing=self.sharing, **kwargs)
+
+    def query(self, **query):
+        """
+        Gets the results of query, with optional parameters sort, limit, skip, and fields.
+
+        :param query: Optional parameters. Valid options are sort, limit, skip, and fields
+        :type query: ``dict``
+
+        :return: Array of documents retrieved by query.
+        :rtype: ``array``
+        """
+        return json.loads(self._get('', **query).body.read())
+
+    def query_by_id(self, id):
+        """
+        Returns object with _id = id.
+
+        :param id: Value for ID. If not a string will be coerced to string.
+        :type id: ``string``
+
+        :return: Document with id
+        :rtype: ``dict``
+        """
+        return json.loads(self._get(UrlEncoded(str(id))).body.read())
+
+    def insert(self, data):
+        """
+        Inserts item into this collection. An _id field will be generated if not assigned in the data.
+
+        :param data: Document to insert
+        :type data: ``string``
+
+        :return: _id of inserted object
+        :rtype: ``dict``
+        """
+        return json.loads(self._post('', headers=KVStoreCollectionData.JSON_HEADER, body=data).body.read())
+
+    def delete(self, query=None):
+        """
+        Deletes all data in collection if query is absent. Otherwise, deletes all data matched by query.
+
+        :param query: Query to select documents to delete
+        :type query: ``string``
+
+        :return: Result of DELETE request
+        """
+        return self._delete('', **({'query': query}) if query else {})
+
+    def delete_by_id(self, id):
+        """
+        Deletes document that has _id = id.
+
+        :param id: id of document to delete
+        :type id: ``string``
+
+        :return: Result of DELETE request
+        """
+        return self._delete(UrlEncoded(str(id)))
+
+    def update(self, id, data):
+        """
+        Replaces document with _id = id with data.
+
+        :param id: _id of document to update
+        :type id: ``string``
+        :param data: the new document to insert
+        :type data: ``string``
+
+        :return: id of replaced document
+        :rtype: ``dict``
+        """
+        return json.loads(self._post(UrlEncoded(str(id)), headers=KVStoreCollectionData.JSON_HEADER, body=data).body.read())
+
+    def batch_find(self, *dbqueries):
+        """
+        Returns array of results from queries dbqueries.
+
+        :param dbqueries: Array of individual queries as dictionaries
+        :type dbqueries: ``array`` of ``dict``
+        
+        :return: Results of each query
+        :rtype: ``array`` of ``array``
+        """
+        if len(dbqueries) < 1: 
+            raise Exception('Must have at least one query.')
+        
+        data = json.dumps(dbqueries)
+
+        return json.loads(self._post('batch_find', headers=KVStoreCollectionData.JSON_HEADER, body=data).body.read())
+
+    def batch_save(self, *documents):
+        """
+        Inserts or updates every document specified in documents.
+
+        :param documents: Array of documents to save as dictionaries
+        :type documents: ``array`` of ``dict``
+        
+        :return: Results of update operation as overall stats
+        :rtype: ``dict``
+        """
+        if len(documents) < 1: 
+            raise Exception('Must have at least one document.')
+        
+        data = json.dumps(documents)
+
+        return json.loads(self._post('batch_save', headers=KVStoreCollectionData.JSON_HEADER, body=data).body.read())
diff --git a/tests/test_examples.py b/tests/test_examples.py
@@ -182,6 +182,11 @@ def test_job(self):
             "job.py",
             "job.py list",
             "job.py list @0")
+
+    def test_kvstore(self):
+        self.check_commands(
+            "kvstore.py --help",
+            "kvstore.py")
         
     def test_loggers(self):
         self.check_commands(
diff --git a/tests/test_kvstore_batch.py b/tests/test_kvstore_batch.py
diff --git a/tests/test_kvstore_conf.py b/tests/test_kvstore_conf.py
diff --git a/tests/test_kvstore_data.py b/tests/test_kvstore_data.py
