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