ripple-sql 0.1.0__py3-none-any.whl

ripple/engine/pipeline.py

@@ -0,0 +1,265 @@
+"""Top-level extraction pipeline.
+
+extract_column_lineage_with_ctes wraps the statement extractor with CTE
+mapping collection, primary-source detection, dialect fallback, and warning
+aggregation. extract_lineage_complete is the unified API: C_con
+(contributing columns) plus C_ref (filter columns and join keys).
+"""
+
+import logging
+from collections.abc import Callable
+from typing import Any
+
+from ripple.engine.cte_tracing import extract_cte_column_lineage
+from ripple.engine.macro_source import extract_primary_source_from_macro
+from ripple.engine.statement import extract_column_lineage_fast
+from ripple.engine.types import (
+    META_CTE_COLLISIONS,
+    META_EXTRACTION_ERROR,
+    META_EXTRACTION_ERROR_TYPE,
+    META_FALLBACK_DIALECT,
+    META_PREFIX,
+    META_QUALIFIED_VIA_SCHEMA,
+    META_WARNINGS,
+    UnifiedLineageResult,
+    WarehouseColumns,
+)
+from ripple.engine.validation import get_trust_level
+
+logger = logging.getLogger(__name__)
+
+
+def extract_column_lineage_with_ctes(
+    raw_sql: str,
+    dialect: str = "snowflake",
+    cte_mappings_func: Callable[[str], dict] | None = None,
+    join_mappings_func: Callable[[str], dict] | None = None,
+    clean_jinja_func: Callable[[str], str] | None = None,
+    clean_jinja_minimal_func: Callable[[str], str] | None = None,
+    warehouse_columns: WarehouseColumns | None = None,
+    real_tables: set[str] | None = None,
+    self_names: set[str] | None = None,
+) -> dict[str, Any]:
+    """Extract column lineage AND CTE/alias->ref mappings from SQL.
+
+    Returns:
+        {
+            "column_lineage": {column_name: [sources]},
+            "cte_mappings": {cte_name: {type, ref/source/table}},
+            "cte_column_lineage": {cte_name: {column: sources}},
+            "primary_source": {type, source, table} (optional),
+            "warnings": [...] (if any issues detected),
+            "qualified_via_schema": bool (if warehouse schema used),
+        }
+
+    Args:
+        raw_sql: Raw SQL (possibly with Jinja templates)
+        dialect: SQL dialect (snowflake, bigquery, databricks, etc.)
+        cte_mappings_func: Function to extract CTE mappings (injected)
+        join_mappings_func: Function to extract JOIN mappings (injected)
+        clean_jinja_func: Function to clean Jinja from SQL (injected)
+        clean_jinja_minimal_func: Function for minimal Jinja cleaning (injected)
+        warehouse_columns: Schema info from warehouse {table: {column: data_type}}
+        real_tables: Set of real table names (for CTE collision detection)
+    """
+    # Use injected functions or return empty dicts
+    if cte_mappings_func is None:
+
+        def cte_mappings_func(sql):
+            return {}
+
+    if join_mappings_func is None:
+
+        def join_mappings_func(sql):
+            return {}
+
+    cte_mappings = cte_mappings_func(raw_sql)
+    join_mappings = join_mappings_func(raw_sql)
+    cte_column_lineage = extract_cte_column_lineage(
+        raw_sql,
+        dialect=dialect,
+        clean_jinja_func=clean_jinja_func,
+        real_tables=real_tables,
+        self_names=self_names,
+    )
+
+    # Collect warnings from CTE extraction
+    all_warnings: list[dict[str, Any]] = []
+    if META_CTE_COLLISIONS in cte_column_lineage:
+        collisions = cte_column_lineage.pop(META_CTE_COLLISIONS)
+        if isinstance(collisions, list):
+            all_warnings.extend(collisions)
+
+    # Merge: CTE mappings take precedence over join mappings
+    combined_mappings = {**join_mappings, **cte_mappings}
+
+    primary_source = extract_primary_source_from_macro(raw_sql)
+
+    column_lineage = extract_column_lineage_fast(
+        raw_sql,
+        dialect=dialect,
+        clean_jinja_func=clean_jinja_func,
+        clean_jinja_minimal_func=clean_jinja_minimal_func,
+        warehouse_columns=warehouse_columns,
+        real_tables=real_tables,
+        self_names=self_names,
+    )
+
+    # Collect warnings from column lineage extraction
+    if META_WARNINGS in column_lineage:
+        warnings = column_lineage.pop(META_WARNINGS)
+        if isinstance(warnings, list):
+            all_warnings.extend(warnings)
+
+    # Track if schema was used for qualification
+    qualified_via_schema = bool(column_lineage.pop(META_QUALIFIED_VIA_SCHEMA, False))
+
+    has_error = META_EXTRACTION_ERROR in column_lineage
+    has_columns = any(k for k in column_lineage if not k.startswith(META_PREFIX))
+
+    if has_error or not has_columns:
+        fallback_dialects = ["snowflake", "postgres", "bigquery", "duckdb"]
+        last_error = column_lineage.get(META_EXTRACTION_ERROR) if has_error else None
+        last_error_type = column_lineage.get(META_EXTRACTION_ERROR_TYPE) if has_error else None
+
+        for fallback in fallback_dialects:
+            if fallback != dialect:
+                fallback_result = extract_column_lineage_fast(
+                    raw_sql,
+                    dialect=fallback,
+                    clean_jinja_func=clean_jinja_func,
+                    clean_jinja_minimal_func=clean_jinja_minimal_func,
+                    warehouse_columns=warehouse_columns,
+                    real_tables=real_tables,
+                )
+                fallback_has_error = META_EXTRACTION_ERROR in fallback_result
+                fallback_has_columns = any(
+                    k for k in fallback_result if not k.startswith(META_PREFIX)
+                )
+
+                if fallback_has_columns and not fallback_has_error:
+                    # Degrade confidence for fallback-parsed lineage
+                    for col_key, col_sources in fallback_result.items():
+                        if col_key.startswith(META_PREFIX):
+                            continue
+                        if isinstance(col_sources, list):
+                            for src in col_sources:
+                                original = src.get("confidence", 1.0)
+                                src["confidence"] = round(original * 0.7, 3)
+                                src["trust_level"] = get_trust_level(src["confidence"])
+                    fallback_result[META_FALLBACK_DIALECT] = fallback
+                    column_lineage = fallback_result
+                    logger.debug(f"Column lineage extracted with fallback dialect: {fallback}")
+                    break
+                elif fallback_has_error:
+                    last_error = fallback_result.get(META_EXTRACTION_ERROR)
+                    last_error_type = fallback_result.get(META_EXTRACTION_ERROR_TYPE)
+        else:
+            if last_error:
+                column_lineage = {
+                    META_EXTRACTION_ERROR: last_error,
+                    META_EXTRACTION_ERROR_TYPE: last_error_type,
+                }
+
+    fallback_dialect = column_lineage.pop(META_FALLBACK_DIALECT, None)
+
+    result = {
+        "column_lineage": column_lineage,
+        "cte_mappings": combined_mappings,
+        "cte_column_lineage": cte_column_lineage,
+    }
+
+    if fallback_dialect:
+        result["fallback_dialect"] = fallback_dialect
+
+    if primary_source:
+        result["primary_source"] = primary_source
+
+    if all_warnings:
+        result["warnings"] = all_warnings
+
+    if qualified_via_schema:
+        result["qualified_via_schema"] = True
+
+    return result
+
+
+def extract_lineage_complete(
+    sql: str,
+    dialect: str = "snowflake",
+    warehouse_columns: WarehouseColumns | None = None,
+    real_tables: set[str] | None = None,
+    clean_jinja_func: Callable[[str], str] | None = None,
+    function_relations: set[str] | None = None,
+    self_names: set[str] | None = None,
+) -> UnifiedLineageResult:
+    """Extract complete lineage: contributing + referenced columns.
+
+    Returns both:
+    - C_con: Columns that contribute to output values
+    - C_ref: Columns referenced in WHERE/JOIN (affect row selection)
+
+    This follows the LineageX paper's model, giving complete visibility
+    into which columns affect the query results.
+
+    Args:
+        sql: SQL query to analyze
+        dialect: SQL dialect
+        warehouse_columns: Schema for disambiguation
+        real_tables: Real table names for CTE collision detection
+        clean_jinja_func: Optional Jinja cleaner
+
+    Returns:
+        UnifiedLineageResult with contributing, filter_columns, join_keys
+    """
+    from ripple.engine.dependencies import extract_column_dependencies
+
+    # Clean SQL if cleaner provided
+    cleaned_sql = clean_jinja_func(sql) if clean_jinja_func else sql
+
+    contributing = extract_column_lineage_fast(
+        sql,
+        dialect=dialect,
+        warehouse_columns=warehouse_columns,
+        real_tables=real_tables,
+        clean_jinja_func=clean_jinja_func,
+        function_relations=function_relations,
+        self_names=self_names,
+    )
+
+    dependencies = extract_column_dependencies(
+        cleaned_sql, dialect, warehouse_columns=warehouse_columns
+    )
+
+    # Combine warnings
+    raw_warnings = contributing.pop(META_WARNINGS, [])
+    all_warnings: list[dict[str, Any]] = raw_warnings if isinstance(raw_warnings, list) else []
+    all_warnings.extend([{"warning": w} for w in dependencies.warnings])
+
+    # A statement the extractor couldn't analyze must say so here, not
+    # vanish when the underscore keys are stripped (MAP/REDUCE script
+    # transforms, unparseable dialect corners, future failures).
+    if META_EXTRACTION_ERROR in contributing:
+        all_warnings.append(
+            {
+                "warning_type": "extraction_error",
+                "trust_level": "review_required",
+                "message": str(contributing.get(META_EXTRACTION_ERROR)),
+                "error_type": str(contributing.get(META_EXTRACTION_ERROR_TYPE, "Unknown")),
+            }
+        )
+
+    qualified = bool(contributing.pop(META_QUALIFIED_VIA_SCHEMA, False))
+
+    clean_contributing: dict[str, list[dict[str, Any]]] = {
+        k: v for k, v in contributing.items() if not k.startswith(META_PREFIX)
+    }
+
+    return UnifiedLineageResult(
+        contributing=clean_contributing,
+        filter_columns=dependencies.filter_columns,
+        join_keys=dependencies.join_keys,
+        window_keys=dependencies.window_keys,
+        warnings=all_warnings,
+        qualified_via_schema=qualified,
+    )

ripple/engine/preprocess.py

@@ -0,0 +1,174 @@
+"""Pre-parse rewrites for constructs sqlglot can't carry through.
+
+IDENTIFIER('literal') (Databricks/Spark parameterized identifiers) either
+fails to parse or parses to a nameless node that would emit lineage to an
+empty table. A single string-literal argument is statically resolvable, so
+rewrite it to the identifier before parsing and record that we did.
+Non-literal arguments (:params, || concat) are left alone; the extractor
+flags them as unresolvable instead of guessing.
+
+Single-brace str.format placeholders in relation position ({hscic}.t) are
+the same family: statically rewritable before parsing, with the braced
+spelling kept verbatim in the emitted name.
+"""
+
+from __future__ import annotations
+
+import re
+from collections.abc import Callable
+
+_IDENTIFIER_LITERAL = re.compile(r"\bIDENTIFIER\s*\(\s*'([A-Za-z0-9_.`\"$]+)'\s*\)", re.IGNORECASE)
+
+
+def rewrite_identifier_literals(sql: str) -> tuple[str, int]:
+    """Replace IDENTIFIER('name') with name. Returns (sql, replacements)."""
+    rewritten, count = _IDENTIFIER_LITERAL.subn(r"\1", sql)
+    return rewritten, count
+
+
+_IDENTIFIER_ANY = re.compile(r"\bIDENTIFIER\s*\(", re.IGNORECASE)
+
+
+def has_unresolved_identifier(sql: str) -> bool:
+    """True when an IDENTIFIER(...) call with a non-literal argument remains."""
+    return bool(_IDENTIFIER_ANY.search(sql))
+
+
+# a surviving pure name/attribute chain ({{params.x}}, the double-brace
+# sibling of the single-brace class) is quotable like any placeholder;
+# jinja spans with anything else in them stay protected verbatim
+_JINJA_CHAIN_BODY = r"[A-Za-z_][A-Za-z0-9_]*(?:\.[A-Za-z_][A-Za-z0-9_]*)*"
+_JINJA_CHAIN = re.compile(rf"\{{\{{{_JINJA_CHAIN_BODY}\}}\}}")
+
+# Dots INSIDE a quoted chain still split the identifier into table parts in
+# BigQuery; encoded before quoting, restored to "." at the identifier
+# accessors (column_ref.column_parts / qualified_table_name and the CTE map
+# keys) so the emitted spelling is byte-for-byte as written. A long ASCII
+# token, not a private-use codepoint: a real identifier containing U+E000
+# must round-trip untouched (cycle-12 review, F3).
+JINJA_DOT_SENTINEL = "__ripple_jinja_dot__"
+
+# A raw chain already containing the token must not read as an encoded dot:
+# the accessors would keep the qualifier and mutate its spelling (cycle-13
+# review, F1). Encoding bumps any pre-existing __ripple_jinja_dot[N]__ one
+# level up and restoration bumps it back down, so a bare token in a parsed
+# name always means a dot we encoded and raw spellings round-trip. Applied
+# to gap-text chains only: re-prepared model SQL carries its chains inside
+# quoted spans, which keeps the bump single-shot.
+_SENTINEL_ANY_LEVEL = re.compile(r"__ripple_jinja_dot(\d*)__")
+_SENTINEL_ESCAPED = re.compile(r"__ripple_jinja_dot(\d+)__")
+
+
+def _encode_chain_dots(chain: str) -> str:
+    escaped = _SENTINEL_ANY_LEVEL.sub(
+        lambda m: f"__ripple_jinja_dot{int(m.group(1) or -1) + 1}__", chain
+    )
+    return escaped.replace(".", JINJA_DOT_SENTINEL)
+
+
+_PROTECTED = re.compile(
+    rf"""'(?:\\.|''|[^'\\])*'
+      | "(?:\\.|""|[^"\\])*"
+      | `[^`]*`
+      | \[(?:\]\]|[^\]])*\]
+      | \{{\{{(?!{_JINJA_CHAIN_BODY}\}}\}}).*?\}}\}}
+      | \{{%.*?%\}}
+      | \{{\#.*?\#\}}
+      | --[^\n]*
+      | /\*.*?\*/
+    """,
+    re.VERBOSE | re.DOTALL,
+)
+
+# a hyphen joins atoms only when tightly adjacent, so BigQuery's
+# my-{env}.dataset.t quotes as one segment while spaced `a - {x}` stays two
+_ENCODED_CHAIN = rf"\{{\{{[A-Za-z_](?:[A-Za-z0-9_]|{JINJA_DOT_SENTINEL})*\}}\}}"
+_PLACEHOLDER_ATOM = rf"(?:{_ENCODED_CHAIN}|\{{[A-Za-z_][A-Za-z0-9_]*\}}|[A-Za-z0-9_])+"
+_PLACEHOLDER_SEGMENT = re.compile(rf"{_PLACEHOLDER_ATOM}(?:-{_PLACEHOLDER_ATOM})*")
+
+
+def restore_jinja_dots(name: str) -> str:
+    """Undo the in-chain dot encoding; identity for ordinary names."""
+    if "__ripple_jinja_dot" not in name:
+        return name
+    name = name.replace(JINJA_DOT_SENTINEL, ".")
+    return _SENTINEL_ESCAPED.sub(
+        lambda m: f"__ripple_jinja_dot{int(m.group(1)) - 1 if m.group(1) != '0' else ''}__",
+        name,
+    )
+
+
+def _identifier_quotes(dialect: str) -> tuple[str, str]:
+    try:
+        from sqlglot.dialects.dialect import Dialect
+
+        d = Dialect.get_or_raise(dialect or None)
+        return d.IDENTIFIER_START, d.IDENTIFIER_END
+    except Exception:
+        return '"', '"'
+
+
+def rewrite_brace_placeholders(sql: str, dialect: str = "") -> tuple[str, int]:
+    """Quote name segments containing {placeholder} so they parse as identifiers.
+
+    Python str.format placeholders like {hscic}.practice_statistics fail to
+    parse and zero the whole statement. Each dotted-name segment holding a
+    single-brace {ident} placeholder is wrapped in the dialect's identifier
+    quote, so sqlglot reads a plain identifier whose name keeps the braces
+    exactly as written. Ground truth cites these relations with the braces,
+    so the spelling must survive into the emitted table name. String
+    literals, jinja double braces, and already-quoted content are left
+    alone. Returns (sql, segments quoted).
+    """
+    if "{" not in sql:
+        return sql, 0
+    start_q, end_q = _identifier_quotes(dialect)
+    # BigQuery splits a backticked name on its dots, shredding a jinja
+    # chain across catalog/db/table ("destination_dataset_project_id}}"
+    # became a relation part, bitcoin_etl round 12); encoding in-chain
+    # dots inside backticks keeps the split on real path separators
+    encode_quoted = (dialect or "").lower() == "bigquery"
+    count = 0
+
+    def quote_segment(m: re.Match[str]) -> str:
+        nonlocal count
+        seg = m.group(0)
+        if "{" not in seg:
+            return seg
+        count += 1
+        return f"{start_q}{seg}{end_q}"
+
+    def rewrite_gap(text: str) -> str:
+        encoded = _JINJA_CHAIN.sub(lambda m: _encode_chain_dots(m.group(0)), text)
+        return _PLACEHOLDER_SEGMENT.sub(quote_segment, encoded)
+
+    out: list[str] = []
+    pos = 0
+    for m in _PROTECTED.finditer(sql):
+        out.append(rewrite_gap(sql[pos : m.start()]))
+        span = m.group(0)
+        if encode_quoted and span.startswith("`"):
+            span = _JINJA_CHAIN.sub(lambda c: c.group(0).replace(".", JINJA_DOT_SENTINEL), span)
+        out.append(span)
+        pos = m.end()
+    out.append(rewrite_gap(sql[pos:]))
+    return "".join(out), count
+
+
+def prepare_sql_for_parse(
+    sql: str,
+    dialect: str = "",
+    clean_jinja_func: Callable[[str], str] | None = None,
+) -> tuple[str, int]:
+    """The one pre-parse pipeline for every site that parses user SQL.
+
+    Jinja cleaning runs first, then the single-brace rewrite: rewriting
+    first quotes across a rendered boundary and corrupts adjacent
+    spellings like {prefix}_{{suffix}} (cycle-10 review, F9). Any
+    parse of user SQL that skips this drifts from the lineage parse and
+    loses edges the way graph._alias_map did (F8). Returns
+    (sql, brace segments quoted).
+    """
+    if clean_jinja_func is not None:
+        sql = clean_jinja_func(sql)
+    return rewrite_brace_placeholders(sql, dialect)

ripple/engine/safe_gen.py

@@ -0,0 +1,21 @@
+"""Crash-safe SQL generation for user-derived ASTs.
+
+sqlglot's generator can raise on nodes its own parser produced: T-SQL
+`SET @x -= 1` parses to an empty Neg() that neg_sql cannot render
+(IndexError), and one such statement crashed a whole project load
+(fhir_server, holdout round 11). Every site that re-generates user SQL
+during load or analysis goes through here and degrades to skipping that
+statement, never aborting the file.
+"""
+
+from __future__ import annotations
+
+from sqlglot import exp
+
+
+def safe_sql(node: exp.Expression, dialect: str | None) -> str | None:
+    """node.sql(dialect=...) or None when the generator refuses the node."""
+    try:
+        return node.sql(dialect=dialect)
+    except Exception:
+        return None

ripple/engine/schema_qualification.py

@@ -0,0 +1,151 @@
+"""Schema-aware column qualification using warehouse metadata.
+
+This module provides disambiguation of unqualified column references
+using actual warehouse schema information. This is key for achieving
+high-confidence lineage: it resolves ambiguity instead of only flagging it.
+
+Key Functions:
+- qualify_sql_with_schema: Qualify SQL AST using warehouse schema
+- resolve_column_table: Find which table a column belongs to
+- build_sqlglot_schema: Convert warehouse columns to SQLGlot format
+
+Example:
+    Given SQL: SELECT col FROM a, b
+    And warehouse schema showing 'col' only exists in table 'a'
+    Result: SELECT a.col FROM a, b (ambiguity resolved)
+"""
+
+import logging
+from typing import TYPE_CHECKING
+
+from ripple.engine.types import WarehouseColumns
+
+if TYPE_CHECKING:
+    from sqlglot import exp
+
+logger = logging.getLogger(__name__)
+
+
+def build_sqlglot_schema(
+    warehouse_columns: WarehouseColumns,
+) -> dict[str, dict[str, str]]:
+    """Convert warehouse columns to SQLGlot schema format.
+
+    SQLGlot's qualify() expects a schema dict in the format:
+    {table_name: {column_name: data_type}}
+
+    Args:
+        warehouse_columns: Dict of {table: {column: data_type}}
+
+    Returns:
+        Schema dict compatible with SQLGlot's qualify()
+    """
+    # SQLGlot schema format is same as our WarehouseColumns format
+    return warehouse_columns
+
+
+def resolve_column_table(
+    column: str,
+    candidate_tables: list[str],
+    warehouse_columns: WarehouseColumns,
+) -> tuple[str | None, bool]:
+    """Resolve which table a column belongs to using warehouse schema.
+
+    This is the key function for disambiguating unqualified columns.
+    Returns the table name if the column exists in exactly one table,
+    otherwise returns None.
+
+    Args:
+        column: Column name to resolve (case-insensitive)
+        candidate_tables: List of possible source tables
+        warehouse_columns: Schema info from warehouse
+
+    Returns:
+        Tuple of (resolved_table_name, was_resolved)
+        - (table_name, True) if column exists in exactly one table
+        - (None, False) if column doesn't exist in any known table
+        - (None, False) if column exists in multiple tables (still ambiguous)
+    """
+    column_lower = column.lower()
+    matching_tables = []
+
+    for table in candidate_tables:
+        table_lower = table.lower()
+        # Try both exact table name and lowercase version
+        if table in warehouse_columns:
+            cols = warehouse_columns[table]
+        elif table_lower in warehouse_columns:
+            cols = warehouse_columns[table_lower]
+        else:
+            continue
+
+        col_names_lower = {c.lower() for c in cols}
+        if column_lower in col_names_lower:
+            matching_tables.append(table)
+
+    if len(matching_tables) == 1:
+        return matching_tables[0], True
+    return None, False
+
+
+def qualify_sql_with_schema(
+    sql: str,
+    dialect: str,
+    warehouse_columns: WarehouseColumns,
+) -> tuple["exp.Expression | None", bool]:
+    """Qualify unqualified columns in SQL using warehouse schema.
+
+    Uses SQLGlot's optimizer.qualify to add table qualifiers to columns.
+    This resolves ambiguous column references like SELECT col FROM a, b
+    into SELECT a.col FROM a, b when 'col' only exists in table 'a'.
+
+    Returns the qualified AST directly so the caller doesn't re-parse.
+    The caller should use this AST instead of re-parsing a SQL string.
+
+    Args:
+        sql: SQL to qualify
+        dialect: SQL dialect
+        warehouse_columns: Schema info from warehouse
+
+    Returns:
+        Tuple of (qualified_ast, was_qualified)
+        - Returns (AST, True) if qualification succeeded
+        - Returns (None, False) if qualification failed
+    """
+    try:
+        import sqlglot
+        from sqlglot.optimizer.qualify import qualify
+    except Exception as e:
+        logger.debug(f"SQLGlot import failed: {e}")
+        return None, False
+
+    schema_dict = build_sqlglot_schema(warehouse_columns)
+    try:
+        parsed = sqlglot.parse_one(sql, read=dialect)
+        qualified = qualify(parsed, schema=schema_dict, dialect=dialect)
+        return qualified, True
+    except Exception as e:
+        logger.debug(f"Strict qualification failed, retrying best-effort: {e}")
+    # Partial parent schemas are the norm (star-passthrough parents): qualify
+    # what the schema resolves, leave the rest as written instead of throwing
+    # away every resolved column with it.
+    try:
+        parsed = sqlglot.parse_one(sql, read=dialect)
+        qualified = qualify(
+            parsed,
+            schema=schema_dict,
+            dialect=dialect,
+            validate_qualify_columns=False,
+            allow_partial_qualification=True,
+            infer_schema=True,
+        )
+        return qualified, True
+    except Exception as e:
+        logger.debug(f"SQLGlot qualification failed, using heuristics: {e}")
+        return None, False
+
+
+# Underscore-prefixed aliases
+_build_sqlglot_schema = build_sqlglot_schema
+_resolve_column_table = resolve_column_table
+_qualify_sql_with_schema = qualify_sql_with_schema