sqlleaf 0.0.1__py3-none-any.whl

sqlleaf/dialects/plpgsql/p_tokenizer.py

@@ -0,0 +1,41 @@
+from enum import IntEnum
+
+from sqlglot.tokens import TokenType
+from sqlglot import tokenizer_core
+
+# Monkey patch sqlglot's tokenizer with PL/pgSQL keywords
+tt = tokenizer_core.TokenType
+PLPGSQL_CUSTOM_TOKEN_NAMES = (
+    "ASSERT",
+    "BY",
+    "CLOSE",
+    "CONTINUE",
+    # Synthetic tokens (not keywords)
+    "LABEL_BEGIN",  # <<
+    "LABEL_END",    # >>
+    "DDOT",
+    "EXIT",
+    "FOREACH",
+    "IF",
+    "LOOP",
+    "MOVE",
+    "OPEN",
+    "PERFORM",
+    "RAISE",
+    "RETURN",
+    "REVERSE",
+    "WHILE",
+)
+PLPGSQL_KEYWORD_TOKEN_NAMES = (
+    *(token for token in PLPGSQL_CUSTOM_TOKEN_NAMES if token not in {"DDOT", "LABEL_BEGIN", "LABEL_END"}),
+    "DECLARE",
+    "GET",
+)
+token_dict = {m.name: m.value for m in TokenType}
+
+next_value = max(token_dict.values()) + 1
+for token in PLPGSQL_CUSTOM_TOKEN_NAMES:
+    token_dict[token] = next_value
+    next_value += 1
+
+TokenType = IntEnum('TokenType', token_dict)

sqlleaf/exception.py

@@ -0,0 +1,67 @@
+import logging
+
+ERROR_LEVEL = "stop"
+
+
+logger = logging.getLogger("sqlleaf")
+
+
+def raise_error(clazz: type, message: str):
+    """
+    Raise a custom exception based on the error level.
+    """
+    if ERROR_LEVEL == "stop":
+        raise clazz(message)
+    elif ERROR_LEVEL == "continue":
+        logger.error(message)
+
+
+def set_error_action(action: str) -> None:
+    """
+    Set the error level for exceptions.
+    """
+    if action not in ["continue", "stop"]:
+        raise ValueError("The error action must be 'continue' or 'stop'.")
+
+    global ERROR_LEVEL
+    ERROR_LEVEL = action
+
+
+class SqlGlotException(Exception):
+    def __init__(self, message, table=""):
+        super().__init__(message)
+
+        self.message = message
+
+    def __str__(self):
+        return "%s" % (self.message,)
+
+
+class SqlLeafException(Exception):
+    def __init__(self, message):
+        super().__init__(message)
+
+        self.message = message
+
+    def __str__(self):
+        return "%s" % (self.message,)
+
+
+class InvalidQueryError(SqlLeafException):
+    pass
+
+
+class UnsupportedFeatureError(SqlLeafException):
+    pass
+
+
+class MappingError(SqlLeafException):
+    pass
+
+
+class GeneratorError(SqlLeafException):
+    pass
+
+
+class GraphError(SqlLeafException):
+    pass

sqlleaf/lineage.py

@@ -0,0 +1,326 @@
+import json
+import logging
+import typing as t
+
+import networkx as nx
+
+from sqlleaf import exception, mappings, path, util
+from sqlleaf.models.hook import Hook
+from sqlleaf.models.node import EdgeAttributes, GraphAttributes, N
+from sqlleaf.models.query import (
+    CopyQuery,
+    CTASQuery,
+    InsertQuery,
+    PutQuery,
+    Q,
+    TableQuery,
+    UnloadQuery,
+    UpdateQuery,
+    ViewQuery,
+)
+from sqlleaf.path import LineagePath
+from sqlleaf.processors.collector import collector
+from sqlleaf.processors.generator import generator
+from sqlleaf.processors.transformer import transformer
+from sqlleaf.typing import SqlObjectType
+
+logging.getLogger("sqlglot").setLevel(logging.WARNING)
+logger = logging.getLogger("sqlleaf")
+
+
+class Lineage:
+    """
+    Holds the lineage as a networkx graph.
+    """
+
+    def __init__(self):
+        self.graph = new_graph()  # The graph that contains all lineage
+        self.subgraphs: t.List[nx.MultiDiGraph] = []  # The subgraphs that make up the main graph
+        self.paths: t.Dict[str, t.List[LineagePath]] = {}  # The paths throughout the graph
+        self.collected_queries: collector.CollectQueryResult | None = None
+        self.user_defined_hooks: list[Hook] = []
+        self.object_mapping: mappings.ObjectMapping | None = None
+
+    def generate(self, sql: str, dialect: str, on_error: str = "stop") -> None:
+        """
+        Generate lineage for one or more SQL statements.
+        """
+        exception.set_error_action(action=on_error)
+
+        if not self.object_mapping:
+            self.object_mapping = mappings.ObjectMapping(dialect=dialect)
+
+        self.collected_queries = collector.collect_queries(
+            text=sql, dialect=dialect, object_mapping=self.object_mapping
+        )
+
+        for parent_holder in self.collected_queries.queries:
+            # A parent holder wraps a top-level query, possibly containing downstream holders
+            graph = new_graph()
+            query_holders = parent_holder.get_all_holders()
+
+            for holder in query_holders:
+                transformer.transform_query(holder)
+
+            for holder in query_holders:
+                # Transform and produce lineage only for certain queries
+                if query_has_lineage(holder.original):
+                    generator.generate_lineage_for_query(holder, graph, self.user_defined_hooks)
+
+            # Associate the query with the graph even if it has no lineage
+            graph.graph["attrs"].add_query_to_graph(parent_holder)
+            self.merge_graph(graph)
+            self.graph.graph["attrs"].add_query_to_graph(parent_holder)
+
+    def merge_graph(self, subgraph: nx.MultiDiGraph) -> None:
+        """
+        Merge the subgraph into the main graph, and also track the individual subgraphs.
+        """
+        self.subgraphs.append(subgraph)
+
+        for n, data in subgraph.nodes(data=True):
+            if self.graph.has_node(n):
+                old_node_attrs = self.graph.nodes[n]["attrs"]
+
+                # The incoming graph's edges must have their NodeAttributes updated to
+                # match the existing graph's NodeAttributes.
+                # This is because different graphs with identical Nodes will have different NodeAttributes models.
+                for par, chi, edge_data in subgraph.edges(data=True):
+                    # Overwrite the new edge's Node to be the old Node
+                    if edge_data["attrs"].parent.full_name == n:
+                        edge_data["attrs"].parent = old_node_attrs
+                    if edge_data["attrs"].child.full_name == n:
+                        edge_data["attrs"].child = old_node_attrs
+            else:
+                self.graph.add_node(n, **data)
+
+        self.graph.add_edges_from(subgraph.edges(data=True))
+
+    def get_edges(self) -> t.List[EdgeAttributes]:
+        edges = [data["attrs"] for par, chi, data in self.graph.edges(data=True)]
+        edges = sorted(edges, key=lambda e: (e.parent.full_name, e.child.full_name))
+        for edge in edges:
+            logger.debug(f"Found edge: {edge.parent.friendly_name} -> {edge.child.friendly_name}")
+        return edges
+
+    def get_nodes(self) -> t.List[N]:
+        nodes = [data["attrs"] for n, data in self.graph.nodes(data=True)]
+        # TODO: sort on selected index?
+        nodes = sorted(
+            nodes, key=lambda e: (getattr(e, "catalog", ""), getattr(e, "schema", ""), getattr(e, "table", ""), e.name)
+        )
+        return nodes
+
+    def get_original_queries(self) -> t.List[Q]:
+        """
+        Get the original queries from the graph.
+        """
+        return [q.original for q in self.graph.graph["attrs"].queries]
+
+    def get_transformed_queries(self) -> t.List[Q]:
+        """
+        Get the transformed queries from the graph.
+        """
+        return [q.transformed for q in self.graph.graph["attrs"].queries]
+
+    def get_paths(self) -> t.Generator[LineagePath]:
+        """
+        paths: [
+            {
+                "id": "",
+                "length": 2,
+                "hops": [
+                    "edge1",
+                    "edge2
+                ]
+            }
+        ]
+        """
+        for p in path.find_all_paths(graph=self.graph):
+            yield p
+
+    def print_json(self):
+        """
+        Print the graph as JSON.
+        """
+        nodes = self.get_nodes()
+        edges = self.get_edges()
+        queries = self.get_original_queries()
+        paths = self.get_paths()
+
+        _nodes = [n.to_dict() for n in nodes]
+        _edges = [e.to_dict() for e in edges]
+        _queries = [q.to_dict() for q in queries]
+        _paths = [p.to_dict() for p in paths]
+
+        print(
+            json.dumps(
+                {
+                    "node": _nodes,
+                    "edges": _edges,
+                    "queries": _queries,
+                    "paths": _paths,
+                },
+                indent=2,
+            )
+        )
+
+    def print_tree(self, full_name=False):
+        """
+        Print from the leaves to the root (as left to right) so that the tree is displayed correctly.
+        For example:
+
+        INSERT INTO fruit.processed
+        SELECT SUBSTRING(name, 2, 4) AS name
+        FROM fruit.raw
+
+        Output:
+
+        column[fruit.processed.name]
+        └── function[SUBSTRING]
+            ├── column[fruit.raw.name]
+            ├── literal[2]
+            └── literal[4]
+        """
+        g = self.graph.reverse()  # We print from the leaves to the roots
+        root_columns = util.get_root_nodes(g)
+        seen = set()
+        symbol = "└──"
+
+        attr = "full_name" if full_name else "friendly_name"
+
+        # TODO: this may not be needed since the NodeAttributes are node
+        for i, root in enumerate(root_columns):
+            for depth, edge_attrs in util.find_edges_downward(
+                g, root
+            ):  # TODO: fetch edges in order of function argument index
+                # Swap the parent and child
+                parent_name = edge_attrs.child.full_name
+                child_name = edge_attrs.parent.full_name
+                num_descendents_of_child = len(nx.descendants(g, child_name))
+
+                parent_node = g.nodes[parent_name]["attrs"]
+                child_node = g.nodes[child_name]["attrs"]
+
+                if parent_name not in seen:
+                    symbol = "└──"
+
+                    # Print arrows
+                    if depth > 0:
+                        prefix = ((depth - 1) * 4 * " ") + f"{symbol} "
+                    else:
+                        prefix = ""
+                    print("%s%s" % (prefix, getattr(parent_node, attr)))
+                seen.add(parent_name)
+
+                # Print the child if we're at the end of the path
+                if num_descendents_of_child == 0:
+                    if depth == 0:
+                        # Direct load (source -> target)
+                        prefix = "└── "
+                    else:
+                        prefix = (depth * 4 * " ") + f"{symbol} "
+                    print("%s%s" % (prefix, getattr(child_node, attr)))
+                    seen.add(child_name)
+
+    def print_paths(self):
+        """
+        Iterate over all the paths in the graph and print each one.
+
+        Example output:
+          column[fruit.raw.apple] -> function[UPPER] - column[fruit.processed.apple]
+        """
+        for _path in self.get_paths():
+            nodes = _path.node_hops()
+            path = ""
+
+            for i, node in enumerate(nodes):
+                path += node.friendly_name
+                if i < len(nodes) - 1:
+                    path += " -> "
+            print(path)
+
+    def register_hooks(self, hooks: list[Hook]) -> None:
+        """
+        Register user-defined hooks.
+
+        A hook allows a user to provide their own logic to control how nodes are created.
+        Each hook is represented by a Hook object which specifies:
+        - name: a unique identifier for the hook
+        - kind: the Node class the hook applies to
+        - func: a callable that accepts a Node instance and returns the Node to use, or None to skip it
+
+        A hook runs after a node has been created, but before the node is added to the graph.
+        This gives the user maximum control over the node's properties, exposing attributes
+        such as the object's name and its underlying SQL expression for inspection.
+
+        When we attempt to create the path [A -> B -> C] in the graph, if we have a hook that
+        prevents the creation of B, then we set A is the previous node, and C as the next,
+        resulting in the path [A -> C].
+
+        Examples:
+            ### 1. Never create a FunctionNode.
+            hooks = [Hook(name="skip_all_functions", kind=FunctionNode, func=lambda n: None)]
+
+            ### 2. Only create a FunctionNode if its name is 'SUM'.
+            def my_func(node):
+                return node if node.name == "SUM" else None
+            hooks = [Hook(name="only_sum", kind=FunctionNode, func=my_func)]
+        """
+        self.user_defined_hooks.extend(hooks)
+
+    def deregister_hooks(self, hooks: list[Hook]) -> None:
+        """
+        The inverse of `register_hooks()`: remove a hook (by name) from the list of user-defined hooks.
+        """
+        self.user_defined_hooks = [udh for udh in self.user_defined_hooks if udh not in hooks]
+
+def new_graph() -> nx.MultiDiGraph:
+    """
+    A graph has attributes along with its node and edges.
+    """
+    return nx.MultiDiGraph(attrs=GraphAttributes())
+
+
+QUERIES_WITH_LINEAGE = (
+    CTASQuery,
+    CopyQuery,
+    InsertQuery,
+    PutQuery,
+    TableQuery,
+    UnloadQuery,
+    UpdateQuery,
+    ViewQuery,
+)
+
+
+def query_has_lineage(query: Q) -> bool:
+    """
+    Check if a query has lineage within its expressions.
+
+    We distinguish between a query's definition and a query that is called.
+    That is, in order for a `CREATE FUNCTION` or `CREATE PROCEDURE` to have
+    lineage, it must be executed by an invoking statement, e.g. `CALL()` or `SELECT UDF()`
+    Simply having statements inside its definition is not sufficient to produce lineage.
+    """
+    has_lineage = True
+    if not isinstance(query, QUERIES_WITH_LINEAGE):
+        has_lineage = False
+    elif isinstance(query, CopyQuery) and query.source_info.type == SqlObjectType.VALUES:
+        # COPY TO STDOUT VALUES (..) # TODO: this should have lineage
+        has_lineage = False
+    elif isinstance(query, CopyQuery) and not query.is_query_active():
+        has_lineage = False
+    elif isinstance(query, CTASQuery) and not query.load_data:
+        # CREATE TABLE WITH NO DATA
+        has_lineage = False
+    elif isinstance(query, CTASQuery) and query.source_info.type == SqlObjectType.PREPARED_STATEMENT:
+        # CREATE TABLE AS EXECUTE
+        has_lineage = False
+    elif isinstance(query, TableQuery) and query.property != "external":
+        # CREATE EXTERNAL TABLE
+        has_lineage = False
+
+    if not has_lineage:
+        logger.debug(f"Query type '{query.__class__.__name__}' does NOT have lineage. Skipping.")
+    return has_lineage