Fix ES proxy to handle additional attribute (#35)

* Fix ES proxy to handle additional attribute

* tmp remove test

* Fix

search_service/__init__.py

@@ -9,6 +9,7 @@
 from typing import Dict, Any    # noqa: F401
 
 from search_service.api.table import SearchTableAPI, SearchTableFieldAPI
+from search_service.api.user import SearchUserAPI
 from search_service.api.healthcheck import healthcheck
 
 # For customized flask use below arguments to override.
@@ -69,6 +70,9 @@ def create_app(*, config_module_class: str) -> Flask:
     api.add_resource(SearchTableFieldAPI,
                      '/search/field/<field_name>/field_val/<field_value>')
 
+    # User Search API
+    api.add_resource(SearchUserAPI, '/search_user')
+
     app.register_blueprint(api_bp)
 
     return app

search_service/api/table.py

@@ -99,11 +99,11 @@ def get(self, *, field_name: str,
 
         try:
             results = self.proxy.fetch_table_search_results_with_field(
-                query_term=args.get('query_term'),
+                query_term=args.get('query_term', ''),
                 field_name=field_name,
                 field_value=field_value,
                 page_index=args['page_index'],
-                index=args['index']
+                index=args.get('index')
             )
 
             return results, HTTPStatus.OK

search_service/api/user.py

@@ -55,7 +55,7 @@ def get(self) -> Iterable[Any]:
             results = self.proxy.fetch_user_search_results(
                 query_term=args['query_term'],
                 page_index=args['page_index'],
-                index=args.get(args['index'])
+                index=args.get('index')
             )
 
             return results, HTTPStatus.OK

search_service/models/base.py

@@ -0,0 +1,13 @@
+from abc import ABCMeta, abstractmethod
+from typing import Set
+
+
+class Base(metaclass=ABCMeta):
+    """
+    A base class for ES model
+    """
+
+    @abstractmethod
+    def get_attrs(cls) -> Set:
+        # return a set of attributes for the class
+        ...

search_service/models/table.py

@@ -1,7 +1,8 @@
-from typing import Iterable
+from typing import Iterable, Set
+from .base import Base
 
 
-class Table:
+class Table(Base):
     def __init__(self, *,
                  name: str,
                  key: str,
@@ -22,6 +23,20 @@ def __init__(self, *,
         self.tags = tags
         self.last_updated_epoch = last_updated_epoch
 
+    @classmethod
+    def get_attrs(cls) -> Set:
+        return {
+            'name',
+            'key',
+            'description',
+            'cluster',
+            'database',
+            'schema_name',
+            'column_names',
+            'tags',
+            'last_updated_epoch'
+        }
+
     def __repr__(self) -> str:
         return 'Table(name={!r}, key={!r}, description={!r}, ' \
                'cluster={!r} database={!r}, schema_name={!r}, column_names={!r}, ' \

search_service/models/user.py

@@ -1,4 +1,8 @@
-class User:
+from typing import Set
+from .base import Base
+
+
+class User(Base):
     def __init__(self, *,
                  first_name: str,
                  last_name: str,
@@ -19,6 +23,20 @@ def __init__(self, *,
         self.is_active = is_active
         self.employee_type = employee_type
 
+    @classmethod
+    def get_attrs(cls) -> Set:
+        return {
+            'name',
+            'first_name',
+            'last_name',
+            'team_name',
+            'email',
+            'manager_email',
+            'github_username',
+            'is_active',
+            'employee_type'
+        }
+
     def __repr__(self) -> str:
         return 'User(name={!r}, first_name={!r}, last_name={!r}, ' \
                'team_name={!r} email={!r}, manager_email={!r}, github_username={!r}, ' \

search_service/proxy/elasticsearch.py

@@ -72,15 +72,23 @@ def _get_search_result(self, page_index: int,
         for hit in response:
 
             try:
-                results.append(model(**vars(hit)))
+                # ES hit: {'_d_': {'key': xxx...}
+                es_payload = hit.__dict__.get('_d_', {})
+                if not es_payload:
+                    raise Exception('The ES doc not contain required field')
+                result = {}
+
+                for attr, val in es_payload.items():
+                    if attr in model.get_attrs():
+                        result[attr] = val
+                results.append(model(**result))
             except Exception:
                 LOGGING.exception('The record doesnt contain specified field.')
 
         return SearchResult(total_results=response.hits.total,
                             results=results)
 
-    def _search_helper(self, query_term: str,
-                       page_index: int,
+    def _search_helper(self, page_index: int,
                        client: Search,
                        query_name: dict,
                        model: Any) -> SearchResult:
@@ -91,21 +99,19 @@ def _search_helper(self, query_term: str,
           2. Uses multi match query to search term in multiple fields.
           `Link https://www.elastic.co/guide/en/elasticsearch/reference/current/query-dsl-multi-match-query.html`_
 
-        :param query_term:
         :param page_index:
         :param client:
         :param query_name: name of query to query the ES
         :return:
         """
 
-        if query_term and query_name:
+        if query_name:
             q = query.Q(query_name)
             client = client.query(q)
-            return self._get_search_result(page_index=page_index,
-                                           client=client,
-                                           model=model)
-        else:
-            raise Exception('Query term and query name must be provided!')
+
+        return self._get_search_result(page_index=page_index,
+                                       client=client,
+                                       model=model)
 
     def _search_wildcard_helper(self, field_value: str,
                                 page_index: int,
@@ -118,6 +124,7 @@ def _search_wildcard_helper(self, field_value: str,
         :param page_index:
         :param client:
         :param field_name
+        :param query_name: name of query
         :return:
         """
         if field_value and field_name:
@@ -166,17 +173,7 @@ def fetch_table_search_results_with_field(self, *,
             'column': 'column_names.raw'
         }
 
-        # Convert field name to actual type in ES doc
-        field_name = mapping[field_name]
-
-        # We allow user to use ? * for wildcard support
-        m = re.search('[\?\*]', field_value)
-        if m:
-            return self._search_wildcard_helper(field_value=field_value,
-                                                page_index=page_index,
-                                                client=s,
-                                                field_name=field_name)
-        else:
+        if query_term:
             query_name = {
                 "function_score": {
                     "query": {
@@ -196,9 +193,23 @@ def fetch_table_search_results_with_field(self, *,
                     }
                 }
             }
-            s = s.filter('term', **{field_name: field_value})
-            return self._search_helper(query_term=query_term,
-                                       page_index=page_index,
+        else:
+            query_name = {}
+
+        # Convert field name to actual type in ES doc
+        new_field_name = mapping[field_name]
+
+        # We allow user to use ? * for wildcard support
+        m = re.search('[\?\*]', field_value)
+        if m:
+            return self._search_wildcard_helper(field_value=field_value,
+                                                page_index=page_index,
+                                                client=s,
+                                                field_name=new_field_name)
+        else:
+
+            s = s.filter('term', **{new_field_name: field_value})
+            return self._search_helper(page_index=page_index,
                                        client=s,
                                        query_name=query_name,
                                        model=Table)
@@ -243,8 +254,7 @@ def fetch_table_search_results(self, *,
             }
         }
 
-        return self._search_helper(query_term=query_term,
-                                   page_index=page_index,
+        return self._search_helper(page_index=page_index,
                                    client=s,
                                    query_name=query_name,
                                    model=Table)
@@ -280,8 +290,7 @@ def fetch_user_search_results(self, *,
             }
         }
 
-        return self._search_helper(query_term=query_term,
-                                   page_index=page_index,
+        return self._search_helper(page_index=page_index,
                                    client=s,
                                    query_name=query_name,
                                    model=User)

setup.py

@@ -1,6 +1,6 @@
 from setuptools import setup, find_packages
 
-__version__ = '1.1.2rc1'
+__version__ = '1.1.2'
 
 
 setup(

tests/unit/proxy/test_elasticsearch.py

@@ -1,3 +1,4 @@
+from typing import Any
 import unittest
 from unittest.mock import patch, MagicMock
 
@@ -33,6 +34,36 @@ def __init__(self, *,
         self.last_updated_epoch = last_updated_epoch
 
 
+class MockKVSearchResult:
+    def __init__(self, *,
+                 first_name: str,
+                 last_name: str,
+                 name: str,
+                 team_name: str,
+                 email: str,
+                 manager_email: str,
+                 github_username: str,
+                 is_active: bool,
+                 employee_type: str,
+                 new_attr: str) -> None:
+        self.name = name
+        self.first_name = first_name
+        self.last_name = last_name
+        self.team_name = team_name
+        self.email = email
+        self.manager_email = manager_email
+        self.github_username = github_username
+        self.is_active = is_active
+        self.employee_type = employee_type
+        self.new_attr = new_attr
+
+
+class Response:
+    def __init__(self,
+                 result: Any):
+        self._d_ = result
+
+
 class TestElasticsearchProxy(unittest.TestCase):
 
     def setUp(self) -> None:
@@ -73,15 +104,16 @@ def setUp(self) -> None:
                                   tags=['match'],
                                   last_updated_epoch=1527283287)
 
-        self.mock_result4 = User(name='First Last',
-                                 first_name='First',
-                                 last_name='Last',
-                                 team_name='Test team',
-                                 email='[email protected]',
-                                 github_username='ghub',
-                                 manager_email='[email protected]',
-                                 is_active=True,
-                                 employee_type='FTE')
+        self.mock_result4 = MockKVSearchResult(name='First Last',
+                                               first_name='First',
+                                               last_name='Last',
+                                               team_name='Test team',
+                                               email='[email protected]',
+                                               github_username='ghub',
+                                               manager_email='[email protected]',
+                                               is_active=True,
+                                               employee_type='FTE',
+                                               new_attr='aaa')
 
     def test_setup_client(self) -> None:
         self.es_proxy = ElasticsearchProxy(
@@ -135,7 +167,7 @@ def test_search_with_one_table_result(self,
 
         mock_results = MagicMock()
         mock_results.hits.total = 1
-        mock_results.__iter__.return_value = [self.mock_result1]
+        mock_results.__iter__.return_value = [Response(result=vars(self.mock_result1))]
         mock_search.return_value = mock_results
 
         expected = SearchResult(total_results=1,
@@ -165,7 +197,8 @@ def test_search_with_multiple_result(self,
 
         mock_results = MagicMock()
         mock_results.hits.total = 2
-        mock_results.__iter__.return_value = [self.mock_result1, self.mock_result2]
+        mock_results.__iter__.return_value = [Response(result=vars(self.mock_result1)),
+                                              Response(result=vars(self.mock_result2))]
         mock_search.return_value = mock_results
 
         expected = SearchResult(total_results=2,
@@ -264,38 +297,13 @@ def test_search_regex_match_field(self,
                              vars(expected.results[0]),
                              "Search result doesn't match with expected result!")
 
-    @patch('search_service.proxy.elasticsearch.ElasticsearchProxy._search_helper')
-    def test_search_user_match(self, mock_search: MagicMock) -> None:
-
-        mock_search.return_value = SearchResult(total_results=1,
-                                                results=[self.mock_result4])
-
-        expected = SearchResult(total_results=1,
-                                results=[User(name='First Last',
-                                              first_name='First',
-                                              last_name='Last',
-                                              team_name='Test team',
-                                              email='[email protected]',
-                                              github_username='ghub',
-                                              manager_email='[email protected]',
-                                              is_active=True,
-                                              employee_type='FTE')])
-
-        resp = self.es_proxy.fetch_user_search_results(query_term='First',
-                                                       index='user_search_index')
-        self.assertEquals(resp.total_results, expected.total_results)
-
-        self.assertDictEqual(vars(resp.results[0]),
-                             vars(expected.results[0]),
-                             "Search result doesn't match with expected result!")
-
     @patch('elasticsearch_dsl.Search.execute')
     def test_search_with_one_user_result(self,
                                          mock_search: MagicMock) -> None:
 
         mock_results = MagicMock()
         mock_results.hits.total = 1
-        mock_results.__iter__.return_value = [self.mock_result4]
+        mock_results.__iter__.return_value = [Response(result=vars(self.mock_result4))]
         mock_search.return_value = mock_results
 
         expected = SearchResult(total_results=1,