feat: Column Lineage API (#280)

* starting commit for column lineage API Signed-off-by: Allison Suarez Miranda <[email protected]> * flake8 fix Signed-off-by: Allison Suarez Miranda <[email protected]> * table uri and column name isntead of id Signed-off-by: Allison Suarez Miranda <[email protected]> * sort imports Signed-off-by: Allison Suarez Miranda <[email protected]> * more import sorting Signed-off-by: Allison Suarez Miranda <[email protected]>
amundsen-io · Mar 12, 2021 · 681893f · 681893f
1 parent fc6c0f8
commit 681893f
Show file tree

Hide file tree

Showing 5 changed files with 75 additions and 4 deletions.
diff --git a/metadata_service/__init__.py b/metadata_service/__init__.py
@@ -15,7 +15,7 @@
 from flask_restful import Api
 
 from metadata_service.api.badge import BadgeAPI
-from metadata_service.api.column import ColumnDescriptionAPI
+from metadata_service.api.column import ColumnDescriptionAPI, ColumnLineageAPI
 from metadata_service.api.dashboard import (DashboardBadgeAPI,
                                             DashboardDescriptionAPI,
                                             DashboardDetailAPI,
@@ -110,6 +110,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(ColumnLineageAPI,
+                     '/table/<path:table_uri>/column/<column_name>/lineage')
     api.add_resource(Neo4jDetailAPI,
                      '/latest_updated_ts')
     api.add_resource(TagAPI,

diff --git a/metadata_service/api/column.py b/metadata_service/api/column.py
@@ -3,16 +3,45 @@
 
 import json
 from http import HTTPStatus
-from typing import Iterable, Union
+from typing import Iterable, Mapping, Union
 
+from amundsen_common.models.lineage import LineageSchema
 from flasgger import swag_from
 from flask import request
-from flask_restful import Resource
+from flask_restful import Resource, reqparse
 
+from metadata_service.entity.resource_type import ResourceType
 from metadata_service.exception import NotFoundException
 from metadata_service.proxy import get_proxy_client
 
 
+class ColumnLineageAPI(Resource):
+    """
+    ColumnLineageAPI supports GET operation to get column lineage
+    """
+    def __init__(self) -> None:
+        self.client = get_proxy_client()
+        self.parser = reqparse.RequestParser()
+        self.parser.add_argument('direction', type=str, required=False)
+        self.parser.add_argument('depth', type=int, required=False)
+        super(ColumnLineageAPI, self).__init__()
+
+    @swag_from('swagger_doc/column/lineage_get.yml')
+    def get(self, table_uri: str, column_name: str) -> Iterable[Union[Mapping, int, None]]:
+        args = self.parser.parse_args()
+        direction = args.get('direction', 'both')
+        depth = args.get('depth', 0)
+        try:
+            lineage = self.client.get_lineage(id=f"{table_uri}/{column_name}",
+                                              resource_type=ResourceType.Column,
+                                              direction=direction,
+                                              depth=depth)
+            schema = LineageSchema()
+            return schema.dump(lineage), HTTPStatus.OK
+        except Exception as e:
+            return {'message': f'Exception raised when getting lineage: {e}'}, HTTPStatus.NOT_FOUND
+
+
 class ColumnDescriptionAPI(Resource):
     """
     ColumnDescriptionAPI supports PUT and GET operations to upsert column description

diff --git a/metadata_service/api/swagger_doc/column/lineage_get.yml b/metadata_service/api/swagger_doc/column/lineage_get.yml
@@ -0,0 +1,40 @@
+Get lineage
+---
+tags:
+  - 'lineage'
+parameters:
+  - name: table_uri
+    in: path
+    type: string
+    schema:
+      type: string
+    required: true
+    example: 'hive://gold.test_schema/test_table1'
+  - name: column_name
+    in: path
+    type: string
+    schema:
+      type: string
+    required: true
+    example: 'test_col_name'
+  - name: direction
+    in: query
+    type: string
+    schema:
+      type: string
+    required: false
+    example: 'upstream'
+  - name: depth
+    in: query
+    type: integer
+    schema:
+      type: integer
+    required: false
+    example: 0
+responses:
+  200:
+    description: 'Lineage for requested direction and depth'
+    content:
+      application/json:
+        schema:
+          $ref: '#/components/schemas/Lineage'
diff --git a/...e/api/swagger_doc/lineage/lineage_get.yml → ...ice/api/swagger_doc/table/lineage_get.yml b/...e/api/swagger_doc/lineage/lineage_get.yml → ...ice/api/swagger_doc/table/lineage_get.yml
diff --git a/metadata_service/api/table.py b/metadata_service/api/table.py
@@ -47,7 +47,7 @@ def __init__(self) -> None:
         self.parser.add_argument('depth', type=int, required=False)
         super(TableLineageAPI, self).__init__()
 
-    @swag_from('swagger_doc/lineage/lineage_get.yml')
+    @swag_from('swagger_doc/table/lineage_get.yml')
     def get(self, id: str) -> Iterable[Union[Mapping, int, None]]:
         args = self.parser.parse_args()
         direction = args.get('direction', 'both')