refactor: use swagger for user model (#111)

* refactor: use swagger for user model * fix lint * fix lint * bump deps * remove deps bump
amundsen-io · Jun 16, 2020 · ca57770 · ca57770
1 parent 17c6739
commit ca57770
Show file tree

Hide file tree

Showing 5 changed files with 25 additions and 29 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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, *,