Merge pull request #43 from postatum/94698540_fields_arg_error

Fix `fields` param usage in ES.get_collection

nefertari/elasticsearch.py

@@ -59,6 +59,28 @@ def _bulk_body(body):
     return ES.api.bulk(body=body)
 
 
+def process_fields_param(fields):
+    """ Process 'fields' ES param.
+
+    * Fields list is split if needed
+    * '_type' field is added, if not present, so the actual value is
+      displayed instead of 'None'
+    * '_source=False' is returned as well, so document source is not
+      loaded from ES. This is done because source is not used when
+      'fields' param is provided
+    """
+    if not fields:
+        return fields
+    if isinstance(fields, basestring):
+        fields = split_strip(fields)
+    if '_type' not in fields:
+        fields.append('_type')
+    return {
+        'fields': fields,
+        '_source': False,
+    }
+
+
 def apply_sort(_sort):
     _sort_param = []
 
@@ -315,7 +337,9 @@ def get_by_ids(self, ids, **params):
             body=dict(docs=docs)
         )
         if fields:
-            params['fields'] = fields
+            fields_params = process_fields_param(fields)
+            params.update(fields_params)
+
         documents = _ESDocs()
         documents._nefertari_meta = dict(
             start=_start,
@@ -419,13 +443,15 @@ def get_collection(self, **params):
         if '_count' in params:
             return self.do_count(_params)
 
-        # pop the fields before passing to search.
-        # ES does not support passing names of nested structures
-        _fields = _params.pop('fields', '')
+        fields = _params.pop('fields', '')
+        if fields:
+            fields_params = process_fields_param(fields)
+            _params.update(fields_params)
+
         documents = _ESDocs()
         documents._nefertari_meta = dict(
             start=_params['from_'],
-            fields=_fields)
+            fields=fields)
 
         try:
             data = ES.api.search(**_params)
@@ -439,7 +465,7 @@ def get_collection(self, **params):
             return documents
 
         for da in data['hits']['hits']:
-            _d = da['fields'] if _fields else da['_source']
+            _d = da['fields'] if fields else da['_source']
             _d['_score'] = da['_score']
             documents.append(dict2obj(_d))
 

tests/test_elasticsearch.py

@@ -40,6 +40,21 @@ def test_perform_request_no_index(self, mock_log):
 
 
 class TestHelperFunctions(object):
+    def test_process_fields_param_no_fields(self):
+        assert es.process_fields_param(None) is None
+
+    def test_process_fields_param_string(self):
+        assert es.process_fields_param('foo,bar') == {
+            'fields': ['foo', 'bar', '_type'],
+            '_source': False
+        }
+
+    def test_process_fields_param_list(self):
+        assert es.process_fields_param(['foo', 'bar']) == {
+            'fields': ['foo', 'bar', '_type'],
+            '_source': False
+        }
+
     @patch('nefertari.elasticsearch.ES')
     def test_includeme(self, mock_es):
         config = Mock()
@@ -368,15 +383,15 @@ def test_get_by_ids_fields(self, mock_mget):
         docs = obj.get_by_ids(documents, _limit=1, _fields=['name'])
         mock_mget.assert_called_once_with(
             body={'docs': [{'_index': 'foondex', '_type': 'story', '_id': 1}]},
-            fields=['name']
+            fields=['name', '_type'], _source=False
         )
         assert len(docs) == 1
         assert not hasattr(docs[0], '_id')
         assert not hasattr(docs[0], '_type')
         assert docs[0].name == 'bar'
         assert docs._nefertari_meta['total'] == 1
         assert docs._nefertari_meta['start'] == 0
-        assert docs._nefertari_meta['fields'] == ['name']
+        assert docs._nefertari_meta['fields'] == ['name', '_type']
 
     @patch('nefertari.elasticsearch.ES.api.mget')
     def test_get_by_ids_no_index_raise(self, mock_mget):
@@ -525,14 +540,16 @@ def test_get_collection_fields(self, mock_search):
         }
         docs = obj.get_collection(
             fields=['foo'], body={'foo': 'bar'}, from_=0)
-        mock_search.assert_called_once_with(body={'foo': 'bar'}, from_=0)
+        mock_search.assert_called_once_with(
+            body={'foo': 'bar'}, fields=['foo', '_type'], from_=0,
+            _source=False)
         assert len(docs) == 1
         assert docs[0].id == 1
         assert docs[0]._score == 2
         assert docs[0].foo == 'bar'
         assert docs._nefertari_meta['total'] == 4
         assert docs._nefertari_meta['start'] == 0
-        assert docs._nefertari_meta['fields'] == ['foo']
+        assert docs._nefertari_meta['fields'] == ['foo', '_type']
         assert docs._nefertari_meta['took'] == 2.8
 
     @patch('nefertari.elasticsearch.ES.api.search')