sqlglot.lineage

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