working new changes

Signed-off-by: adoami <[email protected]>
amundsen-io · Mar 16, 2021 · e8e27b2 · e8e27b2
1 parent 669df4c
commit e8e27b2
Show file tree

Hide file tree

Showing 7 changed files with 106 additions and 14 deletions.
diff --git a/metadata_service/__init__.py b/metadata_service/__init__.py
@@ -112,7 +112,7 @@ def create_app(*, config_module_class: str) -> Flask:
     api.add_resource(ColumnDescriptionAPI,
                      '/table/<path:table_uri>/column/<column_name>/description')
     api.add_resource(ColumnBadgeAPI,
-                     '/table/<path:id>/column/badge/<badge>')
+                     '/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
@@ -101,22 +101,22 @@ def __init__(self) -> None:
 
         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]]:
+    @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=id,
+        return self._badge_common.put(id=f"{table_uri}/{column_name}",
                                       resource_type=ResourceType.Column,
                                       badge_name=badge,
                                       category=category)
 
-    @swag_from('swagger_doc/badge/badge_delete.yml')
-    def delete(self, id: str, badge: str) -> Iterable[Union[Mapping, int, None]]:
+    @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=id,
+        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/badge/badge_delete.yml b/metadata_service/api/swagger_doc/badge/badge_delete.yml
@@ -10,8 +10,7 @@ parameters:
     schema:
       type: string
     required: true
-    table_example: 'hive://gold.test_schema/test_table1'
-    column_example: 'hive://gold.test_schema/test_table1/test_column'
+    example: 'hive://gold.test_schema/test_table1'
   - name: badge_name
     in: path
     type: string

diff --git a/metadata_service/api/swagger_doc/badge/badge_put.yml b/metadata_service/api/swagger_doc/badge/badge_put.yml
@@ -10,8 +10,7 @@ parameters:
     schema:
       type: string
     required: true
-    table_example: 'hive://gold.test_schema/test_table1'
-    column_example: 'hive://gold.test_schema/test_table1/test_column'
+    example: 'hive://gold.test_schema/test_table1'
   - name: badge
     in: path
     type: string

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,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'
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,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'
diff --git a/tests/unit/api/column/test_column_badge_api.py b/tests/unit/api/column/test_column_badge_api.py
@@ -27,23 +27,23 @@ def tearDown(self) -> None:
 
     def test_block_bad_badge_name(self) -> None:
         self.app.config['WHITELIST_BADGES'] = []
-        response = self.app.test_client().put(f'/table/{TABLE_NAME}/{COLUMN_NAME}/column'
+        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_NAME}/column'
+        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_NAME}/column'
+        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)