sqlglot.lineage API documentation

Edit on GitHub

# sqlglot.lineage

View Source

      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__()

@dataclass(frozen=True)
class Node: View Source

    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
expression: sqlglot.expressions.Expression
source: sqlglot.expressions.Expression
downstream: List[Node]
alias: str = ''
def walk(self) -> Iterator[Node]: View Source

    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: View Source

    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: View Source

     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                    (
    117                        i
    118                        for i, select in enumerate(scope.expression.selects)
    119                        if select.alias_or_name == column or select.is_star
    120                    ),
    121                    -1,  # mypy will not allow a None here, but a negative index should never be returned
    122                )
    123            )
    124
    125            if index == -1:
    126                raise ValueError(f"Could not find {column} in {scope.expression}")
    127
    128            for s in scope.union_scopes:
    129                to_node(index, scope=s, upstream=upstream)
    130
    131            return upstream
    132
    133        subquery = select.unalias()
    134
    135        if isinstance(subquery, exp.Subquery):
    136            upstream = upstream or Node(name="SUBQUERY", source=scope.expression, expression=select)
    137            scope = t.cast(Scope, build_scope(subquery.unnest()))
    138
    139            for select in subquery.named_selects:
    140                to_node(select, scope=scope, upstream=upstream)
    141
    142            return upstream
    143
    144        if isinstance(scope.expression, exp.Select):
    145            # For better ergonomics in our node labels, replace the full select with
    146            # a version that has only the column we care about.
    147            #   "x", SELECT x, y FROM foo
    148            #     => "x", SELECT x FROM foo
    149            source = t.cast(exp.Expression, scope.expression.select(select, append=False))
    150        else:
    151            source = scope.expression
    152
    153        # Create the node for this step in the lineage chain, and attach it to the previous one.
    154        node = Node(
    155            name=f"{scope_name}.{column}" if scope_name else str(column),
    156            source=source,
    157            expression=select,
    158            alias=alias or "",
    159        )
    160        if upstream:
    161            upstream.downstream.append(node)
    162
    163        # if the select is a star add all scope sources as downstreams
    164        if select.is_star:
    165            for source in scope.sources.values():
    166                node.downstream.append(Node(name=select.sql(), source=source, expression=source))
    167
    168        # Find all columns that went into creating this one to list their lineage nodes.
    169        source_columns = set(select.find_all(exp.Column))
    170
    171        # If the source is a UDTF find columns used in the UTDF to generate the table
    172        if isinstance(source, exp.UDTF):
    173            source_columns |= set(source.find_all(exp.Column))
    174
    175        for c in source_columns:
    176            table = c.table
    177            source = scope.sources.get(table)
    178
    179            if isinstance(source, Scope):
    180                # The table itself came from a more specific scope. Recurse into that one using the unaliased column name.
    181                to_node(
    182                    c.name, scope=source, scope_name=table, upstream=node, alias=aliases.get(table)
    183                )
    184            else:
    185                # The source is not a scope - we've reached the end of the line. At this point, if a source is not found
    186                # it means this column's lineage is unknown. This can happen if the definition of a source used in a query
    187                # is not passed into the `sources` map.
    188                source = source or exp.Placeholder()
    189                node.downstream.append(Node(name=c.sql(), source=source, expression=source))
    190
    191        return node
    192
    193    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: View Source

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

    202    def __init__(
    203        self,
    204        node: Node,
    205        dialect: DialectType = None,
    206        imports: bool = True,
    207        **opts: t.Any,
    208    ):
    209        self.node = node
    210        self.imports = imports
    211
    212        self.options = {
    213            "height": "500px",
    214            "width": "100%",
    215            "layout": {
    216                "hierarchical": {
    217                    "enabled": True,
    218                    "nodeSpacing": 200,
    219                    "sortMethod": "directed",
    220                },
    221            },
    222            "interaction": {
    223                "dragNodes": False,
    224                "selectable": False,
    225            },
    226            "physics": {
    227                "enabled": False,
    228            },
    229            "edges": {
    230                "arrows": "to",
    231            },
    232            "nodes": {
    233                "font": "20px monaco",
    234                "shape": "box",
    235                "widthConstraint": {
    236                    "maximum": 300,
    237                },
    238            },
    239            **opts,
    240        }
    241
    242        self.nodes = {}
    243        self.edges = []
    244
    245        for node in node.walk():
    246            if isinstance(node.expression, exp.Table):
    247                label = f"FROM {node.expression.this}"
    248                title = f"<pre>SELECT {node.name} FROM {node.expression.this}</pre>"
    249                group = 1
    250            else:
    251                label = node.expression.sql(pretty=True, dialect=dialect)
    252                source = node.source.transform(
    253                    lambda n: exp.Tag(this=n, prefix="<b>", postfix="</b>")
    254                    if n is node.expression
    255                    else n,
    256                    copy=False,
    257                ).sql(pretty=True, dialect=dialect)
    258                title = f"<pre>{source}</pre>"
    259                group = 0
    260
    261            node_id = id(node)
    262
    263            self.nodes[node_id] = {
    264                "id": node_id,
    265                "label": label,
    266                "title": title,
    267                "group": group,
    268            }
    269
    270            for d in node.downstream:
    271                self.edges.append({"from": node_id, "to": id(d)})

node
imports
options
nodes
edges