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>)
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, 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
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], 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.
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