Expand ORQuery to support fields, ranges, and count
authorJonathan Griffin <jgriffin@mozilla.com>
Mon, 10 Oct 2011 11:03:45 -0700
changeset 40 103772f1ae22cea33e36520b01e58b50da552f17
parent 39 07a51ccc9cd10de9c4ebed949ef294057a602813
child 41 96db675fdbdcb51409cb9d0cd9299d0746c07b02
push id41
push userjgriffin@mozilla.com
push dateMon, 10 Oct 2011 18:03:48 +0000
Expand ORQuery to support fields, ranges, and count
mozautoeslib/eslib.py
--- a/mozautoeslib/eslib.py
+++ b/mozautoeslib/eslib.py
@@ -88,17 +88,18 @@ class ESLib(object):
       self._add_fieldlist_to_boolquery(boolquery, include, True)
     if exclude:
       self._add_fieldlist_to_boolquery(boolquery, exclude, False)
     if sort:
       boolquery.sort = sort
 
     return boolquery
 
-  def ORQuery(self, ORItems, size=10000, doc_type=None, useFieldQueries=False):
+  def ORQuery(self, ORItems, size=10000, doc_type=None, useFieldQueries=False,
+              fields=None, count=False):
     """Return a list of hits that match any of the combination of terms
        specified in the ORItems list of dicts.
 
        Example:
          return hits that match any of the following mahine/starttime
          combinations:
 
        result = eslib.ORQuery([
@@ -111,35 +112,47 @@ class ESLib(object):
       self.doc_type = doc_type
 
     resultlist = []
 
     orList = []
     for item in ORItems:
       andList = []
       for key in item:
-        if useFieldQueries:
+        if isinstance(item[key], list):
+          andList.append(RangeQuery(ESRange(key, item[key][0], item[key][1])))
+        elif useFieldQueries:
           andList.append(QueryFilter(FieldQuery(FieldParameter(key, item[key]))))
         else:
           andList.append(TermFilter(key, item[key]))
       orList.append(ANDFilter(andList))
     orq = ORFilter(orList)
 
     q = FilteredQuery(MatchAllQuery(), orq)
-    result = self.connection.search(query=q,
+    s = Search(query=q, fields=fields)
+
+    if count:
+      result = self.connection.count(query=q,
+                                     indexes=[self.read_index],
+                                     doc_types=self.doc_type)
+      if not 'count' in result:
+        raise Exception("'count' not found in response data")
+      return result['count']
+
+    result = self.connection.search(query=s,
                                     size=size,
                                     indexes=[self.read_index],
                                     doc_types=self.doc_type)
 
     if result and result['hits'] and result['hits']['hits']:
       # partially flatten the data
       for hit in result['hits']['hits']:
-        if not '_source' in hit:
-          raise Exception("Key ['_source'] not found in response hit")
-        resultlist.append(hit['_source'])
+        if not '_source' in hit and not 'fields' in hit:
+          raise Exception("Neither '_source' nor 'fields' found in response hit")
+        resultlist.append(hit['_source'] if '_source' in hit else hit['fields'])
 
     return resultlist
 
   def query(self, include={}, exclude={}, size=None, doc_type=None, sort=None,
             withSource=False):
     """Return a list of hits which match all the fields in 'include',
        but none of the fields in 'exclude', up to a maximum of 'size' hits,
        or all hits when 'size' is None.