Edit on GitHub

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

Build the lineage graph for a column of a SQL query.

Arguments:
  • column: The column to build the lineage for.
  • 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.
  • **kwargs: Qualification optimizer kwargs.
Returns:

A lineage node.

class LineageHTML:
174class LineageHTML:
175    """Node to HTML generator using vis.js.
176
177    https://visjs.github.io/vis-network/docs/network/
178    """
179
180    def __init__(
181        self,
182        node: Node,
183        dialect: DialectType = None,
184        imports: bool = True,
185        **opts: t.Any,
186    ):
187        self.node = node
188        self.imports = imports
189
190        self.options = {
191            "height": "500px",
192            "width": "100%",
193            "layout": {
194                "hierarchical": {
195                    "enabled": True,
196                    "nodeSpacing": 200,
197                    "sortMethod": "directed",
198                },
199            },
200            "interaction": {
201                "dragNodes": False,
202                "selectable": False,
203            },
204            "physics": {
205                "enabled": False,
206            },
207            "edges": {
208                "arrows": "to",
209            },
210            "nodes": {
211                "font": "20px monaco",
212                "shape": "box",
213                "widthConstraint": {
214                    "maximum": 300,
215                },
216            },
217            **opts,
218        }
219
220        self.nodes = {}
221        self.edges = []
222
223        for node in node.walk():
224            if isinstance(node.expression, exp.Table):
225                label = f"FROM {node.expression.this}"
226                title = f"<pre>SELECT {node.name} FROM {node.expression.this}</pre>"
227                group = 1
228            else:
229                label = node.expression.sql(pretty=True, dialect=dialect)
230                source = node.source.transform(
231                    lambda n: exp.Tag(this=n, prefix="<b>", postfix="</b>")
232                    if n is node.expression
233                    else n,
234                    copy=False,
235                ).sql(pretty=True, dialect=dialect)
236                title = f"<pre>{source}</pre>"
237                group = 0
238
239            node_id = id(node)
240
241            self.nodes[node_id] = {
242                "id": node_id,
243                "label": label,
244                "title": title,
245                "group": group,
246            }
247
248            for d in node.downstream:
249                self.edges.append({"from": node_id, "to": id(d)})
250
251    def __str__(self):
252        nodes = json.dumps(list(self.nodes.values()))
253        edges = json.dumps(self.edges)
254        options = json.dumps(self.options)
255        imports = (
256            """<script type="text/javascript" src="https://unpkg.com/vis-data@latest/peer/umd/vis-data.min.js"></script>
257  <script type="text/javascript" src="https://unpkg.com/vis-network@latest/peer/umd/vis-network.min.js"></script>
258  <link rel="stylesheet" type="text/css" href="https://unpkg.com/vis-network/styles/vis-network.min.css" />"""
259            if self.imports
260            else ""
261        )
262
263        return f"""<div>
264  <div id="sqlglot-lineage"></div>
265  {imports}
266  <script type="text/javascript">
267    var nodes = new vis.DataSet({nodes})
268    nodes.forEach(row => row["title"] = new DOMParser().parseFromString(row["title"], "text/html").body.childNodes[0])
269
270    new vis.Network(
271        document.getElementById("sqlglot-lineage"),
272        {{
273            nodes: nodes,
274            edges: new vis.DataSet({edges})
275        }},
276        {options},
277    )
278  </script>
279</div>"""
280
281    def _repr_html_(self) -> str:
282        return self.__str__()

Node to HTML generator using vis.js.

https://visjs.github.io/vis-network/docs/network/

LineageHTML( node: Node, dialect: Union[str, sqlglot.dialects.dialect.Dialect, Type[sqlglot.dialects.dialect.Dialect], NoneType] = None, imports: bool = True, **opts: Any)
180    def __init__(
181        self,
182        node: Node,
183        dialect: DialectType = None,
184        imports: bool = True,
185        **opts: t.Any,
186    ):
187        self.node = node
188        self.imports = imports
189
190        self.options = {
191            "height": "500px",
192            "width": "100%",
193            "layout": {
194                "hierarchical": {
195                    "enabled": True,
196                    "nodeSpacing": 200,
197                    "sortMethod": "directed",
198                },
199            },
200            "interaction": {
201                "dragNodes": False,
202                "selectable": False,
203            },
204            "physics": {
205                "enabled": False,
206            },
207            "edges": {
208                "arrows": "to",
209            },
210            "nodes": {
211                "font": "20px monaco",
212                "shape": "box",
213                "widthConstraint": {
214                    "maximum": 300,
215                },
216            },
217            **opts,
218        }
219
220        self.nodes = {}
221        self.edges = []
222
223        for node in node.walk():
224            if isinstance(node.expression, exp.Table):
225                label = f"FROM {node.expression.this}"
226                title = f"<pre>SELECT {node.name} FROM {node.expression.this}</pre>"
227                group = 1
228            else:
229                label = node.expression.sql(pretty=True, dialect=dialect)
230                source = node.source.transform(
231                    lambda n: exp.Tag(this=n, prefix="<b>", postfix="</b>")
232                    if n is node.expression
233                    else n,
234                    copy=False,
235                ).sql(pretty=True, dialect=dialect)
236                title = f"<pre>{source}</pre>"
237                group = 0
238
239            node_id = id(node)
240
241            self.nodes[node_id] = {
242                "id": node_id,
243                "label": label,
244                "title": title,
245                "group": group,
246            }
247
248            for d in node.downstream:
249                self.edges.append({"from": node_id, "to": id(d)})
node
imports
options
nodes
edges