pyDKB/storages: add ES implementation for queries.

Added concept of "raw" queries: it is useful when query must be stored,
yet can not be converted to the format acceptable by storage client.

Here it is JSON with %(parameter)s values: it may fail to be parsed as
proper JSON if there is something like `"taskname": %(name)s`.

Author: mgolosova

Utils/Dataflow/pyDKB/storages/Storage.py

@@ -86,21 +86,42 @@ def read_query(self, fname, qname=None):
         """
         raise NotImplementedError
 
-    def save_query(self, query, qname=None):
+    def query_is_raw(self, query):
+        """ Check if given query is not compiled ("raw").
+
+        :param query: query body
+        :type query: obj
+
+        :return: True/False
+        :rtype: bool
+        """
+        raise NotImplementedError
+
+    def save_query(self, query, qname=None, raw=False):
         """ Save query for further usage.
 
         :param query: query content
         :type query: object
         :param qname: query name (must not start with '__')
         :type qname: str
+        :param raw: store "raw" (not compiled) version of query
+        :type raw: bool
         """
         if qname and qname.startswith('__'):
             raise ValueError("Query name must not start with '__'"
                              " (reserved for service needs).")
+        if not raw:
+            try:
+                raw = self.query_is_raw(query)
+            except NotImplementedError:
+                pass
+        prefix = ''
         if not qname:
             qname = '__last'
-        self.stored_queries[qname] = query
-        self.stored_queries[__last'] = query
+        if raw:
+            prefix = '__raw'
+        self.stored_queries[prefix+qname] = query
+        self.stored_queries[prefix+'__last'] = query
 
     def get_query(self, qname):
         """ Get query by name.
@@ -119,8 +140,12 @@ def get_query(self, qname):
             q = self.stored_queries[qname]
             self.stored_queries['__last'] = q
         except KeyError:
-            raise QueryError("Query used before saving: '%s'"
-                             % qname)
+            # There still may be raw version of the query
+            try:
+                q = self.stored_queries['__raw'+qname]
+            except KeyError:
+                raise QueryError("Query used before saving: '%s'"
+                                 % qname)
         self.stored_queries['__last'] = q
         return q
 

Utils/Dataflow/pyDKB/storages/es.py

@@ -2,11 +2,14 @@
 pyDKB.storages.es
 """
 
+import json
+
 from Storage import Storage
 from . import storageType
 from exceptions import (StorageException,
                         NotFound,
-                        InvalidRequest)
+                        InvalidRequest,
+                        MissedParameter)
 
 
 try:
@@ -29,6 +32,9 @@ class ES(Storage):
     # Default index
     index = None
 
+    # Default datetime format
+    datetime_fmt = '%d-%m-%Y %H:%M:%S'
+
     type = storageType.ES
 
     def __init__(self, name):
@@ -106,3 +112,104 @@ def get(self, id, fields=None, index=None, doc_type='_all', parent=None):
                 raise NotFound(self.name, id=id, index=index)
             raise InvalidRequest(err)
         return r.get('_source', {})
+
+
+    def read_query(self, fname, qname=None):
+        """ Read query from file and save it.
+
+        Raise ``QueryNotFound`` in case of failure.
+
+        :param fname: file name
+        :type fname: str
+        :param qname: query name (for futher usage)
+        :type qname: str
+        """
+        raw = False
+        try:
+            with open(fname, 'r') as f:
+                query = f.read()
+            query = json.loads(query)
+        except IOError:
+            raise QueryNotFound(qname, fname)
+        except ValueError:
+            # Query with parameters may fail when try to parse as JSON
+            # In this case we just store it as "raw" version
+            raw = True
+        self.save_query(query, qname, raw)
+
+
+    def query_is_raw(self, query):
+        """ Check if given query is not compiled ("raw").
+
+        :param query: query body
+        :type query: str, dict
+
+        :return: True/False
+        :rtype: bool
+        """
+        return not isinstance(query, dict)
+
+
+    def exec_query(self, qname=None, **kwargs):
+        """ Execute stored query with given parameters.
+
+        :param qname: query name (if None, last used/read
+                      one will be used)
+        :type qname: str, NoneType
+        :param kwargs: query parameters (applied with old-style
+                       string formatting operator '%'). Parameter
+                       name, started with '_', is treated as special
+                       one:
+                        * _size  -- for ES request "size";
+                        * _type  -- for ES request "doc_type";
+                        * _index -- for ES index to use.
+        :type kwargs: dict
+
+        :return: storage response
+        :rtype: object
+        """
+        query = self.get_query(qname)
+        raw = self.query_is_raw(query)
+        params = {}
+        for key in kwargs:
+            if key.startswith('_'):
+                continue
+            try:
+                params[key] = json.dumps(kwargs[key])
+            except TypeError, err:
+                if 'datetime' in str(err):
+                    val = json.dumps(kwargs[key].strftime(self.datetime_fmt))
+                    params[key] = val
+                else:
+                    raise
+        q = {}
+        q['index'] = kwargs.get('_index', self.index)
+        q['size'] = kwargs.get('_size')
+        q['doc_type'] = kwargs.get('_type')
+        if params:
+            try:
+                if not raw:
+                    query = json.dumps(query)
+                    raw = True
+                query = query % params
+            except KeyError, err:
+                raise MissedParameter(qname, str(err))
+        if raw:
+            try:
+                query = json.loads(query)
+            except ValueError, err:
+                msg = "Failed to parse query"
+                if qname:
+                    msg += " (%r)" % qname
+                msg += ": %s" % err
+                raise QueryError(msg)
+        q['body'] = query
+        try:
+            r = self.client().search(**q)
+        except RequestError, err:
+            msg = "Query failed"
+            if qname:
+                msg += ": (%r)" % qname
+            msg += ": %s" % err
+            raise QueryError(msg)
+        return r

Utils/Dataflow/pyDKB/storages/exceptions.py

@@ -80,3 +80,26 @@ def __init__(self, message, *args, **kwargs):
 class QueryError(StorageException):
     """ Exception indicating issues with stored queries. """
     pass
+
+
+class MissedParameter(QueryError):
+    """ Exception indicating that some query parameters are missed. """
+
+    def __init__(self, qname=None, param=None):
+        """ Initialize exception.
+
+        :param qname: query name
+        :type qname: str, NoneType
+        :param param: parameter name(s)
+        :type param: str, list(str)
+        """
+        message = 'Missed query parameters'
+        if param:
+            if isinstance(param, list):
+                p = ', '.join(param)
+            else:
+                p = param
+            message += ": %s" % p
+        if qname:
+            message += " ('%s')" % qname
+        super(MissedParameter, self).__init__(message)