From 8e979a205e46a3ac983a95fc9ede0c7c84e86f55 Mon Sep 17 00:00:00 2001 From: tianru zhou Date: Thu, 21 Jan 2021 12:05:13 -0800 Subject: [PATCH 1/4] fix: fix table post/put api bug Signed-off-by: tianru zhou --- search_service/api/document.py | 13 +++-- .../api/swagger_doc/document/table_post.yml | 25 ++++----- .../api/swagger_doc/document/table_put.yml | 25 ++++----- .../api/swagger_doc/document/user_post.yml | 25 ++++----- .../api/swagger_doc/document/user_put.yml | 25 ++++----- search_service/api/swagger_doc/template.yml | 50 ++++++++++++++++- search_service/models/table.py | 30 ++++++---- search_service/models/tag.py | 3 + search_service/models/user.py | 4 +- search_service/proxy/atlas.py | 3 +- search_service/proxy/elasticsearch.py | 5 +- .../api/document/test_document_tables_api.py | 2 +- .../api/document/test_document_users_api.py | 2 +- tests/unit/api/table/fixtures.py | 8 ++- tests/unit/proxy/test_atlas.py | 5 +- tests/unit/proxy/test_elasticsearch.py | 56 +++++++++++++------ 16 files changed, 187 insertions(+), 94 deletions(-) diff --git a/search_service/api/document.py b/search_service/api/document.py index 8fbdfe71..9479eb04 100644 --- a/search_service/api/document.py +++ b/search_service/api/document.py @@ -2,10 +2,11 @@ # SPDX-License-Identifier: Apache-2.0 import logging +import json from http import HTTPStatus from typing import Tuple, Any - +from ast import literal_eval from flasgger import swag_from from flask_restful import Resource, reqparse from search_service.proxy import get_proxy_client @@ -58,7 +59,7 @@ def post(self) -> Tuple[Any, int]: :param data: list of data objects to be indexed in Elasticsearch :return: name of new index """ - self.parser.add_argument('data', required=True) + self.parser.add_argument('data', required=True, action='append') args = self.parser.parse_args() try: @@ -78,11 +79,15 @@ def put(self) -> Tuple[Any, int]: :param data: list of data objects to be indexed in Elasticsearch :return: name of index """ - self.parser.add_argument('data', required=True) + self.parser.add_argument('data', required=True, action='append') args = self.parser.parse_args() try: - data = self.schema(many=True, strict=False).loads(args.get('data')).data + table_dict_list = [] + for table_str in args.get('data'): + table_dict_list.append(literal_eval(table_str)) + table_list_json = json.dumps(table_dict_list) + data = self.schema(many=True, strict=False).loads(table_list_json).data results = self.proxy.update_document(data=data, index=args.get('index')) return results, HTTPStatus.OK except RuntimeError as e: diff --git a/search_service/api/swagger_doc/document/table_post.yml b/search_service/api/swagger_doc/document/table_post.yml index 971e610b..ca6e086a 100644 --- a/search_service/api/swagger_doc/document/table_post.yml +++ b/search_service/api/swagger_doc/document/table_post.yml @@ -11,19 +11,18 @@ parameters: type: string default: 'table_search_index' required: false - - name: body - in: body - schema: - type: object - name: data - properties: - data: - type: array - description: 'List of tables' - items: - $ref: '#/components/schemas/TableFields' - description: 'Tables to create' - required: true +requestBody: + content: + 'application/json': + schema: + type: object + properties: + data: + type: array + items: + $ref: '#/components/schemas/TableFields' + description: 'Tables to create' + required: true responses: 200: description: Empty json response diff --git a/search_service/api/swagger_doc/document/table_put.yml b/search_service/api/swagger_doc/document/table_put.yml index a2c3abfe..ba039a3f 100644 --- a/search_service/api/swagger_doc/document/table_put.yml +++ b/search_service/api/swagger_doc/document/table_put.yml @@ -11,19 +11,18 @@ parameters: type: string default: 'table_search_index' required: false - - name: body - in: body - schema: - type: object - name: data - properties: - data: - type: array - description: 'List of tables' - items: - $ref: '#/components/schemas/TableFields' - description: 'Tables to update' - required: true +requestBody: + content: + 'application/json': + schema: + type: object + properties: + data: + type: array + items: + $ref: '#/components/schemas/TableFields' + description: 'Tables to update' + required: true responses: 200: description: Empty json response diff --git a/search_service/api/swagger_doc/document/user_post.yml b/search_service/api/swagger_doc/document/user_post.yml index f1249d43..f950cec6 100644 --- a/search_service/api/swagger_doc/document/user_post.yml +++ b/search_service/api/swagger_doc/document/user_post.yml @@ -11,19 +11,18 @@ parameters: type: string default: user_search_index required: false - - name: body - in: body - schema: - type: object - name: data - properties: - data: - type: array - description: 'List of users' - items: - $ref: '#/components/schemas/UserFields' - description: 'Users to create' - required: true +requestBody: + content: + 'application/json': + schema: + type: object + properties: + data: + type: array + items: + $ref: '#/components/schemas/UserFields' + description: 'Users to create' + required: true responses: 200: description: Empty json response diff --git a/search_service/api/swagger_doc/document/user_put.yml b/search_service/api/swagger_doc/document/user_put.yml index 89216b33..2aae4971 100644 --- a/search_service/api/swagger_doc/document/user_put.yml +++ b/search_service/api/swagger_doc/document/user_put.yml @@ -11,19 +11,18 @@ parameters: type: string default: user_search_index required: false - - name: body - in: body - schema: - type: object - name: data - properties: - data: - type: array - description: 'List of users' - items: - $ref: '#/components/schemas/UserFields' - description: 'Users to update' - required: true +requestBody: + content: + 'application/json': + schema: + type: object + properties: + data: + type: array + items: + $ref: '#/components/schemas/UserFields' + description: 'Users to update' + required: true responses: 200: description: Empty json response diff --git a/search_service/api/swagger_doc/template.yml b/search_service/api/swagger_doc/template.yml index 8484c7b7..e7503b6c 100644 --- a/search_service/api/swagger_doc/template.yml +++ b/search_service/api/swagger_doc/template.yml @@ -52,6 +52,10 @@ components: TableFields: type: object properties: + id: + type: string + description: 'elasticsearch doc id' + example: 'M81jD3cBdULZTSY96PSh' name: type: string description: 'name of table' @@ -82,16 +86,54 @@ components: type: string description: 'list of column names' example: ['col1', 'col2'] + column_descriptions: + type: array + items: + type: string + description: 'list of column descriptions' + example: ['column description1', 'column description2'] + programmatic_descriptions: + type: array + items: + type: string + description: 'list of programmatic descriptions' + example: ['programmatic description1', 'programmatic description2'] tags: type: array items: - type: string + type: object + properties: + tag_name: + type: string description: 'list of table tags' - example: ['tag2', 'tag1'] + example: [{'tag_name': 'tag1'}, {'tag_name': 'tag2'}] + badges: + type: array + items: + type: object + properties: + tag_name: + type: string + description: 'list of table badges' + example: [{'tag_name': 'badge1'}, {'tag_name': 'badge2'}] last_updated_timestamp: type: integer description: 'table last updated time' example: 1568814420 + display_name: + type: string + description: 'table display name' + example: 'display_name' + total_usage: + type: int + description: 'total usage' + example: 0 + schema_description: + type: array + items: + type: string + description: 'list of schema descriptions' + example: ['schema description1', 'schema description2'] DashboardFields: type: object properties: @@ -130,6 +172,10 @@ components: UserFields: type: object properties: + id: + type: string + description: 'elasticsearch doc id' + example: 'M81jD3cBdULZTSY96PSh' name: type: string description: 'user name' diff --git a/search_service/models/table.py b/search_service/models/table.py index 185e341f..c5f2408c 100644 --- a/search_service/models/table.py +++ b/search_service/models/table.py @@ -8,6 +8,7 @@ from .base import Base from search_service.models.tag import Tag +import time @attr.s(auto_attribs=True, kw_only=True) @@ -15,18 +16,19 @@ class Table(Base): """ This represents the part of a table stored in the search proxy """ - database: str - cluster: str - schema: str - name: str - key: str + id: str + database: Optional[str] = None + cluster: Optional[str] = None + schema: Optional[str] = None + name: Optional[str] = None + key: Optional[str] = None display_name: Optional[str] = None - tags: List[Tag] - badges: List[Tag] + tags: List[Tag] = [] + badges: List[Tag] = [] description: Optional[str] = None - last_updated_timestamp: int + last_updated_timestamp: int = int(time.time()) # The following properties are lightly-transformed properties from the normal table object: - column_names: List[str] + column_names: List[str] = [] column_descriptions: List[str] = [] programmatic_descriptions: List[str] = [] # The following are search-only properties: @@ -34,12 +36,18 @@ class Table(Base): schema_description: Optional[str] = attr.ib(default=None) def get_id(self) -> str: - # uses the table key as the document id in ES - return self.key + return self.id + + def get_attrs_dict(self) -> dict: + attrs_dict = self.__dict__.copy() + attrs_dict['tags'] = [str(tag) for tag in self.tags] + attrs_dict['badges'] = [str(tag) for tag in self.badges] + return attrs_dict @classmethod def get_attrs(cls) -> Set: return { + 'id', 'name', 'key', 'description', diff --git a/search_service/models/tag.py b/search_service/models/tag.py index 41067be7..8c12145e 100644 --- a/search_service/models/tag.py +++ b/search_service/models/tag.py @@ -13,6 +13,9 @@ class Tag: def __init__(self, tag_name: str): self.tag_name = tag_name + def __str__(self): + return self.tag_name + class TagSchema(AttrsSchema): class Meta: diff --git a/search_service/models/user.py b/search_service/models/user.py index 8f38a317..63b03f54 100644 --- a/search_service/models/user.py +++ b/search_service/models/user.py @@ -17,14 +17,16 @@ class User(Base, CommonUser): This represents the part of a user stored in the search proxy """ manager_email: Optional[str] = None + id: str def get_id(self) -> str: # uses the user email as the document id in ES - return self.email if self.email else '' + return self.id @classmethod def get_attrs(cls) -> Set: return { + 'id', 'full_name', 'first_name', 'last_name', diff --git a/search_service/proxy/atlas.py b/search_service/proxy/atlas.py index a6ba4065..593a9e0e 100644 --- a/search_service/proxy/atlas.py +++ b/search_service/proxy/atlas.py @@ -92,7 +92,8 @@ def _prepare_tables(self, response: EntityCollection, enhance_metadata: bool = F badges: List[Tag] = tags - table = Table(name=entity_name, + table = Table(id=f"{entity.typeName}://{db_cluster}.{db_name}/{entity_name}", + name=entity_name, key=f"{entity.typeName}://{db_cluster}.{db_name}/{entity_name}", description=entity_attrs.get('description'), cluster=db_cluster, diff --git a/search_service/proxy/elasticsearch.py b/search_service/proxy/elasticsearch.py index 55a31381..d112126b 100644 --- a/search_service/proxy/elasticsearch.py +++ b/search_service/proxy/elasticsearch.py @@ -116,6 +116,8 @@ def _get_search_result(self, page_index: int, for hit in response: try: + es_metadata = hit.__dict__.get('meta', {}) + # ES hit: {'_d_': {'key': xxx...} es_payload = hit.__dict__.get('_d_', {}) if not es_payload: @@ -124,6 +126,7 @@ def _get_search_result(self, page_index: int, for attr, val in es_payload.items(): if attr in model.get_attrs(): result[attr] = self._get_instance(attr=attr, val=val) + result['id'] = self._get_instance(attr='id', val=es_metadata['id']) results.append(model(**result)) except Exception: @@ -590,7 +593,7 @@ def _build_update_actions(self, data: List[Table], index_key: str) -> List[Dict[ for item in data: actions.append({'update': {'_index': index_key, '_type': item.get_type(), '_id': item.get_id()}}) - actions.append({'doc': item.__dict__}) + actions.append({'doc': item.get_attrs_dict()}) return actions def _build_delete_actions(self, data: List[str], index_key: str, type: str) -> List[Dict[str, Any]]: diff --git a/tests/unit/api/document/test_document_tables_api.py b/tests/unit/api/document/test_document_tables_api.py index 1ef39f46..d22b9e78 100644 --- a/tests/unit/api/document/test_document_tables_api.py +++ b/tests/unit/api/document/test_document_tables_api.py @@ -33,7 +33,7 @@ def test_post(self, get_proxy: MagicMock, RequestParser: MagicMock) -> None: @patch('search_service.api.document.get_proxy_client') def test_put(self, get_proxy: MagicMock, RequestParser: MagicMock) -> None: mock_proxy = get_proxy.return_value = Mock() - RequestParser().parse_args.return_value = dict(data='{}', index='fake_index') + RequestParser().parse_args.return_value = dict(data=[], index='fake_index') response = DocumentTablesAPI().put() self.assertEqual(list(response)[1], HTTPStatus.OK) diff --git a/tests/unit/api/document/test_document_users_api.py b/tests/unit/api/document/test_document_users_api.py index 4d20bf4a..bad5da06 100644 --- a/tests/unit/api/document/test_document_users_api.py +++ b/tests/unit/api/document/test_document_users_api.py @@ -33,7 +33,7 @@ def test_post(self, get_proxy: MagicMock, RequestParser: MagicMock) -> None: @patch('search_service.api.document.get_proxy_client') def test_put(self, get_proxy: MagicMock, RequestParser: MagicMock) -> None: mock_proxy = get_proxy.return_value = Mock() - RequestParser().parse_args.return_value = dict(data='{}', index='fake_index') + RequestParser().parse_args.return_value = dict(data=[], index='fake_index') response = DocumentUsersAPI().put() self.assertEqual(list(response)[1], HTTPStatus.OK) diff --git a/tests/unit/api/table/fixtures.py b/tests/unit/api/table/fixtures.py index 50f71aca..fd914c3b 100644 --- a/tests/unit/api/table/fixtures.py +++ b/tests/unit/api/table/fixtures.py @@ -6,7 +6,8 @@ def mock_proxy_results() -> Table: - return Table(name='hello', + return Table(id='world', + name='hello', key='world', description='des1', cluster='clust', @@ -22,7 +23,8 @@ def mock_proxy_results() -> Table: def mock_default_proxy_results() -> Table: - return Table(name='', + return Table(id='', + name='', key='', description='', cluster='', @@ -39,6 +41,7 @@ def mock_default_proxy_results() -> Table: def mock_json_response() -> dict: return { + "id": "world", "name": "hello", "key": "world", "description": "des1", @@ -59,6 +62,7 @@ def mock_json_response() -> dict: def default_json_response() -> dict: return { + "id": '', "name": '', "key": '', "description": '', diff --git a/tests/unit/proxy/test_atlas.py b/tests/unit/proxy/test_atlas.py index 995d831e..4d808c05 100644 --- a/tests/unit/proxy/test_atlas.py +++ b/tests/unit/proxy/test_atlas.py @@ -221,7 +221,10 @@ def test_setup_config(self) -> None: def test_search_normal(self) -> None: expected = SearchTableResult(total_results=2, - results=[Table(name=self.entity1_name, + results=[Table(id=f"{self.entity_type}://" + f"{self.cluster}.{self.db}/" + f"{self.entity1_name}", + name=self.entity1_name, key=f"{self.entity_type}://" f"{self.cluster}.{self.db}/" f"{self.entity1_name}", diff --git a/tests/unit/proxy/test_elasticsearch.py b/tests/unit/proxy/test_elasticsearch.py index afe0e6de..f13149dc 100644 --- a/tests/unit/proxy/test_elasticsearch.py +++ b/tests/unit/proxy/test_elasticsearch.py @@ -76,6 +76,20 @@ def __init__(self, self._d_ = result +class TableResponse: + def __init__(self, + result: Any): + self._d_ = result + self.meta = {'id': result['key']} + + +class UserResponse: + def __init__(self, + result: Any): + self._d_ = result + self.meta = {'id': result['email']} + + class TestElasticsearchProxy(unittest.TestCase): def setUp(self) -> None: @@ -112,7 +126,8 @@ def setUp(self) -> None: badges=self.mock_empty_badge, last_updated_timestamp=1527283287) - self.mock_result3 = Table(name='test_table3', + self.mock_result3 = Table(id='test_key3', + name='test_table3', key='test_key3', description='test_description3', cluster='gold', @@ -221,11 +236,12 @@ def test_search_with_one_table_result(self, mock_results = MagicMock() mock_results.hits.total = 1 - mock_results.__iter__.return_value = [Response(result=vars(self.mock_result1))] + mock_results.__iter__.return_value = [TableResponse(result=vars(self.mock_result1))] mock_search.return_value = mock_results expected = SearchResult(total_results=1, - results=[Table(name='test_table', + results=[Table(id='test_key', + name='test_table', key='test_key', description='test_description', cluster='gold', @@ -252,12 +268,13 @@ def test_search_with_multiple_result(self, mock_results = MagicMock() mock_results.hits.total = 2 - mock_results.__iter__.return_value = [Response(result=vars(self.mock_result1)), - Response(result=vars(self.mock_result2))] + mock_results.__iter__.return_value = [TableResponse(result=vars(self.mock_result1)), + TableResponse(result=vars(self.mock_result2))] mock_search.return_value = mock_results expected = SearchResult(total_results=2, - results=[Table(name='test_table', + results=[Table(id='test_key', + name='test_table', key='test_key', description='test_description', cluster='gold', @@ -267,7 +284,8 @@ def test_search_with_multiple_result(self, tags=[], badges=self.mock_empty_badge, last_updated_timestamp=1527283287), - Table(name='test_table2', + Table(id='test_key2', + name='test_table2', key='test_key2', description='test_description2', cluster='gold', @@ -294,11 +312,12 @@ def test_search_with_multiple_result(self, def test_search_table_filter(self, mock_search: MagicMock) -> None: mock_results = MagicMock() mock_results.hits.total = 1 - mock_results.__iter__.return_value = [Response(result=vars(self.mock_result1))] + mock_results.__iter__.return_value = [TableResponse(result=vars(self.mock_result1))] mock_search.return_value = mock_results expected = SearchResult(total_results=1, - results=[Table(name='test_table', + results=[Table(id='test_key', + name='test_table', key='test_key', description='test_description', cluster='gold', @@ -319,7 +338,6 @@ def test_search_table_filter(self, mock_search: MagicMock) -> None: } } resp = self.es_proxy.fetch_search_results_with_filter(search_request=search_request, query_term='test') - self.assertEqual(resp.total_results, expected.total_results) self.assertIsInstance(resp.results[0], Table) self.assertDictEqual(vars(resp.results[0]), vars(expected.results[0])) @@ -472,11 +490,12 @@ def test_search_with_one_user_result(self, mock_results = MagicMock() mock_results.hits.total = 1 - mock_results.__iter__.return_value = [Response(result=vars(self.mock_result4))] + mock_results.__iter__.return_value = [UserResponse(result=vars(self.mock_result4))] mock_search.return_value = mock_results expected = SearchResult(total_results=1, - results=[User(full_name='First Last', + results=[User(id='test@email.com', + full_name='First Last', first_name='First', last_name='Last', team_name='Test team', @@ -511,13 +530,13 @@ def test_create_document(self, mock_uuid: MagicMock) -> None: mock_uuid.return_value = new_index_name mock_elasticsearch.indices.get_alias.return_value = dict([(new_index_name, {})]) start_data = [ - Table(cluster='blue', column_names=['1', '2'], database='snowflake', - schema='test_schema', description='A table for something', + Table(id='snowflake://blue.test_schema/bank_accounts', cluster='blue', column_names=['1', '2'], + database='snowflake', schema='test_schema', description='A table for something', key='snowflake://blue.test_schema/bank_accounts', last_updated_timestamp=0, name='bank_accounts', tags=[], badges=self.mock_empty_badge, column_descriptions=['desc'], schema_description='schema description 1'), - Table(cluster='blue', column_names=['5', '6'], database='snowflake', - schema='test_schema', description='A table for lots of things!', + Table(id='snowflake://blue.test_schema/bitcoin_wallets', cluster='blue', column_names=['5', '6'], + database='snowflake', schema='test_schema', description='A table for lots of things!', key='snowflake://blue.test_schema/bitcoin_wallets', last_updated_timestamp=0, name='bitcoin_wallets', tags=[], badges=self.mock_empty_badge, schema_description='schema description 2', programmatic_descriptions=["test"]) @@ -531,6 +550,7 @@ def test_create_document(self, mock_uuid: MagicMock) -> None: } }, { + 'id': 'snowflake://blue.test_schema/bank_accounts', 'cluster': 'blue', 'column_names': ['1', '2'], 'column_descriptions': ['desc'], @@ -555,6 +575,7 @@ def test_create_document(self, mock_uuid: MagicMock) -> None: } }, { + 'id': 'snowflake://blue.test_schema/bitcoin_wallets', 'cluster': 'blue', 'column_names': ['5', '6'], 'column_descriptions': [], @@ -593,7 +614,7 @@ def test_update_document(self, mock_uuid: MagicMock) -> None: table_key = 'snowflake://blue.test_schema/bitcoin_wallets' expected_alias = 'table_search_index' data = [ - Table(cluster='blue', column_names=['5', '6'], database='snowflake', + Table(id=table_key, cluster='blue', column_names=['5', '6'], database='snowflake', schema='test_schema', description='A table for lots of things!', key=table_key, last_updated_timestamp=0, name='bitcoin_wallets', tags=[], column_descriptions=['hello'], badges=self.mock_empty_badge, @@ -609,6 +630,7 @@ def test_update_document(self, mock_uuid: MagicMock) -> None: }, { 'doc': { + 'id': table_key, 'cluster': 'blue', 'column_names': ['5', '6'], 'column_descriptions': ['hello'], From 24693e84ae9eefc7e9be863276a8ac193b5fc712 Mon Sep 17 00:00:00 2001 From: tianru zhou Date: Thu, 21 Jan 2021 18:15:19 -0800 Subject: [PATCH 2/4] fix syntax issue v1 Signed-off-by: tianru zhou --- search_service/models/tag.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/search_service/models/tag.py b/search_service/models/tag.py index 8c12145e..178240a4 100644 --- a/search_service/models/tag.py +++ b/search_service/models/tag.py @@ -13,7 +13,7 @@ class Tag: def __init__(self, tag_name: str): self.tag_name = tag_name - def __str__(self): + def __str__(self) -> str: return self.tag_name From 01e0999e8438a729e5627aaff5638893121ee584 Mon Sep 17 00:00:00 2001 From: tianru zhou Date: Tue, 26 Jan 2021 17:32:30 -0800 Subject: [PATCH 3/4] address PR comments Signed-off-by: tianru zhou --- search_service/api/document.py | 4 +- search_service/models/table.py | 30 +++++++------ search_service/proxy/elasticsearch.py | 29 ++++++++++++- .../api/document/test_document_tables_api.py | 43 ++++++++++++++++++- tests/unit/api/table/fixtures.py | 4 +- tests/unit/proxy/test_elasticsearch.py | 17 +++++--- 6 files changed, 100 insertions(+), 27 deletions(-) diff --git a/search_service/api/document.py b/search_service/api/document.py index 9479eb04..78a5c089 100644 --- a/search_service/api/document.py +++ b/search_service/api/document.py @@ -83,9 +83,7 @@ def put(self) -> Tuple[Any, int]: args = self.parser.parse_args() try: - table_dict_list = [] - for table_str in args.get('data'): - table_dict_list.append(literal_eval(table_str)) + table_dict_list = [literal_eval(table_str) for table_str in args.get('data')] table_list_json = json.dumps(table_dict_list) data = self.schema(many=True, strict=False).loads(table_list_json).data results = self.proxy.update_document(data=data, index=args.get('index')) diff --git a/search_service/models/table.py b/search_service/models/table.py index c5f2408c..f8da6645 100644 --- a/search_service/models/table.py +++ b/search_service/models/table.py @@ -17,20 +17,20 @@ class Table(Base): This represents the part of a table stored in the search proxy """ id: str - database: Optional[str] = None - cluster: Optional[str] = None - schema: Optional[str] = None - name: Optional[str] = None - key: Optional[str] = None + database: str + cluster: str + schema: str + name: str + key: str display_name: Optional[str] = None - tags: List[Tag] = [] - badges: List[Tag] = [] + tags: Optional[List[Tag]] = None + badges: Optional[List[Tag]] = None description: Optional[str] = None last_updated_timestamp: int = int(time.time()) # The following properties are lightly-transformed properties from the normal table object: - column_names: List[str] = [] - column_descriptions: List[str] = [] - programmatic_descriptions: List[str] = [] + column_names: Optional[List[str]] = None + column_descriptions: Optional[List[str]] = None + programmatic_descriptions: Optional[List[str]] = None # The following are search-only properties: total_usage: int = 0 schema_description: Optional[str] = attr.ib(default=None) @@ -40,8 +40,14 @@ def get_id(self) -> str: def get_attrs_dict(self) -> dict: attrs_dict = self.__dict__.copy() - attrs_dict['tags'] = [str(tag) for tag in self.tags] - attrs_dict['badges'] = [str(tag) for tag in self.badges] + if self.tags is not None: + attrs_dict['tags'] = [str(tag) for tag in self.tags] + else: + attrs_dict['tags'] = None + if self.badges is not None: + attrs_dict['badges'] = [str(badge) for badge in self.badges] + else: + attrs_dict['badges'] = None return attrs_dict @classmethod diff --git a/search_service/proxy/elasticsearch.py b/search_service/proxy/elasticsearch.py index d112126b..1fd4525a 100644 --- a/search_service/proxy/elasticsearch.py +++ b/search_service/proxy/elasticsearch.py @@ -117,8 +117,33 @@ def _get_search_result(self, page_index: int, for hit in response: try: es_metadata = hit.__dict__.get('meta', {}) - - # ES hit: {'_d_': {'key': xxx...} + """ + ES hit example: + { + '_d_': { + 'name': 'name', + 'database': 'database', + 'schema': 'schema', + 'key': 'database://cluster.schema/name', + 'cluster': 'cluster', + 'column_descriptions': ['description1', 'description2'], + 'column_names': ['colname1', 'colname2'], + 'description': None, + 'display_name': 'display name', + 'last_updated_timestamp': 12345678, + 'programmatic_descriptions': [], + 'schema_description': None, + 'tags': ['tag1', 'tag2'], + 'badges': [], + 'total_usage': 0 + }, + 'mata': { + 'index': 'table index', + 'id': 'table id', + 'type': 'type' + } + } + """ es_payload = hit.__dict__.get('_d_', {}) if not es_payload: raise Exception('The ES doc not contain required field') diff --git a/tests/unit/api/document/test_document_tables_api.py b/tests/unit/api/document/test_document_tables_api.py index d22b9e78..44f4be3c 100644 --- a/tests/unit/api/document/test_document_tables_api.py +++ b/tests/unit/api/document/test_document_tables_api.py @@ -8,7 +8,9 @@ from search_service.api.document import DocumentTablesAPI from search_service import create_app - +from search_service.models.table import Table +from search_service.models.tag import Tag +import json class TestDocumentTablesAPI(unittest.TestCase): def setUp(self) -> None: @@ -39,6 +41,45 @@ def test_put(self, get_proxy: MagicMock, RequestParser: MagicMock) -> None: self.assertEqual(list(response)[1], HTTPStatus.OK) mock_proxy.update_document.assert_called_with(data=[], index='fake_index') + @patch('search_service.api.document.reqparse.RequestParser') + @patch('search_service.api.document.get_proxy_client') + def test_put_multiple_tables(self, get_proxy: MagicMock, RequestParser: MagicMock) -> None: + mock_proxy = get_proxy.return_value = Mock() + input_data = [ + json.dumps({ + 'id': 'table1', + 'key': 'table1', + 'cluster': 'cluster1', + 'database': 'database1', + 'name': 'name1', + 'schema': 'schema1', + 'last_updated_timestamp': 12345678, + 'tags': [{'tag_name': 'tag1'}, {'tag_name': 'tag2'}] + }), + json.dumps({ + 'id': 'table2', + 'key': 'table2', + 'cluster': 'cluster2', + 'database': 'database2', + 'name': 'name2', + 'schema': 'schema2', + 'last_updated_timestamp': 12345678, + 'tags': [{'tag_name': 'tag3'}, {'tag_name': 'tag4'}] + }) + ] + RequestParser().parse_args.return_value = dict(data=input_data, index='fake_index') + + expected_data = [Table(id='table1', database='database1', cluster='cluster1', schema='schema1', name='name1', + key='table1', tags=[Tag(tag_name='tag1'), Tag(tag_name='tag2')], + last_updated_timestamp=12345678), + Table(id='table2', database='database2', cluster='cluster2', schema='schema2', name='name2', + key='table2', tags=[Tag(tag_name='tag3'), Tag(tag_name='tag4')], + last_updated_timestamp=12345678)] + + response = DocumentTablesAPI().put() + self.assertEqual(list(response)[1], HTTPStatus.OK) + mock_proxy.update_document.assert_called_with(data=expected_data, index='fake_index') + def test_should_not_reach_create_with_id(self) -> None: response = self.app.test_client().post('/document_table/1') diff --git a/tests/unit/api/table/fixtures.py b/tests/unit/api/table/fixtures.py index fd914c3b..045ebfa9 100644 --- a/tests/unit/api/table/fixtures.py +++ b/tests/unit/api/table/fixtures.py @@ -56,7 +56,7 @@ def mock_json_response() -> dict: "schema_description": 'schema description', 'programmatic_descriptions': [], 'total_usage': 0, - 'column_descriptions': [] + 'column_descriptions': None } @@ -77,5 +77,5 @@ def default_json_response() -> dict: "schema_description": '', 'programmatic_descriptions': [], 'total_usage': 0, - 'column_descriptions': [] + 'column_descriptions': None } diff --git a/tests/unit/proxy/test_elasticsearch.py b/tests/unit/proxy/test_elasticsearch.py index f13149dc..b8813db6 100644 --- a/tests/unit/proxy/test_elasticsearch.py +++ b/tests/unit/proxy/test_elasticsearch.py @@ -30,7 +30,7 @@ def __init__(self, *, tags: Iterable[Tag], badges: Iterable[Tag], last_updated_timestamp: int, - programmatic_descriptions: List[str] = []) -> None: + programmatic_descriptions: List[str] = None) -> None: self.name = name self.key = key self.description = description @@ -250,7 +250,8 @@ def test_search_with_one_table_result(self, column_names=['test_col1', 'test_col2'], tags=[], badges=self.mock_empty_badge, - last_updated_timestamp=1527283287)]) + last_updated_timestamp=1527283287, + programmatic_descriptions=[])]) resp = self.es_proxy.fetch_table_search_results(query_term='test_query_term') @@ -283,7 +284,8 @@ def test_search_with_multiple_result(self, column_names=['test_col1', 'test_col2'], tags=[], badges=self.mock_empty_badge, - last_updated_timestamp=1527283287), + last_updated_timestamp=1527283287, + programmatic_descriptions=[]), Table(id='test_key2', name='test_table2', key='test_key2', @@ -326,7 +328,8 @@ def test_search_table_filter(self, mock_search: MagicMock) -> None: column_names=['test_col1', 'test_col2'], tags=self.mock_empty_tag, badges=self.mock_empty_badge, - last_updated_timestamp=1527283287)]) + last_updated_timestamp=1527283287, + programmatic_descriptions=[])]) search_request = { 'type': 'AND', 'filters': { @@ -564,7 +567,7 @@ def test_create_document(self, mock_uuid: MagicMock) -> None: 'tags': [], 'badges': [], 'total_usage': 0, - 'programmatic_descriptions': [], + 'programmatic_descriptions': None, 'schema_description': 'schema description 1', }, { @@ -578,7 +581,7 @@ def test_create_document(self, mock_uuid: MagicMock) -> None: 'id': 'snowflake://blue.test_schema/bitcoin_wallets', 'cluster': 'blue', 'column_names': ['5', '6'], - 'column_descriptions': [], + 'column_descriptions': None, 'database': 'snowflake', 'schema': 'test_schema', 'description': 'A table for lots of things!', @@ -644,7 +647,7 @@ def test_update_document(self, mock_uuid: MagicMock) -> None: 'tags': [], 'badges': [], 'total_usage': 0, - 'programmatic_descriptions': [], + 'programmatic_descriptions': None, 'schema_description': 'schema description 1', } } From 56ac659d90add75e94ddb521e5a977b1913fd57a Mon Sep 17 00:00:00 2001 From: tianru zhou Date: Tue, 26 Jan 2021 17:53:07 -0800 Subject: [PATCH 4/4] fix lint test failure Signed-off-by: tianru zhou --- tests/unit/api/document/test_document_tables_api.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/unit/api/document/test_document_tables_api.py b/tests/unit/api/document/test_document_tables_api.py index 44f4be3c..e6827dc9 100644 --- a/tests/unit/api/document/test_document_tables_api.py +++ b/tests/unit/api/document/test_document_tables_api.py @@ -12,6 +12,7 @@ from search_service.models.tag import Tag import json + class TestDocumentTablesAPI(unittest.TestCase): def setUp(self) -> None: self.app = create_app(config_module_class='search_service.config.Config')