api: save original index when prefixing

invenio_search/__init__.py

@@ -287,7 +287,7 @@ def index():
 
 from __future__ import absolute_import, print_function
 
-from .api import RecordsSearch
+from .api import RecordsSearch, UnPrefixedRecordsSearch
 from .ext import InvenioSearch
 from .proxies import current_search, current_search_client
 from .version import __version__
@@ -296,6 +296,7 @@ def index():
     '__version__',
     'InvenioSearch',
     'RecordsSearch',
+    'UnPrefixedRecordsSearch',
     'current_search',
     'current_search_client',
 )

invenio_search/api.py

@@ -11,6 +11,7 @@
 import hashlib
 from functools import partial
 
+import six
 from elasticsearch import VERSION as ES_VERSION
 from elasticsearch_dsl import FacetedSearch, Search
 from elasticsearch_dsl.faceted_search import FacetedResponse
@@ -65,7 +66,7 @@ def __ge__(self, other):
         return False
 
 
-class RecordsSearch(Search):
+class BaseRecordsSearch(Search):
     """Example subclass for searching records using Elastic DSL."""
 
     class Meta:
@@ -84,15 +85,7 @@ class Meta:
 
     def __init__(self, **kwargs):
         """Use Meta to set kwargs defaults."""
-        # at object instantiation, kwargs['index'] is not defined.
-        # Elasticsearch-dsl-py re-instantiated the object at each search
-        # by cloning it and passing as kwargs the list of indices
-        # kwargs['index'] = ['index-name1', 'index-name2']
-        if not kwargs.get('index') and getattr(self.Meta, 'index', None):
-            _index_name = prefix_index(app=current_app,
-                                       index=getattr(self.Meta, 'index', None))
-            kwargs.setdefault('index', _index_name)
-
+        kwargs.setdefault('index', getattr(self.Meta, 'index', None))
         kwargs.setdefault('doc_type', getattr(self.Meta, 'doc_types', None))
         kwargs.setdefault('using', current_search_client)
         kwargs.setdefault('extra', {})
@@ -101,7 +94,7 @@ def __init__(self, **kwargs):
         if min_score:
             kwargs['extra'].update(min_score=min_score)
 
-        super(RecordsSearch, self).__init__(**kwargs)
+        super(BaseRecordsSearch, self).__init__(**kwargs)
 
         default_filter = getattr(self.Meta, 'default_filter', None)
         if default_filter:
@@ -186,3 +179,60 @@ def _get_user_hash(self):
             alg.update(user_hash.encode('utf8'))
             return alg.hexdigest()
         return None
+
+
+class PrefixedIndexList(list):
+    """Custom list type for avoiding double prefixing."""
+
+    pass
+
+
+class RecordsSearch(BaseRecordsSearch):
+    """Example subclass for searching records using index prefixing."""
+
+    def __init__(self, **kwargs):
+        """Using PrefixedIndexList type to avoid double prefixing."""
+        # at object instantiation, kwargs['index'] is not defined.
+        # Elasticsearch-dsl-py re-instantiated the object at each search
+        # by cloning it and passing as kwargs the list of indices
+        # kwargs['index'] = ['index-name1', 'index-name2']
+        _index_param = kwargs.get('index', getattr(self.Meta, 'index', None))
+        if not isinstance(_index_param, PrefixedIndexList):
+            if isinstance(_index_param, (tuple, list)):
+                _prefixed_index_list = [
+                    prefix_index(app=current_app,
+                                 index=_index) for _index in _index_param]
+                kwargs.update({'index': _prefixed_index_list})
+            elif isinstance(_index_param, six.string_types):
+                _splitted_index = _index_param.strip().split(',')
+                if len(_splitted_index) > 1:
+                    _prefix_index_list = [
+                        prefix_index(current_app, _index)
+                        for _index in _splitted_index]
+                    _prefix_index_param = ','.join(_prefix_index_list)
+                    kwargs.update({
+                        'index': _prefix_index_param})
+                else:
+                    kwargs.update({
+                        'index': prefix_index(app=current_app,
+                                              index=_index_param)})
+                _index_param = [_index_param]
+            self._original_index = _index_param
+
+        super(RecordsSearch, self).__init__(**kwargs)
+        if self._index:
+            self._index = PrefixedIndexList(self._index)
+
+    def _clone(self):
+        """Clone `_original_index` attribute.
+
+        During re-instantiation Elasticsearch-dsl-py calls `self._clone`
+        to copy over the search object. We override the method so we can
+        copy the `_original_index` attribute.
+        """
+        s = super(RecordsSearch, self)._clone()
+        s._original_index = self._original_index
+        return s
+
+
+UnPrefixedRecordsSearch = BaseRecordsSearch

tests/test_query.py

@@ -16,7 +16,7 @@
 from elasticsearch_dsl import Q, Search
 from flask import request
 
-from invenio_search.api import DefaultFilter, RecordsSearch
+from invenio_search.api import BaseRecordsSearch, DefaultFilter, RecordsSearch
 
 
 def test_empty_query(app):
@@ -106,7 +106,7 @@ def params(self, **kwargs):
 
 def test_es_preference_param_no_request(app):
     """Test that the preference param is not added when not in a request."""
-    RecordsSearch.__bases__ = (SpySearch,)
+    BaseRecordsSearch.__bases__ = (SpySearch,)
 
     rs = RecordsSearch()
     new_rs = rs.with_preference_param()
@@ -115,7 +115,7 @@ def test_es_preference_param_no_request(app):
 
 def test_es_preference_param(app):
     """Test the preference param is correctly added in a request."""
-    RecordsSearch.__bases__ = (SpySearch,)
+    BaseRecordsSearch.__bases__ = (SpySearch,)
 
     with app.test_request_context('/', headers={'User-Agent': 'Chrome'},
                                   environ_base={'REMOTE_ADDR': '212.54.1.8'}):
@@ -142,3 +142,70 @@ def test_elasticsearch_query_min_score(app):
     search_dict = q.to_dict()
     assert 'min_score' in search_dict
     assert search_dict['min_score'] == app.config['SEARCH_RESULTS_MIN_SCORE']
+
+
+def _test_original_index_is_stored_when_prefixing(q, prefixed_index,
+                                                  original_index):
+    """Test original index is stored."""
+
+    q = q.query(Q('match', title='Higgs'))
+    q = q.sort()
+    search_dict = q.to_dict()
+    assert 'query' in search_dict
+    assert q._index == prefixed_index
+    assert q._original_index == original_index
+
+
+def test_prefix_index_from_meta(app):
+    """Test that index is prefixed when you use `Meta.index` field."""
+    class TestSearch(RecordsSearch):
+        class Meta:
+            index = 'myindex'
+
+    prefix_value = 'myprefix-'
+    index_value = TestSearch.Meta.index
+    app.config['SEARCH_INDEX_PREFIX'] = prefix_value
+
+    prefixed_index = ['{}{}'.format(prefix_value, index_value)]
+    q = TestSearch()
+    _test_original_index_is_stored_when_prefixing(q, prefixed_index,
+                                                  [index_value])
+
+
+def test_prefix_index_from_kwargs(app):
+    """Test that index is prefixed when pass it through kwargs."""
+    prefix_value = 'myprefix-'
+    index_value = 'myindex'
+    app.config['SEARCH_INDEX_PREFIX'] = prefix_value
+
+    prefixed_index = ['{}{}'.format(prefix_value, index_value)]
+    q = RecordsSearch(index=index_value)
+    _test_original_index_is_stored_when_prefixing(q, prefixed_index,
+                                                  [index_value])
+
+
+def test_prefix_index_list(app):
+    """Test that index is prefixed when pass it through kwargs."""
+    prefix_value = 'myprefix-'
+    index_value = ['myindex', 'myanotherindex']
+    app.config['SEARCH_INDEX_PREFIX'] = prefix_value
+
+    prefixed_index = ['{}{}'.format(prefix_value, _index)
+                      for _index in index_value]
+
+    q = RecordsSearch(index=index_value)
+    _test_original_index_is_stored_when_prefixing(q, prefixed_index,
+                                                  index_value)
+
+
+def test_prefix_multi_index_string(app):
+    """Test that index is prefixed when pass it through kwargs."""
+    prefix_value = 'myprefix-'
+    index_value = 'myindex,myanotherindex'
+    app.config['SEARCH_INDEX_PREFIX'] = prefix_value
+
+    prefixed_index = [','.join(['{}{}'.format(prefix_value, _index)
+                                for _index in index_value.split(',')])]
+    q = RecordsSearch(index=index_value)
+    _test_original_index_is_stored_when_prefixing(q, prefixed_index,
+                                                  [index_value])