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
 11from sqlglot.optimizer.scope import ScopeType
 12
 13if t.TYPE_CHECKING:
 14    from sqlglot.dialects.dialect import DialectType
 15    from collections.abc import Iterator, Mapping, Sequence
 16    from sqlglot._typing import GraphHTMLArgs
 17    from typing_extensions import Unpack
 18
 19logger = logging.getLogger("sqlglot")
 20
 21
 22@dataclass(frozen=True)
 23class Node:
 24    name: str
 25    expression: exp.Expr
 26    source: exp.Expr
 27    downstream: list[Node] = field(default_factory=list)
 28    source_name: str = ""
 29    reference_node_name: str = ""
 30
 31    def walk(self) -> Iterator[Node]:
 32        visited: set[int] = set()
 33        queue = [self]
 34        while queue:
 35            node = queue.pop()
 36            node_id = id(node)
 37            if node_id in visited:
 38                continue
 39            visited.add(node_id)
 40            yield node
 41            queue.extend(reversed(node.downstream))
 42
 43    def to_html(self, dialect: DialectType = None, **opts: Unpack[GraphHTMLArgs]) -> GraphHTML:
 44        nodes = {}
 45        edges = []
 46
 47        for node in self.walk():
 48            if isinstance(node.expression, exp.Table):
 49                label = f"FROM {node.expression.this}"
 50                title = f"<pre>SELECT {node.name} FROM {node.expression.this}</pre>"
 51                group = 1
 52            else:
 53                label = node.expression.sql(pretty=True, dialect=dialect)
 54                source = node.source.transform(
 55                    lambda n: (
 56                        exp.Tag(this=n, prefix="<b>", postfix="</b>") if n is node.expression else n
 57                    ),
 58                    copy=False,
 59                ).sql(pretty=True, dialect=dialect)
 60                title = f"<pre>{source}</pre>"
 61                group = 0
 62
 63            node_id = id(node)
 64
 65            nodes[node_id] = {
 66                "id": node_id,
 67                "label": label,
 68                "title": title,
 69                "group": group,
 70            }
 71
 72            for d in node.downstream:
 73                edges.append({"from": node_id, "to": id(d)})
 74        return GraphHTML(nodes, edges, **opts)
 75
 76
 77def lineage(
 78    column: str | exp.Column,
 79    sql: str | exp.Expr,
 80    schema: dict | Schema | None = None,
 81    sources: Mapping[str, str | exp.Query] | None = None,
 82    dialect: DialectType = None,
 83    scope: Scope | None = None,
 84    trim_selects: bool = True,
 85    copy: bool = True,
 86    **kwargs,
 87) -> Node:
 88    """Build the lineage graph for a column of a SQL query.
 89
 90    Args:
 91        column: The column to build the lineage for.
 92        sql: The SQL string or expression.
 93        schema: The schema of tables.
 94        sources: A mapping of queries which will be used to continue building lineage.
 95        dialect: The dialect of input SQL.
 96        scope: A pre-created scope to use instead.
 97        trim_selects: Whether to clean up selects by trimming to only relevant columns.
 98        copy: Whether to copy the Expr arguments.
 99        **kwargs: Qualification optimizer kwargs.
100
101    Returns:
102        A lineage node.
103    """
104
105    expression = maybe_parse(sql, copy=copy, dialect=dialect)
106    column = normalize_identifiers.normalize_identifiers(column, dialect=dialect).name
107
108    if sources:
109        expression = exp.expand(
110            expression,
111            {
112                k: t.cast(exp.Query, maybe_parse(v, copy=copy, dialect=dialect))
113                for k, v in sources.items()
114            },
115            dialect=dialect,
116            copy=copy,
117        )
118
119    if not scope:
120        expression = qualify.qualify(
121            expression,
122            dialect=dialect,
123            schema=schema,
124            **{"validate_qualify_columns": False, "identify": False, **kwargs},  # type: ignore
125        )
126
127        scope = build_scope(expression)
128
129    if not scope:
130        raise SqlglotError("Cannot build lineage, sql must be SELECT")
131
132    selectable = scope.expression
133    if not isinstance(selectable, exp.Selectable) or not any(
134        select.alias_or_name == column for select in selectable.selects
135    ):
136        raise SqlglotError(f"Cannot find column '{column}' in query.")
137
138    return to_node(column, scope, dialect, trim_selects=trim_selects, _cache={})
139
140
141def to_node(
142    column: str | int,
143    scope: Scope,
144    dialect: DialectType,
145    scope_name: str | None = None,
146    upstream: Node | None = None,
147    source_name: str | None = None,
148    reference_node_name: str | None = None,
149    trim_selects: bool = True,
150    _cache: dict[tuple, Node] | None = None,
151) -> Node:
152    cache_key = (column, id(scope), scope_name, source_name, reference_node_name)
153
154    if _cache is not None and cache_key in _cache:
155        cached_node = _cache[cache_key]
156        if upstream:
157            upstream.downstream.append(cached_node)
158        return cached_node
159
160    # Find the specific select clause that is the source of the column we want.
161    # This can either be a specific, named select or a generic `*` clause.
162    selectable = t.cast(exp.Selectable, scope.expression)
163    if isinstance(column, int):
164        if column >= len(selectable.selects):
165            raise SqlglotError(
166                f"Cannot find column's source with index {column} in query: {selectable.sql(dialect=dialect)}"
167            )
168        select = selectable.selects[column]
169    else:
170        select = next(
171            (select for select in selectable.selects if select.alias_or_name == column),
172            exp.Star() if selectable.is_star else scope.expression,
173        )
174
175    if isinstance(scope.expression, exp.Subquery):
176        for inner_scope in scope.subquery_scopes:
177            result = to_node(
178                column,
179                scope=inner_scope,
180                dialect=dialect,
181                upstream=upstream,
182                source_name=source_name,
183                reference_node_name=reference_node_name,
184                trim_selects=trim_selects,
185                _cache=_cache,
186            )
187            if _cache is not None:
188                _cache[cache_key] = result
189            return result
190    if isinstance(scope.expression, exp.SetOperation):
191        name = type(scope.expression).__name__.upper()
192        upstream = upstream or Node(name=name, source=scope.expression, expression=select)
193
194        index = (
195            column
196            if isinstance(column, int)
197            else next(
198                (
199                    i
200                    for i, select in enumerate(selectable.selects)
201                    if select.alias_or_name == column or select.is_star
202                ),
203                -1,  # mypy will not allow a None here, but a negative index should never be returned
204            )
205        )
206
207        if index == -1:
208            raise ValueError(f"Could not find {column} in {scope.expression}")
209
210        for s in scope.union_scopes:
211            to_node(
212                index,
213                scope=s,
214                dialect=dialect,
215                upstream=upstream,
216                source_name=source_name,
217                reference_node_name=reference_node_name,
218                trim_selects=trim_selects,
219                _cache=_cache,
220            )
221
222        if _cache is not None:
223            _cache[cache_key] = upstream
224        return upstream
225
226    if trim_selects and isinstance(scope.expression, exp.Select):
227        # For better ergonomics in our node labels, replace the full select with
228        # a version that has only the column we care about.
229        #   "x", SELECT x, y FROM foo
230        #     => "x", SELECT x FROM foo
231        source: exp.Expr = scope.expression.select(select, append=False)
232    else:
233        source = scope.expression
234
235    # Create the node for this step in the lineage chain, and attach it to the previous one.
236    node = Node(
237        name=f"{scope_name}.{column}" if scope_name else str(column),
238        source=source,
239        expression=select,
240        source_name=source_name or "",
241        reference_node_name=reference_node_name or "",
242    )
243
244    if upstream:
245        upstream.downstream.append(node)
246
247    subquery_scopes = {
248        id(subquery_scope.expression): subquery_scope for subquery_scope in scope.subquery_scopes
249    }
250
251    for subquery in find_all_in_scope(select, *exp.UNWRAPPED_QUERIES):
252        subquery_scope: Scope | None = subquery_scopes.get(id(subquery))
253        if not subquery_scope:
254            logger.warning(f"Unknown subquery scope: {subquery.sql(dialect=dialect)}")
255            continue
256
257        for name in subquery.named_selects:
258            to_node(
259                name,
260                scope=subquery_scope,
261                dialect=dialect,
262                upstream=node,
263                trim_selects=trim_selects,
264                _cache=_cache,
265            )
266
267    # if the select is a star add all scope sources as downstreams
268    if isinstance(select, exp.Star):
269        for src in scope.sources.values():
270            src_expr = src.expression if isinstance(src, Scope) else src
271            node.downstream.append(
272                Node(name=select.sql(comments=False), source=src_expr, expression=src_expr)
273            )
274
275    # Find all columns that went into creating this one to list their lineage nodes.
276    source_columns = set(find_all_in_scope(select, exp.Column))
277
278    # If the source is a UDTF find columns used in the UDTF to generate the table
279    if isinstance(source, exp.UDTF):
280        source_columns |= set(source.find_all(exp.Column))
281        derived_tables: Sequence[exp.Expr] = [
282            src.expression.parent
283            for src in scope.sources.values()
284            if isinstance(src, Scope) and src.is_derived_table and src.expression.parent
285        ]
286    else:
287        derived_tables = scope.derived_tables
288
289    source_names = {
290        dt.alias: dt.comments[0].split()[1]
291        for dt in derived_tables
292        if dt.comments and dt.comments[0].startswith("source: ")
293    }
294
295    pivots = scope.pivots
296    pivot = pivots[0] if len(pivots) == 1 and not pivots[0].unpivot else None
297    if pivot:
298        # For each aggregation function, the pivot creates a new column for each field in category
299        # combined with the aggfunc. So the columns parsed have this order: cat_a_value_sum, cat_a,
300        # b_value_sum, b. Because of this step wise manner the aggfunc 'sum(value) as value_sum'
301        # belongs to the column indices 0, 2, and the aggfunc 'max(price)' without an alias belongs
302        # to the column indices 1, 3. Here, only the columns used in the aggregations are of interest
303        # in the lineage, so lookup the pivot column name by index and map that with the columns used
304        # in the aggregation.
305        #
306        # Example: PIVOT (SUM(value) AS value_sum, MAX(price)) FOR category IN ('a' AS cat_a, 'b')
307        pivot_columns = pivot.args["columns"]
308        pivot_aggs_count = len(pivot.expressions)
309
310        pivot_column_mapping = {}
311        for i, agg in enumerate(pivot.expressions):
312            agg_cols = list(agg.find_all(exp.Column))
313            for col_index in range(i, len(pivot_columns), pivot_aggs_count):
314                pivot_column_mapping[pivot_columns[col_index].name] = agg_cols
315
316    for c in source_columns:
317        table = c.table
318        col_source: exp.Table | Scope | None = scope.sources.get(table)
319
320        if isinstance(col_source, Scope):
321            reference_node_name = None
322            if col_source.scope_type == ScopeType.DERIVED_TABLE and table not in source_names:
323                reference_node_name = table
324            elif col_source.scope_type == ScopeType.CTE:
325                selected_node, _ = scope.selected_sources.get(table, (None, None))
326                reference_node_name = selected_node.name if selected_node else None
327
328            # The table itself came from a more specific scope. Recurse into that one using the unaliased column name.
329            to_node(
330                c.name,
331                scope=col_source,
332                dialect=dialect,
333                scope_name=table,
334                upstream=node,
335                source_name=source_names.get(table) or source_name,
336                reference_node_name=reference_node_name,
337                trim_selects=trim_selects,
338                _cache=_cache,
339            )
340        elif pivot and pivot.alias_or_name == c.table:
341            downstream_columns = []
342
343            column_name = c.name
344            if any(column_name == pivot_column.name for pivot_column in pivot_columns):
345                downstream_columns.extend(pivot_column_mapping[column_name])
346            else:
347                # The column is not in the pivot, so it must be an implicit column of the
348                # pivoted source -- adapt column to be from the implicit pivoted source.
349                pivot_parent = pivot.parent
350                downstream_columns.append(
351                    exp.column(c.this, table=pivot_parent.alias_or_name if pivot_parent else "")
352                )
353
354            for downstream_column in downstream_columns:
355                table = downstream_column.table
356                col_source = scope.sources.get(table)
357                if isinstance(col_source, Scope):
358                    to_node(
359                        downstream_column.name,
360                        scope=col_source,
361                        scope_name=table,
362                        dialect=dialect,
363                        upstream=node,
364                        source_name=source_names.get(table) or source_name,
365                        reference_node_name=reference_node_name,
366                        trim_selects=trim_selects,
367                        _cache=_cache,
368                    )
369                else:
370                    col_expr = col_source or exp.Placeholder()
371                    node.downstream.append(
372                        Node(
373                            name=downstream_column.sql(comments=False),
374                            source=col_expr,
375                            expression=col_expr,
376                        )
377                    )
378        else:
379            # The source is not a scope and the column is not in any pivot - we've reached the end
380            # of the line. At this point, if a source is not found it means this column's lineage
381            # is unknown. This can happen if the definition of a source used in a query is not
382            # passed into the `sources` map.
383            col_expr = col_source or exp.Placeholder()
384            node.downstream.append(
385                Node(name=c.sql(comments=False), source=col_expr, expression=col_expr)
386            )
387
388    if _cache is not None:
389        _cache[cache_key] = node
390
391    return node
392
393
394class GraphHTML:
395    """Node to HTML generator using vis.js.
396
397    https://visjs.github.io/vis-network/docs/network/
398    """
399
400    def __init__(
401        self,
402        nodes: dict,
403        edges: list,
404        imports: bool = True,
405        options: Mapping[str, object] | None = None,
406    ):
407        self.imports = imports
408
409        self.options = {
410            "height": "500px",
411            "width": "100%",
412            "layout": {
413                "hierarchical": {
414                    "enabled": True,
415                    "nodeSpacing": 200,
416                    "sortMethod": "directed",
417                },
418            },
419            "interaction": {
420                "dragNodes": False,
421                "selectable": False,
422            },
423            "physics": {
424                "enabled": False,
425            },
426            "edges": {
427                "arrows": "to",
428            },
429            "nodes": {
430                "font": "20px monaco",
431                "shape": "box",
432                "widthConstraint": {
433                    "maximum": 300,
434                },
435            },
436            **(options or {}),
437        }
438
439        self.nodes = nodes
440        self.edges = edges
441
442    def __str__(self):
443        nodes = json.dumps(list(self.nodes.values()))
444        edges = json.dumps(self.edges)
445        options = json.dumps(self.options)
446        imports = (
447            """<script type="text/javascript" src="https://unpkg.com/vis-data@latest/peer/umd/vis-data.min.js"></script>
448  <script type="text/javascript" src="https://unpkg.com/vis-network@latest/peer/umd/vis-network.min.js"></script>
449  <link rel="stylesheet" type="text/css" href="https://unpkg.com/vis-network/styles/vis-network.min.css" />"""
450            if self.imports
451            else ""
452        )
453
454        return f"""<div>
455  <div id="sqlglot-lineage"></div>
456  {imports}
457  <script type="text/javascript">
458    var nodes = new vis.DataSet({nodes})
459    nodes.forEach(row => row["title"] = new DOMParser().parseFromString(row["title"], "text/html").body.childNodes[0])
460
461    new vis.Network(
462        document.getElementById("sqlglot-lineage"),
463        {{
464            nodes: nodes,
465            edges: new vis.DataSet({edges})
466        }},
467        {options},
468    )
469  </script>
470</div>"""
471
472    def _repr_html_(self) -> str:
473        return self.__str__()