ripple-sql 0.1.0__py3-none-any.whl

ripple/engine/dependencies.py

@@ -0,0 +1,466 @@
+"""Column dependency extraction - WHERE clauses and JOIN keys.
+
+This module extracts C_ref columns (referenced but not contributing to output)
+following the LineageX paper's terminology:
+
+- C_con: Columns that contribute to output values (handled by extraction.py)
+- C_ref: Columns referenced in WHERE, HAVING, JOIN ON (handled here)
+
+Knowing which columns gate rows completes the picture for impact analysis.
+
+Example:
+    SELECT a.name, a.value
+    FROM orders a
+    JOIN customers b ON a.customer_id = b.id
+    WHERE b.status = 'active' AND a.created_at > '2024-01-01'
+
+    C_con: name, value (output columns)
+    C_ref: customer_id, id, status, created_at (filter/join columns)
+"""
+
+import logging
+from typing import TYPE_CHECKING, Literal
+
+from ripple.engine.column_ref import table_and_column
+from ripple.engine.types import (
+    ColumnDependencies,
+    FilterDependency,
+    JoinKeyPair,
+    WarehouseColumns,
+)
+
+if TYPE_CHECKING:
+    from sqlglot import exp
+
+logger = logging.getLogger(__name__)
+
+
+def extract_column_dependencies(
+    sql: str,
+    dialect: str = "snowflake",
+    warehouse_columns: "WarehouseColumns | None" = None,
+) -> ColumnDependencies:
+    """Extract columns used in WHERE clauses and JOIN conditions.
+
+    This complements extract_column_lineage_fast() by tracking C_ref columns
+    that affect which rows are returned, not the output values.
+
+    Args:
+        sql: SQL query to analyze
+        dialect: SQL dialect for parsing
+        warehouse_columns: optional schemas; when given, an unqualified key
+            resolves to its owning table instead of the value's table
+            (cycle-10 review, F2)
+
+    Returns:
+        ColumnDependencies with filter_columns and join_keys
+    """
+    import sqlglot
+
+    result = ColumnDependencies()
+
+    try:
+        # this reparse must see the same pre-parse pipeline the lineage
+        # parse got, or a {placeholder} estate zeroes dependency extraction
+        from ripple.engine.preprocess import prepare_sql_for_parse
+
+        sql, _ = prepare_sql_for_parse(sql, dialect)
+        parsed = None
+        if warehouse_columns:
+            from ripple.engine.schema_qualification import qualify_sql_with_schema
+
+            qualified_ast, was_qualified = qualify_sql_with_schema(sql, dialect, warehouse_columns)
+            if was_qualified and qualified_ast is not None:
+                parsed = qualified_ast
+        if parsed is None:
+            parsed = sqlglot.parse_one(sql, read=dialect)
+        if (dialect or "").lower() == "tsql":
+            # this reparse never saw the XML-method rewrite the lineage AST
+            # got, so a .exist() receiver in a predicate stayed invisible
+            # (cycle-9 review, F6)
+            from ripple.engine.tsql_xml import rewrite_xml_method_calls
+
+            parsed = rewrite_xml_method_calls(parsed)
+
+        result.filter_columns.extend(_extract_where_columns(parsed))
+
+        result.filter_columns.extend(_extract_having_columns(parsed))
+
+        result.filter_columns.extend(_extract_qualify_columns(parsed))
+
+        result.join_keys.extend(_extract_join_keys(parsed))
+
+        result.window_keys.extend(_extract_window_keys(parsed, dialect))
+
+    except Exception as e:
+        logger.warning(f"Dependency extraction failed: {type(e).__name__}: {e}")
+        # Sanitize error message - don't expose internal details to users
+        result.warnings.append(f"Dependency extraction failed: {type(e).__name__}")
+
+    return result
+
+
+def _extract_where_columns(
+    parsed: "exp.Expression",
+) -> list[FilterDependency]:
+    """Extract columns from WHERE clauses."""
+    from sqlglot import exp
+
+    results = []
+
+    for where in parsed.find_all(exp.Where):
+        for col in where.find_all(exp.Column):
+            table, column = table_and_column(col)
+            results.append(
+                FilterDependency(
+                    column=column,
+                    table=table,
+                    predicate_type="where",
+                    operator=_get_parent_operator(col),
+                )
+            )
+
+    return results
+
+
+def _keyed_relation_names(parsed: "exp.Expression") -> set[str]:
+    """Lowercased names a keyed column's qualifier could bind: table names,
+    dotted qualified names, and query aliases."""
+    from sqlglot import exp
+
+    from ripple.engine.column_ref import qualified_table_name
+
+    known: set[str] = set()
+    for table in parsed.find_all(exp.Table):
+        if not table.name:
+            continue
+        known.add(table.name.lower())
+        qualified = qualified_table_name(table).lower()
+        if qualified:
+            known.add(qualified)
+        if table.alias:
+            known.add(table.alias.lower())
+    for cte in parsed.find_all(exp.CTE):
+        if cte.alias_or_name:
+            known.add(cte.alias_or_name.lower())
+    return known
+
+
+def _resolve_keyed_column(col: "exp.Column", known: set[str]) -> tuple[str | None, str]:
+    """(table, column) for a window/selector key, qualification-aware.
+
+    table_and_column's struct guess read db.sch.t.s as table db, column
+    sch, a confident edge to a relation the query never reads (cycle-10
+    review, F1). A dotted prefix matching a relation the statement
+    names is a qualification, not a struct path.
+    """
+    from ripple.engine.column_ref import column_parts
+
+    parts = column_parts(col)
+    if len(parts) <= 2 or parts[0].lower() in known:
+        return table_and_column(col)
+    prefix = ".".join(parts[:-1])
+    if prefix.lower() in known or parts[-2].lower() in known:
+        return prefix, parts[-1]
+    return table_and_column(col)
+
+
+def _lateral_alias_columns(
+    col: "exp.Column", known: set[str]
+) -> "list[tuple[str | None, str]] | None":
+    """The real columns behind a keyed name that is a same-select alias.
+
+    rank_col AS selector, then max_by(v, selector): the key is the aliased
+    expression's own columns, never a column named selector (cycle-10
+    review, F3). A key sitting inside the expression its own alias
+    names is the real column, so it stays textual. None when the name
+    matches no alias.
+    """
+    from sqlglot import exp
+
+    from ripple.engine.column_ref import owning_select
+
+    owner = owning_select(col)
+    if owner is None:
+        return None
+    for item in owner.selects:
+        if not isinstance(item, exp.Alias) or (item.alias or "").lower() != col.name.lower():
+            continue
+        if any(inner is col for inner in item.this.find_all(exp.Column)):
+            return None
+        return [_resolve_keyed_column(inner, known) for inner in item.this.find_all(exp.Column)]
+    return None
+
+
+def _extract_window_keys(
+    parsed: "exp.Expression",
+    dialect: str = "",
+) -> list[FilterDependency]:
+    """Extract the PARTITION BY and ORDER BY columns of every window.
+
+    They frame the window rather than feed its value, so they are C_ref like a
+    WHERE column. Kept because renaming a partition key really does break the
+    window; dropping them would trade one wrong answer for a blind spot.
+    The selector argument of arg_max/arg_min plays the same role as a
+    window's ORDER BY, so it rides along here.
+    """
+    from sqlglot import exp
+
+    from ripple.engine.column_ref import column_parts
+    from ripple.engine.types import LATERAL_ALIAS_DIALECTS
+
+    keyed: list[exp.Expression] = []
+    for window in parsed.find_all(exp.Window):
+        keyed.extend(window.args.get("partition_by") or [])
+        order = window.args.get("order")
+        if order is not None:
+            keyed.append(order)
+    for agg in parsed.find_all(exp.ArgMax, exp.ArgMin):
+        keyed.extend(
+            arg for arg in (agg.args.get("expression"), agg.args.get("count")) if arg is not None
+        )
+
+    known = _keyed_relation_names(parsed)
+    allow_lateral = (dialect or "").lower() in LATERAL_ALIAS_DIALECTS
+
+    results = []
+    seen: set[tuple[str | None, str]] = set()
+    for part in keyed:
+        for col in part.find_all(exp.Column):
+            resolved = None
+            if allow_lateral and len(column_parts(col)) == 1:
+                resolved = _lateral_alias_columns(col, known)
+            if resolved is None:
+                resolved = [_resolve_keyed_column(col, known)]
+            for table, column in resolved:
+                if not column or (table, column) in seen:
+                    continue
+                seen.add((table, column))
+                results.append(
+                    FilterDependency(
+                        column=column,
+                        table=table,
+                        predicate_type="window",
+                    )
+                )
+
+    return results
+
+
+def _extract_having_columns(
+    parsed: "exp.Expression",
+) -> list[FilterDependency]:
+    """Extract columns from HAVING clauses."""
+    from sqlglot import exp
+
+    results = []
+
+    for having in parsed.find_all(exp.Having):
+        for col in having.find_all(exp.Column):
+            table, column = table_and_column(col)
+            results.append(
+                FilterDependency(
+                    column=column,
+                    table=table,
+                    predicate_type="having",
+                    operator=_get_parent_operator(col),
+                )
+            )
+
+    return results
+
+
+def _extract_qualify_columns(
+    parsed: "exp.Expression",
+) -> list[FilterDependency]:
+    """Extract columns from QUALIFY clauses (Snowflake/Databricks)."""
+    from sqlglot import exp
+
+    results = []
+
+    for qualify in parsed.find_all(exp.Qualify):
+        for col in qualify.find_all(exp.Column):
+            table, column = table_and_column(col)
+            results.append(
+                FilterDependency(
+                    column=column,
+                    table=table,
+                    predicate_type="qualify",
+                    operator=_get_parent_operator(col),
+                )
+            )
+
+    return results
+
+
+def _extract_join_keys(
+    parsed: "exp.Expression",
+) -> list[JoinKeyPair]:
+    """Extract column pairs from JOIN ON conditions."""
+    from sqlglot import exp
+
+    results = []
+
+    for join in parsed.find_all(exp.Join):
+        join_type = _get_join_type(join)
+        on_clause = join.args.get("on")
+
+        if on_clause:
+            # Find equality conditions (a.col = b.col)
+            for eq in on_clause.find_all(exp.EQ):
+                left = eq.left
+                right = eq.right
+
+                if isinstance(left, exp.Column) and isinstance(right, exp.Column):
+                    left_table, left_column = table_and_column(left)
+                    right_table, right_column = table_and_column(right)
+                    results.append(
+                        JoinKeyPair(
+                            left_column=left_column,
+                            left_table=left_table or "",
+                            right_column=right_column,
+                            right_table=right_table or "",
+                            join_type=join_type,
+                        )
+                    )
+
+    return results
+
+
+def _get_parent_operator(col: "exp.Column") -> str | None:
+    """Get the operator containing this column (=, IN, BETWEEN, etc.)."""
+    from sqlglot import exp
+
+    parent = col.parent
+    while parent:
+        if isinstance(parent, exp.EQ):
+            return "="
+        elif isinstance(parent, exp.NEQ):
+            return "!="
+        elif isinstance(parent, exp.In):
+            return "IN"
+        elif isinstance(parent, exp.Between):
+            return "BETWEEN"
+        elif isinstance(parent, exp.GT):
+            return ">"
+        elif isinstance(parent, exp.GTE):
+            return ">="
+        elif isinstance(parent, exp.LT):
+            return "<"
+        elif isinstance(parent, exp.LTE):
+            return "<="
+        elif isinstance(parent, exp.Like):
+            return "LIKE"
+        elif isinstance(parent, exp.ILike):
+            return "ILIKE"
+        elif isinstance(parent, exp.Is):
+            return "IS"
+        elif isinstance(parent, exp.RegexpLike):
+            return "REGEXP"
+        parent = parent.parent
+    return None
+
+
+def _get_join_type(join: "exp.Join") -> Literal["inner", "left", "right", "full", "cross"]:
+    """Determine JOIN type from expression.
+
+    Handles: INNER, LEFT, RIGHT, FULL, CROSS, NATURAL, SEMI, ANTI joins.
+    NATURAL/SEMI/ANTI are mapped to "inner" as they don't have a distinct type in our model.
+    """
+    from typing import Literal as Lit
+    from typing import cast
+
+    if join.side:
+        side = join.side.lower()
+        if side in ("left", "right", "full"):
+            return cast(Lit["left", "right", "full"], side)
+    if join.kind:
+        kind = join.kind.lower()
+        if kind == "cross":
+            return "cross"
+        # NATURAL, SEMI, ANTI joins default to "inner" behavior
+    return "inner"
+
+
+def extract_all_referenced_columns(
+    sql: str,
+    dialect: str = "snowflake",
+) -> list[dict[str, str | None]]:
+    """Extract ALL column references from SQL (SELECT, WHERE, JOIN, GROUP BY, etc.).
+
+    This is a convenience function that returns every column reference found,
+    tagged with its context (select, where, join, group_by, order_by, etc.).
+
+    Useful for impact analysis: "which columns does this query touch?"
+
+    Args:
+        sql: SQL query to analyze
+        dialect: SQL dialect
+
+    Returns:
+        List of column references with context
+    """
+    import sqlglot
+    from sqlglot import exp
+
+    results = []
+
+    try:
+        from ripple.engine.preprocess import prepare_sql_for_parse
+
+        sql, _ = prepare_sql_for_parse(sql, dialect)
+        parsed = sqlglot.parse_one(sql, read=dialect)
+
+        # Track which columns we've seen to avoid duplicates
+        seen: set[tuple[str, str, str]] = set()
+
+        def add_column(col: exp.Column, context: str) -> None:
+            table, column = table_and_column(col)
+            key = (table or "", column, context)
+            if key not in seen:
+                seen.add(key)
+                results.append(
+                    {
+                        "column": column,
+                        "table": table,
+                        "context": context,
+                    }
+                )
+
+        # SELECT columns
+        for select in parsed.find_all(exp.Select):
+            for expr in select.selects:
+                for col in expr.find_all(exp.Column):
+                    add_column(col, "select")
+
+        # WHERE columns
+        for where in parsed.find_all(exp.Where):
+            for col in where.find_all(exp.Column):
+                add_column(col, "where")
+
+        # HAVING columns
+        for having in parsed.find_all(exp.Having):
+            for col in having.find_all(exp.Column):
+                add_column(col, "having")
+
+        # JOIN ON columns
+        for join in parsed.find_all(exp.Join):
+            on_clause = join.args.get("on")
+            if on_clause:
+                for col in on_clause.find_all(exp.Column):
+                    add_column(col, "join")
+
+        # GROUP BY columns
+        for group in parsed.find_all(exp.Group):
+            for col in group.find_all(exp.Column):
+                add_column(col, "group_by")
+
+        # ORDER BY columns
+        for order in parsed.find_all(exp.Order):
+            for col in order.find_all(exp.Column):
+                add_column(col, "order_by")
+
+    except Exception as e:
+        logger.warning(f"Failed to extract all column references: {e}")
+
+    return results

ripple/engine/dialect.py

@@ -0,0 +1,132 @@
+"""Map dbt adapter types to SQLGlot dialect names.
+
+dbt manifests include metadata.adapter_type which tells us what database
+the project targets. This module maps those adapter names to SQLGlot's
+dialect identifiers for correct SQL parsing.
+
+Example:
+    manifest["metadata"]["adapter_type"] = "snowflake"
+    → get_sqlglot_dialect(manifest) returns "snowflake"
+
+    manifest["metadata"]["adapter_type"] = "athena"
+    → get_sqlglot_dialect(manifest) returns "trino" (Athena uses Presto/Trino SQL)
+"""
+
+import logging
+from typing import Any, Final
+
+logger = logging.getLogger(__name__)
+
+# Mapping from dbt adapter_type → SQLGlot dialect name
+#
+# Why some mappings aren't 1:1:
+# - athena → trino: AWS Athena uses Presto/Trino SQL syntax
+# - synapse → tsql: Azure Synapse uses T-SQL (SQL Server syntax)
+# - sqlserver → tsql: SQL Server uses T-SQL
+#
+ADAPTER_TO_SQLGLOT_DIALECT: Final[dict[str, str]] = {
+    # Cloud data warehouses
+    "snowflake": "snowflake",
+    "bigquery": "bigquery",
+    "databricks": "databricks",
+    "redshift": "redshift",
+    # Open source / on-prem
+    "postgres": "postgres",
+    "trino": "trino",
+    "presto": "trino",  # Presto uses Trino SQL syntax
+    "starburst": "trino",  # Starburst is enterprise Trino
+    "spark": "spark",
+    "duckdb": "duckdb",
+    "clickhouse": "clickhouse",
+    # AWS
+    "athena": "trino",  # Athena uses Presto/Trino SQL syntax
+    "glue": "spark",  # AWS Glue uses Spark SQL
+    # Microsoft
+    "synapse": "tsql",  # Azure Synapse uses T-SQL
+    "sqlserver": "tsql",  # SQL Server uses T-SQL
+    "fabric": "tsql",  # Microsoft Fabric uses T-SQL
+}
+
+# Default dialect when adapter is unknown or not specified
+DEFAULT_DIALECT: Final[str] = "snowflake"
+
+
+def get_sqlglot_dialect(manifest: dict[str, Any]) -> str:
+    """Get SQLGlot dialect from dbt manifest metadata.
+
+    Extracts the adapter_type from manifest.metadata and maps it
+    to the corresponding SQLGlot dialect for SQL parsing.
+
+    Args:
+        manifest: Parsed dbt manifest.json dictionary
+
+    Returns:
+        SQLGlot dialect string (e.g., "snowflake", "bigquery", "trino")
+        Defaults to "snowflake" if adapter is unknown.
+
+    Example:
+        >>> manifest = {"metadata": {"adapter_type": "databricks"}}
+        >>> get_sqlglot_dialect(manifest)
+        'databricks'
+    """
+    adapter_type = manifest.get("metadata", {}).get("adapter_type", "")
+    dialect = ADAPTER_TO_SQLGLOT_DIALECT.get(adapter_type.lower(), DEFAULT_DIALECT)
+
+    if adapter_type and adapter_type.lower() not in ADAPTER_TO_SQLGLOT_DIALECT:
+        logger.warning(
+            f"Unknown dbt adapter '{adapter_type}', defaulting to '{DEFAULT_DIALECT}' dialect"
+        )
+
+    logger.debug(f"dbt adapter '{adapter_type}' → SQLGlot dialect '{dialect}'")
+    return dialect
+
+
+def get_dialect_from_adapter_type(adapter_type: str | None) -> str:
+    """Get SQLGlot dialect from a dbt adapter type string.
+
+    Args:
+        adapter_type: dbt adapter type (e.g., "snowflake", "bigquery", "databricks")
+
+    Returns:
+        SQLGlot dialect string. Defaults to "snowflake" if unknown.
+
+    Example:
+        >>> get_dialect_from_adapter_type("databricks")
+        'databricks'
+        >>> get_dialect_from_adapter_type("athena")
+        'trino'
+    """
+    if not adapter_type:
+        return DEFAULT_DIALECT
+
+    dialect = ADAPTER_TO_SQLGLOT_DIALECT.get(adapter_type.lower(), DEFAULT_DIALECT)
+
+    if adapter_type.lower() not in ADAPTER_TO_SQLGLOT_DIALECT:
+        logger.warning(
+            f"Unknown dbt adapter '{adapter_type}', defaulting to '{DEFAULT_DIALECT}' dialect"
+        )
+
+    return dialect
+
+
+def get_dialect_from_config(config: dict[str, Any] | None) -> str:
+    """Get the dialect from a config dict that may carry adapter_type.
+
+    Args:
+        config: dict that may contain "adapter_type".
+
+    Returns:
+        SQLGlot dialect string, or DEFAULT_DIALECT if absent.
+    """
+    if not config:
+        return DEFAULT_DIALECT
+    return get_dialect_from_adapter_type(config.get("adapter_type"))
+
+
+def get_supported_adapters() -> list[str]:
+    """Get list of supported dbt adapter types.
+
+    Returns:
+        List of adapter type strings that have SQLGlot dialect mappings.
+    """
+    return list(ADAPTER_TO_SQLGLOT_DIALECT.keys())

ripple/engine/dispatch.py

@@ -0,0 +1,12 @@
+"""Engine entry point.
+
+Every caller goes through this module for lineage extraction, so a different
+engine could slot in behind it later without touching the call sites. Today it
+calls the Python engine directly.
+"""
+
+from __future__ import annotations
+
+from ripple.engine.extraction import extract_lineage_complete
+
+__all__ = ["extract_lineage_complete"]

ripple/engine/extraction.py

@@ -0,0 +1,27 @@
+"""Column lineage extraction - public entry points (facade).
+
+The implementation is split by responsibility:
+- macro_source.py: primary-source detection from dbt macro patterns
+- scope.py: FROM/JOIN/LATERAL/UNNEST alias and expansion registration
+- json_sources.py: JSON path and bracket access source extraction
+- select_sources.py: per-select-item source classification
+- statement.py: statement-level orchestration (extract_column_lineage_fast)
+- pipeline.py: CTE mappings, dialect fallback, unified C_con + C_ref API
+
+This module only re-exports the public functions so existing imports
+(`from ripple.engine.extraction import ...`) keep working.
+"""
+
+from ripple.engine.macro_source import extract_primary_source_from_macro
+from ripple.engine.pipeline import (
+    extract_column_lineage_with_ctes,
+    extract_lineage_complete,
+)
+from ripple.engine.statement import extract_column_lineage_fast
+
+__all__ = [
+    "extract_column_lineage_fast",
+    "extract_column_lineage_with_ctes",
+    "extract_lineage_complete",
+    "extract_primary_source_from_macro",
+]