diff --git a/search_service/api/user.py b/search_service/api/user.py index ee179498..720222e7 100644 --- a/search_service/api/user.py +++ b/search_service/api/user.py @@ -2,29 +2,12 @@ from typing import Iterable, Any from flasgger import swag_from -from flask_restful import Resource, fields, marshal_with, reqparse +from flask_restful import Resource, reqparse from search_service.proxy import get_proxy_client +from search_service.models.user import SearchUserResultSchema -user_fields = { - "name": fields.String, - "first_name": fields.String, - "last_name": fields.String, - "team_name": fields.String, - "email": fields.String, - "manager_email": fields.String, - "github_username": fields.String, - "is_active": fields.Boolean, - "employee_type": fields.String, - "role_name": fields.String -} - -search_user_results = { - "total_results": fields.Integer, - "results": fields.Nested(user_fields, default=[]) -} - USER_INDEX = 'user_search_index' @@ -45,7 +28,6 @@ def __init__(self) -> None: super(SearchUserAPI, self).__init__() - @marshal_with(search_user_results) @swag_from('swagger_doc/user.yml') def get(self) -> Iterable[Any]: """ @@ -63,7 +45,7 @@ def get(self) -> Iterable[Any]: index=args.get('index') ) - return results, HTTPStatus.OK + return SearchUserResultSchema().dump(results).data, HTTPStatus.OK except RuntimeError: diff --git a/search_service/models/user.py b/search_service/models/user.py index d64cd09e..b24464e6 100644 --- a/search_service/models/user.py +++ b/search_service/models/user.py @@ -1,5 +1,5 @@ -from typing import Optional, Set +from typing import Optional, Set, List import attr from amundsen_common.models.user import User as CommonUser @@ -43,3 +43,15 @@ class UserSchema(AttrsSchema): class Meta: target = User register_as_scheme = True + + +@attr.s(auto_attribs=True, kw_only=True) +class SearchUserResult: + total_results: int = attr.ib() + results: List[User] = attr.ib(factory=list) + + +class SearchUserResultSchema(AttrsSchema): + class Meta: + target = SearchUserResult + register_as_scheme = True diff --git a/search_service/proxy/atlas.py b/search_service/proxy/atlas.py index 08d0501a..7ee18b61 100644 --- a/search_service/proxy/atlas.py +++ b/search_service/proxy/atlas.py @@ -10,7 +10,7 @@ from search_service.models.dashboard import SearchDashboardResult from search_service.models.table import SearchTableResult -from search_service.models.search_result import SearchResult +from search_service.models.user import SearchUserResult from search_service.models.table import Table from search_service.models.tag import Tag from search_service.proxy import BaseProxy @@ -266,7 +266,7 @@ def fetch_search_results_with_filter(self, *, def fetch_user_search_results(self, *, query_term: str, page_index: int = 0, - index: str = '') -> SearchResult: + index: str = '') -> SearchUserResult: pass def update_document(self, *, data: List[Dict[str, Any]], index: str = '') -> str: diff --git a/search_service/proxy/base.py b/search_service/proxy/base.py index 4a48949f..efb23c96 100644 --- a/search_service/proxy/base.py +++ b/search_service/proxy/base.py @@ -3,7 +3,7 @@ from search_service.models.dashboard import SearchDashboardResult from search_service.models.table import SearchTableResult -from search_service.models.search_result import SearchResult +from search_service.models.user import SearchUserResult class BaseProxy(metaclass=ABCMeta): @@ -23,7 +23,7 @@ def fetch_table_search_results(self, *, def fetch_user_search_results(self, *, query_term: str, page_index: int = 0, - index: str = '') -> SearchResult: + index: str = '') -> SearchUserResult: pass @abstractmethod diff --git a/search_service/proxy/elasticsearch.py b/search_service/proxy/elasticsearch.py index e2a098c5..84729724 100644 --- a/search_service/proxy/elasticsearch.py +++ b/search_service/proxy/elasticsearch.py @@ -15,6 +15,7 @@ from search_service.api.table import TABLE_INDEX from search_service.models.search_result import SearchResult from search_service.models.table import Table, SearchTableResult +from search_service.models.user import SearchUserResult from search_service.models.user import User from search_service.models.dashboard import Dashboard, SearchDashboardResult from search_service.models.tag import Tag @@ -409,12 +410,12 @@ def fetch_search_results_with_filter(self, *, def fetch_user_search_results(self, *, query_term: str, page_index: int = 0, - index: str = '') -> SearchResult: + index: str = '') -> SearchUserResult: if not index: raise Exception('Index cant be empty for user search') if not query_term: # return empty result for blank query term - return SearchResult(total_results=0, results=[]) + return SearchUserResult(total_results=0, results=[]) s = Search(using=self.elasticsearch, index=index) @@ -440,7 +441,8 @@ def fetch_user_search_results(self, *, return self._search_helper(page_index=page_index, client=s, query_name=query_name, - model=User) + model=User, + search_result_model=SearchUserResult) @timer_with_counter def fetch_dashboard_search_results(self, *,