sqlglot.lineage

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