sqlleaf 0.0.1__py3-none-any.whl

sqlleaf/processors/transformer/expressions/values.py

@@ -0,0 +1,318 @@
+from __future__ import annotations
+
+import logging
+import typing as t
+
+from sqlglot import exp
+
+from sqlleaf import exception, util
+from sqlleaf.models.query import Q
+from sqlleaf.typing import E
+from sqlleaf.util import default_column_index_iterator
+
+logger = logging.getLogger("sqlleaf")
+
+
+def normalize_values(query: Q, expr: exp.Expr) -> exp.Expr:
+    """
+    Convert a VALUES() statement into a SELECT .. UNION statement.
+    This handles two cases:
+    - When it is a top-level query (e.g. INSERT .. VALUES, CTAS .. VALUES)
+    - When it is inside another query, e.g. nested in an expression
+    """
+    if isinstance(expr, exp.Values):
+        return _rewrite_values_statement(query, expression=expr, statement=expr)
+
+    unresolved_ids: t.Set[int] = set()
+
+    # Walk the subtree and rewrite all Values occurrences
+    prev_value = None
+    while True:
+        values = _pick_next_values_node(expr, unresolved_ids)
+        if values and prev_value == values:
+            raise exception.InvalidQueryError("Infinite loop detected while searching for the next VALUES() expression.")
+
+        if values is None:
+            break
+
+        # Determine if there is a CTE ancestor within the current statement scope
+        cte = _cte_ancestor_in_scope(expr, values)
+        if cte is not None:
+            if cte.this is values:
+                _rewrite_values_statement(query, values, cte)
+            else:
+                unresolved_ids.add(id(values))
+        else:
+            # VALUES is directly the expression of an INSERT
+            parent = values.parent
+            if isinstance(parent, exp.Insert) and parent.expression is values:
+                # Inline of _handle_values_in_insert
+                converted = _rewrite_values_statement(query, values, parent)
+                if isinstance(converted, exp.Insert) and (parent is expr):
+                    expr = converted
+
+            # VALUES is directly the expression of a CREATE ... AS
+            elif isinstance(parent, exp.Create) and parent.expression is values:
+                _rewrite_values_statement(query, values, parent)
+
+            # VALUES is one of the sides of a UNION
+            elif isinstance(parent, exp.SetOperation):
+                _rewrite_values_statement(query, values, parent)
+
+            else:
+                # VALUES in a table position (wrapped by Subquery or direct FROM on SELECT/UPDATE)
+                from_ancestor = values.find_ancestor(exp.From)
+                if (
+                    values.find_ancestor(exp.Subquery) is not None
+                    or values.parent_select
+                    or (from_ancestor is not None and from_ancestor.this is values)
+                ):
+                    _rewrite_values_in_table_position(values, query.dialect)
+
+        prev_value = values
+    return expr
+
+
+def _pick_next_values_node(statement: exp.Expr, unresolved: set[int]) -> exp.Values | None:
+    """
+    Return the next exp.Values node in the statement not present in `unresolved`.
+    """
+    for candidate in statement.find_all(exp.Values):
+        if id(candidate) not in unresolved:
+            return candidate
+    return None
+
+
+def _cte_ancestor_in_scope(statement: exp.Expr, values: exp.Values) -> exp.CTE | None:
+    """
+    Return the enclosing CTE only if it belongs to the current statement's subtree.
+    This is to prevent us from using a CTE that does not enclose us.
+    """
+    cte = values.find_ancestor(exp.CTE)
+    if cte is None:
+        return None
+
+    ctes_in_scope = {id(c) for c in statement.find_all(exp.CTE)}
+    if id(cte) not in ctes_in_scope:
+        return None
+
+    return cte
+
+
+def _rewrite_values_in_table_position(values: exp.Values, dialect: str) -> None:
+    """
+    Rewrite VALUES used in FROM/JOIN/LATERAL positions.
+
+    Supports both:
+    - Subquery(Values), which updates the inner expression and preserves the outer alias/columns.
+    - Direct 'From(Values)', which converts and transfers aliases.
+    """
+    converted = _values_to_select_expr(
+        values=values,
+        dialect=dialect,
+    )
+
+    # SELECT (VALUES ...)
+    outer_subquery = values.find_ancestor(exp.Subquery)
+    if outer_subquery is not None and outer_subquery.this is values:
+        if isinstance(converted, exp.Subquery):
+            converted = converted.this  # unwrap nested Subquery
+        outer_subquery.set("this", converted)
+        return
+
+    # SELECT/UPDATE/DELETE FROM (VALUES ...)
+    # Case A: VALUES appears in a SELECT's FROM
+    parent_select = values.parent_select
+    from_ = parent_select and parent_select.args.get("from_")
+    if from_ and from_.this is values:
+        # Ensure we have a Subquery in FROM and preserve the alias/column names
+        original_alias = values.args.get("alias")
+        if original_alias and not converted.args.get("alias"):
+            converted.set("alias", original_alias)
+        from_.set("this", converted)
+        return
+
+    # Case B: VALUES appears directly under a FROM of non-SELECT (e.g., UPDATE ... FROM (VALUES ...))
+    from_ancestor = values.find_ancestor(exp.From)
+    if from_ancestor is not None and from_ancestor.this is values:
+        original_alias = values.args.get("alias")
+        if not isinstance(converted, exp.Subquery):
+            converted = converted.subquery()
+        if original_alias and not converted.args.get("alias"):
+            converted.set("alias", original_alias)
+        from_ancestor.set("this", converted)
+
+
+def _resolve_values_column_names(values: exp.Values, parent: exp.Expr, query: Q) -> list[str]:
+    """
+    Resolve column names for the expressions inside VALUES() based on its surrounding expressions.
+    """
+    columns = []
+    if isinstance(parent, exp.CTE):
+        columns = parent.alias_column_names
+    elif not isinstance(parent, exp.Values):
+        columns = [e.name for e in parent.this.expressions]
+
+    if not columns:
+        # Fall back to the mapping
+        # TODO: remove after downstream bug found
+        child_table = query.target_info.expression
+        if isinstance(child_table, exp.Table):
+            cols = query.object_mapping.find_columns_for_table(child_table)
+            values_lists: t.List[exp.Tuple] = values.expressions
+            columns = list(cols)[: len(values_lists[0].expressions)]
+
+    # Returning empty will populate dialect defaults
+    return list(columns)
+
+
+def _substitute_default_column_references(query: Q, values: exp.Values, target_table: exp.Table | None) -> None:
+    """
+    Replace the column names used inside `INSERT ... VALUES (...)` statements
+    with the referenced column's DEFAULT expression (or NULL), per MySQL's semantics for
+    VALUES(column_name).
+
+    This also applies when the column reference is nested inside another expression,
+    e.g. VALUES(UPPER(name), age * 2).
+
+    Example [given name's default is 5]:
+        INSERT INTO (name) VALUES (name)
+    ->
+        INSERT INTO (name) SELECT 5 AS name
+    """
+    if target_table is None:
+        return
+
+    table_query = query.object_mapping.lookup_table_query(table=target_table, raise_on_missing=False)
+    if not table_query:
+        return
+
+    table_columns = table_query.get_column_defs()
+    table_column_names = {col.name for col in table_columns}
+
+    for value_expr in values.expressions:
+        if not isinstance(value_expr, exp.Tuple):
+            continue
+
+        # Replace each column with its default/generated expression
+        for tuple_expr in value_expr.expressions:
+            for column_ref in list(tuple_expr.find_all(exp.Column)):
+                if column_ref.name in table_column_names:
+                    col_def = next(c for c in table_columns if c.name == column_ref.name)
+                    if default_constraint := util.get_column_constraint_expression(col_def):
+                        column_ref.replace(default_constraint.this.copy())
+                    else:
+                        column_ref.replace(exp.Null())
+
+
+def _rewrite_empty_values_or_values_with_column_names(
+    query: Q, expression: exp.Values, statement: E, columns: t.List[str]
+) -> t.List[str]:
+    """
+    MySQL-only handling for INSERT ... VALUES lists:
+    - Expand empty VALUES() to refer to all target columns so DEFAULTs can be applied.
+    - Replace target column-name references in VALUES(...) with that column's DEFAULT (or NULL).
+
+    Example:
+      CREATE TABLE t(a INT DEFAULT 5, b INT);
+      INSERT INTO t () VALUES();            -> SELECT 5 AS a, NULL AS b
+      INSERT INTO t (a,b) VALUES(a, UPPER(b)) -> SELECT 5 AS a, UPPER(NULL) AS b
+    """
+    if isinstance(statement, exp.Insert) and query.dialect == "mysql":
+        into_table = statement.find(exp.Table)
+
+        # Expand an empty tuple to reference all target columns so substitution can occur
+        if (
+            expression.expressions
+            and isinstance(expression.expressions[0], exp.Tuple)
+            and len(expression.expressions[0].expressions) == 0
+        ):
+            table_query = (
+                query.object_mapping.lookup_table_query(table=into_table, raise_on_missing=False)
+                if into_table is not None
+                else None
+            )
+            if table_query:
+                # Replace the empty tuple values with references to each column name
+                columns = [c.name for c in table_query.get_column_defs()]
+                expression.expressions[0].set("expressions", [exp.column(c) for c in columns])
+
+        _substitute_default_column_references(query, expression, into_table)
+
+    return columns
+
+
+def _rewrite_values_statement(query: Q, expression: exp.Values, statement: E) -> E:
+    """
+    Convert a `VALUES(...)` statement into a `SELECT ... UNION ALL SELECT ...` statement
+    and rewrite the parent statement in-place.
+    """
+    if not isinstance(expression, exp.Values):
+        return statement
+
+    columns = _resolve_values_column_names(expression, statement, query)
+    columns = _rewrite_empty_values_or_values_with_column_names(query, expression, statement, columns)
+
+    # Build the 'SELECT ... UNION ALL SELECT ...'
+    new_statement = _values_to_select_expr(
+        values=expression,
+        dialect=query.dialect,
+        column_names=columns,
+    )
+
+    # Rewrite the parent statement with the existing column list
+    if isinstance(statement, exp.Insert):
+        into_table = statement.find(exp.Table)
+        insert_columns = statement.this.expressions
+
+        insert_expr = exp.insert(
+            expression=new_statement,
+            columns=insert_columns,
+            into=into_table or statement.this,
+            returning=statement.args.get("returning"),
+        )
+        insert_expr.set("conflict", statement.args.get("conflict"))
+        statement.replace(insert_expr)
+        statement = insert_expr
+    elif isinstance(statement, exp.Create):
+        expression.replace(new_statement)
+    elif isinstance(statement, exp.CTE):
+        if isinstance(new_statement, exp.Subquery):
+            new_statement = new_statement.this
+        expression.replace(new_statement)
+    elif isinstance(statement, exp.Values):
+        statement = new_statement
+    elif isinstance(statement, exp.SetOperation):
+        expression.replace(new_statement)
+    else:
+        raise exception.InvalidQueryError(message=f"Unknown statement type: {statement.__class__}")
+
+    return statement
+
+
+def _values_to_select_expr(
+    values: exp.Values,
+    dialect: str,
+    column_names: t.Optional[t.List[str]] = None,
+) -> exp.Expr:
+    """
+    Convert an exp.Values into an exp.Select or exp.Union.
+    """
+    values_lists: t.List[exp.Tuple] = values.expressions
+    if not column_names:
+        column_names = list(default_column_index_iterator(dialect, values_lists[0].expressions))
+
+    selects = []
+    for val_list in values_lists:
+        row_vals = val_list.expressions
+        cols = [exp.alias_(val, str(col)) for col, val in zip(column_names, row_vals)]
+        selects.append(cols)
+
+    if len(selects) > 1:
+        new_selects = [exp.select(*select) for select in selects]
+        return_expr = exp.union(*new_selects, distinct=False)
+    else:
+        return_expr = exp.select(*selects[0])
+
+    # Wrap the query in a subquery if it's not a top-level statement
+    return return_expr.subquery() if values.parent_select else return_expr

sqlleaf/processors/transformer/insert.py

@@ -0,0 +1,125 @@
+import typing as t
+
+from sqlglot import exp
+from sqlglot.optimizer.annotate_types import annotate_types
+
+from sqlleaf.processors.transformer.base import BaseQueryTransformer
+from sqlleaf.typing import E
+
+
+class InsertTransformer(BaseQueryTransformer):
+    """Transformer for INSERT statements."""
+
+    def transform(self, statement: exp.Insert) -> exp.Insert:
+        statement = self._add_information_from_merge(statement)
+        statement = self._add_information_from_multitable_insert(statement)
+        statement = self._process_inner_ctes(statement)
+        statement = self._remove_on_conflict(statement)
+
+        return statement
+
+    def postprocess(self, statement: E) -> E:
+        # TODO: this is a hack to get the types annotated for UDF overloading to work for Inserts.
+        #  Revisit proper type annotation so that every query (not just Inserts) has its types populated to enable UDF resolution.
+        statement = annotate_types(statement, dialect=self.query.dialect, schema=self.query.object_mapping)
+        return super().postprocess(statement)
+
+    def _remove_on_conflict(self, statement: exp.Insert) -> exp.Insert:
+        """
+        Remove the 'ON CONFLICT ...' clause from the INSERT statement.
+        """
+        if conflict := statement.args.get("conflict"):
+            conflict.pop()
+        return statement
+
+    @BaseQueryTransformer._validate_syntax
+    def _add_information_from_merge(self, statement: exp.Insert) -> exp.Insert:
+        """
+        INSERT branch of _add_information_from_merge.
+
+        Transform a nested INSERT (from a MERGE WHEN NOT MATCHED branch) into a
+        fully qualified INSERT … SELECT so it can be processed independently.
+        """
+        ctx = self._extract_merge_context(statement)
+        if ctx is None:
+            return statement
+
+        using = ctx["using"]
+        returning = ctx["returning"]
+        target = self.query.target_info.expression
+
+        # Add the missing information to the INSERT statement
+        insert_columns = self._extract_insert_columns(statement, target, include_system=False)
+        values_lists = self._extract_value_lists(statement.expression)
+
+        if not values_lists:
+            return statement
+
+        # Build a new SELECT
+        values = values_lists[0]
+        aliases = [exp.alias_(val, str(col)) for col, val in zip(insert_columns, values)]
+        new_select = exp.select(*aliases).from_(using)
+
+        # 3. Build the standalone INSERT statement
+        insert_expr = exp.insert(
+            expression=new_select,
+            columns=insert_columns,
+            into=target,
+            dialect=self.query.dialect,
+            returning=returning,
+        )
+
+        # Add the CTEs
+        for cte in ctx["ctes"]:
+            insert_expr = insert_expr.with_(alias=cte.alias_or_name, as_=cte.this)
+
+        statement.replace(insert_expr)
+        return insert_expr
+
+    @BaseQueryTransformer._validate_syntax
+    def _add_information_from_multitable_insert(self, statement: exp.Insert) -> exp.Insert:
+        """
+        Reconstruct a standalone INSERT .. SELECT from a MultitableInsert branch.
+        """
+        ctx = self._extract_multitable_insert_context(statement)
+        if ctx is None:
+            return statement
+
+        source = ctx["source"]
+        target = self.query.target_info.expression
+
+        insert_columns = self._extract_insert_columns(statement, target, include_system=False)
+        values_lists = self._extract_value_lists(statement.expression)
+
+        if not values_lists:
+            return statement
+
+        selects = []
+        for values in values_lists:
+            aliases = [exp.alias_(val, str(col)) for col, val in zip(insert_columns, values)]
+            new_select = exp.select(*aliases).from_(source.subquery())
+            selects.append(new_select)
+
+        new_expression = exp.union(*selects, distinct=False) if len(selects) > 1 else selects[0]
+
+        insert_expr = exp.insert(
+            expression=new_expression,
+            columns=insert_columns,
+            into=target,
+            dialect=self.query.dialect,
+        )
+
+        statement.replace(insert_expr)
+        return insert_expr
+
+    @staticmethod
+    def _extract_value_lists(expression: exp.Expression) -> t.List[t.List[exp.Expression]]:
+        """
+        Handles exp.Tuple, exp.Select, and exp.Union.
+        """
+        values_lists = []
+        if isinstance(expression, exp.Tuple):
+            values_lists = [expression.expressions]
+        elif isinstance(expression, exp.Select):
+            values_lists = [[s.unalias() for s in expression.expressions]]
+        return values_lists

sqlleaf/processors/transformer/merge.py

@@ -0,0 +1,14 @@
+"""
+MergeTransformer — handles MERGE statement transformations.
+"""
+
+from sqlglot import exp
+
+from sqlleaf.processors.transformer.base import BaseQueryTransformer
+
+
+class MergeTransformer(BaseQueryTransformer):
+    """Transformer for MERGE statements."""
+
+    def transform(self, statement: exp.Merge) -> exp.Merge:
+        return self._process_inner_ctes(statement)

sqlleaf/processors/transformer/multitable_insert.py

@@ -0,0 +1,14 @@
+"""
+MultitableInsertTransformer — handles MultitableInserts statement transformations.
+"""
+
+from sqlglot import exp
+
+from sqlleaf.processors.transformer.base import BaseQueryTransformer
+
+
+class MultitableInsertTransformer(BaseQueryTransformer):
+    """Transformer for MultitableInserts statements."""
+
+    def transform(self, statement: exp.MultitableInserts) -> exp.MultitableInserts:
+        return statement

sqlleaf/processors/transformer/replace.py

@@ -0,0 +1,23 @@
+import sqlglot
+from sqlglot import exp
+
+from sqlleaf.processors.transformer.expressions import normalize_values
+from sqlleaf.processors.transformer.insert import InsertTransformer
+
+
+class ReplaceTransformer(InsertTransformer):
+    """Transformer for REPLACE statements."""
+
+    def transform(self, statement: exp.Insert) -> exp.Insert:
+        # Transform REPLACE to INSERT
+        expression = statement.args.get("expression")
+        new_sql = f"INSERT {expression.this}" if expression else "INSERT"
+        statement = sqlglot.parse_one(new_sql, dialect=self.query.dialect)
+
+        # This newly-parsed statement never went through preprocess(), so it hasn't
+        # had DEFAULT VALUES expansion or VALUES->SELECT normalization applied yet.
+        if isinstance(statement, exp.Insert):
+            statement = self._convert_insert_defaults_to_values(statement)
+        statement = normalize_values(self.query, statement)
+
+        return super().transform(statement)

sqlleaf/processors/transformer/transformer.py

@@ -0,0 +1,141 @@
+import logging
+
+from sqlglot import exp
+
+from sqlleaf import util
+from sqlleaf.models.query import (
+    CallQuery,
+    CopyQuery,
+    CTASQuery,
+    DeleteQuery,
+    ExecuteQuery,
+    InsertQuery,
+    MergeQuery,
+    MultitableInsertQuery,
+    Q,
+    QueryHolder,
+    ReplaceQuery,
+    TableQuery,
+    UnloadQuery,
+    UpdateQuery,
+    ValuesQuery,
+)
+from sqlleaf.processors.transformer import (
+    BaseQueryTransformer,
+    CallTransformer,
+    CopyTransformer,
+    CTASTransformer,
+    DeleteTransformer,
+    ExecuteTransformer,
+    InsertTransformer,
+    MergeTransformer,
+    MultitableInsertTransformer,
+    ReplaceTransformer,
+    UnloadTransformer,
+    UpdateTransformer,
+    ValuesTransformer,
+)
+from sqlleaf.typing import E
+
+_TRANSFORMER_MAP: dict[type, type[BaseQueryTransformer]] = {
+    CallQuery: CallTransformer,
+    CTASQuery: CTASTransformer,
+    CopyQuery: CopyTransformer,
+    DeleteQuery: DeleteTransformer,
+    ExecuteQuery: ExecuteTransformer,
+    InsertQuery: InsertTransformer,
+    MergeQuery: MergeTransformer,
+    MultitableInsertQuery: MultitableInsertTransformer,
+    ReplaceQuery: ReplaceTransformer,
+    TableQuery: BaseQueryTransformer,  # pass-through
+    UnloadQuery: UnloadTransformer,
+    UpdateQuery: UpdateTransformer,
+    ValuesQuery: ValuesTransformer,
+}
+
+logger = logging.getLogger("sqlleaf")
+
+
+def transform_query(query_holder: QueryHolder) -> None:
+    """
+    Transform an SQL query into a single, normalized form that we can easily generate lineage from.
+
+    This runs a series of transformations specific to each query type. Queries register their own
+    transformation functions to ensure that the query is fully qualified, with all columns added and
+    functions substituted so that they can be extracted into a graph by the generator.
+    """
+    transformed_query = _transform_query_instance(query=query_holder.original)
+    query_holder.set_transformed_query(query=transformed_query)
+
+
+def _transform_query_instance(query: Q) -> Q:
+    """
+    Helper to transform a Query instance.
+    """
+    statement_to_transform = util.copy_expression(query.statement)
+    transformed_statement = _transform_statement(statement=statement_to_transform, query=query)
+
+    return _build_transformed_query(
+        original_query=query,
+        transformed_statement=transformed_statement,
+    )
+
+
+def _build_transformed_query(
+    original_query: Q,
+    transformed_statement: exp.Expr,
+) -> Q:
+    """
+    Create a new Query instance whose statement is the transformed expression.
+    The Query subclass is selected based on the statement type.
+    """
+    if isinstance(transformed_statement, exp.Insert):
+        new_query = InsertQuery(
+            expr=transformed_statement,
+            dialect=original_query.dialect,
+            object_mapping=original_query.object_mapping,
+            statement_index=original_query.statement_index,
+            skip_annotate=True,
+        )
+        # TODO: everything below this in this function this should not occur
+        #  - requires big refactor
+        # CopyQuery special case: preserve source_info/target_info so that
+        # _apply_optimizations can still read the STREAM/FILE/STAGE type.
+        if isinstance(original_query, (CopyQuery, UnloadQuery)):
+            new_query.source_info = original_query.source_info
+            new_query.target_info = original_query.target_info
+    else:
+        # For statements not converted to INSERT, keep the same Query subclass
+        # but with the new statement.
+        new_query = original_query.__class__.__new__(original_query.__class__)
+        new_query.__dict__.update(original_query.__dict__)
+        new_query.statement = transformed_statement
+
+    # Propagate shared metadata
+    new_query.column_defs = original_query.column_defs
+    new_query.parent_query = original_query.parent_query
+    # Store a reference to the original query so that type-based checks in the
+    # generator (e.g. isinstance(query, UpdateQuery)) can inspect the original class.
+    new_query.original_query = original_query
+    return new_query
+
+
+def _transform_statement(statement: E, query: Q) -> exp.Expr:
+    """
+    Perform a series of transformations against an SQL statement.
+    Dispatches to the appropriate transformer class based on the query type.
+    """
+    logger.debug("---- Transformer ---")
+    logger.debug(f"Query: {statement.sql(dialect=query.dialect)}")
+    logger.debug(f"Transforming: {query.__class__.__name__} - {statement.__class__}")
+
+    transformer_cls = _TRANSFORMER_MAP.get(type(query), BaseQueryTransformer)
+    transformer = transformer_cls(statement, query)
+    stmt = transformer.preprocess(statement)
+    logger.debug(f"After pre-process: {type(stmt)} - {stmt.sql(dialect=query.dialect)}")
+    stmt = transformer.transform(stmt)
+    logger.debug(f"After process: {type(stmt)} - {stmt.sql(dialect=query.dialect)}")
+    stmt = transformer.postprocess(stmt)
+    logger.debug(f"After post-process: {type(stmt)} - {stmt.sql(dialect=query.dialect)}")
+    logger.debug("---- Transformer End ---")
+    return stmt