add highlighting options

Makefile

@@ -12,7 +12,7 @@ endif
 test: clean
 	./manage.py test --settings=$(test_settings) --with-coverage --with-ignore-docstrings \
 		--exclude-dir=notesserver/settings --cover-inclusive --cover-branches \
-		--ignore-files=search_indexes.py \
+		--ignore-files=search_indexes.py --ignore-files=highlight.py\
 		--cover-html --cover-html-dir=build/coverage/html/ \
 		--cover-xml --cover-xml-file=build/coverage/coverage.xml \
 		$(foreach package,$(PACKAGES),--cover-package=$(package)) \

notesapi/v1/tests/test_views.py

@@ -401,12 +401,17 @@ class AnnotationSearchlViewTests(BaseAnnotationViewTests):
         results = self._get_search_results()
         self.assertEqual(results['total'], 2)
 
-        # FIXME class and tag
-        results = self._get_search_results(text="first", highlight=True, highlight_class='class', highlight_tag='tag')
+        results = self._get_search_results(text="first", highlight=True)
         self.assertEqual(results['total'], 1)
         self.assertEqual(len(results['rows']), 1)
         self.assertEqual(results['rows'][0]['text'], '<em>First</em> note')
 
+        results = self._get_search_results(text="first", highlight=True, highlight_tag='tag')
+        self.assertEqual(results['rows'][0]['text'], '<tag>First</tag> note')
+
+        results = self._get_search_results(text="first", highlight=True, highlight_tag='tag', highlight_class='klass')
+        self.assertEqual(results['rows'][0]['text'], '<tag class="klass">First</tag> note')
+
     def test_search_ordering(self):
         """
         Tests ordering of search results.

notesapi/v1/views.py

@@ -13,7 +13,7 @@ from rest_framework.views import APIView
 from notesapi.v1.models import Note
 
 if not settings.ES_DISABLED:
-    from haystack.query import SearchQuerySet
+    from notesserver.highlight import SearchQuerySet
 
 log = logging.getLogger(__name__)
 
@@ -60,7 +60,16 @@ class AnnotationSearchView(APIView):
         ).order_by('-updated')
 
         if params.get('highlight'):
-            query = query.highlight()
+            tag = params.get('highlight_tag', 'em')
+            klass = params.get('highlight_class')
+            opts = {
+                'pre_tags': ['<{tag}{klass_str}>'.format(
+                    tag=tag,
+                    klass_str=' class="{}"'.format(klass) if klass else ''
+                )],
+                'post_tags': ['</{tag}>'.format(tag=tag)],
+            }
+            query = query.highlight(**opts)
 
         results = []
         for item in query:

notesserver/highlight.py

+"""
+django-haystack does not support passing additional highlighting parameters
+to backends, so we use our subclassed one which does.
+"""
+
+import haystack
+from haystack.backends.elasticsearch_backend import (
+    ElasticsearchSearchEngine as OrigElasticsearchSearchEngine,
+    ElasticsearchSearchQuery as OrigElasticsearchSearchQuery,
+    ElasticsearchSearchBackend as OrigElasticsearchSearchBackend)
+from haystack.query import SearchQuerySet as OrigSearchQuerySet
+
+
+class SearchQuerySet(OrigSearchQuerySet):
+    def highlight(self, **kwargs):
+        """Adds highlighting to the results."""
+        clone = self._clone()
+        clone.query.add_highlight(**kwargs)
+        return clone
+
+
+class ElasticsearchSearchQuery(OrigElasticsearchSearchQuery):
+    def add_highlight(self, **kwargs):
+        """Adds highlighting to the search results."""
+        self.highlight = kwargs or True
+
+
+class ElasticsearchSearchBackend(OrigElasticsearchSearchBackend):
+    """
+    Subclassed backend that lets user modify highlighting options
+    """
+    def build_search_kwargs(self, *args, **kwargs):
+        res = super(ElasticsearchSearchBackend, self).build_search_kwargs(*args, **kwargs)
+        index = haystack.connections[self.connection_alias].get_unified_index()
+        content_field = index.document_field
+        highlight = kwargs.get('highlight')
+        if highlight:
+            highlight_options = {
+                'fields': {
+                    content_field: {'store': 'yes'},
+                }
+            }
+            if isinstance(highlight, dict):
+                highlight_options.update(highlight)
+            res['highlight'] = highlight_options
+        return res
+
+
+class ElasticsearchSearchEngine(OrigElasticsearchSearchEngine):
+    backend = ElasticsearchSearchBackend
+    query = ElasticsearchSearchQuery

notesserver/settings/common.py

@@ -21,7 +21,7 @@ CLIENT_SECRET = 'edx-notes-secret'
 ES_DISABLED = False
 HAYSTACK_CONNECTIONS = {
     'default': {
-        'ENGINE': 'haystack.backends.elasticsearch_backend.ElasticsearchSearchEngine',
+        'ENGINE': 'notesserver.highlight.ElasticsearchSearchEngine',
         'URL': 'http://127.0.0.1:9200/',
         'INDEX_NAME': 'notes_index',
     },

notesserver/settings/test.py

@@ -12,7 +12,7 @@ INSTALLED_APPS += ('django_nose',)
 
 HAYSTACK_CONNECTIONS = {
     'default': {
-        'ENGINE': 'haystack.backends.elasticsearch_backend.ElasticsearchSearchEngine',
+        'ENGINE': 'notesserver.highlight.ElasticsearchSearchEngine',
         'URL': 'http://127.0.0.1:9200/',
         'INDEX_NAME': 'notes_index_test',
     },