feat: Search Service Highlighting (#1856)

* feat: Search Service Highlighting

Signed-off-by: Allison Suarez Miranda <[email protected]>

* fixed unit test

Signed-off-by: Allison Suarez Miranda <[email protected]>

* put attrdict access in try in case of keyerror

Signed-off-by: Allison Suarez Miranda <[email protected]>

* stupid

Signed-off-by: Allison Suarez Miranda <[email protected]>

* ???

Signed-off-by: Allison Suarez Miranda <[email protected]>

* mock mathc

Signed-off-by: Allison Suarez Miranda <[email protected]>

* debug

Signed-off-by: Allison Suarez Miranda <[email protected]>

* all unit tests pass

Signed-off-by: Allison Suarez Miranda <[email protected]>

* lint

Signed-off-by: Allison Suarez Miranda <[email protected]>

* cinf fix

Signed-off-by: Allison Suarez Miranda <[email protected]>

* check that request has highlights

Signed-off-by: Allison Suarez Miranda <[email protected]>

* implemented feedback

Signed-off-by: Allison Suarez Miranda <[email protected]>

* empty checks

Signed-off-by: Allison Suarez Miranda <[email protected]>

* updated format response tests anhd put them in utils test file

Signed-off-by: Allison Suarez Miranda <[email protected]>

* fixed all tests

Signed-off-by: Allison Suarez Miranda <[email protected]>

* lint

Signed-off-by: Allison Suarez Miranda <[email protected]>

* mypy

Signed-off-by: Allison Suarez Miranda <[email protected]>

* isort

Signed-off-by: Allison Suarez Miranda <[email protected]>

* added another unit test

Signed-off-by: Allison Suarez Miranda <[email protected]>

* isort

Signed-off-by: Allison Suarez Miranda <[email protected]>

* major oopsie

Signed-off-by: Allison Suarez Miranda <[email protected]>

* some changes

Signed-off-by: Allison Suarez Miranda <[email protected]>

* updated test

Signed-off-by: Allison Suarez Miranda <[email protected]>

* latets pending marshmallow3-annotations bump

Signed-off-by: Allison Suarez Miranda <[email protected]>

* remove debugging line

Signed-off-by: Allison Suarez Miranda <[email protected]>

* isort

Signed-off-by: Allison Suarez Miranda <[email protected]>

* bumped marshmallow3-annotation to 1.1.0

Signed-off-by: Allison Suarez Miranda <[email protected]>

* typing extensions issue attempt to fix

Signed-off-by: Allison Suarez Miranda <[email protected]>

* had to convert sttrdict and lsit in the end

Signed-off-by: Allison Suarez Miranda <[email protected]>

* sort imports

Signed-off-by: Allison Suarez Miranda <[email protected]>

Author: allisonsuarez

common/setup.py

@@ -41,7 +41,7 @@
         'Flask>=1.0.2',
         'attrs>=19.0.0',
         'marshmallow>=3.0',
-        'marshmallow3-annotations>=1.0.0'
+        'marshmallow3-annotations>=1.1.0'
     ],
     extras_require={
         'all': requirements_dev

databuilder/setup.py

@@ -60,7 +60,7 @@
     'Flask==1.0.2',
     'gremlinpython==3.4.3',
     'requests-aws4auth==1.1.0',
-    'typing-extensions==3.7.4',
+    'typing-extensions==4.0.0',
     'overrides==2.5',
     'boto3==1.17.23'
 ]

requirements-common.txt

@@ -17,7 +17,7 @@ itsdangerous<=2.0.1
 Jinja2>=2.10.1,<3.1
 jsonschema>=3.0.1,<4.0
 marshmallow>=3.0,<=3.6
-marshmallow3-annotations>=1.0.0
+marshmallow3-annotations>=1.1.0
 pytz==2021.1
 requests>=2.25.0
 requests-aws4auth==1.1.0

search/search_service/api/search.py

@@ -1,13 +1,14 @@
 # Copyright Contributors to the Amundsen project.
 # SPDX-License-Identifier: Apache-2.0
 
-import json
 from http import HTTPStatus
 from typing import (  # noqa: F401
-    Any, Iterable, List,
+    Any, Dict, Iterable, List,
 )
 
-from amundsen_common.models.search import SearchRequestSchema, SearchResponseSchema
+from amundsen_common.models.search import (
+    HighlightOptions, SearchRequestSchema, SearchResponseSchema,
+)
 from flasgger import swag_from
 from flask_restful import Resource, request
 
@@ -29,14 +30,17 @@ def post(self) -> Iterable[Any]:
         Fetch search results
         :return: json payload of schema
         """
-
-        request_data = SearchRequestSchema().loads(json.dumps(request.get_json()))
+        request_data = SearchRequestSchema().load(request.json, partial=False)
 
         resources: List[AmundsenResource] = []
-        for r in request.get_json().get('resource_types'):
+        highlight_options: Dict[AmundsenResource, HighlightOptions] = {}
+
+        for r in request_data.resource_types:
             resource = RESOURCE_STR_MAPPING.get(r)
             if resource:
                 resources.append(resource)
+                if request_data.highlight_options.get(r):
+                    highlight_options[resource] = request_data.highlight_options.get(r)
             else:
                 err_msg = f'Search for invalid resource "{r}" requested'
                 return {'message': err_msg}, HTTPStatus.BAD_REQUEST
@@ -46,7 +50,8 @@ def post(self) -> Iterable[Any]:
                                                       page_index=request_data.page_index,
                                                       results_per_page=request_data.results_per_page,
                                                       resource_types=resources,
-                                                      filters=request_data.filters)
+                                                      filters=request_data.filters,
+                                                      highlight_options=highlight_options)
             return SearchResponseSchema().dump(search_results), HTTPStatus.OK
 
         except RuntimeError as e:

search/search_service/config.py

@@ -125,4 +125,4 @@ class AwsSearchConfig(LocalConfig):
             connection_class=RequestsHttpConnection
         )
 
-        PROXY_CLIENT_KEY = client
+        ELASTICSEARCH_CLIENT = client

search/search_service/proxy/base.py

@@ -7,7 +7,9 @@
 )
 
 from amundsen_common.models.api.health_check import HealthCheck
-from amundsen_common.models.search import Filter, SearchResponse
+from amundsen_common.models.search import (
+    Filter, HighlightOptions, SearchResponse,
+)
 
 from search_service.models.dashboard import SearchDashboardResult
 from search_service.models.feature import SearchFeatureResult
@@ -36,7 +38,8 @@ def search(self, *,
                page_index: int,
                results_per_page: int,
                resource_types: List[Resource],
-               filters: List[Filter]) -> SearchResponse:
+               filters: List[Filter],
+               highlight_options: Dict[Resource, HighlightOptions]) -> SearchResponse:
         pass
 
     @abstractmethod

search/search_service/proxy/elasticsearch.py

@@ -11,7 +11,9 @@
 from amundsen_common.models.index_map import (
     FEATURE_INDEX_MAP, TABLE_INDEX_MAP, USER_INDEX_MAP,
 )
-from amundsen_common.models.search import Filter, SearchResponse
+from amundsen_common.models.search import (
+    Filter, HighlightOptions, SearchResponse,
+)
 from elasticsearch import Elasticsearch
 from elasticsearch.exceptions import ConnectionError as ElasticConnectionError, NotFoundError
 from elasticsearch_dsl import Search, query
@@ -812,7 +814,8 @@ def search(self, *,
                page_index: int,
                results_per_page: int,
                resource_types: List[Resource],
-               filters: List[Filter]) -> SearchResponse:
+               filters: List[Filter],
+               highlight_options: Dict[Resource, HighlightOptions]) -> SearchResponse:
         LOGGING.warn(DEPRECATION_MSG)
         return SearchResponse(msg=DEPRECATION_MSG,
                               page_index=0,

search/search_service/proxy/es_proxy_utils.py

@@ -1,7 +1,16 @@
 # Copyright Contributors to the Amundsen project.
 # SPDX-License-Identifier: Apache-2.0
 
+import logging
 from enum import Enum
+from typing import (
+    Any, Dict, List,
+)
+
+from amundsen_common.models.search import SearchResponse
+from elasticsearch_dsl.response import Response
+from elasticsearch_dsl.response.hit import Hit
+from elasticsearch_dsl.utils import AttrDict, AttrList
 
 
 class Resource(Enum):
@@ -12,13 +21,97 @@ class Resource(Enum):
 
 
 RESOURCE_STR_MAPPING = {
-    'table': Resource.TABLE,
-    'dashboard': Resource.DASHBOARD,
-    'feature': Resource.FEATURE,
-    'user': Resource.USER,
+    "table": Resource.TABLE,
+    "dashboard": Resource.DASHBOARD,
+    "feature": Resource.FEATURE,
+    "user": Resource.USER,
 }
 
 
 def get_index_for_resource(resource_type: Resource) -> str:
     resource_str = resource_type.name.lower()
     return f"{resource_str}_search_index"
+
+
+class SearchHit():
+    # custom wrapper for elasticsearch_dsl Hit
+    def __init__(self, hit: Hit, fields_mapping: Dict):
+        self.hit = hit
+        self.fields_mapping = fields_mapping
+
+    def _convert_attr_value_to_native(self, attr_value: Any) -> Any:
+        if type(attr_value) is AttrDict:
+            return attr_value.to_dict()
+        elif type(attr_value) is AttrList:
+            return list(attr_value)
+        return attr_value
+
+    def to_search_result(self) -> Dict:
+        result = {}
+        for field, mapped_field in self.fields_mapping.items():
+            field_value = None
+            # get field name instead of subfield
+            mapped_field = mapped_field.split('.')[0]
+
+            if field != mapped_field and hasattr(self.hit, mapped_field):
+                # if the field name doesn't already match get mapped one
+                field_value = getattr(self.hit, mapped_field)
+            elif hasattr(self.hit, field):
+                field_value = getattr(self.hit, field)
+
+            result[field] = self._convert_attr_value_to_native(field_value)
+        result["search_score"] = self.hit.meta.score
+        return result
+
+    def get_highlights(self) -> Dict:
+        highlights = {}
+        try:
+            for highlighted_field, highlighted_value in self.hit.meta.highlight.to_dict().items():
+                parent_field_name = highlighted_field.split('.')[0]
+                highlights[parent_field_name] = highlighted_value
+            return highlights
+
+        except AttributeError:
+            # response doesn't have highlights
+            return highlights
+
+
+def format_resource_response(response: Response, fields_mapping: Dict) -> Dict:
+    results = []
+    for hit in response.hits:
+        search_hit = SearchHit(hit=hit, fields_mapping=fields_mapping)
+        results.append({
+            **search_hit.to_search_result(),
+            'highlight': search_hit.get_highlights()
+        })
+    return {
+        "results": results,
+        "total_results": response.hits.total.value,
+    }
+
+
+def create_search_response(page_index: int,  # noqa: C901
+                           results_per_page: int,
+                           responses: List[Response],
+                           resource_types: List[Resource],
+                           resource_to_field_mapping: Dict) -> SearchResponse:
+    results_per_resource = {}
+    # responses are returned in the order in which the searches appear in msearch request
+    for resource, response in zip(resource_types, responses):
+        msg = ''
+        status_code = 200
+        if response.success():
+            msg = 'Success'
+            results_per_resource[resource.name.lower()] = \
+                format_resource_response(response=response,
+                                         fields_mapping=resource_to_field_mapping[resource])
+        else:
+            msg = f'Query response for {resource} returned an error: {response.to_dict()}'
+            status_code = 500
+            logging.error(msg)
+
+    return SearchResponse(msg=msg,
+                          page_index=page_index,
+                          results_per_page=results_per_page,
+                          results=results_per_resource,
+                          status_code=status_code)

search/search_service/proxy/es_proxy_v2.py

@@ -8,18 +8,20 @@
 )
 
 from amundsen_common.models.api import health_check
-from amundsen_common.models.search import Filter, SearchResponse
+from amundsen_common.models.search import (
+    Filter, HighlightOptions, SearchResponse,
+)
 from elasticsearch import Elasticsearch
 from elasticsearch.exceptions import ConnectionError as ElasticConnectionError, ElasticsearchException
 from elasticsearch_dsl import (
     MultiSearch, Q, Search,
 )
 from elasticsearch_dsl.query import MultiMatch
 from elasticsearch_dsl.response import Response
-from elasticsearch_dsl.utils import AttrDict, AttrList
+from elasticsearch_dsl.utils import AttrList
 from werkzeug.exceptions import InternalServerError
 
-from search_service.proxy.es_proxy_utils import Resource
+from search_service.proxy.es_proxy_utils import Resource, create_search_response
 
 LOGGER = logging.getLogger(__name__)
 
@@ -237,56 +239,6 @@ def _build_elasticsearch_query(self, *,
 
         return es_query
 
-    def _format_response(self, page_index: int,
-                         results_per_page: int,
-                         responses: List[Response],
-                         resource_types: List[Resource]) -> SearchResponse:
-        resource_types_str = [r.name.lower() for r in resource_types]
-        no_results_for_resource = {
-            "results": [],
-            "total_results": 0
-        }
-        results_per_resource = {resource: no_results_for_resource for resource in resource_types_str}
-
-        for r in responses:
-            if r.success():
-                if len(r.hits.hits) > 0:
-                    resource_type = r.hits.hits[0]._source['resource_type']
-                    results = []
-                    for search_result in r.hits.hits:
-                        # mapping gives all the fields in the response
-                        result = {}
-                        fields = self.RESOUCE_TO_MAPPING[Resource[resource_type.upper()]]
-                        for f in fields.keys():
-                            # remove "keyword" from mapping value
-                            field = fields[f].split('.')[0]
-                            try:
-                                result_for_field = search_result._source[field]
-                                # AttrList and AttrDict are not json serializable
-                                if type(result_for_field) is AttrList:
-                                    result_for_field = list(result_for_field)
-                                elif type(result_for_field) is AttrDict:
-                                    result_for_field = result_for_field.to_dict()
-                                result[f] = result_for_field
-                            except KeyError:
-                                logging.debug(f'Field: {field} missing in search response.')
-                                pass
-                        result["search_score"] = search_result._score
-                        results.append(result)
-                    # replace empty results with actual results
-                    results_per_resource[resource_type] = {
-                        "results": results,
-                        "total_results": r.hits.total.value
-                    }
-            else:
-                raise InternalServerError(f"Request to Elasticsearch failed: {r.failures}")
-
-        return SearchResponse(msg="Success",
-                              page_index=page_index,
-                              results_per_page=results_per_page,
-                              results=results_per_resource,
-                              status_code=200)
-
     def execute_queries(self, queries: Dict[Resource, Q],
                         page_index: int,
                         results_per_page: int) -> List[Response]:
@@ -316,8 +268,9 @@ def search(self, *,
                page_index: int,
                results_per_page: int,
                resource_types: List[Resource],
-               filters: List[Filter]) -> SearchResponse:
-        if resource_types == []:
+               filters: List[Filter],
+               highlight_options: Dict[Resource, HighlightOptions]) -> SearchResponse:
+        if not resource_types:
             # if resource types are not defined then search all resources
             resource_types = self.PRIMARY_ENTITIES
 
@@ -332,10 +285,11 @@ def search(self, *,
                                          page_index=page_index,
                                          results_per_page=results_per_page)
 
-        formatted_response = self._format_response(page_index=page_index,
-                                                   results_per_page=results_per_page,
-                                                   responses=responses,
-                                                   resource_types=resource_types)
+        formatted_response = create_search_response(page_index=page_index,
+                                                    results_per_page=results_per_page,
+                                                    responses=responses,
+                                                    resource_types=resource_types,
+                                                    resource_to_field_mapping=self.RESOUCE_TO_MAPPING)
 
         return formatted_response