ripple-sql 0.1.0__py3-none-any.whl

ripple/engine/scope.py

@@ -0,0 +1,488 @@
+"""FROM-scope registration for a single statement.
+
+Builds, for the select that owns the FROM clause:
+- alias_map: alias -> actual table name
+- all_tables: every relation the scope reads
+- array_expansion_sources: FLATTEN/UNNEST/EXPLODE alias -> source array info
+
+Covers plain FROM tables, JOINed tables, LATERAL FLATTEN/EXPLODE (including
+Spark's LATERAL VIEW slot), BigQuery UNNEST, and the JOIN-side variants
+(joined subqueries, CROSS JOIN LATERAL, CROSS JOIN UNNEST AS t(item)).
+"""
+
+from sqlglot import exp
+
+from ripple.engine.column_ref import (
+    extract_array_source as _extract_array_source,
+)
+from ripple.engine.column_ref import (
+    owning_select,
+    qualified_table_name,
+    resolve_expansion_extras,
+    resolve_field_sources,
+    table_function_relation_name,
+)
+from ripple.engine.column_ref import table_function_kind as _table_function_kind
+
+# TODO: the UNNEST sweep below near-duplicates column_ref.collect_unnest_expansions but also mutates alias_map/all_tables; merging needs judgment.
+
+
+def register_table(table: "exp.Table", alias_map: dict[str, str], all_tables: list[str]) -> None:
+    """Enter a FROM/JOIN table into scope under every name a column reference
+    could anchor to, while keeping the value the SQL actually wrote.
+
+    A Table wrapping a function call (CROSS JOIN jsonb_array_elements_text(x)
+    AS e) is a set-returning function, not a relation: registering it
+    polluted the scope with an empty-named table and broke unqualified
+    resolution (kingfisher-summarize, holdout round 7). Those are handled
+    by the expansion machinery.
+
+    The value keeps its qualifiers (pg_catalog.pg_attribute, not pg_attribute):
+    stripping them here is what let a model named after the table it wraps
+    swallow the external relation into itself and lose all lineage
+    (dbt-labs/redshift, holdout round 2). The bare name stays a lookup KEY,
+    because column references anchor by the name part (pg_attribute.attname)
+    even when the FROM clause is qualified.
+    """
+    kind = _table_function_kind(table)
+    if kind == "srf":
+        return  # its alias is an expansion, handled by the expansion machinery
+    if kind == "relation":
+        # opaque table function: a real unknown-schema relation. A
+        # qualified one (sys.fn_xe_file_target_read_file) is a catalog
+        # object cited by its written name; a bare one (read_parquet) is
+        # known only by its alias. Either way it must keep blocking
+        # unique-ownership claims (the cycle-8 review)
+        written = table_function_relation_name(table)
+        name = written or table.alias
+        if not name:
+            return
+        if table.alias:
+            alias_map.setdefault(table.alias, name)
+        if written:
+            alias_map.setdefault(written, written)
+        all_tables.append(name)
+        return
+    qualified = qualified_table_name(table)
+    if table.alias:
+        alias_map[table.alias] = qualified
+    else:
+        claimed = alias_map.get(table.name)
+        if claimed is not None and claimed not in (qualified, table.name):
+            # two qualified relations share this bare name in one scope
+            # (raw.orders JOIN staging.orders): neither may own the bare key,
+            # so it degrades to the ambiguous merged name and resolves at
+            # review trust instead of silently binding the later table
+            alias_map[table.name] = table.name
+        else:
+            alias_map[table.name] = qualified
+        if qualified != table.name:
+            alias_map[qualified] = qualified
+    all_tables.append(qualified)
+
+
+def _fold_struct_parts(source_info: dict, alias_map: dict[str, str]) -> dict:
+    """Decide which identifiers of a dotted expansion argument name the table.
+
+    extract_array_source sees a.b.c.d without scope knowledge. The longest
+    dotted prefix that names a relation in THIS scope (case-insensitively,
+    since SQL identifiers fold) is the table; the next part is the column
+    and the remainder is the path inside it. When no prefix matches, the
+    first part is the root column; reading it any other way invented a
+    `histogram.bin` relation out of
+    UNNEST(first_contentful_paint.histogram.bin) (crux, holdout round 4),
+    and a fully qualified `proj.ds.events`.items prefix must fold to items,
+    not to a phantom `proj` column (review).
+    """
+    parts = source_info.get("struct_parts")
+    if not parts:
+        return source_info
+    lower_keys = {k.lower(): k for k in alias_map}
+    folded = None
+    for i in range(len(parts) - 1, 0, -1):
+        prefix = ".".join(parts[:i]).lower()
+        if prefix in lower_keys:
+            folded = {"source_table": lower_keys[prefix], "source_column": parts[i]}
+            rest = parts[i + 1 :]
+            break
+    if folded is None:
+        folded = {"source_table": "", "source_column": parts[0]}
+        rest = parts[1:]
+    if rest:
+        folded["json_path"] = ".".join(rest)
+    return folded
+
+
+def _register_function_expansion(
+    lateral_alias: str,
+    func_args: "list[exp.Expression]",
+    alias_map: dict[str, str],
+    all_tables: list[str],
+    array_expansion_sources: dict[str, dict[str, str]],
+    alias_columns: list[str] | None = None,
+) -> None:
+    """A LATERAL over a set-returning function (jsonb_array_elements,
+    regexp_split_to_table, ...) expands its argument columns; every
+    reference to the alias (or to its declared output columns) derives
+    from them. FAC's findings_text scored three empty cases because this
+    shape registered nothing (holdout round 4). the review added: chain
+    through a prior expansion when the argument is itself an expansion
+    alias, keep EVERY argument (regexp_split_to_table reads two), and
+    register AS f(col) output-column names."""
+    if not lateral_alias or not func_args:
+        return
+    source_info = _extract_array_source(func_args[0])
+    if not source_info:
+        return
+    source_info = _fold_struct_parts(source_info, alias_map)
+    src_table = source_info.get("source_table", "")
+    src_col = source_info.get("source_column", "")
+    prior = array_expansion_sources.get(src_table) or (
+        array_expansion_sources.get(src_col) if not src_table else None
+    )
+    if prior is not None:
+        # nested lateral: the argument is a previous expansion's output, so
+        # the true source is that expansion's own source
+        path_bits = [p for p in (prior.get("json_path"), source_info.get("json_path")) if p]
+        entry = {
+            "source_table": prior.get("source_table", ""),
+            "source_column": prior.get("source_column", ""),
+            "expansion_type": "LATERAL",
+        }
+        if path_bits:
+            entry["json_path"] = ".".join(path_bits)
+        resolved_table = entry["source_table"]
+    else:
+        resolved_table = alias_map.get(src_table, src_table)
+        if not resolved_table:
+            real = [t for t in all_tables if not t.startswith("(")]
+            if len(real) == 1:
+                resolved_table = real[0]
+        entry = {
+            "source_table": resolved_table,
+            "source_column": src_col,
+            "expansion_type": "LATERAL",
+        }
+        if source_info.get("json_path"):
+            entry["json_path"] = source_info["json_path"]
+    extras = []
+    for arg in func_args[1:]:
+        extra_info = _extract_array_source(arg)
+        if not extra_info:
+            continue
+        extra_info = _fold_struct_parts(extra_info, alias_map)
+        if extra_info.get("source_column"):
+            extras.append(
+                {
+                    "source_table": extra_info.get("source_table", ""),
+                    "source_column": extra_info["source_column"],
+                }
+            )
+    if extras:
+        entry["extra_columns"] = resolve_expansion_extras(
+            {"extra_columns": extras}, alias_map, all_tables
+        )
+    array_expansion_sources[lateral_alias] = entry
+    for col_name in alias_columns or []:
+        # AS f(elem): the declared output column is how the value is read
+        array_expansion_sources.setdefault(col_name, {**entry, "is_unnest_column": True})
+    alias_map[lateral_alias] = resolved_table
+    if resolved_table and resolved_table not in all_tables:
+        all_tables.append(resolved_table)
+
+
+def _lateral_function_args(inner: "exp.Expression") -> "list[exp.Expression]":
+    """The arguments of a lateral set-returning function call, if that is
+    what the LATERAL wraps (subqueries and explode-family nodes are handled
+    by their own branches)."""
+    if isinstance(inner, (exp.Subquery, exp.Select, exp.Explode, exp.Unnest)):
+        return []
+    if isinstance(inner, exp.Func):
+        if getattr(inner, "expressions", None):
+            return list(inner.expressions)
+        this = inner.args.get("this")
+        if isinstance(this, exp.Expression):
+            return [this]
+    return []
+
+
+def _alias_column_names(alias_node) -> list[str]:
+    """Output-column names from a TableAlias like f(elem, idx)."""
+    if alias_node is None:
+        return []
+    return [c.name for c in getattr(alias_node, "columns", []) or [] if getattr(c, "name", "")]
+
+
+def build_scope(
+    select_for_from: "exp.Expression",
+) -> tuple[dict[str, str], list[str], dict[str, dict[str, str]]]:
+    """Register this select's FROM/JOIN/LATERAL/UNNEST relations."""
+    alias_map: dict[str, str] = {}
+    all_tables: list[str] = []
+
+    # Track array expansion sources (FLATTEN, UNNEST, EXPLODE)
+    # Maps alias -> {source_table, source_column, expansion_type}
+    array_expansion_sources: dict[str, dict[str, str]] = {}
+
+    from_clause = select_for_from.args.get("from_") or select_for_from.args.get("from")
+    if from_clause:
+        for table in from_clause.find_all(exp.Table):
+            register_table(table, alias_map, all_tables)
+
+    # Joins inside CTE bodies or nested subqueries belong to their own
+    # select; letting them in shadows this scope's aliases.
+    for join in select_for_from.find_all(exp.Join):
+        if owning_select(join) is not select_for_from:
+            continue
+        join_table = join.this
+        if isinstance(join_table, exp.Table):
+            register_table(join_table, alias_map, all_tables)
+
+    # Set-returning functions in table position (CROSS JOIN
+    # jsonb_array_elements_text(roles) AS role, or the FROM-comma form):
+    # register_table skips them; their alias is an expansion of the
+    # argument (kingfisher-summarize, holdout round 7). Runs after the
+    # relation loops so argument aliases resolve.
+    for table in select_for_from.find_all(exp.Table):
+        if owning_select(table) is not select_for_from:
+            continue
+        if _table_function_kind(table) != "srf":
+            continue
+        func = table.this
+        func_args = list(getattr(func, "expressions", None) or [])
+        if not func_args and isinstance(func.args.get("this"), exp.Expression):
+            func_args = [func.args["this"]]
+        alias = table.alias
+        if alias and func_args:
+            _register_function_expansion(
+                alias,
+                func_args,
+                alias_map,
+                all_tables,
+                array_expansion_sources,
+                alias_columns=_alias_column_names(table.args.get("alias")),
+            )
+
+    # Runs after alias_map is populated so source table aliases resolve.
+    # Spark/Databricks LATERAL VIEW lives in select.args["laterals"], not
+    # under FROM, so both locations are swept.
+    laterals: list[exp.Lateral] = []
+    if from_clause:
+        laterals.extend(
+            lat
+            for lat in from_clause.find_all(exp.Lateral)
+            if owning_select(lat) is select_for_from
+        )
+    laterals.extend(select_for_from.args.get("laterals") or [])
+    if laterals:
+        for lateral in laterals:
+            lateral_alias = lateral.alias
+            inner = lateral.this
+
+            # FLATTEN/EXPLODE/POSEXPLODE parse as Lateral containing
+            # Explode; OUTER and INLINE variants are separate classes
+            expansion_types = (exp.Explode,)
+            for extra in ("ExplodeOuter", "Inline"):
+                if hasattr(exp, extra):
+                    expansion_types = (*expansion_types, getattr(exp, extra))
+            if isinstance(inner, expansion_types):
+                source_col = inner.this
+                if source_col and lateral_alias:
+                    source_info = _extract_array_source(source_col)
+                    if source_info:
+                        source_info = _fold_struct_parts(source_info, alias_map)
+                        # Resolve source table through alias_map; a source
+                        # that is itself an expansion alias chains through
+                        src_table = source_info.get("source_table", "")
+                        if src_table in array_expansion_sources:
+                            resolved_table = array_expansion_sources[src_table].get(
+                                "source_table", ""
+                            )
+                        else:
+                            resolved_table = alias_map.get(src_table, src_table)
+                        if not resolved_table:
+                            real = [t for t in all_tables if not t.startswith("(")]
+                            if len(real) == 1:
+                                resolved_table = real[0]
+                        expansion_entry = {
+                            "source_table": resolved_table,
+                            "source_column": source_info.get("source_column", ""),
+                            "expansion_type": "FLATTEN",
+                        }
+                        if source_info.get("json_path"):
+                            expansion_entry["json_path"] = source_info["json_path"]
+                        array_expansion_sources[lateral_alias] = expansion_entry
+                        # Spark allows the exploded output columns bare
+                        # (item, not just ex.item): register each alias
+                        # column name too
+                        table_alias = lateral.args.get("alias")
+                        for alias_col in table_alias.columns if table_alias else []:
+                            array_expansion_sources.setdefault(
+                                alias_col.name,
+                                {**expansion_entry, "is_unnest_column": True},
+                            )
+                        alias_map[lateral_alias] = resolved_table
+                        if resolved_table and resolved_table not in all_tables:
+                            all_tables.append(resolved_table)
+            else:
+                func_args = _lateral_function_args(inner)
+                if func_args:
+                    _register_function_expansion(
+                        lateral_alias,
+                        func_args,
+                        alias_map,
+                        all_tables,
+                        array_expansion_sources,
+                        alias_columns=_alias_column_names(lateral.args.get("alias")),
+                    )
+
+    # BigQuery UNNEST, scoped to this select's own FROM:
+    # an UNNEST inside a select-item subquery anchors only that
+    # subquery's references (resolved per column below), never sibling
+    # select items. Anonymous UNNESTs anchor unknown struct-root
+    # references to their source array.
+    for unnest in select_for_from.find_all(exp.Unnest):
+        if owning_select(unnest) is not select_for_from:
+            continue
+        unnest_alias = unnest.alias
+        if unnest.expressions:
+            source_col = unnest.expressions[0]
+            source_info = _extract_array_source(source_col)
+            if source_info:
+                source_info = _fold_struct_parts(source_info, alias_map)
+                src_table = source_info.get("source_table", "")
+                resolved_table = alias_map.get(src_table, src_table)
+                # a literal-only named-struct array reads nothing from the
+                # outer relation; anchoring it to the sole table fabricated
+                # struct-guess edges (cycle-12, F17)
+                if not resolved_table and source_info.get("source_column"):
+                    real = [t for t in all_tables if not t.startswith("(")]
+                    if len(real) == 1:
+                        resolved_table = real[0]
+                entry = {
+                    "source_table": resolved_table,
+                    "source_column": source_info.get("source_column", ""),
+                    "expansion_type": "UNNEST",
+                }
+                extras = resolve_expansion_extras(source_info, alias_map, all_tables)
+                if extras:
+                    entry["extra_columns"] = extras
+                field_map = resolve_field_sources(source_info, alias_map, all_tables)
+                if field_map:
+                    entry["field_sources"] = field_map
+                if unnest_alias:
+                    array_expansion_sources[unnest_alias] = entry
+                    alias_map[unnest_alias] = resolved_table
+                else:
+                    array_expansion_sources.setdefault("__anonymous_unnest__", entry)
+                if resolved_table and resolved_table not in all_tables:
+                    all_tables.append(resolved_table)
+
+    # plain joined tables were registered above; only the exotic join shapes remain
+    for join in select_for_from.find_all(exp.Join):
+        if owning_select(join) is not select_for_from:
+            continue
+        join_table = join.this
+        if isinstance(join_table, exp.Subquery) and join_table.alias:
+            alias_map[join_table.alias] = f"({join_table.alias})"
+            all_tables.append(f"({join_table.alias})")
+        elif isinstance(join_table, exp.Lateral):
+            lateral_alias = join_table.alias
+            inner = join_table.this
+            if isinstance(inner, exp.Explode) and lateral_alias:
+                source_col = inner.this
+                source_info = _extract_array_source(source_col)
+                if source_info:
+                    source_info = _fold_struct_parts(source_info, alias_map)
+                    # Resolve source table through alias_map
+                    src_table = source_info.get("source_table", "")
+                    resolved_table = alias_map.get(src_table, src_table)
+                    array_expansion_sources[lateral_alias] = {
+                        "source_table": resolved_table,
+                        "source_column": source_info.get("source_column", ""),
+                        "expansion_type": "FLATTEN",
+                    }
+                    if source_info.get("json_path"):
+                        array_expansion_sources[lateral_alias]["json_path"] = source_info[
+                            "json_path"
+                        ]
+                    alias_map[lateral_alias] = resolved_table
+                    if resolved_table and resolved_table not in all_tables:
+                        all_tables.append(resolved_table)
+            elif lateral_alias:
+                func_args = _lateral_function_args(inner)
+                if func_args:
+                    _register_function_expansion(
+                        lateral_alias,
+                        func_args,
+                        alias_map,
+                        all_tables,
+                        array_expansion_sources,
+                        alias_columns=_alias_column_names(join_table.args.get("alias")),
+                    )
+        elif isinstance(join_table, exp.Unnest):
+            unnest_alias_node = join_table.args.get("alias")
+            if unnest_alias_node:
+                if hasattr(unnest_alias_node, "this"):
+                    unnest_alias = (
+                        unnest_alias_node.this.name
+                        if hasattr(unnest_alias_node.this, "name")
+                        else str(unnest_alias_node.this)
+                    )
+                else:
+                    unnest_alias = str(unnest_alias_node)
+
+                # UNNEST(...) AS t(item): t(item) declares the output columns
+                column_aliases = []
+                if hasattr(unnest_alias_node, "columns"):
+                    column_aliases = [
+                        c.name if hasattr(c, "name") else str(c) for c in unnest_alias_node.columns
+                    ]
+
+                if join_table.expressions:
+                    source_col = join_table.expressions[0]
+                    source_info = _extract_array_source(source_col)
+                    if source_info:
+                        source_info = _fold_struct_parts(source_info, alias_map)
+                        # Resolve the source table through alias_map; the
+                        # sole-table fallback is skipped for literal-only
+                        # struct arrays, same reason as above (F17)
+                        src_table = source_info.get("source_table", "")
+                        resolved_table = alias_map.get(src_table, src_table)
+                        if not resolved_table and source_info.get("source_column"):
+                            real = [t for t in all_tables if not t.startswith("(")]
+                            if len(real) == 1:
+                                resolved_table = real[0]
+
+                        join_extras = resolve_expansion_extras(source_info, alias_map, all_tables)
+                        join_fields = resolve_field_sources(source_info, alias_map, all_tables)
+                        array_expansion_sources[unnest_alias] = {
+                            "source_table": resolved_table,
+                            "source_column": source_info.get("source_column", ""),
+                            "expansion_type": "UNNEST",
+                            "column_aliases": column_aliases,
+                        }
+                        if join_extras:
+                            array_expansion_sources[unnest_alias]["extra_columns"] = join_extras
+                        if join_fields:
+                            array_expansion_sources[unnest_alias]["field_sources"] = join_fields
+                        if source_info.get("json_path"):
+                            array_expansion_sources[unnest_alias]["json_path"] = source_info[
+                                "json_path"
+                            ]
+
+                        alias_map[unnest_alias] = resolved_table
+                        # Also map each column alias to the source
+                        for col_alias in column_aliases:
+                            array_expansion_sources[col_alias] = {
+                                "source_table": resolved_table,
+                                "source_column": source_info.get("source_column", ""),
+                                "expansion_type": "UNNEST",
+                                "is_unnest_column": True,
+                            }
+                            if join_extras:
+                                array_expansion_sources[col_alias]["extra_columns"] = join_extras
+                            if join_fields:
+                                array_expansion_sources[col_alias]["field_sources"] = join_fields
+
+    return alias_map, all_tables, array_expansion_sources