Edit on GitHub

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
 12from sqlglot.schema import ensure_schema
 13
 14if t.TYPE_CHECKING:
 15    from sqlglot.dialects.dialect import DialectType
 16    from collections.abc import Iterator, Mapping, Sequence
 17    from sqlglot._typing import GraphHTMLArgs
 18    from typing_extensions import Unpack
 19
 20logger = logging.getLogger("sqlglot")
 21
 22
 23@dataclass(frozen=True)
 24class Node:
 25    name: str
 26    expression: exp.Expr
 27    source: exp.Expr
 28    downstream: list[Node] = field(default_factory=list)
 29    source_name: str = ""
 30    reference_node_name: str = ""
 31
 32    # Caller-injected per-node data, populated via the `on_node` hook on lineage()
 33    payload: dict[str, t.Any] = field(default_factory=dict)
 34
 35    def walk(self) -> Iterator[Node]:
 36        visited: set[int] = set()
 37        queue = [self]
 38        while queue:
 39            node = queue.pop()
 40            node_id = id(node)
 41            if node_id in visited:
 42                continue
 43            visited.add(node_id)
 44            yield node
 45            queue.extend(reversed(node.downstream))
 46
 47    def to_html(self, dialect: DialectType = None, **opts: Unpack[GraphHTMLArgs]) -> GraphHTML:
 48        nodes = {}
 49        edges = []
 50
 51        for node in self.walk():
 52            if isinstance(node.expression, exp.Table):
 53                label = f"FROM {node.expression.this}"
 54                title = f"<pre>SELECT {node.name} FROM {node.expression.this}</pre>"
 55                group = 1
 56            else:
 57                label = node.expression.sql(pretty=True, dialect=dialect)
 58                source = node.source.transform(
 59                    lambda n: (
 60                        exp.Tag(this=n, prefix="<b>", postfix="</b>") if n is node.expression else n
 61                    ),
 62                    copy=False,
 63                ).sql(pretty=True, dialect=dialect)
 64                title = f"<pre>{source}</pre>"
 65                group = 0
 66
 67            node_id = id(node)
 68
 69            nodes[node_id] = {
 70                "id": node_id,
 71                "label": label,
 72                "title": title,
 73                "group": group,
 74            }
 75
 76            for d in node.downstream:
 77                edges.append({"from": node_id, "to": id(d)})
 78        return GraphHTML(nodes, edges, **opts)
 79
 80
 81@t.overload
 82def lineage(column: str | exp.Column, sql: str | exp.Expr, **kwargs: t.Any) -> Node: ...
 83
 84
 85@t.overload
 86def lineage(column: None, sql: str | exp.Expr, **kwargs: t.Any) -> dict[str, Node]: ...
 87
 88
 89def lineage(
 90    column: str | exp.Column | None,
 91    sql: str | exp.Expr,
 92    schema: dict | Schema | None = None,
 93    sources: Mapping[str, str | exp.Query] | None = None,
 94    dialect: DialectType = None,
 95    scope: Scope | None = None,
 96    trim_selects: bool = True,
 97    copy: bool = True,
 98    on_node: t.Callable[[Node], None] | None = None,
 99    **kwargs,
100) -> Node | dict[str, Node]:
101    """Build the lineage graph for a SQL query.
102
103    If `column` is given, returns the lineage Node for that single output column.
104    If `column` is None, returns a dict mapping every top-level output column name
105    to its lineage Node (with a shared cache so cross-column work is deduplicated).
106
107    Args:
108        column: The column to build the lineage for. Pass None to get all output columns.
109        sql: The SQL string or expression.
110        schema: The schema of tables.
111        sources: A mapping of queries which will be used to continue building lineage.
112        dialect: The dialect of input SQL.
113        scope: A pre-created scope to use instead.
114        trim_selects: Whether to clean up selects by trimming to only relevant columns.
115        copy: Whether to copy the Expr arguments.
116        on_node: Optional callback invoked for every Node created during the walk,
117            after the Node's downstream is populated. Useful for injecting
118            caller-managed data into Node.payload during the walk.
119        **kwargs: Qualification optimizer kwargs.
120
121    Returns:
122        A Node when `column` is provided, or a dict[str, Node] when `column` is None.
123    """
124    expression = maybe_parse(sql, copy=copy, dialect=dialect)
125
126    if sources:
127        expression = exp.expand(
128            expression,
129            {
130                k: t.cast(exp.Query, maybe_parse(v, copy=copy, dialect=dialect))
131                for k, v in sources.items()
132            },
133            dialect=dialect,
134            copy=copy,
135        )
136
137    schema = ensure_schema(schema, dialect=dialect)
138
139    if not scope:
140        expression = qualify.qualify(
141            expression,
142            dialect=dialect,
143            schema=schema,
144            **{"validate_qualify_columns": False, "identify": False, **kwargs},  # type: ignore
145        )
146        scope = build_scope(expression)
147
148    if not scope:
149        raise SqlglotError("Cannot build lineage, sql must be SELECT")
150
151    selectable = scope.expression
152    if not isinstance(selectable, exp.Selectable):
153        raise SqlglotError("Cannot build lineage, sql must be a query")
154
155    cache: dict[tuple, Node] = {}
156    scope_meta: dict[int, tuple[bool, dict[str, exp.Expr]]] = {}
157
158    if column is not None:
159        column_name = normalize_identifiers.normalize_identifiers(column, dialect=dialect).name
160        if not any(select.alias_or_name == column_name for select in selectable.selects):
161            raise SqlglotError(f"Cannot find column '{column_name}' in query.")
162
163        return to_node(
164            column_name,
165            scope,
166            dialect,
167            trim_selects=trim_selects,
168            schema=schema,
169            _cache=cache,
170            _scope_meta=scope_meta,
171            on_node=on_node,
172        )
173
174    result: dict[str, Node] = {}
175    for sel in selectable.selects:
176        name = sel.alias_or_name
177        if not name:
178            raise SqlglotError(
179                f"Cannot fetch lineage for unnamed projection: {sel.sql(dialect=dialect)}."
180            )
181
182        result[name] = to_node(
183            name,
184            scope,
185            dialect,
186            trim_selects=trim_selects,
187            schema=schema,
188            _cache=cache,
189            _scope_meta=scope_meta,
190            on_node=on_node,
191        )
192
193    return result
194
195
196def to_node(
197    column: str | int,
198    scope: Scope,
199    dialect: DialectType,
200    scope_name: str | None = None,
201    upstream: Node | None = None,
202    source_name: str | None = None,
203    reference_node_name: str | None = None,
204    trim_selects: bool = True,
205    schema: Schema | None = None,
206    _cache: dict[tuple, Node] | None = None,
207    _scope_meta: dict[int, tuple[bool, dict[str, exp.Expr]]] | None = None,
208    on_node: t.Callable[[Node], None] | None = None,
209) -> Node:
210    cache_key = (column, id(scope), scope_name, source_name, reference_node_name)
211
212    if _cache is not None and cache_key in _cache:
213        cached_node = _cache[cache_key]
214        if upstream:
215            upstream.downstream.append(cached_node)
216        return cached_node
217
218    # Find the specific select clause that is the source of the column we want.
219    # This can either be a specific, named select or a generic `*` clause.
220    selectable = t.cast(exp.Selectable, scope.expression)
221    if isinstance(column, int):
222        if column >= len(selectable.selects):
223            raise SqlglotError(
224                f"Cannot find column's source with index {column} in query: {selectable.sql(dialect=dialect)}"
225            )
226        select = selectable.selects[column]
227    else:
228        # Resolving a column to its select scans selectable.selects on every call;
229        # memoize a per-scope {name: select} map and is_star bit instead.
230        if _scope_meta is None:
231            select = next(
232                (s for s in selectable.selects if s.alias_or_name == column),
233                exp.Star() if selectable.is_star else scope.expression,
234            )
235        else:
236            scope_id = id(scope)
237            meta = _scope_meta.get(scope_id)
238            if meta is None:
239                select_by_name: dict[str, exp.Expr] = {}
240                for sel in selectable.selects:
241                    select_by_name.setdefault(sel.alias_or_name, sel)
242                meta = (selectable.is_star, select_by_name)
243                _scope_meta[scope_id] = meta
244            is_star, select_by_name = meta
245            select = select_by_name.get(column, exp.Star() if is_star else scope.expression)
246
247    if isinstance(scope.expression, exp.Subquery):
248        for inner_scope in scope.subquery_scopes:
249            result = to_node(
250                column,
251                scope=inner_scope,
252                dialect=dialect,
253                upstream=upstream,
254                source_name=source_name,
255                reference_node_name=reference_node_name,
256                trim_selects=trim_selects,
257                schema=schema,
258                _cache=_cache,
259                _scope_meta=_scope_meta,
260                on_node=on_node,
261            )
262            # Skip caching a passed-in upstream returned by an inner SetOp:
263            # a sibling call at the same key with that node as its upstream
264            # would otherwise self-loop on the cache hit.
265            if _cache is not None and result is not upstream:
266                _cache[cache_key] = result
267            return result
268    if isinstance(scope.expression, exp.SetOperation):
269        name = type(scope.expression).__name__.upper()
270        created_setop = upstream is None
271        upstream = upstream or Node(name=name, source=scope.expression, expression=select)
272
273        index = (
274            column
275            if isinstance(column, int)
276            else next(
277                (
278                    i
279                    for i, select in enumerate(selectable.selects)
280                    if select.alias_or_name == column or select.is_star
281                ),
282                -1,  # mypy will not allow a None here, but a negative index should never be returned
283            )
284        )
285
286        if index == -1:
287            raise ValueError(f"Could not find {column} in {scope.expression}")
288
289        for s in scope.union_scopes:
290            to_node(
291                index,
292                scope=s,
293                dialect=dialect,
294                upstream=upstream,
295                source_name=source_name,
296                reference_node_name=reference_node_name,
297                trim_selects=trim_selects,
298                schema=schema,
299                _cache=_cache,
300                _scope_meta=_scope_meta,
301                on_node=on_node,
302            )
303
304        if _cache is not None and created_setop:
305            _cache[cache_key] = upstream
306        if created_setop and on_node:
307            on_node(upstream)
308        return upstream
309
310    if trim_selects and isinstance(scope.expression, exp.Select):
311        # For better ergonomics in our node labels, replace the full select with
312        # a version that has only the column we care about.
313        #   "x", SELECT x, y FROM foo
314        #     => "x", SELECT x FROM foo
315        source: exp.Expr = scope.expression.select(select, append=False)
316    else:
317        source = scope.expression
318
319    # Create the node for this step in the lineage chain, and attach it to the previous one.
320    node = Node(
321        name=f"{scope_name}.{column}" if scope_name else str(column),
322        source=source,
323        expression=select,
324        source_name=source_name or "",
325        reference_node_name=reference_node_name or "",
326    )
327
328    if upstream:
329        upstream.downstream.append(node)
330
331    subquery_scopes = {
332        id(subquery_scope.expression): subquery_scope for subquery_scope in scope.subquery_scopes
333    }
334
335    for subquery in find_all_in_scope(select, *exp.UNWRAPPED_QUERIES):
336        subquery_scope: Scope | None = subquery_scopes.get(id(subquery))
337        if not subquery_scope:
338            logger.warning(f"Unknown subquery scope: {subquery.sql(dialect=dialect)}")
339            continue
340
341        for name in subquery.named_selects:
342            to_node(
343                name,
344                scope=subquery_scope,
345                dialect=dialect,
346                upstream=node,
347                trim_selects=trim_selects,
348                schema=schema,
349                _cache=_cache,
350                _scope_meta=_scope_meta,
351                on_node=on_node,
352            )
353
354    # if the select is a star add all scope sources as downstreams
355    if isinstance(select, exp.Star):
356        for src in scope.sources.values():
357            src_expr = src.expression if isinstance(src, Scope) else src
358            star_node = Node(name=select.sql(comments=False), source=src_expr, expression=src_expr)
359            node.downstream.append(star_node)
360            if on_node:
361                on_node(star_node)
362
363    # Find all columns that went into creating this one to list their lineage nodes.
364    source_columns = set(find_all_in_scope(select, exp.Column))
365
366    # If the source is a UDTF find columns used in the UDTF to generate the table
367    if isinstance(source, exp.UDTF):
368        source_columns |= set(source.find_all(exp.Column))
369        derived_tables: Sequence[exp.Expr] = [
370            src.expression.parent
371            for src in scope.sources.values()
372            if isinstance(src, Scope) and src.is_derived_table and src.expression.parent
373        ]
374    else:
375        derived_tables = scope.derived_tables
376
377    source_names = {
378        dt.alias: dt.comments[0].split()[1]
379        for dt in derived_tables
380        if dt.comments and dt.comments[0].startswith("source: ")
381    }
382
383    pivots = scope.pivots
384    if pivots and pivots[0].parent is not pivots[-1].parent:
385        # The scope's pivots only form a chain when they all hang off of the same source;
386        # otherwise they can't be folded, so their columns degrade to unresolved leaves
387        pivots = []
388
389    pivot_renames: dict[str, str] = {}
390    pivot_column_mapping: dict[str, list[exp.Column]] = {}
391
392    if pivots:
393        pivot_renames, pivot_column_mapping = _pivot_chain_mapping(pivots, scope, schema)
394
395    for c in source_columns:
396        table = c.table
397        col_source: exp.Table | Scope | None = scope.sources.get(table)
398
399        if isinstance(col_source, Scope):
400            reference_node_name = None
401            if col_source.scope_type == ScopeType.DERIVED_TABLE and table not in source_names:
402                reference_node_name = table
403            elif col_source.scope_type == ScopeType.CTE:
404                selected_node, _ = scope.selected_sources.get(table, (None, None))
405                reference_node_name = selected_node.name if selected_node else None
406
407            # The table itself came from a more specific scope. Recurse into that one using the unaliased column name.
408            to_node(
409                c.name,
410                scope=col_source,
411                dialect=dialect,
412                scope_name=table,
413                upstream=node,
414                source_name=source_names.get(table) or source_name,
415                reference_node_name=reference_node_name,
416                trim_selects=trim_selects,
417                schema=schema,
418                _cache=_cache,
419                _scope_meta=_scope_meta,
420                on_node=on_node,
421            )
422        elif pivots and pivots[-1].alias_or_name == c.table:
423            # Only the last operator in a chain names the resulting source
424            pivot_parent = pivots[-1].parent
425            downstream_columns = []
426
427            column_name = c.name
428            if column_name in pivot_column_mapping:
429                downstream_columns.extend(pivot_column_mapping[column_name])
430            else:
431                # The column is not in the pivot, so it must be an implicit column of the
432                # pivoted source -- adapt column to be from the implicit pivoted source.
433                downstream_columns.append(
434                    exp.column(
435                        pivot_renames.get(c.name, c.this),
436                        table=pivot_parent.alias_or_name if pivot_parent else None,
437                    )
438                )
439
440            for downstream_column in downstream_columns:
441                if not downstream_column.table:
442                    # Some dialects (e.g. bigquery) don't qualify the IN-list columns,
443                    # but they can only come from the pivoted source
444                    downstream_column = exp.column(
445                        downstream_column.this,
446                        table=pivot_parent.alias_or_name if pivot_parent else None,
447                    )
448
449                table = downstream_column.table
450                col_source = scope.sources.get(table)
451                if isinstance(col_source, exp.Table) and not col_source.db:
452                    # A pivoted CTE reference maps to the raw table in `scope.sources`,
453                    # so recover the CTE's scope to keep tracing through it
454                    col_source = scope.cte_sources.get(col_source.name, col_source)
455                if isinstance(col_source, Scope):
456                    to_node(
457                        downstream_column.name,
458                        scope=col_source,
459                        scope_name=table,
460                        dialect=dialect,
461                        upstream=node,
462                        source_name=source_names.get(table) or source_name,
463                        reference_node_name=reference_node_name,
464                        trim_selects=trim_selects,
465                        schema=schema,
466                        _cache=_cache,
467                        _scope_meta=_scope_meta,
468                        on_node=on_node,
469                    )
470                else:
471                    col_expr = col_source or exp.Placeholder()
472                    pivot_leaf = Node(
473                        name=downstream_column.sql(comments=False),
474                        source=col_expr,
475                        expression=col_expr,
476                    )
477                    node.downstream.append(pivot_leaf)
478                    if on_node:
479                        on_node(pivot_leaf)
480        else:
481            # The source is not a scope and the column is not in any pivot - we've reached the end
482            # of the line. At this point, if a source is not found it means this column's lineage
483            # is unknown. This can happen if the definition of a source used in a query is not
484            # passed into the `sources` map.
485            col_expr = col_source or exp.Placeholder()
486            leaf = Node(name=c.sql(comments=False), source=col_expr, expression=col_expr)
487            node.downstream.append(leaf)
488            if on_node:
489                on_node(leaf)
490
491    if _cache is not None:
492        _cache[cache_key] = node
493
494    if on_node:
495        on_node(node)
496
497    return node
498
499
500def _pre_pivot_columns(pivot: exp.Pivot, scope: Scope, schema: Schema | None = None) -> list[str]:
501    """
502    The columns the first operator of a chain sees, taken from the projections of a
503    derived table or CTE source, or from the schema for a physical table. Returns an
504    empty list when they can't be determined (e.g. an unexpanded star), since anything
505    positional over them would silently shift.
506    """
507    parent = pivot.parent
508    columns: list[str] = []
509    if isinstance(parent, exp.DerivedTable) and isinstance(parent.this, exp.Query):
510        columns = parent.this.named_selects
511    elif isinstance(parent, exp.Table):
512        cte_source = scope.cte_sources.get(parent.name) if not parent.db else None
513        if isinstance(cte_source, Scope) and isinstance(cte_source.expression, exp.Query):
514            columns = cte_source.expression.named_selects
515        elif schema is not None:
516            columns = list(schema.column_names(parent, only_visible=True))
517
518    return [] if "*" in columns else columns
519
520
521def _pivot_chain_mapping(
522    pivots: list[exp.Pivot], scope: Scope, schema: Schema | None = None
523) -> tuple[dict[str, str], dict[str, list[exp.Column]]]:
524    """
525    Fold a chain of (UN)PIVOT operators into a single view of its output, since each one
526    consumes the previous one's columns rather than the pivoted source's.
527
528    Returns the composed output-name -> pre-chain-name renames (from alias column lists),
529    and the composed output-name -> source columns it derives from.
530    """
531    available = _pre_pivot_columns(pivots[0], scope, schema)
532    renames: dict[str, str] = {}
533    mapping: dict[str, list[exp.Column]] = {}
534
535    for pivot in pivots:
536        # Renames are positional over the operator's full output, so they can only be
537        # applied when the columns going into it are known
538        step_renames = (
539            pivot.output_columns(available) if pivot.alias_column_names and available else {}
540        )
541        step_mapping = _pivot_column_mapping(pivot)
542        if step_renames:
543            step_mapping = {
544                post: step_mapping[pre] for post, pre in step_renames.items() if pre in step_mapping
545            }
546
547        # Columns this operator consumed may have been produced by an earlier one, or be
548        # alias-list renames of passthroughs; resolve through what we've folded so far
549        def resolve(col: exp.Column) -> list[exp.Column]:
550            if col.name in mapping:
551                return mapping[col.name]
552            if col.name in renames:
553                return [exp.column(renames[col.name])]
554            return [col]
555
556        composed = {
557            out: [resolved for col in cols for resolved in resolve(col)]
558            for out, cols in step_mapping.items()
559        }
560
561        # Whatever an earlier operator produced and this one didn't consume passes through,
562        # under whatever name this operator's alias column list gives it
563        consumed = {col.name for cols in step_mapping.values() for col in cols}
564        pre_to_post = {pre: post for post, pre in step_renames.items()}
565        for out, cols in mapping.items():
566            if out not in consumed:
567                composed.setdefault(pre_to_post.get(out, out), cols)
568
569        renames = (
570            {post: renames.get(pre, pre) for post, pre in step_renames.items()}
571            if step_renames
572            else renames
573        )
574        mapping = composed
575        available = list(pivot.output_columns(available)) if available else []
576
577    return renames, mapping
578
579
580def _pivot_column_mapping(pivot: exp.Pivot) -> dict[str, list[exp.Column]]:
581    """Map each (UN)PIVOT output column name to the source columns it's derived from."""
582    mapping: dict[str, list[exp.Column]] = {}
583
584    if pivot.unpivot:
585        # UNPIVOT((v1, v2) FOR name IN ((a1, a2), (b1, b2))): each value column is derived
586        # positionally from the IN-list entries, and the name column from all of them
587        value_columns = [
588            identifier for e in pivot.expressions for identifier in e.find_all(exp.Identifier)
589        ]
590        for value_column in value_columns:
591            mapping[value_column.name] = []
592
593        for field in pivot.fields:
594            if not isinstance(field, exp.In):
595                continue
596
597            name_columns = mapping.setdefault(field.this.name, [])
598            for entry in field.expressions:
599                entry_columns = list(entry.find_all(exp.Column))
600                name_columns.extend(entry_columns)
601
602                if len(entry_columns) == len(value_columns):
603                    for value_column, column in zip(value_columns, entry_columns):
604                        mapping[value_column.name].append(column)
605                else:
606                    for value_column in value_columns:
607                        mapping[value_column.name].extend(entry_columns)
608
609        return mapping
610
611    # For each aggregation function, the pivot creates a new column for each field in category
612    # combined with the aggfunc. So the columns parsed have this order: cat_a_value_sum, cat_a,
613    # b_value_sum, b. Because of this step wise manner the aggfunc 'sum(value) as value_sum'
614    # belongs to the column indices 0, 2, and the aggfunc 'max(price)' without an alias belongs
615    # to the column indices 1, 3. Here, only the columns used in the aggregations are of interest
616    # in the lineage, so lookup the pivot column name by index and map that with the columns used
617    # in the aggregation.
618    #
619    # Example: PIVOT (SUM(value) AS value_sum, MAX(price)) FOR category IN ('a' AS cat_a, 'b')
620    pivot_columns = pivot.args["columns"]
621    pivot_aggs_count = len(pivot.expressions)
622
623    mapping = {}
624    for i, agg in enumerate(pivot.expressions):
625        agg_cols = list(agg.find_all(exp.Column))
626        for col_index in range(i, len(pivot_columns), pivot_aggs_count):
627            mapping[pivot_columns[col_index].name] = agg_cols
628    return mapping
629
630
631class GraphHTML:
632    """Node to HTML generator using vis.js.
633
634    https://visjs.github.io/vis-network/docs/network/
635    """
636
637    def __init__(
638        self,
639        nodes: dict,
640        edges: list,
641        imports: bool = True,
642        options: Mapping[str, object] | None = None,
643    ):
644        self.imports = imports
645
646        self.options = {
647            "height": "500px",
648            "width": "100%",
649            "layout": {
650                "hierarchical": {
651                    "enabled": True,
652                    "nodeSpacing": 200,
653                    "sortMethod": "directed",
654                },
655            },
656            "interaction": {
657                "dragNodes": False,
658                "selectable": False,
659            },
660            "physics": {
661                "enabled": False,
662            },
663            "edges": {
664                "arrows": "to",
665            },
666            "nodes": {
667                "font": "20px monaco",
668                "shape": "box",
669                "widthConstraint": {
670                    "maximum": 300,
671                },
672            },
673            **(options or {}),
674        }
675
676        self.nodes = nodes
677        self.edges = edges
678
679    def __str__(self):
680        nodes = json.dumps(list(self.nodes.values()))
681        edges = json.dumps(self.edges)
682        options = json.dumps(self.options)
683        imports = (
684            """<script type="text/javascript" src="https://unpkg.com/vis-data@latest/peer/umd/vis-data.min.js"></script>
685  <script type="text/javascript" src="https://unpkg.com/vis-network@latest/peer/umd/vis-network.min.js"></script>
686  <link rel="stylesheet" type="text/css" href="https://unpkg.com/vis-network/styles/vis-network.min.css" />"""
687            if self.imports
688            else ""
689        )
690
691        return f"""<div>
692  <div id="sqlglot-lineage"></div>
693  {imports}
694  <script type="text/javascript">
695    var nodes = new vis.DataSet({nodes})
696    nodes.forEach(row => row["title"] = new DOMParser().parseFromString(row["title"], "text/html").body.childNodes[0])
697
698    new vis.Network(
699        document.getElementById("sqlglot-lineage"),
700        {{
701            nodes: nodes,
702            edges: new vis.DataSet({edges})
703        }},
704        {options},
705    )
706  </script>
707</div>"""
708
709    def _repr_html_(self) -> str:
710        return self.__str__()
logger = <Logger sqlglot (WARNING)>
@dataclass(frozen=True)
class Node:
24@dataclass(frozen=True)
25class Node:
26    name: str
27    expression: exp.Expr
28    source: exp.Expr
29    downstream: list[Node] = field(default_factory=list)
30    source_name: str = ""
31    reference_node_name: str = ""
32
33    # Caller-injected per-node data, populated via the `on_node` hook on lineage()
34    payload: dict[str, t.Any] = field(default_factory=dict)
35
36    def walk(self) -> Iterator[Node]:
37        visited: set[int] = set()
38        queue = [self]
39        while queue:
40            node = queue.pop()
41            node_id = id(node)
42            if node_id in visited:
43                continue
44            visited.add(node_id)
45            yield node
46            queue.extend(reversed(node.downstream))
47
48    def to_html(self, dialect: DialectType = None, **opts: Unpack[GraphHTMLArgs]) -> GraphHTML:
49        nodes = {}
50        edges = []
51
52        for node in self.walk():
53            if isinstance(node.expression, exp.Table):
54                label = f"FROM {node.expression.this}"
55                title = f"<pre>SELECT {node.name} FROM {node.expression.this}</pre>"
56                group = 1
57            else:
58                label = node.expression.sql(pretty=True, dialect=dialect)
59                source = node.source.transform(
60                    lambda n: (
61                        exp.Tag(this=n, prefix="<b>", postfix="</b>") if n is node.expression else n
62                    ),
63                    copy=False,
64                ).sql(pretty=True, dialect=dialect)
65                title = f"<pre>{source}</pre>"
66                group = 0
67
68            node_id = id(node)
69
70            nodes[node_id] = {
71                "id": node_id,
72                "label": label,
73                "title": title,
74                "group": group,
75            }
76
77            for d in node.downstream:
78                edges.append({"from": node_id, "to": id(d)})
79        return GraphHTML(nodes, edges, **opts)
Node( name: str, expression: sqlglot.expressions.core.Expr, source: sqlglot.expressions.core.Expr, downstream: list[Node] = <factory>, source_name: str = '', reference_node_name: str = '', payload: dict[str, typing.Any] = <factory>)
name: str
downstream: list[Node]
source_name: str = ''
reference_node_name: str = ''
payload: dict[str, typing.Any]
def walk(self) -> Iterator[Node]:
36    def walk(self) -> Iterator[Node]:
37        visited: set[int] = set()
38        queue = [self]
39        while queue:
40            node = queue.pop()
41            node_id = id(node)
42            if node_id in visited:
43                continue
44            visited.add(node_id)
45            yield node
46            queue.extend(reversed(node.downstream))
def to_html( self, dialect: Union[str, sqlglot.dialects.Dialect, type[sqlglot.dialects.Dialect], NoneType] = None, **opts: typing_extensions.Unpack[sqlglot._typing.GraphHTMLArgs]) -> GraphHTML:
48    def to_html(self, dialect: DialectType = None, **opts: Unpack[GraphHTMLArgs]) -> GraphHTML:
49        nodes = {}
50        edges = []
51
52        for node in self.walk():
53            if isinstance(node.expression, exp.Table):
54                label = f"FROM {node.expression.this}"
55                title = f"<pre>SELECT {node.name} FROM {node.expression.this}</pre>"
56                group = 1
57            else:
58                label = node.expression.sql(pretty=True, dialect=dialect)
59                source = node.source.transform(
60                    lambda n: (
61                        exp.Tag(this=n, prefix="<b>", postfix="</b>") if n is node.expression else n
62                    ),
63                    copy=False,
64                ).sql(pretty=True, dialect=dialect)
65                title = f"<pre>{source}</pre>"
66                group = 0
67
68            node_id = id(node)
69
70            nodes[node_id] = {
71                "id": node_id,
72                "label": label,
73                "title": title,
74                "group": group,
75            }
76
77            for d in node.downstream:
78                edges.append({"from": node_id, "to": id(d)})
79        return GraphHTML(nodes, edges, **opts)
def lineage( column: str | sqlglot.expressions.core.Column | None, sql: str | sqlglot.expressions.core.Expr, schema: dict | sqlglot.schema.Schema | None = None, sources: Mapping[str, str | sqlglot.expressions.query.Query] | None = None, dialect: Union[str, sqlglot.dialects.Dialect, type[sqlglot.dialects.Dialect], NoneType] = None, scope: sqlglot.optimizer.scope.Scope | None = None, trim_selects: bool = True, copy: bool = True, on_node: Optional[Callable[[Node], NoneType]] = None, **kwargs) -> Node | dict[str, Node]:
 90def lineage(
 91    column: str | exp.Column | None,
 92    sql: str | exp.Expr,
 93    schema: dict | Schema | None = None,
 94    sources: Mapping[str, str | exp.Query] | None = None,
 95    dialect: DialectType = None,
 96    scope: Scope | None = None,
 97    trim_selects: bool = True,
 98    copy: bool = True,
 99    on_node: t.Callable[[Node], None] | None = None,
100    **kwargs,
101) -> Node | dict[str, Node]:
102    """Build the lineage graph for a SQL query.
103
104    If `column` is given, returns the lineage Node for that single output column.
105    If `column` is None, returns a dict mapping every top-level output column name
106    to its lineage Node (with a shared cache so cross-column work is deduplicated).
107
108    Args:
109        column: The column to build the lineage for. Pass None to get all output columns.
110        sql: The SQL string or expression.
111        schema: The schema of tables.
112        sources: A mapping of queries which will be used to continue building lineage.
113        dialect: The dialect of input SQL.
114        scope: A pre-created scope to use instead.
115        trim_selects: Whether to clean up selects by trimming to only relevant columns.
116        copy: Whether to copy the Expr arguments.
117        on_node: Optional callback invoked for every Node created during the walk,
118            after the Node's downstream is populated. Useful for injecting
119            caller-managed data into Node.payload during the walk.
120        **kwargs: Qualification optimizer kwargs.
121
122    Returns:
123        A Node when `column` is provided, or a dict[str, Node] when `column` is None.
124    """
125    expression = maybe_parse(sql, copy=copy, dialect=dialect)
126
127    if sources:
128        expression = exp.expand(
129            expression,
130            {
131                k: t.cast(exp.Query, maybe_parse(v, copy=copy, dialect=dialect))
132                for k, v in sources.items()
133            },
134            dialect=dialect,
135            copy=copy,
136        )
137
138    schema = ensure_schema(schema, dialect=dialect)
139
140    if not scope:
141        expression = qualify.qualify(
142            expression,
143            dialect=dialect,
144            schema=schema,
145            **{"validate_qualify_columns": False, "identify": False, **kwargs},  # type: ignore
146        )
147        scope = build_scope(expression)
148
149    if not scope:
150        raise SqlglotError("Cannot build lineage, sql must be SELECT")
151
152    selectable = scope.expression
153    if not isinstance(selectable, exp.Selectable):
154        raise SqlglotError("Cannot build lineage, sql must be a query")
155
156    cache: dict[tuple, Node] = {}
157    scope_meta: dict[int, tuple[bool, dict[str, exp.Expr]]] = {}
158
159    if column is not None:
160        column_name = normalize_identifiers.normalize_identifiers(column, dialect=dialect).name
161        if not any(select.alias_or_name == column_name for select in selectable.selects):
162            raise SqlglotError(f"Cannot find column '{column_name}' in query.")
163
164        return to_node(
165            column_name,
166            scope,
167            dialect,
168            trim_selects=trim_selects,
169            schema=schema,
170            _cache=cache,
171            _scope_meta=scope_meta,
172            on_node=on_node,
173        )
174
175    result: dict[str, Node] = {}
176    for sel in selectable.selects:
177        name = sel.alias_or_name
178        if not name:
179            raise SqlglotError(
180                f"Cannot fetch lineage for unnamed projection: {sel.sql(dialect=dialect)}."
181            )
182
183        result[name] = to_node(
184            name,
185            scope,
186            dialect,
187            trim_selects=trim_selects,
188            schema=schema,
189            _cache=cache,
190            _scope_meta=scope_meta,
191            on_node=on_node,
192        )
193
194    return result

Build the lineage graph for a SQL query.

If column is given, returns the lineage Node for that single output column. If column is None, returns a dict mapping every top-level output column name to its lineage Node (with a shared cache so cross-column work is deduplicated).

Arguments:
  • column: The column to build the lineage for. Pass None to get all output columns.
  • sql: The SQL string or expression.
  • schema: The schema of tables.
  • sources: A mapping of queries which will be used to continue building lineage.
  • dialect: The dialect of input SQL.
  • scope: A pre-created scope to use instead.
  • trim_selects: Whether to clean up selects by trimming to only relevant columns.
  • copy: Whether to copy the Expr arguments.
  • on_node: Optional callback invoked for every Node created during the walk, after the Node's downstream is populated. Useful for injecting caller-managed data into Node.payload during the walk.
  • **kwargs: Qualification optimizer kwargs.
Returns:

A Node when column is provided, or a dict[str, Node] when column is None.

def to_node( column: str | int, scope: sqlglot.optimizer.scope.Scope, dialect: Union[str, sqlglot.dialects.Dialect, type[sqlglot.dialects.Dialect], NoneType], scope_name: str | None = None, upstream: Node | None = None, source_name: str | None = None, reference_node_name: str | None = None, trim_selects: bool = True, schema: sqlglot.schema.Schema | None = None, _cache: dict[tuple, Node] | None = None, _scope_meta: dict[int, tuple[bool, dict[str, sqlglot.expressions.core.Expr]]] | None = None, on_node: Optional[Callable[[Node], NoneType]] = None) -> Node:
197def to_node(
198    column: str | int,
199    scope: Scope,
200    dialect: DialectType,
201    scope_name: str | None = None,
202    upstream: Node | None = None,
203    source_name: str | None = None,
204    reference_node_name: str | None = None,
205    trim_selects: bool = True,
206    schema: Schema | None = None,
207    _cache: dict[tuple, Node] | None = None,
208    _scope_meta: dict[int, tuple[bool, dict[str, exp.Expr]]] | None = None,
209    on_node: t.Callable[[Node], None] | None = None,
210) -> Node:
211    cache_key = (column, id(scope), scope_name, source_name, reference_node_name)
212
213    if _cache is not None and cache_key in _cache:
214        cached_node = _cache[cache_key]
215        if upstream:
216            upstream.downstream.append(cached_node)
217        return cached_node
218
219    # Find the specific select clause that is the source of the column we want.
220    # This can either be a specific, named select or a generic `*` clause.
221    selectable = t.cast(exp.Selectable, scope.expression)
222    if isinstance(column, int):
223        if column >= len(selectable.selects):
224            raise SqlglotError(
225                f"Cannot find column's source with index {column} in query: {selectable.sql(dialect=dialect)}"
226            )
227        select = selectable.selects[column]
228    else:
229        # Resolving a column to its select scans selectable.selects on every call;
230        # memoize a per-scope {name: select} map and is_star bit instead.
231        if _scope_meta is None:
232            select = next(
233                (s for s in selectable.selects if s.alias_or_name == column),
234                exp.Star() if selectable.is_star else scope.expression,
235            )
236        else:
237            scope_id = id(scope)
238            meta = _scope_meta.get(scope_id)
239            if meta is None:
240                select_by_name: dict[str, exp.Expr] = {}
241                for sel in selectable.selects:
242                    select_by_name.setdefault(sel.alias_or_name, sel)
243                meta = (selectable.is_star, select_by_name)
244                _scope_meta[scope_id] = meta
245            is_star, select_by_name = meta
246            select = select_by_name.get(column, exp.Star() if is_star else scope.expression)
247
248    if isinstance(scope.expression, exp.Subquery):
249        for inner_scope in scope.subquery_scopes:
250            result = to_node(
251                column,
252                scope=inner_scope,
253                dialect=dialect,
254                upstream=upstream,
255                source_name=source_name,
256                reference_node_name=reference_node_name,
257                trim_selects=trim_selects,
258                schema=schema,
259                _cache=_cache,
260                _scope_meta=_scope_meta,
261                on_node=on_node,
262            )
263            # Skip caching a passed-in upstream returned by an inner SetOp:
264            # a sibling call at the same key with that node as its upstream
265            # would otherwise self-loop on the cache hit.
266            if _cache is not None and result is not upstream:
267                _cache[cache_key] = result
268            return result
269    if isinstance(scope.expression, exp.SetOperation):
270        name = type(scope.expression).__name__.upper()
271        created_setop = upstream is None
272        upstream = upstream or Node(name=name, source=scope.expression, expression=select)
273
274        index = (
275            column
276            if isinstance(column, int)
277            else next(
278                (
279                    i
280                    for i, select in enumerate(selectable.selects)
281                    if select.alias_or_name == column or select.is_star
282                ),
283                -1,  # mypy will not allow a None here, but a negative index should never be returned
284            )
285        )
286
287        if index == -1:
288            raise ValueError(f"Could not find {column} in {scope.expression}")
289
290        for s in scope.union_scopes:
291            to_node(
292                index,
293                scope=s,
294                dialect=dialect,
295                upstream=upstream,
296                source_name=source_name,
297                reference_node_name=reference_node_name,
298                trim_selects=trim_selects,
299                schema=schema,
300                _cache=_cache,
301                _scope_meta=_scope_meta,
302                on_node=on_node,
303            )
304
305        if _cache is not None and created_setop:
306            _cache[cache_key] = upstream
307        if created_setop and on_node:
308            on_node(upstream)
309        return upstream
310
311    if trim_selects and isinstance(scope.expression, exp.Select):
312        # For better ergonomics in our node labels, replace the full select with
313        # a version that has only the column we care about.
314        #   "x", SELECT x, y FROM foo
315        #     => "x", SELECT x FROM foo
316        source: exp.Expr = scope.expression.select(select, append=False)
317    else:
318        source = scope.expression
319
320    # Create the node for this step in the lineage chain, and attach it to the previous one.
321    node = Node(
322        name=f"{scope_name}.{column}" if scope_name else str(column),
323        source=source,
324        expression=select,
325        source_name=source_name or "",
326        reference_node_name=reference_node_name or "",
327    )
328
329    if upstream:
330        upstream.downstream.append(node)
331
332    subquery_scopes = {
333        id(subquery_scope.expression): subquery_scope for subquery_scope in scope.subquery_scopes
334    }
335
336    for subquery in find_all_in_scope(select, *exp.UNWRAPPED_QUERIES):
337        subquery_scope: Scope | None = subquery_scopes.get(id(subquery))
338        if not subquery_scope:
339            logger.warning(f"Unknown subquery scope: {subquery.sql(dialect=dialect)}")
340            continue
341
342        for name in subquery.named_selects:
343            to_node(
344                name,
345                scope=subquery_scope,
346                dialect=dialect,
347                upstream=node,
348                trim_selects=trim_selects,
349                schema=schema,
350                _cache=_cache,
351                _scope_meta=_scope_meta,
352                on_node=on_node,
353            )
354
355    # if the select is a star add all scope sources as downstreams
356    if isinstance(select, exp.Star):
357        for src in scope.sources.values():
358            src_expr = src.expression if isinstance(src, Scope) else src
359            star_node = Node(name=select.sql(comments=False), source=src_expr, expression=src_expr)
360            node.downstream.append(star_node)
361            if on_node:
362                on_node(star_node)
363
364    # Find all columns that went into creating this one to list their lineage nodes.
365    source_columns = set(find_all_in_scope(select, exp.Column))
366
367    # If the source is a UDTF find columns used in the UDTF to generate the table
368    if isinstance(source, exp.UDTF):
369        source_columns |= set(source.find_all(exp.Column))
370        derived_tables: Sequence[exp.Expr] = [
371            src.expression.parent
372            for src in scope.sources.values()
373            if isinstance(src, Scope) and src.is_derived_table and src.expression.parent
374        ]
375    else:
376        derived_tables = scope.derived_tables
377
378    source_names = {
379        dt.alias: dt.comments[0].split()[1]
380        for dt in derived_tables
381        if dt.comments and dt.comments[0].startswith("source: ")
382    }
383
384    pivots = scope.pivots
385    if pivots and pivots[0].parent is not pivots[-1].parent:
386        # The scope's pivots only form a chain when they all hang off of the same source;
387        # otherwise they can't be folded, so their columns degrade to unresolved leaves
388        pivots = []
389
390    pivot_renames: dict[str, str] = {}
391    pivot_column_mapping: dict[str, list[exp.Column]] = {}
392
393    if pivots:
394        pivot_renames, pivot_column_mapping = _pivot_chain_mapping(pivots, scope, schema)
395
396    for c in source_columns:
397        table = c.table
398        col_source: exp.Table | Scope | None = scope.sources.get(table)
399
400        if isinstance(col_source, Scope):
401            reference_node_name = None
402            if col_source.scope_type == ScopeType.DERIVED_TABLE and table not in source_names:
403                reference_node_name = table
404            elif col_source.scope_type == ScopeType.CTE:
405                selected_node, _ = scope.selected_sources.get(table, (None, None))
406                reference_node_name = selected_node.name if selected_node else None
407
408            # The table itself came from a more specific scope. Recurse into that one using the unaliased column name.
409            to_node(
410                c.name,
411                scope=col_source,
412                dialect=dialect,
413                scope_name=table,
414                upstream=node,
415                source_name=source_names.get(table) or source_name,
416                reference_node_name=reference_node_name,
417                trim_selects=trim_selects,
418                schema=schema,
419                _cache=_cache,
420                _scope_meta=_scope_meta,
421                on_node=on_node,
422            )
423        elif pivots and pivots[-1].alias_or_name == c.table:
424            # Only the last operator in a chain names the resulting source
425            pivot_parent = pivots[-1].parent
426            downstream_columns = []
427
428            column_name = c.name
429            if column_name in pivot_column_mapping:
430                downstream_columns.extend(pivot_column_mapping[column_name])
431            else:
432                # The column is not in the pivot, so it must be an implicit column of the
433                # pivoted source -- adapt column to be from the implicit pivoted source.
434                downstream_columns.append(
435                    exp.column(
436                        pivot_renames.get(c.name, c.this),
437                        table=pivot_parent.alias_or_name if pivot_parent else None,
438                    )
439                )
440
441            for downstream_column in downstream_columns:
442                if not downstream_column.table:
443                    # Some dialects (e.g. bigquery) don't qualify the IN-list columns,
444                    # but they can only come from the pivoted source
445                    downstream_column = exp.column(
446                        downstream_column.this,
447                        table=pivot_parent.alias_or_name if pivot_parent else None,
448                    )
449
450                table = downstream_column.table
451                col_source = scope.sources.get(table)
452                if isinstance(col_source, exp.Table) and not col_source.db:
453                    # A pivoted CTE reference maps to the raw table in `scope.sources`,
454                    # so recover the CTE's scope to keep tracing through it
455                    col_source = scope.cte_sources.get(col_source.name, col_source)
456                if isinstance(col_source, Scope):
457                    to_node(
458                        downstream_column.name,
459                        scope=col_source,
460                        scope_name=table,
461                        dialect=dialect,
462                        upstream=node,
463                        source_name=source_names.get(table) or source_name,
464                        reference_node_name=reference_node_name,
465                        trim_selects=trim_selects,
466                        schema=schema,
467                        _cache=_cache,
468                        _scope_meta=_scope_meta,
469                        on_node=on_node,
470                    )
471                else:
472                    col_expr = col_source or exp.Placeholder()
473                    pivot_leaf = Node(
474                        name=downstream_column.sql(comments=False),
475                        source=col_expr,
476                        expression=col_expr,
477                    )
478                    node.downstream.append(pivot_leaf)
479                    if on_node:
480                        on_node(pivot_leaf)
481        else:
482            # The source is not a scope and the column is not in any pivot - we've reached the end
483            # of the line. At this point, if a source is not found it means this column's lineage
484            # is unknown. This can happen if the definition of a source used in a query is not
485            # passed into the `sources` map.
486            col_expr = col_source or exp.Placeholder()
487            leaf = Node(name=c.sql(comments=False), source=col_expr, expression=col_expr)
488            node.downstream.append(leaf)
489            if on_node:
490                on_node(leaf)
491
492    if _cache is not None:
493        _cache[cache_key] = node
494
495    if on_node:
496        on_node(node)
497
498    return node
class GraphHTML:
632class GraphHTML:
633    """Node to HTML generator using vis.js.
634
635    https://visjs.github.io/vis-network/docs/network/
636    """
637
638    def __init__(
639        self,
640        nodes: dict,
641        edges: list,
642        imports: bool = True,
643        options: Mapping[str, object] | None = None,
644    ):
645        self.imports = imports
646
647        self.options = {
648            "height": "500px",
649            "width": "100%",
650            "layout": {
651                "hierarchical": {
652                    "enabled": True,
653                    "nodeSpacing": 200,
654                    "sortMethod": "directed",
655                },
656            },
657            "interaction": {
658                "dragNodes": False,
659                "selectable": False,
660            },
661            "physics": {
662                "enabled": False,
663            },
664            "edges": {
665                "arrows": "to",
666            },
667            "nodes": {
668                "font": "20px monaco",
669                "shape": "box",
670                "widthConstraint": {
671                    "maximum": 300,
672                },
673            },
674            **(options or {}),
675        }
676
677        self.nodes = nodes
678        self.edges = edges
679
680    def __str__(self):
681        nodes = json.dumps(list(self.nodes.values()))
682        edges = json.dumps(self.edges)
683        options = json.dumps(self.options)
684        imports = (
685            """<script type="text/javascript" src="https://unpkg.com/vis-data@latest/peer/umd/vis-data.min.js"></script>
686  <script type="text/javascript" src="https://unpkg.com/vis-network@latest/peer/umd/vis-network.min.js"></script>
687  <link rel="stylesheet" type="text/css" href="https://unpkg.com/vis-network/styles/vis-network.min.css" />"""
688            if self.imports
689            else ""
690        )
691
692        return f"""<div>
693  <div id="sqlglot-lineage"></div>
694  {imports}
695  <script type="text/javascript">
696    var nodes = new vis.DataSet({nodes})
697    nodes.forEach(row => row["title"] = new DOMParser().parseFromString(row["title"], "text/html").body.childNodes[0])
698
699    new vis.Network(
700        document.getElementById("sqlglot-lineage"),
701        {{
702            nodes: nodes,
703            edges: new vis.DataSet({edges})
704        }},
705        {options},
706    )
707  </script>
708</div>"""
709
710    def _repr_html_(self) -> str:
711        return self.__str__()

Node to HTML generator using vis.js.

https://visjs.github.io/vis-network/docs/network/

GraphHTML( nodes: dict, edges: list, imports: bool = True, options: Mapping[str, object] | None = None)
638    def __init__(
639        self,
640        nodes: dict,
641        edges: list,
642        imports: bool = True,
643        options: Mapping[str, object] | None = None,
644    ):
645        self.imports = imports
646
647        self.options = {
648            "height": "500px",
649            "width": "100%",
650            "layout": {
651                "hierarchical": {
652                    "enabled": True,
653                    "nodeSpacing": 200,
654                    "sortMethod": "directed",
655                },
656            },
657            "interaction": {
658                "dragNodes": False,
659                "selectable": False,
660            },
661            "physics": {
662                "enabled": False,
663            },
664            "edges": {
665                "arrows": "to",
666            },
667            "nodes": {
668                "font": "20px monaco",
669                "shape": "box",
670                "widthConstraint": {
671                    "maximum": 300,
672                },
673            },
674            **(options or {}),
675        }
676
677        self.nodes = nodes
678        self.edges = edges
imports
options
nodes
edges