sqlglot.lineage

  1from __future__ import annotations
  2
  3import json
  4import typing as t
  5from dataclasses import dataclass, field
  6
  7from sqlglot import Schema, exp, maybe_parse
  8from sqlglot.errors import SqlglotError
  9from sqlglot.optimizer import Scope, build_scope, qualify
 10
 11if t.TYPE_CHECKING:
 12    from sqlglot.dialects.dialect import DialectType
 13
 14
 15@dataclass(frozen=True)
 16class Node:
 17    name: str
 18    expression: exp.Expression
 19    source: exp.Expression
 20    downstream: t.List[Node] = field(default_factory=list)
 21    alias: str = ""
 22
 23    def walk(self) -> t.Iterator[Node]:
 24        yield self
 25
 26        for d in self.downstream:
 27            if isinstance(d, Node):
 28                yield from d.walk()
 29            else:
 30                yield d
 31
 32    def to_html(self, **opts) -> LineageHTML:
 33        return LineageHTML(self, **opts)
 34
 35
 36def lineage(
 37    column: str | exp.Column,
 38    sql: str | exp.Expression,
 39    schema: t.Optional[t.Dict | Schema] = None,
 40    sources: t.Optional[t.Dict[str, str | exp.Subqueryable]] = None,
 41    dialect: DialectType = None,
 42    **kwargs,
 43) -> Node:
 44    """Build the lineage graph for a column of a SQL query.
 45
 46    Args:
 47        column: The column to build the lineage for.
 48        sql: The SQL string or expression.
 49        schema: The schema of tables.
 50        sources: A mapping of queries which will be used to continue building lineage.
 51        dialect: The dialect of input SQL.
 52        **kwargs: Qualification optimizer kwargs.
 53
 54    Returns:
 55        A lineage node.
 56    """
 57
 58    expression = maybe_parse(sql, dialect=dialect)
 59
 60    if sources:
 61        expression = exp.expand(
 62            expression,
 63            {
 64                k: t.cast(exp.Subqueryable, maybe_parse(v, dialect=dialect))
 65                for k, v in sources.items()
 66            },
 67        )
 68
 69    qualified = qualify.qualify(
 70        expression,
 71        dialect=dialect,
 72        schema=schema,
 73        **{"validate_qualify_columns": False, "identify": False, **kwargs},  # type: ignore
 74    )
 75
 76    scope = build_scope(qualified)
 77
 78    if not scope:
 79        raise SqlglotError("Cannot build lineage, sql must be SELECT")
 80
 81    def to_node(
 82        column: str | int,
 83        scope: Scope,
 84        scope_name: t.Optional[str] = None,
 85        upstream: t.Optional[Node] = None,
 86        alias: t.Optional[str] = None,
 87    ) -> Node:
 88        aliases = {
 89            dt.alias: dt.comments[0].split()[1]
 90            for dt in scope.derived_tables
 91            if dt.comments and dt.comments[0].startswith("source: ")
 92        }
 93
 94        # Find the specific select clause that is the source of the column we want.
 95        # This can either be a specific, named select or a generic `*` clause.
 96        select = (
 97            scope.expression.selects[column]
 98            if isinstance(column, int)
 99            else next(
100                (select for select in scope.expression.selects if select.alias_or_name == column),
101                exp.Star() if scope.expression.is_star else None,
102            )
103        )
104
105        if not select:
106            raise ValueError(f"Could not find {column} in {scope.expression}")
107
108        if isinstance(scope.expression, exp.Union):
109            upstream = upstream or Node(name="UNION", source=scope.expression, expression=select)
110
111            index = (
112                column
113                if isinstance(column, int)
114                else next(
115                    (
116                        i
117                        for i, select in enumerate(scope.expression.selects)
118                        if select.alias_or_name == column or select.is_star
119                    ),
120                    -1,  # mypy will not allow a None here, but a negative index should never be returned
121                )
122            )
123
124            if index == -1:
125                raise ValueError(f"Could not find {column} in {scope.expression}")
126
127            for s in scope.union_scopes:
128                to_node(index, scope=s, upstream=upstream)
129
130            return upstream
131
132        subquery = select.unalias()
133
134        if isinstance(subquery, exp.Subquery):
135            upstream = upstream or Node(name="SUBQUERY", source=scope.expression, expression=select)
136            scope = t.cast(Scope, build_scope(subquery.unnest()))
137
138            for select in subquery.named_selects:
139                to_node(select, scope=scope, upstream=upstream)
140
141            return upstream
142
143        if isinstance(scope.expression, exp.Select):
144            # For better ergonomics in our node labels, replace the full select with
145            # a version that has only the column we care about.
146            #   "x", SELECT x, y FROM foo
147            #     => "x", SELECT x FROM foo
148            source = t.cast(exp.Expression, scope.expression.select(select, append=False))
149        else:
150            source = scope.expression
151
152        # Create the node for this step in the lineage chain, and attach it to the previous one.
153        node = Node(
154            name=f"{scope_name}.{column}" if scope_name else str(column),
155            source=source,
156            expression=select,
157            alias=alias or "",
158        )
159        if upstream:
160            upstream.downstream.append(node)
161
162        # if the select is a star add all scope sources as downstreams
163        if select.is_star:
164            for source in scope.sources.values():
165                node.downstream.append(Node(name=select.sql(), source=source, expression=source))
166
167        # Find all columns that went into creating this one to list their lineage nodes.
168        source_columns = set(select.find_all(exp.Column))
169
170        # If the source is a UDTF find columns used in the UTDF to generate the table
171        if isinstance(source, exp.UDTF):
172            source_columns |= set(source.find_all(exp.Column))
173
174        for c in source_columns:
175            table = c.table
176            source = scope.sources.get(table)
177
178            if isinstance(source, Scope):
179                # The table itself came from a more specific scope. Recurse into that one using the unaliased column name.
180                to_node(
181                    c.name, scope=source, scope_name=table, upstream=node, alias=aliases.get(table)
182                )
183            else:
184                # The source is not a scope - we've reached the end of the line. At this point, if a source is not found
185                # it means this column's lineage is unknown. This can happen if the definition of a source used in a query
186                # is not passed into the `sources` map.
187                source = source or exp.Placeholder()
188                node.downstream.append(Node(name=c.sql(), source=source, expression=source))
189
190        return node
191
192    return to_node(column if isinstance(column, str) else column.name, scope)
193
194
195class LineageHTML:
196    """Node to HTML generator using vis.js.
197
198    https://visjs.github.io/vis-network/docs/network/
199    """
200
201    def __init__(
202        self,
203        node: Node,
204        dialect: DialectType = None,
205        imports: bool = True,
206        **opts: t.Any,
207    ):
208        self.node = node
209        self.imports = imports
210
211        self.options = {
212            "height": "500px",
213            "width": "100%",
214            "layout": {
215                "hierarchical": {
216                    "enabled": True,
217                    "nodeSpacing": 200,
218                    "sortMethod": "directed",
219                },
220            },
221            "interaction": {
222                "dragNodes": False,
223                "selectable": False,
224            },
225            "physics": {
226                "enabled": False,
227            },
228            "edges": {
229                "arrows": "to",
230            },
231            "nodes": {
232                "font": "20px monaco",
233                "shape": "box",
234                "widthConstraint": {
235                    "maximum": 300,
236                },
237            },
238            **opts,
239        }
240
241        self.nodes = {}
242        self.edges = []
243
244        for node in node.walk():
245            if isinstance(node.expression, exp.Table):
246                label = f"FROM {node.expression.this}"
247                title = f"<pre>SELECT {node.name} FROM {node.expression.this}</pre>"
248                group = 1
249            else:
250                label = node.expression.sql(pretty=True, dialect=dialect)
251                source = node.source.transform(
252                    lambda n: exp.Tag(this=n, prefix="<b>", postfix="</b>")
253                    if n is node.expression
254                    else n,
255                    copy=False,
256                ).sql(pretty=True, dialect=dialect)
257                title = f"<pre>{source}</pre>"
258                group = 0
259
260            node_id = id(node)
261
262            self.nodes[node_id] = {
263                "id": node_id,
264                "label": label,
265                "title": title,
266                "group": group,
267            }
268
269            for d in node.downstream:
270                self.edges.append({"from": node_id, "to": id(d)})
271
272    def __str__(self):
273        nodes = json.dumps(list(self.nodes.values()))
274        edges = json.dumps(self.edges)
275        options = json.dumps(self.options)
276        imports = (
277            """<script type="text/javascript" src="https://unpkg.com/vis-data@latest/peer/umd/vis-data.min.js"></script>
278  <script type="text/javascript" src="https://unpkg.com/vis-network@latest/peer/umd/vis-network.min.js"></script>
279  <link rel="stylesheet" type="text/css" href="https://unpkg.com/vis-network/styles/vis-network.min.css" />"""
280            if self.imports
281            else ""
282        )
283
284        return f"""<div>
285  <div id="sqlglot-lineage"></div>
286  {imports}
287  <script type="text/javascript">
288    var nodes = new vis.DataSet({nodes})
289    nodes.forEach(row => row["title"] = new DOMParser().parseFromString(row["title"], "text/html").body.childNodes[0])
290
291    new vis.Network(
292        document.getElementById("sqlglot-lineage"),
293        {{
294            nodes: nodes,
295            edges: new vis.DataSet({edges})
296        }},
297        {options},
298    )
299  </script>
300</div>"""
301
302    def _repr_html_(self) -> str:
303        return self.__str__()