sqlleaf 0.0.1__py3-none-any.whl

sqlleaf/models/query/unload.py

@@ -0,0 +1,64 @@
+from __future__ import annotations
+
+import typing as t
+from dataclasses import dataclass
+
+import sqlglot
+from sqlglot import TokenType, exp
+
+from sqlleaf import exception, mappings, util
+from sqlleaf.models.query.base import Query
+from sqlleaf.typing import SourceInfo, TargetInfo
+
+
+@dataclass(frozen=True)
+class UnloadQueryProperties:
+    file_format: str
+
+
+class UnloadQuery(Query):
+    KIND = "unload"
+
+    def __init__(self, expr: exp.Command, dialect: str, object_mapping: mappings.ObjectMapping, statement_index: int):
+        source, target = self.get_source_and_target_expressions(expr)
+
+        source_type = self._determine_expression_type(source, dialect)
+        target_type = self._determine_expression_type(target, dialect)
+
+        super().__init__(
+            dialect=dialect,
+            statement=source,
+            statement_index=statement_index,
+            object_mapping=object_mapping,
+            source_info=SourceInfo(expression=source, type=source_type),
+            target_info=TargetInfo(expression=target, type=target_type),
+        )
+        self.parameters = UnloadQueryProperties(file_format="UNKNOWN")
+        util.qualify_and_annotate(self.source_info.expression, self.dialect, self.object_mapping)
+
+    def get_source_and_target_expressions(self, statement: exp.Command) -> t.Tuple[exp.Select, exp.Literal]:
+        """
+        Parse an UNLOAD statement for Redshift.
+        We parse this ourselves due to missing support in sqlglot.
+        """
+        # Syntax: "UNLOAD ('SELECT ...') TO ..."
+        expected_tokens = [TokenType.L_PAREN, TokenType.STRING, TokenType.R_PAREN, TokenType.VAR, TokenType.STRING]
+        actual_tokens = sqlglot.tokenize(statement.expression.name, dialect="redshift")
+
+        # Basic validation - ensure the token types match
+        for i in range(len(expected_tokens)):
+            if expected_tokens[i] != actual_tokens[i].token_type:
+                # This may be incorrect! Use the parser instead once available.
+                raise exception.InvalidQueryError(
+                    message=f"Invalid syntax for UNLOAD expression: {statement.sql(dialect='redshift')}"
+                )
+
+        select_expr = sqlglot.parse_one(actual_tokens[1].text, dialect="redshift")
+        if not isinstance(select_expr, exp.Select):
+            raise exception.InvalidQueryError(
+                message=f"Invalid expression inside UNLOAD. Expected SELECT "
+                f"but got: {select_expr.sql(dialect='redshift')}"
+            )
+
+        to_location = actual_tokens[4].text
+        return select_expr, t.cast(exp.Literal, exp.convert(to_location))

sqlleaf/models/query/update.py

@@ -0,0 +1,45 @@
+from __future__ import annotations
+
+from sqlglot import exp
+
+from sqlleaf import mappings, util
+from sqlleaf.models.query.base import Query
+from sqlleaf.typing import SourceInfo, TargetInfo
+
+
+class UpdateQuery(Query):
+    KIND = "update"
+
+    def __init__(
+        self,
+        expr: exp.Update,
+        dialect: str,
+        object_mapping: mappings.ObjectMapping,
+        statement_index: int,
+        table: exp.Table | None = None,
+    ):
+        if not table:
+            table = util.get_table(expr)
+
+        if isinstance(expr, exp.OnConflict):
+            source = expr
+        else:
+            source = expr.args["expressions"]
+        target = table
+
+        source_type = self._determine_expression_type(source, dialect)
+        target_type = self._determine_expression_type(target, dialect)
+
+        super().__init__(
+            dialect=dialect,
+            statement=expr,
+            statement_index=statement_index,
+            object_mapping=object_mapping,
+            source_info=SourceInfo(expression=source, type=source_type),
+            target_info=TargetInfo(expression=target, type=target_type),
+        )
+        self.only = table.args.get("only", False) if table else False  # Not available inside a MERGE
+
+    def get_ctes(self):
+        with_ = self.statement.args.get("with_", None)
+        return with_.expressions if with_ else []

sqlleaf/models/query/user_defined_function.py

@@ -0,0 +1,297 @@
+from __future__ import annotations
+
+import typing as t
+from dataclasses import dataclass
+
+from sqlglot import exp
+
+from sqlleaf import mappings, util
+from sqlleaf.models.query.base import Query
+from sqlleaf.typing import TargetInfo
+
+
+@dataclass(frozen=True)
+class UserDefinedFunctionParameters:
+    schema_name: t.Optional[str]
+    function_name: str
+    return_type: t.Optional[exp.DataType]
+    language: t.Optional[str]
+    parameters: t.List[FunctionParam]
+    return_columns: t.List[exp.ColumnDef]
+    inner_statements: t.List[exp.Expr]
+
+    @classmethod
+    def from_expression(cls, expr: exp.Create, object_mapping: mappings.ObjectMapping) -> UserDefinedFunctionParameters:
+        """
+        Collect the required information from the function DDL.
+        """
+        function_name, schema_name, parameters = _extract_function_info(expr)
+        return_type, return_columns = _extract_return_info(expr, parameters, object_mapping)
+
+        # Filter parameters to only include those that can be passed as input
+        input_parameters = [p for p in parameters if p.is_input or p.is_variadic]
+
+        language = _extract_language(expr)
+
+        body_expr = expr.args.get("expression")
+        inner_statements = []
+
+        if body_expr:
+            inner_statements = util.iter_inner_statements(body_expr, object_mapping.dialect, wrap=True)
+
+        return cls(
+            schema_name=schema_name,
+            function_name=function_name,
+            return_type=return_type,
+            language=language,
+            parameters=input_parameters,
+            return_columns=return_columns,
+            inner_statements=inner_statements,
+        )
+
+
+class UserDefinedFunctionQuery(Query):
+    KIND = "udf"
+
+    def __init__(
+        self,
+        expr: exp.Create,
+        dialect,
+        object_mapping: mappings.ObjectMapping,
+        statement_index: int,
+    ):
+        source = None
+        target = expr.this.this
+
+        target_type = self._determine_expression_type(target, dialect)
+
+        super().__init__(
+            dialect=dialect,
+            statement=expr,
+            statement_index=statement_index,
+            object_mapping=object_mapping,
+            source_info=source,
+            target_info=TargetInfo(expression=target, type=target_type),
+        )
+
+        self.properties = self.get_params()
+        self.column_defs = self.properties.return_columns
+
+    @property
+    def parameters(self) -> t.List[FunctionParam]:
+        return self.properties.parameters
+
+    @property
+    def schema_name(self) -> t.Optional[str]:
+        return self.properties.schema_name
+
+    @property
+    def function_name(self) -> str:
+        return self.properties.function_name
+
+    @property
+    def return_type(self) -> t.Optional[exp.DataType]:
+        return self.properties.return_type
+
+    @property
+    def language(self) -> t.Optional[str]:
+        return self.properties.language
+
+    @property
+    def return_columns(self) -> t.List[exp.ColumnDef]:
+        return self.properties.return_columns
+
+    @property
+    def inner_statements(self) -> t.List[exp.Expr]:
+        return self.properties.inner_statements
+
+    @property
+    def name(self):
+        return ".".join([var for var in [self.schema_name, self.function_name] if var])
+
+    def get_params(self) -> UserDefinedFunctionParameters:
+        return UserDefinedFunctionParameters.from_expression(self.statement, self.object_mapping)
+
+
+def get_user_defined_data_type(kind: t.Optional[str | exp.Identifier | exp.Dot] = None) -> exp.DataType:
+    if kind:
+        if isinstance(kind, str):
+            if "." in kind:
+                parts = kind.split(".")
+                kind = exp.Dot.build([exp.Identifier(this=p, quoted=False) for p in parts])
+            else:
+                kind = exp.Identifier(this=kind, quoted=False)
+        return exp.DataType.build(kind, udt=True, dialect="postgres")
+    return exp.DataType.build("USER-DEFINED", dialect="postgres")
+
+
+@dataclass
+class FunctionParam:
+    """A parameter of a user-defined function."""
+
+    name: str
+    type: exp.DataType
+    default: t.Optional[exp.Expression] = None
+    is_input: bool = True
+    is_output: bool = False
+    is_variadic: bool = False
+
+
+def _extract_function_info(expression: exp.Create) -> t.Tuple[str, t.Optional[str], t.List[FunctionParam]]:
+    """Extracts the name, schema, and parameters from a CREATE FUNCTION expression."""
+    target = expression.this
+    parameters = []
+
+    # Process the function definition which includes name and parameters
+    if isinstance(target, exp.UserDefinedFunction):
+        function_expr = target.this
+        # Iterate through column definitions to find parameters
+        for i, col_def in enumerate(target.expressions):
+            if isinstance(col_def, exp.ColumnDef):
+                param_name = col_def.this.name
+                param_type = col_def.kind
+                param_default = None
+                is_input = True
+                is_output = False
+                is_variadic = False
+
+                # Determine the parameter types (IN, OUT, etc)
+                constraints = col_def.args.get("constraints")
+                if constraints:
+                    for constraint in constraints:
+                        kind = constraint.args.get("kind")
+                        if isinstance(kind, exp.DefaultColumnConstraint):
+                            param_default = kind.this
+                        elif isinstance(constraint, exp.InOutColumnConstraint):
+                            is_input = constraint.args.get("input_")
+                            is_output = constraint.args.get("output")
+                            is_variadic = constraint.args.get("variadic")
+
+                # Unnamed parameter logic inside ColumnDef (e.g., VARIADIC numeric[])
+                if param_name.upper() == "VARIADIC":
+                    is_variadic = True
+                    param_name = f"${i + 1}"
+                    # The type might be stored in a way that includes the rest of the definition
+                    # In Postgres unnamed VARIADIC, sqlglot might parse it such that 'variadic' is the name
+                    # and the type is numeric[]
+
+                # Handle special keywords used as parameter names by assigning positional placeholders
+                if param_name.upper() in ("IN", "OUT", "INOUT"):
+                    if param_name.upper() == "OUT":
+                        is_input = False
+                        is_output = True
+                    elif param_name.upper() == "INOUT":
+                        is_input = True
+                        is_output = True
+                    param_name = f"${i + 1}"
+                parameters.append(
+                    FunctionParam(
+                        name=param_name,
+                        type=param_type,
+                        default=param_default,
+                        is_input=is_input,
+                        is_output=is_output,
+                        is_variadic=is_variadic,
+                    )
+                )
+            elif _is_unnamed_parameter(col_def):
+                # Handle unnamed parameters like CREATE FUNCTION hello(TEXT) or (VARIADIC numeric[])
+                param_name = f"${i + 1}"
+                is_variadic = False
+
+                type_node = col_def
+                if isinstance(col_def, exp.UserDefinedFunction):
+                    is_variadic = col_def.sql().upper().startswith("variadic")
+                    type_node = col_def.this
+
+                type_sql = type_node.this.lower()
+
+                if type_sql in ("anyelement", "anyarray"):
+                    param_type = get_user_defined_data_type(kind=type_sql)
+                else:
+                    try:
+                        param_type = exp.DataType.build(type_node.this, dialect="postgres")
+                    except Exception:
+                        # e.g. for tables as parameters - CREATE FUNCTION hello(people), where 'people' is a table
+                        param_type = get_user_defined_data_type(kind=type_node.this)
+
+                parameters.append(FunctionParam(name=param_name, type=param_type, is_variadic=is_variadic))
+    else:
+        function_expr = target.this
+
+    # Determine function name and schema from the table-like expression structure
+    if isinstance(function_expr, exp.Table):
+        name = function_expr.name
+        schema = function_expr.args.get("db").this if function_expr.args.get("db") else None
+    else:
+        name = function_expr.sql()
+        schema = None
+
+    return name, schema, parameters
+
+
+def _is_unnamed_parameter(col_def: exp.ColumnDef) -> bool:
+    return isinstance(col_def, exp.Identifier) or (
+        isinstance(col_def, exp.UserDefinedFunction) and isinstance(col_def.this, exp.Identifier)
+    )
+
+
+def _extract_return_info(
+    expression: exp.Create, parameters: t.List[FunctionParam], object_mapping: mappings.ObjectMapping
+) -> t.Tuple[t.Optional[exp.DataType], t.List[exp.ColumnDef]]:
+    """Extracts return type and return columns from a CREATE FUNCTION expression."""
+    return_type = None
+    return_columns = []
+
+    # Check for OUT/INOUT parameters that determine the return structure
+    out_params: t.List[FunctionParam] = [p for p in parameters if p.is_output]
+
+    if out_params:
+        if len(out_params) == 1:
+            return_type = out_params[0].type
+        else:
+            return_type = get_user_defined_data_type()
+            for param in out_params:
+                col_name = param.name
+                if col_name.startswith("$"):
+                    col_name = f"column{col_name[1:]}"
+                return_columns.append(exp.ColumnDef(this=exp.to_identifier(col_name), kind=param.type))
+
+    # Loop through properties to find the RETURNS clause
+    for prop in expression.args.get("properties", {}).expressions:
+        if isinstance(prop, exp.ReturnsProperty):
+            if prop.this is None:
+                continue
+
+            # Check if it's a RETURNS TABLE(...) with multiple columns
+            if isinstance(prop.this, exp.Schema):
+                return_type = get_user_defined_data_type()
+                # Extract individual column definitions for table-returning functions
+                for col_def in prop.this.expressions:
+                    if isinstance(col_def, exp.ColumnDef):
+                        return_columns.append(col_def)
+            else:
+                # Standard return type (e.g., RETURNS TEXT, RETURNS person, RETURNS TABLE_NAME)
+                return_type_sql = prop.this.sql().lower()
+                return_type = prop.this
+
+                if return_type_sql == "void":
+                    return_type = exp.DataType.build("NULL")
+                else:
+                    target_table = exp.to_table(return_type_sql)
+                    found_query = object_mapping.lookup_type_query(target_table, raise_on_missing=False)
+                    if not found_query:
+                        found_query = object_mapping.lookup_table_query(target_table, raise_on_missing=False)
+
+                    if found_query:
+                        return_columns = found_query.get_column_defs()
+
+    return return_type, return_columns
+
+
+def _extract_language(expression: exp.Create) -> t.Optional[str]:
+    """Extracts the language from a CREATE FUNCTION expression."""
+    for prop in expression.args.get("properties", {}).expressions:
+        if isinstance(prop, exp.LanguageProperty):
+            return prop.this.sql().lower()
+    return None

sqlleaf/models/query/values.py

@@ -0,0 +1,20 @@
+from sqlglot import exp
+
+from sqlleaf.models.query.base import Query
+from sqlleaf.typing import SourceInfo, SqlObjectType, TargetInfo
+
+
+class ValuesQuery(Query):
+    KIND = "values"
+
+    def __init__(self, expr: exp.Values, dialect: str, object_mapping, statement_index: int):
+        source_info = SourceInfo(expression=expr, type=SqlObjectType.VALUES)
+        target_info = TargetInfo(expression=None, type=SqlObjectType.NONE)
+        super().__init__(
+            dialect=dialect,
+            statement=expr,
+            statement_index=statement_index,
+            object_mapping=object_mapping,
+            source_info=source_info,
+            target_info=target_info,
+        )

sqlleaf/models/query/view.py

@@ -0,0 +1,41 @@
+from __future__ import annotations
+
+import typing as t
+
+from sqlglot import exp
+
+from sqlleaf import mappings, util
+from sqlleaf.models.query.base import Query
+from sqlleaf.models.query.table import TableQuery
+from sqlleaf.typing import SourceInfo, TargetInfo
+
+
+class ViewQuery(Query):
+    KIND = "view"
+
+    def __init__(
+        self,
+        expr: exp.Create,
+        dialect: str,
+        object_mapping: mappings.ObjectMapping,
+        columns: t.List[exp.ColumnDef],
+        statement_index: int,
+    ):
+        source = expr.expression
+        target = util.get_table(expr)
+
+        source_type = self._determine_expression_type(source, dialect)
+        target_type = self._determine_expression_type(target, dialect)
+
+        super().__init__(
+            dialect=dialect,
+            statement=expr,
+            statement_index=statement_index,
+            object_mapping=object_mapping,
+            source_info=SourceInfo(expression=source, type=source_type),
+            target_info=TargetInfo(expression=target, type=target_type),
+        )
+        self.column_defs: t.List[exp.ColumnDef] = columns
+        self.inherited_by: t.List[TableQuery] = []
+
+        self.property: str = util.find_property(expr, self.target_info.expression, dialect)

sqlleaf/path.py

@@ -0,0 +1,194 @@
+from __future__ import annotations
+
+import logging
+import typing as t
+
+import networkx as nx
+
+from sqlleaf import exception, util
+from sqlleaf.models.node import EdgeAttributes, N
+
+logger = logging.getLogger("sqlleaf")
+
+
+class LineagePath:
+    def __init__(self, hops: t.List[EdgeAttributes]):
+        self.hops = hops
+        self.length = len(hops)
+
+    @property
+    def id(self) -> str:
+        return "path:" + util.short_sha256_hash(":".join(self.get_edge_ids()))
+
+    def node_hops(self) -> t.List[N]:
+        """
+        Return the list of nodes on this path.
+        """
+        hops = [self.hops[0].parent]
+        for hop in self.hops:
+            hops.append(hop.child)
+        return hops
+
+    def get_edge_ids(self) -> list[str]:
+        """
+        In order to distinguish between multiple edges that are part of the same path,
+        we need to create a unique id based off data that differentiates them.
+        This is done using the edges' "id" attribute.
+        """
+        return [edge.id for edge in self.hops]
+
+    def to_dict(self) -> dict:
+        return {
+            "id": self.id,
+            "length": len(self.hops),
+            "hops": [edge.id for edge in self.hops],
+        }
+
+
+def find_all_paths(graph: nx.MultiDiGraph) -> t.Generator[LineagePath]:
+    """
+    Find all the unique paths in the graph.
+
+    There are two algorithms depending on the graph's structure:
+    - one for those containing root node (a Rooted graph)
+    - one for those containing only cycles (a Circuit graph)
+    """
+
+    # Cycle handling logic.
+    #  [x] Stage 1. Error: Throw an error/warning on cycles that include multiple columns.
+    #  [ ] Stage 2. Recovery: For every cycle, print it to the console, and then
+    #  remove all their edges; then proceed with printing the remainder.
+    cycles = []
+    cycle_errors = 0
+
+    for cycle, errors in util.get_cycles(graph):
+        if errors > 0:
+            cycle_errors += errors
+        else:
+            cycles.append(cycle)
+
+    if cycle_errors:
+        raise exception.GraphError(message=f"Found {cycle_errors} errors with cycles in graph. Remove these.")
+
+    cycles = sorted(cycles)  # simple_cycles() underneath is non-deterministic
+    root_columns = util.get_root_nodes(graph)
+
+    if root_columns:
+        logger.debug(f"Found root columns in graph: {root_columns}")
+        for i, root in enumerate(root_columns):
+            for path in find_edges_from_root(graph, root):
+                if not path:
+                    continue
+
+                if path[0].parent.id == path[0].child.id:
+                    # Skip selfloops; these are covered in the dedicated loop below
+                    continue
+
+                logger.debug(
+                    "Found edge path using root: %s --- %s",
+                    [e.id for e in path],
+                    [(e.parent.friendly_name, e.child.friendly_name) for e in path],
+                )
+                lineage_path = LineagePath(hops=path)
+                yield lineage_path
+
+    # Check for selfloops (e.g. a column that goes through another node, such as a function)
+    for cycle in cycles:
+        for path in find_edges_along_cycle_path(graph, cycle):
+            logger.debug(
+                "Found edge path from cycle: %s --- %s",
+                [e.id for e in path],
+                [(e.parent.friendly_name, e.child.friendly_name) for e in path],
+            )
+            lineage_path = LineagePath(hops=path)
+            yield lineage_path
+
+
+def find_edges_along_cycle_path(
+    g: nx.MultiDiGraph, cycle: t.List[str], path: t.Optional[t.List[EdgeAttributes]] = None
+) -> t.Generator[t.List[EdgeAttributes]]:
+    """
+    Given a cycle, find all the edges along it and return these as the new path.
+    A path must be provided so that iteration doesn't deviate into non-cycle paths.
+    """
+    if path is None:
+        path = []
+
+    depth = len(path)
+    if depth == len(cycle) - 1:
+        yield path
+    else:
+        node_src = cycle[depth]
+        node_dst = cycle[depth + 1]
+        edges = g.get_edge_data(node_src, node_dst)
+        for idx, data in edges.items():
+            hop = data["attrs"]
+            yield from find_edges_along_cycle_path(g, cycle, path + [hop])
+
+
+def find_edges_from_root(
+    g: nx.MultiDiGraph, node: str, path: t.Optional[t.List[EdgeAttributes]] = None, seen: t.Optional[t.Set[str]] = None
+) -> t.Generator[t.List[EdgeAttributes]]:
+    """
+    Find all the complete paths in a graph by traversing the descendants of a node until we find
+    a node without any descendants.
+
+    This is the same as the regular find_paths(), but we iterate over all the edges between any two node
+    so that we include all paths.
+
+    For example, given the graph:
+        A -> B -> C -> D
+
+    If there are two edges between A->B and two edges between C->D:
+           __         __
+          /  \\       /  \
+     --> A    B --> C    D -->
+          \\__/       \\__/
+
+    then traversing all the edges gives:
+        A -> B -> C -> D
+        A -> B -> C -> D
+        A -> B -> C -> D
+        A -> B -> C -> D
+
+    Thus we consider each path during traversal, as each likely has slightly different attributes.
+    """
+    if path is None:
+        path = []
+    if seen is None:
+        seen = set()
+
+    if node in seen:
+        yield path
+    else:
+        # If there's a loop [.., X, .., X], discard everything outside of the Xs as the path.
+        if selfloop_edges := g.get_edge_data(node, node):
+            for idx, data in selfloop_edges.items():
+                hop = data["attrs"]
+                yield [hop]
+
+        # Get direct descendants
+        desc = nx.descendants_at_distance(g, node, 1)
+        if not desc:
+            yield path
+        else:
+            desc = sorted(desc)  # nx.desc() above is non-deterministic
+            for n in desc:
+                yield from _traverse_path_along_edges(g, node, n, path, seen)
+
+
+def _traverse_path_along_edges(
+    g: nx.MultiDiGraph, node_src: str, node_dst: str, path: t.Optional[t.List] = None, seen: t.Optional[t.Set] = None
+) -> t.Generator[t.List[EdgeAttributes]]:
+    """
+    Get the list of edges between two node, and find the paths for each of them.
+    """
+    if path is None:
+        path = []
+    if seen is None:
+        seen = set()
+
+    edges = g.get_edge_data(node_src, node_dst)
+    for idx, data in edges.items():
+        hop = data["attrs"]
+        yield from find_edges_from_root(g, node_dst, path + [hop], seen.union([node_src]))