Feat: Introduce API endpoint for column lineage (#691)

* Introduce API endpoint for column lineage

* PR feedback

* Use to_column instead of parse_one

Author: vchan

setup.py

@@ -43,7 +43,7 @@
         "requests",
         "rich",
         "ruamel.yaml",
-        "sqlglot>=11.4.4",
+        "sqlglot>=11.5.3",
     ],
     extras_require={
         "dev": [
@@ -76,10 +76,10 @@
             "types-requests==2.28.8",
         ],
         "web": [
-            "fastapi==0.89.1",
+            "fastapi==0.95.0",
             "hyperscript==0.0.1",
             "pyarrow==11.0.0",
-            "uvicorn==0.20.0",
+            "uvicorn==0.21.1",
         ],
         "snowflake": [
             "snowflake-connector-python[pandas]",

tests/web/test_main.py

@@ -469,3 +469,18 @@ def test_get_environments(project_context: Context) -> None:
             "expiration_ts": None,
         }
     }
+
+
+def test_get_lineage(web_sushi_context: Context) -> None:
+    response = client.get(
+        "/api/lineage", params={"model": "sushi.top_waiters", "column": "revenue"}
+    )
+    assert response.status_code == 200
+    assert response.json() == {
+        "sushi.top_waiters": {"revenue": {"sushi.waiter_revenue_by_day": ["revenue"]}},
+        "sushi.waiter_revenue_by_day": {
+            "revenue": {"sushi.items": ["price"], "sushi.order_items": ["quantity"]}
+        },
+        "sushi.items": {"price": {}},
+        "sushi.order_items": {"quantity": {}},
+    }

web/server/api/endpoints/__init__.py

@@ -7,6 +7,7 @@
     environments,
     events,
     files,
+    lineage,
     models,
     plan,
 )
@@ -21,5 +22,6 @@
 api_router.include_router(plan.router, prefix="/plan")
 api_router.include_router(environments.router, prefix="/environments")
 api_router.include_router(events.router, prefix="/events")
+api_router.include_router(lineage.router, prefix="/lineage")
 api_router.include_router(models.router, prefix="/models")
 api_router.include_router(context.router, prefix="/context")

web/server/api/endpoints/lineage.py

@@ -0,0 +1,63 @@
+from __future__ import annotations
+
+import collections
+import traceback
+import typing as t
+
+from fastapi import APIRouter, Depends, HTTPException
+from sqlglot import exp
+from sqlglot.lineage import Node, lineage
+from starlette.status import HTTP_422_UNPROCESSABLE_ENTITY
+
+from sqlmesh.core.context import Context
+from web.server.settings import get_loaded_context
+
+router = APIRouter()
+
+
+def _get_table(node: Node) -> str:
+    """Get a node's table/source"""
+    if isinstance(node.expression, exp.Table):
+        return exp.table_name(node.expression)
+    else:
+        return node.alias
+
+
+def _process_downstream(downstream: t.List[Node]) -> t.Dict[str, t.List[str]]:
+    """Aggregate a list of downstream nodes by table/source"""
+    graph = collections.defaultdict(list)
+    for node in downstream:
+        column = exp.to_column(node.name).name
+        table = _get_table(node)
+        graph[table].append(column)
+    return graph
+
+
+@router.get("/")
+async def column_lineage(
+    column: str,
+    model: str,
+    context: Context = Depends(get_loaded_context),
+) -> t.Dict[str, t.Dict[str, t.Dict[str, t.List[str]]]]:
+    """Get a column's lineage"""
+    try:
+        node = lineage(
+            column=column,
+            sql=context.models[model].render_query(),
+            sources={
+                model: context.models[model].render_query() for model in context.dag.upstream(model)
+            },
+        )
+    except Exception:
+        raise HTTPException(
+            status_code=HTTP_422_UNPROCESSABLE_ENTITY, detail=traceback.format_exc()
+        )
+
+    graph = {}
+    table = model
+    for i, node in enumerate(node.walk()):
+        if i > 0:
+            table = _get_table(node)
+            column = exp.to_column(node.name).name
+        graph[table] = {column: _process_downstream(node.downstream)}
+    return graph