diff --git a/search_service/api/table.py b/search_service/api/table.py index 82ad4dfd..5d317992 100644 --- a/search_service/api/table.py +++ b/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 diff --git a/search_service/models/table.py b/search_service/models/table.py index 19c0f3a8..c0d894dd 100644 --- a/search_service/models/table.py +++ b/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 diff --git a/search_service/proxy/atlas.py b/search_service/proxy/atlas.py index ce8ebb39..9566b645 100644 --- a/search_service/proxy/atlas.py +++ b/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, diff --git a/search_service/proxy/base.py b/search_service/proxy/base.py index 9247cd51..5bc32da0 100644 --- a/search_service/proxy/base.py +++ b/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 diff --git a/search_service/proxy/elasticsearch.py b/search_service/proxy/elasticsearch.py index 172ac843..ad6d5af8 100644 --- a/search_service/proxy/elasticsearch.py +++ b/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,7 +460,7 @@ 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, @@ -467,7 +468,7 @@ def fetch_table_search_results_with_filter(self, *, 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, *, diff --git a/setup.py b/setup.py index c98ebbac..03b47aba 100644 --- a/setup.py +++ b/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: diff --git a/tests/unit/api/table/fixtures.py b/tests/unit/api/table/fixtures.py index 50e1fb23..927dcc94 100644 --- a/tests/unit/api/table/fixtures.py +++ b/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': [] } diff --git a/tests/unit/api/table/test_search_table_api.py b/tests/unit/api/table/test_search_table_api.py index f7c5d3d6..e18877ef 100644 --- a/tests/unit/api/table/test_search_table_api.py +++ b/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,7 +25,7 @@ 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') @@ -31,6 +33,7 @@ def test_should_get_result_for_search(self) -> None: "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, @@ -38,7 +41,7 @@ def test_should_get_result_for_search(self) -> None: 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,7 +53,8 @@ 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') @@ -58,6 +62,7 @@ def test_should_get_default_response_values_when_values_not_in_proxy_response(se "total_results": 1, "results": [default_json_response()] } + self.assertEqual(response.json, expected_response) def test_should_fail_without_query_term(self) -> None: diff --git a/tests/unit/proxy/test_atlas.py b/tests/unit/proxy/test_atlas.py index d42a8d58..b396f27d 100644 --- a/tests/unit/proxy/test_atlas.py +++ b/tests/unit/proxy/test_atlas.py @@ -4,8 +4,7 @@ from typing import Any, Callable, Dict, List, Tuple from search_service import create_app, config -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.tag import Tag from search_service.proxy import get_proxy_client @@ -218,19 +217,19 @@ def test_setup_config(self) -> None: self.assertEqual(client.page_size, 1337) # type: ignore def test_search_normal(self) -> None: - expected = SearchResult(total_results=2, - results=[Table(name=self.entity1_name, - key=f"{self.entity_type}://" - f"{self.cluster}.{self.db}/" - f"{self.entity1_name}", - description=self.entity1_description, - cluster=self.cluster, - database=self.entity_type, - schema=self.db, - column_names=[], - tags=[Tag(tag_name='PII_DATA')], - badges=[Tag(tag_name='PII_DATA')], - last_updated_timestamp=123)]) + expected = SearchTableResult(total_results=2, + results=[Table(name=self.entity1_name, + key=f"{self.entity_type}://" + f"{self.cluster}.{self.db}/" + f"{self.entity1_name}", + description=self.entity1_description, + cluster=self.cluster, + database=self.entity_type, + schema=self.db, + column_names=[], + tags=[Tag(tag_name='PII_DATA')], + badges=[Tag(tag_name='PII_DATA')], + last_updated_timestamp=123)]) entity1 = self.to_class(self.entity1) entity_collection = MagicMock() entity_collection.entities = [entity1] @@ -246,8 +245,8 @@ def test_search_normal(self) -> None: "Search Result doesn't match with expected result!") def test_search_empty(self) -> None: - expected = SearchResult(total_results=0, - results=[]) + expected = SearchTableResult(total_results=0, + results=[]) self.proxy.atlas.search_dsl = self.dsl_inject([ (lambda dsl: "select count()" in dsl, {"attributes": {"name": ["count()"], "values": [[0]]}}), @@ -261,13 +260,13 @@ def test_search_empty(self) -> None: ]) resp = self.proxy.fetch_table_search_results(query_term="Table1") self.assertTrue(resp.total_results == 0, "there should no search results") - self.assertIsInstance(resp, SearchResult, "Search result received is not of 'SearchResult' type!") + self.assertIsInstance(resp, SearchTableResult, "Search result received is not of 'SearchResult' type!") self.assertDictEqual(vars(resp), vars(expected), "Search Result doesn't match with expected result!") def test_unknown_field(self) -> None: - expected = SearchResult(total_results=0, - results=[]) + expected = SearchTableResult(total_results=0, + results=[]) self.proxy.atlas.search_dsl = self.dsl_inject([ (lambda dsl: "select count()" in dsl, {"attributes": {"name": ["count()"], "values": [[0]]}}), @@ -281,6 +280,6 @@ def test_unknown_field(self) -> None: ]) resp = self.proxy.fetch_table_search_results(query_term="unknown:Table1") self.assertTrue(resp.total_results == 0, "there should no search results") - self.assertIsInstance(resp, SearchResult, "Search result received is not of 'SearchResult' type!") + self.assertIsInstance(resp, SearchTableResult, "Search result received is not of 'SearchResult' type!") self.assertDictEqual(vars(resp), vars(expected), "Search Result doesn't match with expected result!")