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>)
expression: sqlglot.expressions.core.Expr
source: sqlglot.expressions.core.Expr
downstream: list[Node]
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
columnis provided, or a dict[str, Node] whencolumnis 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.
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