refactor: use swagger for table model in search

search_service/api/table.py

@@ -1,39 +1,12 @@
 from http import HTTPStatus
 from typing import Any, Dict, Iterable  # noqa: F401
 
-from flask_restful import Resource, fields, marshal_with, reqparse
+from flask_restful import Resource, reqparse
 from flasgger import swag_from
 
+from search_service.models.table import SearchTableResultSchema
 from search_service.proxy import get_proxy_client
 
-tag_fields = {
-    "tag_name": fields.String
-}
-
-table_fields = {
-    "name": fields.String,
-    "key": fields.String,
-    # description can be empty, if no description is present in DB
-    "description": fields.String,
-    "cluster": fields.String,
-    "database": fields.String,
-    "schema": fields.String,
-    "column_names": fields.List(fields.String),
-    # tags can be empty list
-    "tags": fields.List(fields.Nested(tag_fields)),
-    # badges can be an empty list
-    "badges": fields.List(fields.Nested(tag_fields)),
-    # last etl timestamp as epoch
-    "last_updated_timestamp": fields.Integer,
-    "display_name": fields.String,
-    "schema_description": fields.String,
-    "programmatic_descriptions": fields.List(fields.String)
-}
-
-search_table_results = {
-    "total_results": fields.Integer,
-    "results": fields.Nested(table_fields, default=[])
-}
 
 TABLE_INDEX = 'table_search_index'
 
@@ -54,7 +27,6 @@ def __init__(self) -> None:
 
         super(SearchTableAPI, self).__init__()
 
-    @marshal_with(search_table_results)
     @swag_from('swagger_doc/table/search_table.yml')
     def get(self) -> Iterable[Any]:
         """
@@ -73,7 +45,7 @@ def get(self) -> Iterable[Any]:
                 index=args.get('index')
             )
 
-            return results, HTTPStatus.OK
+            return SearchTableResultSchema().dump(results).data, HTTPStatus.OK
 
         except RuntimeError:
 
@@ -101,7 +73,6 @@ def __init__(self) -> None:
 
         super(SearchTableFilterAPI, self).__init__()
 
-    @marshal_with(search_table_results)
     @swag_from('swagger_doc/table/search_table_filter.yml')
     def post(self) -> Iterable[Any]:
         """
@@ -130,7 +101,7 @@ def post(self) -> Iterable[Any]:
                 page_index=page_index,
                 index=args['index']
             )
-            return results, HTTPStatus.OK
+            return SearchTableResultSchema().dump(results).data, HTTPStatus.OK
         except RuntimeError:
             err_msg = 'Exception encountered while processing search request'
             return {'message': err_msg}, HTTPStatus.INTERNAL_SERVER_ERROR

search_service/models/table.py

@@ -61,3 +61,15 @@ class TableSchema(AttrsSchema):
     class Meta:
         target = Table
         register_as_scheme = True
+
+
+@attr.s(auto_attribs=True, kw_only=True)
+class SearchTableResult:
+        total_results: int = attr.ib()
+        results: List[Table] = attr.ib(factory=list)
+
+
+class SearchTableResultSchema(AttrsSchema):
+    class Meta:
+        target = SearchTableResult
+        register_as_scheme = True

search_service/proxy/atlas.py

@@ -9,6 +9,7 @@
 from atlasclient.utils import parse_table_qualified_name
 
 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.table import Table
 from search_service.models.tag import Tag
@@ -191,7 +192,7 @@ def _prepare_basic_search_query(self, limit: int, page_index: int, query_term: O
     def fetch_table_search_results(self, *,
                                    query_term: str,
                                    page_index: int = 0,
-                                   index: str = '') -> SearchResult:
+                                   index: str = '') -> SearchTableResult:
         """
         Conduct a 'Basic Search' in Amundsen UI.
 
@@ -202,23 +203,23 @@ def fetch_table_search_results(self, *,
         :param query_term: Search Query Term
         :param page_index: Index of search page user is currently on (for pagination)
         :param index: Search Index (different resource corresponding to different index)
-        :return: SearchResult Object
+        :return: SearchTableResult Object
         """
         if not query_term:
             # return empty result for blank query term
-            return SearchResult(total_results=0, results=[])
+            return SearchTableResult(total_results=0, results=[])
 
         query_params = self._prepare_basic_search_query(self.page_size, page_index, query_term=query_term)
 
         tables, approx_count = self._atlas_basic_search(query_params)
 
-        return SearchResult(total_results=approx_count, results=tables)
+        return SearchTableResult(total_results=approx_count, results=tables)
 
     def fetch_table_search_results_with_filter(self, *,
                                                query_term: str,
                                                search_request: dict,
                                                page_index: int = 0,
-                                               index: str = '') -> SearchResult:
+                                               index: str = '') -> SearchTableResult:
         """
         Conduct an 'Advanced Search' to narrow down search results with a use of filters.
 
@@ -229,7 +230,7 @@ def fetch_table_search_results_with_filter(self, *,
         :param search_request: Values from Filters
         :param page_index: Index of search page user is currently on (for pagination)
         :param index: Search Index (different resource corresponding to different index)
-        :return: SearchResult Object
+        :return: SearchTableResult Object
         """
         _filters = search_request.get('filters', dict())
 
@@ -260,7 +261,7 @@ def fetch_table_search_results_with_filter(self, *,
 
         tables, approx_count = self._atlas_basic_search(query_params)
 
-        return SearchResult(total_results=approx_count, results=tables)
+        return SearchTableResult(total_results=approx_count, results=tables)
 
     def fetch_user_search_results(self, *,
                                   query_term: str,

search_service/proxy/base.py

@@ -2,6 +2,7 @@
 from typing import Any, Dict, List
 
 from search_service.models.dashboard import SearchDashboardResult
+from search_service.models.table import SearchTableResult
 from search_service.models.search_result import SearchResult
 
 
@@ -15,7 +16,7 @@ class BaseProxy(metaclass=ABCMeta):
     def fetch_table_search_results(self, *,
                                    query_term: str,
                                    page_index: int = 0,
-                                   index: str = '') -> SearchResult:
+                                   index: str = '') -> SearchTableResult:
         pass
 
     @abstractmethod
@@ -48,7 +49,7 @@ def fetch_table_search_results_with_filter(self, *,
                                                query_term: str,
                                                search_request: dict,
                                                page_index: int = 0,
-                                               index: str = '') -> SearchResult:
+                                               index: str = '') -> SearchTableResult:
         pass
 
     @abstractmethod

search_service/proxy/elasticsearch.py

@@ -14,7 +14,7 @@
 from search_service.api.user import USER_INDEX
 from search_service.api.table import TABLE_INDEX
 from search_service.models.search_result import SearchResult
-from search_service.models.table import Table
+from search_service.models.table import Table, SearchTableResult
 from search_service.models.user import User
 from search_service.models.dashboard import Dashboard, SearchDashboardResult
 from search_service.models.tag import Tag
@@ -289,7 +289,7 @@ def _search_wildcard_helper(self, field_value: str,
     def fetch_table_search_results(self, *,
                                    query_term: str,
                                    page_index: int = 0,
-                                   index: str = '') -> SearchResult:
+                                   index: str = '') -> SearchTableResult:
         """
         Query Elasticsearch and return results as list of Table objects
 
@@ -302,7 +302,7 @@ def fetch_table_search_results(self, *,
             current_app.config.get(config.ELASTICSEARCH_INDEX_KEY, DEFAULT_ES_INDEX)
         if not query_term:
             # return empty result for blank query term
-            return SearchResult(total_results=0, results=[])
+            return SearchTableResult(total_results=0, results=[])
 
         s = Search(using=self.elasticsearch, index=current_index)
         query_name = {
@@ -332,7 +332,8 @@ def fetch_table_search_results(self, *,
         return self._search_helper(page_index=page_index,
                                    client=s,
                                    query_name=query_name,
-                                   model=Table)
+                                   model=Table,
+                                   search_result_model=SearchTableResult)
 
     @staticmethod
     def get_model_by_index(index: str) -> Any:
@@ -447,7 +448,7 @@ def fetch_table_search_results_with_filter(self, *,
                                                query_term: str,
                                                search_request: dict,
                                                page_index: int = 0,
-                                               index: str = '') -> SearchResult:
+                                               index: str = '') -> SearchTableResult:
         """
         Query Elasticsearch and return results as list of Table objects
         :param search_request: A json representation of search request
@@ -459,15 +460,15 @@ def fetch_table_search_results_with_filter(self, *,
             current_app.config.get(config.ELASTICSEARCH_INDEX_KEY, DEFAULT_ES_INDEX)  # type: str
         if not search_request:
             # return empty result for blank query term
-            return SearchResult(total_results=0, results=[])
+            return SearchTableResult(total_results=0, results=[])
 
         try:
             query_string = self.convert_query_json_to_query_dsl(search_request=search_request,
                                                                 query_term=query_term)  # type: str
         except Exception as e:
             LOGGING.exception(e)
             # return nothing if any exception is thrown under the hood
-            return SearchResult(total_results=0, results=[])
+            return SearchTableResult(total_results=0, results=[])
         s = Search(using=self.elasticsearch, index=current_index)
 
         query_name = {
@@ -488,7 +489,8 @@ def fetch_table_search_results_with_filter(self, *,
         return self._search_helper(page_index=page_index,
                                    client=s,
                                    query_name=query_name,
-                                   model=model)
+                                   model=model,
+                                   search_result_model=SearchTableResult)
 
     @timer_with_counter
     def fetch_user_search_results(self, *,

setup.py

@@ -2,7 +2,7 @@
 
 from setuptools import setup, find_packages
 
-__version__ = '2.3.3'
+__version__ = '2.3.4'
 
 requirements_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), 'requirements.txt')
 with open(requirements_path) as requirements_file:

tests/unit/api/table/fixtures.py

@@ -18,6 +18,22 @@ def mock_proxy_results() -> Table:
                  programmatic_descriptions=[])
 
 
+def mock_default_proxy_results() -> Table:
+    return Table(name='',
+                 key='',
+                 description='',
+                 cluster='',
+                 database='',
+                 display_name='',
+                 schema='',
+                 column_names=[],
+                 tags=[],
+                 badges=[],
+                 last_updated_timestamp=0,
+                 schema_description='',
+                 programmatic_descriptions=[])
+
+
 def mock_json_response() -> dict:
     return {
         "name": "hello",
@@ -33,22 +49,26 @@ def mock_json_response() -> dict:
         "last_updated_timestamp": 1568324871,
         "schema_description": 'schema description',
         'programmatic_descriptions': [],
+        'total_usage': 0,
+        'column_descriptions': []
     }
 
 
 def default_json_response() -> dict:
     return {
-        "name": None,
-        "key": None,
-        "description": None,
-        "cluster": None,
-        "database": None,
-        "display_name": None,
-        "schema": None,
-        "column_names": None,
-        "tags": None,
-        "badges": None,
+        "name": '',
+        "key": '',
+        "description": '',
+        "cluster": '',
+        "database": '',
+        "display_name": '',
+        "schema": '',
+        "column_names": [],
+        "tags": [],
+        "badges": [],
         "last_updated_timestamp": 0,
-        "schema_description": None,
-        'programmatic_descriptions': None,
+        "schema_description": '',
+        'programmatic_descriptions': [],
+        'total_usage': 0,
+        'column_descriptions': []
     }

tests/unit/api/table/test_search_table_api.py

@@ -4,11 +4,13 @@
 from mock import patch, Mock
 
 from search_service import create_app
-from search_service.models.search_result import SearchResult
-from tests.unit.api.table.fixtures import mock_json_response, mock_proxy_results, default_json_response
+from search_service.models.table import SearchTableResult
+from tests.unit.api.table.fixtures import mock_json_response, mock_proxy_results, default_json_response, \
+    mock_default_proxy_results
 
 
 class TestSearchTableAPI(TestCase):
+
     def setUp(self) -> None:
         self.app = create_app(config_module_class='search_service.config.Config')
         self.app_context = self.app.app_context()
@@ -23,22 +25,23 @@ def tear_down(self) -> None:
 
     def test_should_get_result_for_search(self) -> None:
         result = mock_proxy_results()
-        self.mock_proxy.fetch_table_search_results.return_value = SearchResult(total_results=1, results=[result])
+        self.mock_proxy.fetch_table_search_results.return_value = SearchTableResult(total_results=1, results=[result])
 
         response = self.app.test_client().get('/search?query_term=searchterm')
 
         expected_response = {
             "total_results": 1,
             "results": [mock_json_response()]
         }
+
         self.assertEqual(response.json, expected_response)
         self.assertEqual(response.status_code, HTTPStatus.OK)
         self.mock_proxy.fetch_table_search_results.assert_called_with(query_term='searchterm', page_index=0,
                                                                       index='table_search_index')
 
     def test_should_give_empty_result_when_there_are_no_results_from_proxy(self) -> None:
         self.mock_proxy.fetch_table_search_results.return_value = \
-            SearchResult(total_results=0, results=[])
+            SearchTableResult(total_results=0, results=[])
 
         response = self.app.test_client().get('/search?query_term=searchterm')
 
@@ -50,14 +53,16 @@ def test_should_give_empty_result_when_there_are_no_results_from_proxy(self) ->
 
     def test_should_get_default_response_values_when_values_not_in_proxy_response(self) -> None:
         self.mock_proxy.fetch_table_search_results.return_value = \
-            SearchResult(total_results=1, results=[{}])
+            SearchTableResult(total_results=1,
+                              results=[mock_default_proxy_results()])
 
         response = self.app.test_client().get('/search?query_term=searchterm')
 
         expected_response = {
             "total_results": 1,
             "results": [default_json_response()]
         }
+
         self.assertEqual(response.json, expected_response)
 
     def test_should_fail_without_query_term(self) -> None: