feat: Add Column Badge API

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,

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)

metadata_service/api/swagger_doc/column/badge_delete.yml

@@ -0,0 +1,47 @@
+Delete badges of a resource
+---
+tags:
+  - 'table'
+  - 'dashboard'
+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'

metadata_service/api/swagger_doc/column/badge_put.yml

@@ -0,0 +1,47 @@
+Add badge to a resource
+---
+tags:
+  - 'table'
+  - 'dashboard'
+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'

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

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()

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)