sqlleaf 0.0.1__py3-none-any.whl

sqlleaf/mappings.py

@@ -0,0 +1,376 @@
+from __future__ import annotations
+
+import typing as t
+
+from sqlglot import MappingSchema, exp
+from sqlglot.dialects.dialect import DialectType
+from sqlglot.schema import nested_set
+from sqlglot.trie import new_trie
+
+from sqlleaf import exception, util
+from sqlleaf.models.query import (
+    CTASQuery,
+    DatabaseQuery,
+    PrepareQuery,
+    ProcedureQuery,
+    Q,
+    SchemaQuery,
+    SequenceQuery,
+    StageQuery,
+    TableQuery,
+    TriggerQuery,
+    TypeQuery,
+    UserDefinedFunctionQuery,
+    ViewQuery,
+)
+
+ColumnMapping = t.Union[t.Dict, str, t.List]
+
+
+class ObjectMapping(MappingSchema):
+    """
+    Extends sqlglot.MappingSchema to provide additional functionality related to tracking exp.Table
+
+    Specifically, we need to track the exp.Table inside the exp.Create statements, as they contain more information
+    than the exp.Table that we encounter later when parsing INSERT statements.
+    """
+
+    def __init__(self, dialect: str):
+        """
+        Initialize a mapping of tables parts to exp.Table
+        """
+        super().__init__(
+            dialect=dialect, normalize=False
+        )  # Set `normalize=False` to prevent an unnecessary second parse.
+        self.kind_mapping = {}
+        self.kind_mapping_trie = {}
+        self.session_variables: dict[str, exp.Expr] = {}
+
+    def add_database_query(self, query: DatabaseQuery) -> None:
+        self._add_query(kind="database", query=query, dialect=query.dialect)
+
+    def add_prepare_query(self, query: PrepareQuery) -> None:
+        self._add_query(kind="prepare", query=query, dialect=query.dialect)
+
+    def add_procedure_query(self, query: ProcedureQuery) -> None:
+        self._add_query(kind="procedure", query=query, dialect=query.dialect)
+
+    def add_schema_query(self, query: SchemaQuery) -> None:
+        self._add_query(kind="schema", query=query, dialect=query.dialect)
+
+    def add_sequence_query(self, query: SequenceQuery) -> None:
+        self._add_query(kind="sequence", query=query, dialect=query.dialect)
+
+    def add_stage_query(self, query: StageQuery) -> None:
+        self._add_query(kind="stage", query=query, dialect=query.dialect)
+
+    def add_table_query(
+        self, query: TableQuery | ViewQuery | CTASQuery, column_mapping: t.Optional[ColumnMapping] = None
+    ) -> None:
+        self._add_query(kind="table", query=query, column_mapping=column_mapping, dialect=query.dialect)
+
+    def add_trigger_query(self, query: TriggerQuery) -> None:
+        self._add_query(kind="trigger", query=query, dialect=query.dialect)
+
+    def add_type_query(self, query: TypeQuery) -> None:
+        self._add_query(kind="type", query=query, dialect=query.dialect)
+
+    def add_udf_query(self, query: UserDefinedFunctionQuery, column_mapping: t.Optional[ColumnMapping] = None) -> None:
+        self._add_query(kind="udf", query=query, column_mapping=column_mapping, dialect=query.dialect)
+
+    def _add_query(
+        self,
+        kind: str,
+        query: Q,
+        column_mapping: t.Optional[ColumnMapping] = None,
+        dialect: DialectType = None,
+        normalize: t.Optional[bool] = None,
+        match_depth: bool = False,
+    ) -> None:
+        """
+        Register or update a table. Updates are only performed if a new column mapping is provided.
+        The added table must have the necessary number of qualifiers in its path to match the schema's nesting level.
+
+        Args:
+            kind: the expression's kind
+            query: the query to store
+            column_mapping: a column mapping that describes the structure of the table.
+            dialect: the SQL dialect that will be used to parse `table` if it's a string.
+            normalize: whether to normalize identifiers according to the dialect of interest.
+            match_depth: whether to enforce that the table must match the schema's depth or not.
+        """
+        # Initialize the structures. These are required by sqlglot.
+        if kind not in self.kind_mapping:
+            self.kind_mapping[kind] = {}
+            self.kind_mapping_trie[kind] = new_trie({})
+
+        # Normalize the table parts (catalog.db.table.column)
+        table = query.get_target_as_table()
+        normalized_table = self._normalize_table(table, dialect=dialect, normalize=normalize)
+        parts = self.table_parts(normalized_table)
+
+        if kind == "udf":
+            # Store the UDF with other UDFs of the same name
+            udfs_with_same_name = self.lookup_udf_query(table, raise_on_missing=False) or []
+            query = [query] + udfs_with_same_name
+
+        # Store the object
+        nested_set(self.kind_mapping[kind], tuple(reversed(parts)), query)
+        new_trie([parts], self.kind_mapping_trie[kind])
+
+        if kind in ["table", "udf"] and column_mapping is not None:
+            # Track the table's columns. UDFs are supported by sqlglot when qualifying columns
+            # TODO: there is a bug here where multiple UDFs with the same name that have different return columns
+            #  will have their column names overwritten (non-merged)
+            self._add_columns_for_table(
+                table=table,
+                column_mapping=column_mapping,
+                dialect=dialect,
+                normalize=normalize,
+                match_depth=match_depth,
+            )
+
+    def _add_columns_for_table(
+        self,
+        table: exp.Table,
+        column_mapping: t.Optional[ColumnMapping] = None,
+        dialect: DialectType = None,
+        normalize: t.Optional[bool] = None,
+        match_depth: bool = False,
+    ):
+        super().add_table(
+            table=table,
+            column_mapping=column_mapping,
+            dialect=dialect,
+            normalize=normalize,
+            match_depth=match_depth,
+        )
+
+    def find_columns_for_table(
+        self,
+        table: exp.Table,
+        raise_on_missing: bool = True,
+        ensure_data_types: bool = False,
+    ):
+        """
+        A nicer name for the parent's function.
+        """
+        return super().find(
+            table,
+            raise_on_missing=raise_on_missing,
+            ensure_data_types=ensure_data_types,
+        )
+
+    def lookup_prepare_query(self, table: exp.Table, raise_on_missing: bool = True) -> PrepareQuery | None:
+        return self._lookup_query(kind="prepare", object=table, raise_on_missing=raise_on_missing)
+
+    def lookup_procedure_query(self, table: exp.Table, raise_on_missing: bool = True) -> ProcedureQuery | None:
+        return self._lookup_query(kind="procedure", object=table, raise_on_missing=raise_on_missing)
+
+    def lookup_sequence_query(self, table: exp.Table, raise_on_missing: bool = True) -> SequenceQuery | None:
+        return self._lookup_query(kind="sequence", object=table, raise_on_missing=raise_on_missing)
+
+    def lookup_stage_query(self, table: exp.Table, raise_on_missing: bool = True) -> StageQuery | None:
+        return self._lookup_query(kind="stage", object=table, raise_on_missing=raise_on_missing)
+
+    def lookup_table_query(self, table: exp.Table, raise_on_missing: bool = True) -> TableQuery | None:
+        return self._lookup_query(kind="table", object=table, raise_on_missing=raise_on_missing)
+
+    def lookup_trigger_query(self, table: exp.Table, raise_on_missing: bool = True) -> TriggerQuery | None:
+        return self._lookup_query(kind="trigger", object=table, raise_on_missing=raise_on_missing)
+
+    def lookup_type_query(self, table: exp.Table, raise_on_missing: bool = True) -> TypeQuery | None:
+        return self._lookup_query(kind="type", object=table, raise_on_missing=raise_on_missing)
+
+    def lookup_udf_query(
+        self, table: exp.Table, raise_on_missing: bool = True
+    ) -> t.List[UserDefinedFunctionQuery] | None:
+        # A UDF can have multiple names but different types, thus we return a list
+        return self._lookup_query(kind="udf", object=table, raise_on_missing=raise_on_missing)
+
+    def _lookup_query(
+        self,
+        kind: str,
+        object: exp.Table,
+        raise_on_missing: bool,
+    ) -> Q | None:
+        """
+        Returns the Query for a given object kind and exp.Table.
+
+        This is different from the MappingSchema's find(), which returns the column mappings.
+        This returns an exp.Table.
+
+        Args:
+            kind: the expression's kind
+            object: the target object.
+            raise_on_missing: whether to raise in case the schema is not found.
+
+        Returns:
+            The schema of the target object.
+        """
+        if kind not in self.kind_mapping:
+            if raise_on_missing:
+                raise exception.MappingError(f"Could not find '{object.this}' of type '{kind}' in mapping.")
+            return None
+
+        parts = self.table_parts(object)[0: len(self.supported_table_args)]
+        resolved_parts = self._find_in_trie(parts, self.kind_mapping_trie[kind], raise_on_missing)
+
+        if resolved_parts is None:
+            return None
+
+        result = self.nested_get(resolved_parts, self.kind_mapping[kind], raise_on_missing=raise_on_missing)
+        if not result:
+            return None
+        elif isinstance(result, dict):
+            # The mapping object has varying depth if some objects use a catalog and others don't
+            if object.name in result:
+                return result[object.name]
+            else:
+                return None
+        else:
+            # Must be a Query
+            return result
+
+    # Override sqlglot's property. It seems to be buggy when using different dict sizes (catalog, schema, etc)
+    @property
+    def supported_table_args(self) -> t.Tuple[str, ...]:
+        return exp.TABLE_PARTS
+
+    def get_table_or_stage(
+        self, table: exp.Table | exp.Var, raise_on_missing: bool = True
+    ) -> TableQuery | StageQuery | None:
+        """
+        Get the 'CREATE' query for a table or stage.
+        """
+        table_expr = table
+        if isinstance(table_expr, exp.Var):
+            table_expr = exp.Table(this=table)
+
+        if str(table_expr).startswith("@"):
+            child_object_query = self.lookup_stage_query(table=table_expr, raise_on_missing=raise_on_missing)
+        else:
+            child_object_query = self.lookup_table_query(table=table_expr, raise_on_missing=raise_on_missing)
+
+        if not child_object_query and raise_on_missing:
+            raise exception.MappingError(message=f"Unknown table: {str(table_expr)}")
+
+        return child_object_query
+
+    def lookup_udf_call(
+        self,
+        node: exp.Anonymous,
+    ) -> t.Optional[UserDefinedFunctionQuery]:
+        """
+        Looks up the UDF definition for a single exp.Anonymous node.
+        Returns the matched UDF definition, or None if not found.
+        """
+        function_schema, function_name = util.get_udf_name(node)
+        udf_object = exp.table_(table=function_name, db=function_schema)
+        candidates = self.lookup_udf_query(table=udf_object, raise_on_missing=False)
+        if not candidates:
+            return None
+
+        return resolve_overloaded_function(node, candidates)
+
+
+def resolve_overloaded_function(
+    node: exp.Anonymous, candidates: t.List[UserDefinedFunctionQuery]
+) -> t.Optional[UserDefinedFunctionQuery]:
+    """
+    Resolves the best function candidate for an overloaded function call.
+    Applies some precedence rules (e.g., type matching, or non-variadic > variadic) but many are currently
+    excluded due to the complexity of the rules.
+    """
+    if len(candidates) == 1:
+        return candidates[0]
+
+    # Function overloading: find the best match based on arguments
+    args = node.expressions
+    matches = []
+    for candidate in candidates:
+        if match_function_arguments(args, candidate):
+            matches.append(candidate)
+
+    if not matches:
+        raise exception.MappingError(f"No matching function signatures found for args: {args}")
+
+    if len(matches) == 1:
+        return matches[0]
+
+    # Preference rule: non-variadic is preferred over variadic
+    non_variadic_matches = [m for m in matches if not any(p.is_variadic for p in m.parameters)]
+    if non_variadic_matches:
+        # Prefer more specific non-variadic matches over polymorphic ones
+        specific_matches = [
+            m
+            for m in non_variadic_matches
+            if not any(str(p.type).lower() in ("anyelement", "anyarray") for p in m.parameters)
+        ]
+        if specific_matches:
+            return specific_matches[0]
+
+        # If we only have polymorphic matches, prefer anyarray for array arguments
+        if all(arg.type and arg.type.is_type(exp.DataType.Type.ARRAY) for arg in args):
+            anyarray_matches = [
+                m for m in non_variadic_matches if any(str(p.type).lower() == "anyarray" for p in m.parameters)
+            ]
+            if anyarray_matches:
+                return anyarray_matches[0]
+
+        return non_variadic_matches[0]
+
+    return matches[0]
+
+
+def match_function_arguments(args: t.List[exp.Expr], candidate: UserDefinedFunctionQuery) -> bool:
+    """
+    Checks if the provided arguments match the function candidate's parameters.
+    Handles exact types, polymorphic types (anyelement, anyarray), and VARIADIC parameters.
+    """
+    params = candidate.parameters
+    arg_count = len(args)
+    param_count = len(params)
+
+    has_variadic = any(p.is_variadic for p in params)
+
+    if not has_variadic:
+        if arg_count != param_count:
+            return False
+
+        for i, arg in enumerate(args):
+            if not match_type(arg, params[i].type):
+                return False
+        return True
+
+def match_type(arg: exp.Expr, target_type: exp.DataType) -> bool:
+    """Matches an argument expression to a target data type."""
+    arg_type = arg.type
+    if not arg_type:
+        # If type isn't annotated, we can't be sure, but let's try to be lenient or rely on annotate_types
+        return True
+
+    if arg_type.this == exp.DataType.Type.VARCHAR:
+        arg_type = exp.DataType.build("TEXT")
+
+    target_type_name = target_type.sql().lower()
+
+    if target_type_name == "anyelement":
+        return True
+    if target_type_name == "anyarray":
+        return arg_type.is_type(exp.DataType.Type.ARRAY)
+
+    # If target is DECIMAL but arg is TEXT, they don't match
+    if target_type.is_type(exp.DataType.Type.DECIMAL) and arg_type.is_type(exp.DataType.Type.TEXT):
+        return False
+
+    # Handle numeric type matching - Postgres allows implicit cast from literal to numeric
+    if target_type.is_type(exp.DataType.Type.DECIMAL):
+        if arg_type.is_type(exp.DataType.Type.DOUBLE, exp.DataType.Type.FLOAT, exp.DataType.Type.DECIMAL):
+            return True
+        # Literals often come as DOUBLE but should match NUMERIC
+        if isinstance(arg, (exp.Literal, exp.Cast)):
+            return True
+
+    return exp.DataType.is_type(arg_type, target_type)
+

sqlleaf/models/context.py

@@ -0,0 +1,153 @@
+from __future__ import annotations
+
+import dataclasses
+import logging
+import typing as t
+from dataclasses import InitVar, dataclass
+
+import networkx as nx
+from sqlglot import exp
+from sqlglot.optimizer import Scope, traverse_scope
+
+from sqlleaf import util
+from sqlleaf.typing import TableOrScopeType
+
+logger = logging.getLogger("sqlleaf")
+
+
+class ScopePositions:
+    def __init__(self):
+        self.positions: t.Dict[int, t.Dict[int, int]] = {}
+
+    def get_position_for_expr(self, expr: exp.Expr) -> dict[int, int]:
+        return self.positions[id(expr)]
+
+    def calculate(self, scope: Scope) -> None:
+        """
+        Determine the height and width of every scope (SELECT statement) in the query's expression tree.
+        This iterates over every expression in the tree via Depth-First Search, looking for scopes.
+        """
+        root_expr = scope.expression.root()
+        scopes = {id(scope.expression): scope for scope in list(traverse_scope(root_expr))}
+
+        # For each height, map to the current width
+        heights_to_widths = {}
+        expr_ids_to_positions = {}
+        stack = [(root_expr, 1)]
+
+        while stack:
+            node, h = stack.pop()
+            node_id = id(node)
+
+            if node_id in scopes:
+                logger.debug(f"Found scope expr ({node.__class__.__name__}): {node.sql()}")
+
+                if not expr_ids_to_positions:  # Root node
+                    expr_ids_to_positions[node_id] = (0, 0)
+                    heights_to_widths[0] = 0
+                else:
+                    # Track the width across varying heights
+                    w = heights_to_widths.get(h, 0)
+                    expr_ids_to_positions[node_id] = (h, w)
+                    heights_to_widths[h] = w + 1
+                    logger.debug(f"Set height={h} width={w}")
+                    h = h + 1
+                scopes.pop(node_id)
+
+            for v in node.iter_expressions(reverse=True):
+                stack.append((v, h))
+
+        self.positions = expr_ids_to_positions
+
+
+@dataclass(frozen=True)
+class GeneratorContext[Q, N]:
+    graph: nx.MultiDiGraph
+    query: Q
+    expr: exp.Expr
+    scope: TableOrScopeType
+    child_node: N
+    scope_positions: ScopePositions = ScopePositions()
+    data_type: exp.DataType | None = None
+    # Override the data_type if needed
+    new_data_type: InitVar[exp.DataType | None] = None
+
+    def __post_init__(self, new_data_type: exp.DataType | None = None):
+        """
+        Called via replace() or if a new object is instantiated
+        """
+        expr_type = new_data_type if new_data_type else self.get_data_type_of_expr(self.expr)
+        unwrapped_expr = util.unwrap_expression(self.expr)
+
+        object.__setattr__(self, "data_type", expr_type)
+        object.__setattr__(self, "expr", unwrapped_expr)
+
+    def get_child_node(self) -> N:
+        assert self.child_node is not None
+        return self.child_node
+
+    @staticmethod
+    def get_data_type_of_expr(expr: exp.Expr) -> exp.DataType:
+        """
+        Determine the expression's data type. If it's missing, use an ancestor's data type.
+        """
+
+        def is_missing_type(x: exp.Expr) -> bool:
+            return not (x.type or x.is_type(exp.DType.UNKNOWN.into_expr()))
+
+        if isinstance(expr, exp.ColumnDef):
+            return expr.kind or exp.DType.UNKNOWN.into_expr()
+        elif is_missing_type(expr) and expr.parent:
+            # Use an ancestor's type
+            parent = expr.parent
+            while parent:
+                if not is_missing_type(parent):
+                    return t.cast(exp.DataType, parent.type)
+                parent = parent.parent
+
+            return t.cast(exp.DataType, expr.parent.type)
+        return expr.type or exp.DType.UNKNOWN.into_expr()
+
+    def replace(self, **args) -> GeneratorContext[Q, N]:
+        """
+        Convenience method to take us to this class definition.
+        """
+        return dataclasses.replace(self, **args)
+
+
+@dataclass(frozen=True)
+class PositionContext:
+    # The position of this query inside a list of queries, e.g. SELECT 'a'; SELECT 'b' -> a=0, b=1
+    statement_index: str
+
+    # The position of this column inside a set of selected columns (e.g. SELECT 'a', 'b') -> a=0, b=1
+    select_index: int = 0
+
+    # The depth of the function: e.g. SELECT UPPER(LOWER('a')) -> LOWER=0, UPPER=1
+    function_depth: int = 0
+
+    # The argument of a function: e.g. SELECT my.func('a', 'b') -> a=0, b=1
+    function_arg_index: int = 0
+
+    # The depth of a subquery, e.g. WITH cte AS (SELECT 'a') SELECT 'a' -> The first a=1, second a=0
+    query_depth: int = 0
+
+    # The width of a subquery, e.g. SELECT 4 + (SELECT 5) + (SELECT 6) -> depth(4)=0, depth(5)=1, depth(6)=1
+    query_width: int = 0
+
+    def as_str(self) -> str:
+        parts = [
+            f"query_depth={self.query_depth}",
+            f"query_width={self.query_width}",
+            f"statement={self.statement_index}",
+            f"select={self.select_index}",
+            f"func_depth={self.function_depth}",
+            f"func_arg={self.function_arg_index}",
+        ]
+        return " ".join(parts)
+
+    def replace(self, **args) -> PositionContext:
+        """
+        Convenience method to take us to this class definition.
+        """
+        return dataclasses.replace(self, **args)

sqlleaf/models/hook.py

@@ -0,0 +1,9 @@
+import typing as t
+
+from sqlleaf.models.node import N
+
+class Hook:
+    def __init__(self, name: str, kind: N, func: t.Callable[[N], N | None]):
+        self.name = name
+        self.kind = kind
+        self.func = func

sqlleaf/models/node/__init__.py

@@ -0,0 +1,27 @@
+import typing as t
+
+from sqlleaf.models.node.base import EdgeAttributes as EdgeAttributes
+from sqlleaf.models.node.base import GraphAttributes as GraphAttributes
+from sqlleaf.models.node.base import NodeAttributes as NodeAttributes
+from sqlleaf.models.node.column import ColumnNode as ColumnNode
+from sqlleaf.models.node.column import FileColumnNode as FileColumnNode
+from sqlleaf.models.node.column import StageColumnNode as StageColumnNode
+from sqlleaf.models.node.dynamodb import DynamoDbNode as DynamoDbNode
+from sqlleaf.models.node.function import FunctionNode as FunctionNode
+from sqlleaf.models.node.interval import IntervalNode as IntervalNode
+from sqlleaf.models.node.json import JsonPathNode as JsonPathNode
+from sqlleaf.models.node.literal import LiteralNode as LiteralNode
+from sqlleaf.models.node.null import NullNode as NullNode
+from sqlleaf.models.node.pivot import PivotNode as PivotNode
+from sqlleaf.models.node.pivot import UnpivotNode as UnpivotNode
+from sqlleaf.models.node.program import ProgramNode as ProgramNode
+from sqlleaf.models.node.sequence import SequenceNode as SequenceNode
+from sqlleaf.models.node.star import StarNode as StarNode
+from sqlleaf.models.node.stream import StreamNode as StreamNode
+from sqlleaf.models.node.user_defined_function import UserDefinedFunctionNode as UserDefinedFunctionNode
+from sqlleaf.models.node.var import VarNode as VarNode
+from sqlleaf.models.node.variable import VariableNode as VariableNode
+from sqlleaf.models.node.window import WindowNode as WindowNode
+
+N = t.TypeVar("N", bound=NodeAttributes)
+TargetNodeType = ColumnNode | DynamoDbNode | FileColumnNode | StageColumnNode | StreamNode | ProgramNode