sqlglot.lineage

  1from __future__ import annotations
  2
  3import json
  4import typing as t
  5from dataclasses import dataclass, field
  6
  7from sqlglot import Schema, exp, maybe_parse
  8from sqlglot.errors import SqlglotError
  9from sqlglot.optimizer import Scope, build_scope, find_all_in_scope, qualify
 10
 11if t.TYPE_CHECKING:
 12    from sqlglot.dialects.dialect import DialectType
 13
 14
 15@dataclass(frozen=True)
 16class Node:
 17    name: str
 18    expression: exp.Expression
 19    source: exp.Expression
 20    downstream: t.List[Node] = field(default_factory=list)
 21    alias: str = ""
 22
 23    def walk(self) -> t.Iterator[Node]:
 24        yield self
 25
 26        for d in self.downstream:
 27            if isinstance(d, Node):
 28                yield from d.walk()
 29            else:
 30                yield d
 31
 32    def to_html(self, dialect: DialectType = None, **opts) -> GraphHTML:
 33        nodes = {}
 34        edges = []
 35
 36        for node in self.walk():
 37            if isinstance(node.expression, exp.Table):
 38                label = f"FROM {node.expression.this}"
 39                title = f"<pre>SELECT {node.name} FROM {node.expression.this}</pre>"
 40                group = 1
 41            else:
 42                label = node.expression.sql(pretty=True, dialect=dialect)
 43                source = node.source.transform(
 44                    lambda n: exp.Tag(this=n, prefix="<b>", postfix="</b>")
 45                    if n is node.expression
 46                    else n,
 47                    copy=False,
 48                ).sql(pretty=True, dialect=dialect)
 49                title = f"<pre>{source}</pre>"
 50                group = 0
 51
 52            node_id = id(node)
 53
 54            nodes[node_id] = {
 55                "id": node_id,
 56                "label": label,
 57                "title": title,
 58                "group": group,
 59            }
 60
 61            for d in node.downstream:
 62                edges.append({"from": node_id, "to": id(d)})
 63        return GraphHTML(nodes, edges, **opts)
 64
 65
 66def lineage(
 67    column: str | exp.Column,
 68    sql: str | exp.Expression,
 69    schema: t.Optional[t.Dict | Schema] = None,
 70    sources: t.Optional[t.Dict[str, str | exp.Subqueryable]] = None,
 71    dialect: DialectType = None,
 72    **kwargs,
 73) -> Node:
 74    """Build the lineage graph for a column of a SQL query.
 75
 76    Args:
 77        column: The column to build the lineage for.
 78        sql: The SQL string or expression.
 79        schema: The schema of tables.
 80        sources: A mapping of queries which will be used to continue building lineage.
 81        dialect: The dialect of input SQL.
 82        **kwargs: Qualification optimizer kwargs.
 83
 84    Returns:
 85        A lineage node.
 86    """
 87
 88    expression = maybe_parse(sql, dialect=dialect)
 89
 90    if sources:
 91        expression = exp.expand(
 92            expression,
 93            {
 94                k: t.cast(exp.Subqueryable, maybe_parse(v, dialect=dialect))
 95                for k, v in sources.items()
 96            },
 97            dialect=dialect,
 98        )
 99
100    qualified = qualify.qualify(
101        expression,
102        dialect=dialect,
103        schema=schema,
104        **{"validate_qualify_columns": False, "identify": False, **kwargs},  # type: ignore
105    )
106
107    scope = build_scope(qualified)
108
109    if not scope:
110        raise SqlglotError("Cannot build lineage, sql must be SELECT")
111
112    def to_node(
113        column: str | int,
114        scope: Scope,
115        scope_name: t.Optional[str] = None,
116        upstream: t.Optional[Node] = None,
117        alias: t.Optional[str] = None,
118    ) -> Node:
119        aliases = {
120            dt.alias: dt.comments[0].split()[1]
121            for dt in scope.derived_tables
122            if dt.comments and dt.comments[0].startswith("source: ")
123        }
124
125        # Find the specific select clause that is the source of the column we want.
126        # This can either be a specific, named select or a generic `*` clause.
127        select = (
128            scope.expression.selects[column]
129            if isinstance(column, int)
130            else next(
131                (select for select in scope.expression.selects if select.alias_or_name == column),
132                exp.Star() if scope.expression.is_star else None,
133            )
134        )
135
136        if not select:
137            raise ValueError(f"Could not find {column} in {scope.expression}")
138
139        if isinstance(scope.expression, exp.Union):
140            upstream = upstream or Node(name="UNION", source=scope.expression, expression=select)
141
142            index = (
143                column
144                if isinstance(column, int)
145                else next(
146                    (
147                        i
148                        for i, select in enumerate(scope.expression.selects)
149                        if select.alias_or_name == column or select.is_star
150                    ),
151                    -1,  # mypy will not allow a None here, but a negative index should never be returned
152                )
153            )
154
155            if index == -1:
156                raise ValueError(f"Could not find {column} in {scope.expression}")
157
158            for s in scope.union_scopes:
159                to_node(index, scope=s, upstream=upstream)
160
161            return upstream
162
163        if isinstance(scope.expression, exp.Select):
164            # For better ergonomics in our node labels, replace the full select with
165            # a version that has only the column we care about.
166            #   "x", SELECT x, y FROM foo
167            #     => "x", SELECT x FROM foo
168            source = t.cast(exp.Expression, scope.expression.select(select, append=False))
169        else:
170            source = scope.expression
171
172        # Create the node for this step in the lineage chain, and attach it to the previous one.
173        node = Node(
174            name=f"{scope_name}.{column}" if scope_name else str(column),
175            source=source,
176            expression=select,
177            alias=alias or "",
178        )
179
180        if upstream:
181            upstream.downstream.append(node)
182
183        subquery_scopes = {
184            id(subquery_scope.expression): subquery_scope
185            for subquery_scope in scope.subquery_scopes
186        }
187
188        for subquery in find_all_in_scope(select, exp.Subqueryable):
189            subquery_scope = subquery_scopes[id(subquery)]
190
191            for name in subquery.named_selects:
192                to_node(name, scope=subquery_scope, upstream=node)
193
194        # if the select is a star add all scope sources as downstreams
195        if select.is_star:
196            for source in scope.sources.values():
197                node.downstream.append(Node(name=select.sql(), source=source, expression=source))
198
199        # Find all columns that went into creating this one to list their lineage nodes.
200        source_columns = set(find_all_in_scope(select, exp.Column))
201
202        # If the source is a UDTF find columns used in the UTDF to generate the table
203        if isinstance(source, exp.UDTF):
204            source_columns |= set(source.find_all(exp.Column))
205
206        for c in source_columns:
207            table = c.table
208            source = scope.sources.get(table)
209
210            if isinstance(source, Scope):
211                # The table itself came from a more specific scope. Recurse into that one using the unaliased column name.
212                to_node(
213                    c.name, scope=source, scope_name=table, upstream=node, alias=aliases.get(table)
214                )
215            else:
216                # The source is not a scope - we've reached the end of the line. At this point, if a source is not found
217                # it means this column's lineage is unknown. This can happen if the definition of a source used in a query
218                # is not passed into the `sources` map.
219                source = source or exp.Placeholder()
220                node.downstream.append(Node(name=c.sql(), source=source, expression=source))
221
222        return node
223
224    return to_node(column if isinstance(column, str) else column.name, scope)
225
226
227class GraphHTML:
228    """Node to HTML generator using vis.js.
229
230    https://visjs.github.io/vis-network/docs/network/
231    """
232
233    def __init__(
234        self, nodes: t.Dict, edges: t.List, imports: bool = True, options: t.Optional[t.Dict] = None
235    ):
236        self.imports = imports
237
238        self.options = {
239            "height": "500px",
240            "width": "100%",
241            "layout": {
242                "hierarchical": {
243                    "enabled": True,
244                    "nodeSpacing": 200,
245                    "sortMethod": "directed",
246                },
247            },
248            "interaction": {
249                "dragNodes": False,
250                "selectable": False,
251            },
252            "physics": {
253                "enabled": False,
254            },
255            "edges": {
256                "arrows": "to",
257            },
258            "nodes": {
259                "font": "20px monaco",
260                "shape": "box",
261                "widthConstraint": {
262                    "maximum": 300,
263                },
264            },
265            **(options or {}),
266        }
267
268        self.nodes = nodes
269        self.edges = edges
270
271    def __str__(self):
272        nodes = json.dumps(list(self.nodes.values()))
273        edges = json.dumps(self.edges)
274        options = json.dumps(self.options)
275        imports = (
276            """<script type="text/javascript" src="https://unpkg.com/vis-data@latest/peer/umd/vis-data.min.js"></script>
277  <script type="text/javascript" src="https://unpkg.com/vis-network@latest/peer/umd/vis-network.min.js"></script>
278  <link rel="stylesheet" type="text/css" href="https://unpkg.com/vis-network/styles/vis-network.min.css" />"""
279            if self.imports
280            else ""
281        )
282
283        return f"""<div>
284  <div id="sqlglot-lineage"></div>
285  {imports}
286  <script type="text/javascript">
287    var nodes = new vis.DataSet({nodes})
288    nodes.forEach(row => row["title"] = new DOMParser().parseFromString(row["title"], "text/html").body.childNodes[0])
289
290    new vis.Network(
291        document.getElementById("sqlglot-lineage"),
292        {{
293            nodes: nodes,
294            edges: new vis.DataSet({edges})
295        }},
296        {options},
297    )
298  </script>
299</div>"""
300
301    def _repr_html_(self) -> str:
302        return self.__str__()