feat: Add Column Badge API (#273)

* Badge Column API Signed-off-by: adoami <[email protected]> * neo4j query change Signed-off-by: adoami <[email protected]> * delete badge Signed-off-by: adoami <[email protected]> * yaml Signed-off-by: adoami <[email protected]> * log Signed-off-by: adoami <[email protected]> * import Signed-off-by: adoami <[email protected]> * Delete Signed-off-by: adoami <[email protected]> * WORKING Signed-off-by: adoami <[email protected]> * added ResourceType to add_badge test Signed-off-by: adoami <[email protected]> * resource_type Signed-off-by: adoami <[email protected]> * lint Signed-off-by: adoami <[email protected]> * proxy Signed-off-by: adoami <[email protected]> * other proxies Signed-off-by: adoami <[email protected]> * lint Signed-off-by: adoami <[email protected]> * add test Signed-off-by: adoami <[email protected]> * shorten lines Signed-off-by: adoami <[email protected]> * indenting Signed-off-by: adoami <[email protected]> * remove backslashg Signed-off-by: adoami <[email protected]> * updates Signed-off-by: adoami <[email protected]> * add test for column badge Signed-off-by: adoami <[email protected]> * chore: Bumped version for release 3.2.0 (#274) * New release for 3.2.0 Signed-off-by: github-actions <[email protected]> * Bumped version to 3.2.0 Signed-off-by: github-actions <[email protected]> Co-authored-by: github-actions <[email protected]> Signed-off-by: adoami <[email protected]> * chore(deps): bump flask-cors from 3.0.8 to 3.0.10 (#268) Bumps [flask-cors](https://github.com/corydolphin/flask-cors) from 3.0.8 to 3.0.10. - [Release notes](https://github.com/corydolphin/flask-cors/releases) - [Changelog](https://github.com/corydolphin/flask-cors/blob/master/CHANGELOG.md) - [Commits](corydolphin/flask-cors@3.0.8...3.0.10) Signed-off-by: dependabot-preview[bot] <[email protected]> Co-authored-by: dependabot-preview[bot] <27856297+dependabot-preview[bot]@users.noreply.github.com> Signed-off-by: adoami <[email protected]> * added docker publish, need to do for other repos (#275) Signed-off-by: Allison Suarez Miranda <[email protected]> Signed-off-by: adoami <[email protected]> * New release for 3.3.0 (#277) Signed-off-by: github-actions <[email protected]> Co-authored-by: github-actions <[email protected]> Signed-off-by: adoami <[email protected]> * chore(deps): bump boto3 from 1.12.12 to 1.17.23 (#276) Bumps [boto3](https://github.com/boto/boto3) from 1.12.12 to 1.17.23. - [Release notes](https://github.com/boto/boto3/releases) - [Changelog](https://github.com/boto/boto3/blob/develop/CHANGELOG.rst) - [Commits](boto/boto3@1.12.12...1.17.23) Signed-off-by: dependabot-preview[bot] <[email protected]> Co-authored-by: dependabot-preview[bot] <27856297+dependabot-preview[bot]@users.noreply.github.com> Signed-off-by: adoami <[email protected]> * refactor: Remove Global call (#278) * Remove the global call and set it to ignore the shard in the init of the proxy Signed-off-by: Andrew Ciambrone <[email protected]> * clean up spacing Signed-off-by: Andrew Ciambrone <[email protected]> * isort lint Signed-off-by: Andrew Ciambrone <[email protected]> Signed-off-by: adoami <[email protected]> * refactor: Separate configs based on the backends/proxies (#279) * fixes pip install and configs Signed-off-by: verdan <[email protected]> * cleanup Signed-off-by: verdan <[email protected]> * test fix Signed-off-by: verdan <[email protected]> * removes the requirements changes from branch Signed-off-by: verdan <[email protected]> Signed-off-by: adoami <[email protected]> * merge Signed-off-by: adoami <[email protected]> * remove logs Signed-off-by: adoami <[email protected]> * badge Signed-off-by: adoami <[email protected]> * revert resource Signed-off-by: adoami <[email protected]> * Working new method Signed-off-by: adoami <[email protected]> * lint Signed-off-by: adoami <[email protected]> * working new changes Signed-off-by: adoami <[email protected]> * update yaml Signed-off-by: adoami <[email protected]> Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> Co-authored-by: github-actions <[email protected]> Co-authored-by: dependabot-preview[bot] <27856297+dependabot-preview[bot]@users.noreply.github.com> Co-authored-by: Allison Suarez Miranda <[email protected]> Co-authored-by: Andrew Ciambrone <[email protected]> Co-authored-by: Verdan Mahmood <[email protected]>
amundsen-io · Mar 19, 2021 · ee0ac63 · ee0ac63
1 parent db33af9
commit ee0ac63
Show file tree

Hide file tree

Showing 7 changed files with 206 additions and 5 deletions.
diff --git a/metadata_service/__init__.py b/metadata_service/__init__.py
@@ -15,7 +15,8 @@
 from flask_restful import Api
 
 from metadata_service.api.badge import BadgeAPI
-from metadata_service.api.column import ColumnDescriptionAPI, ColumnLineageAPI
+from metadata_service.api.column import (ColumnBadgeAPI, ColumnDescriptionAPI,
+                                         ColumnLineageAPI)
 from metadata_service.api.dashboard import (DashboardBadgeAPI,
                                             DashboardDescriptionAPI,
                                             DashboardDetailAPI,
@@ -110,6 +111,8 @@ def create_app(*, config_module_class: str) -> Flask:
                      '/table/<path:id>/dashboard/')
     api.add_resource(ColumnDescriptionAPI,
                      '/table/<path:table_uri>/column/<column_name>/description')
+    api.add_resource(ColumnBadgeAPI,
+                     '/table/<path:table_uri>/column/<column_name>/badge/<badge>')
     api.add_resource(ColumnLineageAPI,
                      '/table/<path:table_uri>/column/<column_name>/lineage')
     api.add_resource(Neo4jDetailAPI,

diff --git a/metadata_service/api/column.py b/metadata_service/api/column.py
@@ -10,6 +10,7 @@
 from flask import request
 from flask_restful import Resource, reqparse
 
+from metadata_service.api.badge import BadgeCommon
 from metadata_service.entity.resource_type import ResourceType
 from metadata_service.exception import NotFoundException
 from metadata_service.proxy import get_proxy_client
@@ -89,3 +90,33 @@ def get(self, table_uri: str, column_name: str) -> Union[tuple, int, None]:
 
         except Exception:
             return {'message': 'Internal server error!'}, HTTPStatus.INTERNAL_SERVER_ERROR
+
+
+class ColumnBadgeAPI(Resource):
+    def __init__(self) -> None:
+        self.client = get_proxy_client()
+        self.parser = reqparse.RequestParser()
+        self.parser.add_argument('category', type=str, required=True)
+        super(ColumnBadgeAPI, self).__init__()
+
+        self._badge_common = BadgeCommon(client=self.client)
+
+    @swag_from('swagger_doc/column/badge_put.yml')
+    def put(self, table_uri: str, badge: str, column_name: str) -> Iterable[Union[Mapping, int, None]]:
+        args = self.parser.parse_args()
+        category = args.get('category', '')
+
+        return self._badge_common.put(id=f"{table_uri}/{column_name}",
+                                      resource_type=ResourceType.Column,
+                                      badge_name=badge,
+                                      category=category)
+
+    @swag_from('swagger_doc/column/badge_delete.yml')
+    def delete(self, table_uri: str, badge: str, column_name: str) -> Iterable[Union[Mapping, int, None]]:
+        args = self.parser.parse_args()
+        category = args.get('category', '')
+
+        return self._badge_common.delete(id=f"{table_uri}/{column_name}",
+                                         resource_type=ResourceType.Column,
+                                         badge_name=badge,
+                                         category=category)
diff --git a/metadata_service/api/swagger_doc/column/badge_delete.yml b/metadata_service/api/swagger_doc/column/badge_delete.yml
@@ -0,0 +1,46 @@
+Delete badges of a resource
+---
+tags:
+  - 'column'
+parameters:
+  - name: table_uri
+    in: path
+    type: string
+    schema:
+      type: string
+    required: true
+    example: 'hive://gold.test_schema/test_table1'
+  - name: badge_name
+    in: path
+    type: string
+    schema:
+      type: string
+    required: true
+    example: 'beta'
+  - name: category
+    in: query
+    type: string
+    schema:
+      type: string
+    required: true
+    example: 'table_status'
+  - name: column_name
+    in: path
+    type: string
+    schema:
+      type: string
+    required: true
+    example: 'beta'
+responses:
+  200:
+    description: 'The badge was deleted successfully'
+    content:
+      application/json:
+        schema:
+          $ref: '#/components/schemas/MessageResponse'
+  404:
+    description: 'Table or badge not found'
+    content:
+      application/json:
+        schema:
+          $ref: '#/components/schemas/ErrorResponse'
diff --git a/metadata_service/api/swagger_doc/column/badge_put.yml b/metadata_service/api/swagger_doc/column/badge_put.yml
@@ -0,0 +1,46 @@
+Add badge to a resource
+---
+tags:
+  - 'column'
+parameters:
+  - name: table_uri
+    in: path
+    type: string
+    schema:
+      type: string
+    required: true
+    example: 'hive://gold.test_schema/test_table1'
+  - name: badge
+    in: path
+    type: string
+    schema:
+      type: string
+    required: true
+    example: 'beta'
+  - name: category
+    in: query
+    type: string
+    schema:
+      type: string
+    required: true
+    example: 'table_status'
+  - name: column_name
+    in: path
+    type: string
+    schema:
+      type: string
+    required: true
+    example: 'beta'
+responses:
+  200:
+    description: 'The badge was added successfully'
+    content:
+      application/json:
+        schema:
+          $ref: '#/components/schemas/MessageResponse'
+  404:
+    description: 'Table not found, or badge is not whitelisted'
+    content:
+      application/json:
+        schema:
+          $ref: '#/components/schemas/ErrorResponse'
diff --git a/metadata_service/proxy/neo4j_proxy.py b/metadata_service/proxy/neo4j_proxy.py
@@ -592,7 +592,7 @@ def add_badge(self, *,
                   id: str,
                   badge_name: str,
                   category: str = '',
-                  resource_type: ResourceType = ResourceType.Table) -> None:
+                  resource_type: ResourceType) -> None:
 
         LOGGER.info('New badge {} for id {} with category {} '
                     'and resource type {}'.format(badge_name, id, category, resource_type.name))
@@ -643,7 +643,7 @@ def add_badge(self, *,
     def delete_badge(self, id: str,
                      badge_name: str,
                      category: str,
-                     resource_type: ResourceType = ResourceType.Table) -> None:
+                     resource_type: ResourceType) -> None:
 
         # TODO for some reason when deleting it will say it was successful
         # even when the badge never existed to begin with

diff --git a/tests/unit/api/column/test_column_badge_api.py b/tests/unit/api/column/test_column_badge_api.py
@@ -0,0 +1,53 @@
+# Copyright Contributors to the Amundsen project.
+# SPDX-License-Identifier: Apache-2.0
+
+import unittest
+from http import HTTPStatus
+from unittest.mock import Mock, patch
+
+from metadata_service.entity.badge import Badge
+from tests.unit.test_basics import BasicTestCase
+
+TABLE_NAME = 'magic'
+BADGE_NAME = 'alpha'
+COLUMN_NAME = 'beta'
+
+
+class TestColumnBadgeAPI(BasicTestCase):
+    def setUp(self) -> None:
+        super().setUp()
+
+        self.mock_client = patch('metadata_service.api.column.get_proxy_client')
+        self.mock_proxy = self.mock_client.start().return_value = Mock()
+
+    def tearDown(self) -> None:
+        super().tearDown()
+
+        self.mock_client.stop()
+
+    def test_block_bad_badge_name(self) -> None:
+        self.app.config['WHITELIST_BADGES'] = []
+        response = self.app.test_client().put(f'/table/{TABLE_NAME}/column/{COLUMN_NAME}'
+                                              f'/badge/{BADGE_NAME}?category=table_status')
+
+        self.assertEqual(response.status_code, HTTPStatus.NOT_FOUND)
+
+    def test_block_badge_missing_category(self) -> None:
+        self.app.config['WHITELIST_BADGES'] = [Badge(badge_name='alpha',
+                                                     category='table_status')]
+        response = self.app.test_client().put(f'/table/{TABLE_NAME}/column/{COLUMN_NAME}'
+                                              f'/badge/{BADGE_NAME}')
+
+        self.assertEqual(response.status_code, HTTPStatus.BAD_REQUEST)
+
+    def test_badge_with_category(self) -> None:
+        self.app.config['WHITELIST_BADGES'] = [Badge(badge_name='alpha',
+                                                     category='table_status')]
+        response = self.app.test_client().put(f'/table/{TABLE_NAME}/column/{COLUMN_NAME}'
+                                              f'/badge/{BADGE_NAME}?category=table_status')
+
+        self.assertEqual(response.status_code, HTTPStatus.OK)
+
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/tests/unit/proxy/test_neo4j_proxy.py b/tests/unit/proxy/test_neo4j_proxy.py
@@ -412,7 +412,7 @@ def test_delete_owner(self) -> None:
             self.assertEqual(mock_run.call_count, 1)
             self.assertEqual(mock_commit.call_count, 1)
 
-    def test_add_badge(self) -> None:
+    def test_add_table_badge(self) -> None:
         with patch.object(GraphDatabase, 'driver') as mock_driver:
             mock_session = MagicMock()
             mock_driver.return_value.session.return_value = mock_session
@@ -427,7 +427,29 @@ def test_add_badge(self) -> None:
 
             neo4j_proxy = Neo4jProxy(host='DOES_NOT_MATTER', port=0000)
             neo4j_proxy.add_badge(id='dummy_uri',
-                                  badge_name='hive')
+                                  badge_name='hive',
+                                  resource_type=ResourceType.Table)
+            # we call neo4j twice in add_tag call
+            self.assertEqual(mock_run.call_count, 3)
+            self.assertEqual(mock_commit.call_count, 1)
+
+    def test_add_column_badge(self) -> None:
+        with patch.object(GraphDatabase, 'driver') as mock_driver:
+            mock_session = MagicMock()
+            mock_driver.return_value.session.return_value = mock_session
+
+            mock_transaction = MagicMock()
+            mock_session.begin_transaction.return_value = mock_transaction
+
+            mock_run = MagicMock()
+            mock_transaction.run = mock_run
+            mock_commit = MagicMock()
+            mock_transaction.commit = mock_commit
+
+            neo4j_proxy = Neo4jProxy(host='DOES_NOT_MATTER', port=0000)
+            neo4j_proxy.add_badge(id='dummy_uri/dummy_column',
+                                  badge_name='hive',
+                                  resource_type=ResourceType.Column)
             # we call neo4j twice in add_tag call
             self.assertEqual(mock_run.call_count, 3)
             self.assertEqual(mock_commit.call_count, 1)