feat: badges with category and badge_type fields

metadata_service/__init__.py

@@ -15,13 +15,15 @@
 from flask_restful import Api
 
 from metadata_service.api.column import ColumnDescriptionAPI
-from metadata_service.api.dashboard import DashboardDetailAPI, DashboardDescriptionAPI, DashboardTagAPI
+from metadata_service.api.dashboard import (DashboardDetailAPI, DashboardDescriptionAPI,
+                                            DashboardTagAPI, DashboardBadgeAPI)
 from metadata_service.api.healthcheck import healthcheck
 from metadata_service.api.popular_tables import PopularTablesAPI
 from metadata_service.api.system import Neo4jDetailAPI
 from metadata_service.api.table \
-    import TableDetailAPI, TableOwnerAPI, TableTagAPI, TableDescriptionAPI, TableDashboardAPI
+    import TableDetailAPI, TableOwnerAPI, TableTagAPI, TableBadgeAPI, TableDescriptionAPI, TableDashboardAPI
 from metadata_service.api.tag import TagAPI
+from metadata_service.api.badge import BadgeAPI
 from metadata_service.api.user import (UserDetailAPI, UserFollowAPI,
                                        UserFollowsAPI, UserOwnsAPI,
                                        UserOwnAPI, UserReadsAPI)
@@ -92,6 +94,8 @@ def create_app(*, config_module_class: str) -> Flask:
                      '/table/<path:id>/description')
     api.add_resource(TableTagAPI,
                      '/table/<path:id>/tag/<tag>')
+    api.add_resource(TableBadgeAPI,
+                     '/table/<path:id>/badge/<badge>')
     api.add_resource(TableOwnerAPI,
                      '/table/<path:table_uri>/owner/<owner>')
     api.add_resource(TableDashboardAPI,
@@ -102,6 +106,8 @@ def create_app(*, config_module_class: str) -> Flask:
                      '/latest_updated_ts')
     api.add_resource(TagAPI,
                      '/tags/')
+    api.add_resource(BadgeAPI,
+                     '/badges/')
     api.add_resource(UserDetailAPI,
                      '/user',
                      '/user/<path:id>')
@@ -121,6 +127,8 @@ def create_app(*, config_module_class: str) -> Flask:
                      '/dashboard/<path:id>/description')
     api.add_resource(DashboardTagAPI,
                      '/dashboard/<path:id>/tag/<tag>')
+    api.add_resource(DashboardBadgeAPI,
+                     '/dashboard/<path:id>/badge/<badge>')
     app.register_blueprint(api_bp)
 
     if app.config.get('SWAGGER_ENABLED'):

metadata_service/api/badge.py

@@ -0,0 +1,111 @@
+# Copyright Contributors to the Amundsen project.
+# SPDX-License-Identifier: Apache-2.0
+
+from flask_restful import Resource, fields, marshal
+
+from http import HTTPStatus
+from typing import Iterable, Union, Mapping, Tuple, Any
+
+from flasgger import swag_from
+from flask import current_app as app
+
+from metadata_service.entity.resource_type import ResourceType
+from metadata_service.entity.badge import Badge
+from metadata_service.exception import NotFoundException
+from metadata_service.proxy import get_proxy_client
+from metadata_service.proxy.base_proxy import BaseProxy
+
+badge_fields = {
+    'badge_name': fields.String,
+    'category': fields.String,
+    'badge_type': fields.String,
+}
+
+badges_fields = {
+    'badges': fields.List(fields.Nested(badge_fields))
+}
+
+
+class BadgeAPI(Resource):
+    def __init__(self) -> None:
+        self.client = get_proxy_client()
+        super(BadgeAPI, self).__init__()
+
+    @swag_from('swagger_doc/badge/badge_get.yml')
+    def get(self) -> Iterable[Union[Mapping, int, None]]:
+        """
+        API to get all existing badges
+        """
+        badges = self.client.get_badges()
+        return marshal({'badges': badges}, badges_fields), HTTPStatus.OK
+
+
+class BadgeCommon:
+    def __init__(self, client: BaseProxy) -> None:
+        self.client = client
+
+    def put(self, id: str, resource_type: ResourceType,
+            badge_name: str,
+            category: str = '',
+            badge_type: str = '') -> Tuple[Any, HTTPStatus]:
+
+        if badge_type == '' or category == '':
+            return \
+                {'message': f'The badge {badge_name} for resource id {id} is not added successfully because '
+                            f'category `{category}` and badge_type `{badge_type}` parameters are required '
+                            'for badges'}, \
+                HTTPStatus.NOT_FOUND
+
+        # TODO check resource type is column when adding a badge of category column after
+        # implementing column level badges
+        whitelist_badges = app.config.get('WHITELIST_BADGES', [])
+        incomimg_badge = Badge(badge_name=badge_name,
+                               category=category,
+                               badge_type=badge_type)
+        # need to check whether the badge combination is part of the whitelist:
+
+        in_whitelist = False
+        for badge in whitelist_badges:
+            if incomimg_badge.badge_name == badge.badge_name and incomimg_badge.category == badge.category \
+               and incomimg_badge.badge_type == badge.badge_type:
+                in_whitelist = True
+        if not in_whitelist:
+            return \
+                {'message': f'The badge {badge_name} with category {category} badge_type {badge_type} for resource '
+                            f'id {id} and resource_type {resource_type.name} is not added successfully because '
+                            'this combination of values is not part of the whitelist'}, \
+                HTTPStatus.NOT_FOUND
+
+        try:
+            self.client.add_badge(id=id,
+                                  badge_name=badge_name,
+                                  category=category,
+                                  badge_type=badge_type,
+                                  resource_type=resource_type)
+            return {'message': f'The badge {badge_name} with category {category} and type {badge_type} was '
+                               f'added successfully to resurce with id {id}'}, HTTPStatus.OK
+        except Exception as e:
+            return {'message': f'The badge {badge_name} with category {category}, badge type {badge_type} '
+                               f'for resource id {id} and resource_type {resource_type.name} failed to '
+                               'be added'}, \
+                HTTPStatus.NOT_FOUND
+
+    def delete(self, id: str, badge_name: str,
+               category: str,
+               badge_type: str,
+               resource_type: ResourceType) -> Tuple[Any, HTTPStatus]:
+        try:
+            self.client.delete_badge(id=id,
+                                     resource_type=resource_type,
+                                     badge_name=badge_name,
+                                     category=category,
+                                     badge_type=badge_type)
+            return \
+                {'message': f'The badge {badge_name} with category {category}, badge type {badge_type} for resource '
+                            f'id {id} and resource_type {resource_type.name} was deleted successfully'}, \
+                HTTPStatus.OK
+        except NotFoundException:
+            return \
+                {'message': f'The badge {badge_name} with category {category}, badge type {badge_type} for resource '
+                            f'id {id} and resource_type {resource_type.name} was not deleted successfully'}, \
+                HTTPStatus.NOT_FOUND

metadata_service/api/dashboard.py

@@ -11,6 +11,7 @@
 
 from metadata_service.api import BaseAPI
 from metadata_service.api.tag import TagCommon
+from metadata_service.api.badge import BadgeCommon
 from metadata_service.entity.dashboard_detail import DashboardSchema
 from metadata_service.entity.description import DescriptionSchema
 from metadata_service.entity.resource_type import ResourceType
@@ -74,6 +75,46 @@ def put(self, id: str) -> Iterable[Union[Mapping, int, None]]:
             return {'message': 'id {} does not exist'.format(id)}, HTTPStatus.NOT_FOUND
 
 
+class DashboardBadgeAPI(Resource):
+    """
+    DashboardBadgeAPI that supports PUT and DELETE operation to add or delete badges
+    on Dashboard
+    """
+    def __init__(self) -> None:
+        self.client = get_proxy_client()
+        self.parser = reqparse.RequestParser()
+        self.parser.add_argument('category', type=str, required=True)
+        self.parser.add_argument('badge_type', type=str, required=True)
+        super(DashboardBadgeAPI, self).__init__()
+
+        self._badge_common = BadgeCommon(client=self.client)
+
+    @swag_from('swagger_doc/badge/badge_put.yml')
+    def put(self, id: str, badge: str) -> Iterable[Union[Mapping, int, None]]:
+        args = self.parser.parse_args()
+
+        category = args.get('category', '')
+        badge_type = args.get('badge_type', '')
+
+        return self._badge_common.put(id=id,
+                                      resource_type=ResourceType.Dashboard,
+                                      badge_name=badge,
+                                      category=category,
+                                      badge_type=badge_type)
+
+    @swag_from('swagger_doc/badge/badge_delete.yml')
+    def delete(self, id: str, badge: str) -> Iterable[Union[Mapping, int, None]]:
+        args = self.parser.parse_args()
+        category = args.get('category', '')
+        badge_type = args.get('badge_type', '')
+
+        return self._badge_common.delete(id=id,
+                                         resource_type=ResourceType.Dashboard,
+                                         badge_name=badge,
+                                         category=category,
+                                         badge_type=badge_type)
+
+
 class DashboardTagAPI(Resource):
     """
     DashboardTagAPI that supports PUT and DELETE operation to add or delete tag

metadata_service/api/swagger_doc/badge/badge_delete.yml

@@ -0,0 +1,47 @@
+Delete badges of a resource
+---
+tags:
+  - 'table'
+  - 'dashboard'
+parameters:
+  - name: id
+    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: badge_type
+    in: query
+    type: string
+    schema:
+      type: string
+    required: true
+    example: 'neutral'
+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/badge/badge_get.yml

@@ -0,0 +1,16 @@
+Get badges
+---
+tags:
+  - 'badge'
+responses:
+  200:
+    description: 'Badges with sentiment and category'
+    content:
+      application/json:
+        schema:
+          type: object
+          properties:
+            badges:
+              type: array
+              items:
+                $ref: '#/components/schemas/Badge'

metadata_service/api/swagger_doc/badge/badge_put.yml

@@ -0,0 +1,47 @@
+Add badge to a resource
+---
+tags:
+  - 'table'
+  - 'dashboard'
+parameters:
+  - name: id
+    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: 'test_badge'
+  - name: category
+    in: query
+    type: string
+    schema:
+      type: string
+    required: true
+    example: 'table_status'
+  - name: badge_type
+    in: query
+    type: string
+    schema:
+      type: string
+    required: true
+    example: 'positive'
+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/api/swagger_doc/template.yml

@@ -47,6 +47,19 @@ components:
         tag_count:
           type: integer
           description: 'Count of times used'
+    Badge:
+    # TODO add column level badges later, probably will use this same object
+        type: object
+        properties:
+          badge_name:
+            type: string
+            description: 'name of the badge'
+          category:
+            type: string
+            description: 'category that badge belongs to'
+          badge_type:
+            type: string
+            description: 'negative, neutral, positive, or warning badge'
     TableReaderFields:
       type: object
       properties:

metadata_service/api/table.py

@@ -12,6 +12,7 @@
 
 from metadata_service.api import BaseAPI
 from metadata_service.api.tag import TagCommon
+from metadata_service.api.badge import BadgeCommon
 from metadata_service.entity.resource_type import ResourceType
 from metadata_service.entity.dashboard_summary import DashboardSummarySchema
 from metadata_service.exception import NotFoundException
@@ -160,6 +161,42 @@ def delete(self, id: str, tag: str) -> Iterable[Union[Mapping, int, None]]:
                                        tag_type=tag_type)
 
 
+class TableBadgeAPI(Resource):
+    def __init__(self) -> None:
+        self.client = get_proxy_client()
+        self.parser = reqparse.RequestParser()
+        self.parser.add_argument('category', type=str, required=True)
+        self.parser.add_argument('badge_type', type=str, required=True)
+        super(TableBadgeAPI, self).__init__()
+
+        self._badge_common = BadgeCommon(client=self.client)
+
+    @swag_from('swagger_doc/badge/badge_put.yml')
+    def put(self, id: str, badge: str) -> Iterable[Union[Mapping, int, None]]:
+        args = self.parser.parse_args()
+        # TODO should I have default here?
+        category = args.get('category', '')
+        badge_type = args.get('badge_type', '')
+
+        return self._badge_common.put(id=id,
+                                      resource_type=ResourceType.Table,
+                                      badge_name=badge,
+                                      category=category,
+                                      badge_type=badge_type)
+
+    @swag_from('swagger_doc/badge/badge_delete.yml')
+    def delete(self, id: str, badge: str) -> Iterable[Union[Mapping, int, None]]:
+        args = self.parser.parse_args()
+        category = args.get('category', '')
+        badge_type = args.get('badge_type', '')
+
+        return self._badge_common.delete(id=id,
+                                         resource_type=ResourceType.Table,
+                                         badge_name=badge,
+                                         category=category,
+                                         badge_type=badge_type)
+
+
 class TableDashboardAPI(BaseAPI):
     """
     TableDashboard API that supports GET operation providing list of Dashboards using a table.