bqsqlparse 0.1.0__py3-none-any.whl

bqsqlparse/lineage.py

@@ -0,0 +1,902 @@
+"""Column (attribute) level lineage extraction for BigQuery SQL.
+
+Given a statement AST (or SQL text), this module traces every output column
+back to the physical source table columns it derives from — through CTEs,
+subqueries, joins, set operations, UNNEST (including struct field access),
+star expansion, PIVOT/UNPIVOT, and DML (INSERT / UPDATE / MERGE / CTAS).
+
+Usage::
+
+    from bqsqlparse import extract_lineage
+
+    result = extract_lineage(sql, schema={"proj.ds.orders": ["id", "amount"]})
+    for col in result.columns:
+        print(col.name, "<-", [f"{s.table}.{s.column}" for s in col.sources])
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass, field as dc_field
+from typing import Dict, List, Optional, Sequence, Set, Tuple, Union
+
+from . import nodes as n
+from .errors import LineageError
+from .functions import is_aggregate, is_navigation
+from .parser import parse, parse_one
+
+# Transformation kinds
+IDENTITY = "IDENTITY"
+EXPRESSION = "EXPRESSION"
+AGGREGATION = "AGGREGATION"
+WINDOW = "WINDOW"
+CONSTANT = "CONSTANT"
+STAR = "STAR"
+OFFSET = "OFFSET"
+
+
+@dataclass(frozen=True, order=True)
+class SourceColumn:
+    """A physical source attribute: ``table`` is the fully-qualified table
+    name; ``column`` may include struct field paths (``payload.user.id``) or
+    be ``*`` when a star could not be expanded without a schema."""
+
+    table: str
+    column: str
+
+    def __str__(self) -> str:
+        return f"{self.table}.{self.column}"
+
+
+@dataclass
+class ColumnUsage:
+    """Where a physical source column is used across the statement.
+
+    ``contexts`` is a sorted list drawn from: SELECT, JOIN, WHERE, GROUP_BY,
+    HAVING, QUALIFY, ORDER_BY, UNNEST, PIVOT, TABLE_FUNCTION, SET, INSERT.
+    """
+
+    table: str
+    column: str
+    contexts: List[str]
+
+    def to_dict(self) -> dict:
+        return {"table": self.table, "column": self.column, "contexts": self.contexts}
+
+
+@dataclass
+class ColumnLineage:
+    name: str
+    expression: str
+    transformation: str
+    sources: List[SourceColumn] = dc_field(default_factory=list)
+    indirect_sources: List[SourceColumn] = dc_field(default_factory=list)
+
+    def to_dict(self) -> dict:
+        return {
+            "name": self.name,
+            "expression": self.expression,
+            "transformation": self.transformation,
+            "sources": [{"table": s.table, "column": s.column} for s in self.sources],
+            "indirect_sources": [
+                {"table": s.table, "column": s.column} for s in self.indirect_sources
+            ],
+        }
+
+
+@dataclass
+class LineageResult:
+    target: Optional[str]
+    columns: List[ColumnLineage]
+    tables: List[str]
+    ctes: List[str] = dc_field(default_factory=list)
+    usage: List[ColumnUsage] = dc_field(default_factory=list)
+
+    @property
+    def edges(self) -> List[Tuple[str, str]]:
+        """(source ``table.column``, target ``target.column``) pairs."""
+        tgt = self.target or "<query>"
+        out = []
+        for col in self.columns:
+            for s in col.sources:
+                out.append((str(s), f"{tgt}.{col.name}"))
+        return out
+
+    def to_dict(self) -> dict:
+        return {
+            "target": self.target,
+            "tables": self.tables,
+            "ctes": self.ctes,
+            "columns": [c.to_dict() for c in self.columns],
+            "usage": [u.to_dict() for u in self.usage],
+        }
+
+    def to_json(self, **kwargs) -> str:
+        kwargs.setdefault("indent", 2)
+        return json.dumps(self.to_dict(), **kwargs)
+
+
+# ---------------------------------------------------------------------------
+# Internal model
+# ---------------------------------------------------------------------------
+
+@dataclass
+class _Out:
+    name: str
+    sources: Set[SourceColumn]
+    transformation: str
+    expression: str = ""
+    indirect: Set[SourceColumn] = dc_field(default_factory=set)
+
+
+@dataclass
+class _Derived:
+    """The resolved output of a query/subquery/CTE."""
+
+    outputs: "Dict[str, _Out]"  # ordered; keys lowercased
+    order: List[str]
+    star_sources: List[SourceColumn]  # unexpanded pass-through columns
+
+    def get(self, name: str) -> Optional[_Out]:
+        return self.outputs.get(name.lower())
+
+
+@dataclass
+class _Rel:
+    kind: str  # 'table' | 'derived' | 'unnest'
+    alias: Optional[str] = None
+    table_name: Optional[str] = None
+    schema_cols: Optional[List[str]] = None
+    derived: Optional[_Derived] = None
+    unnest_sources: Set[SourceColumn] = dc_field(default_factory=set)
+    offset_alias: Optional[str] = None
+
+    def names(self) -> List[str]:
+        """Qualifier names this relation answers to (lowercased)."""
+        out = []
+        if self.alias:
+            out.append(self.alias.lower())
+        elif self.table_name:
+            out.append(self.table_name.lower())
+            last = self.table_name.split(".")[-1].lower()
+            if last != self.table_name.lower():
+                out.append(last)
+        return out
+
+    def has_column(self, name: str) -> Optional[bool]:
+        """True/False if known, None if unknown schema."""
+        name = name.lower()
+        if self.kind == "table":
+            if self.schema_cols is None:
+                return None
+            return name in (c.lower() for c in self.schema_cols)
+        if self.kind == "derived":
+            if self.derived.get(name) is not None:
+                return True
+            return None if self.derived.star_sources else False
+        if self.kind == "unnest":
+            return name in ((self.alias or "").lower(), (self.offset_alias or "").lower())
+        return None
+
+    def resolve(self, rest: Sequence[str]) -> Set[SourceColumn]:
+        """Resolve a column (+ optional struct field path) inside this rel."""
+        if self.kind == "table":
+            return {SourceColumn(self.table_name, ".".join(rest))}
+        if self.kind == "unnest":
+            if rest and self.offset_alias and rest[0].lower() == self.offset_alias.lower():
+                return set()
+            suffix = list(rest)
+            if suffix and self.alias and suffix[0].lower() == self.alias.lower():
+                suffix = suffix[1:]
+            if not suffix:
+                return set(self.unnest_sources)
+            return {
+                SourceColumn(s.table, f"{s.column}.{'.'.join(suffix)}" if s.column != "*" else ".".join(suffix))
+                for s in self.unnest_sources
+            }
+        # derived
+        out = self.derived.get(rest[0])
+        if out is not None:
+            if len(rest) == 1:
+                return set(out.sources)
+            suffix = ".".join(rest[1:])
+            return {
+                SourceColumn(s.table, f"{s.column}.{suffix}" if s.column != "*" else suffix)
+                for s in out.sources
+            }
+        # pass through unexpanded star
+        return {
+            SourceColumn(s.table, ".".join(rest))
+            for s in self.derived.star_sources
+        }
+
+
+class _Scope:
+    def __init__(self, rels: List[_Rel], parent: Optional["_Scope"] = None,
+                 variables: Optional[Set[str]] = None):
+        self.rels = rels
+        self.parent = parent
+        self.variables = variables if variables is not None else (
+            parent.variables if parent is not None else set())
+
+    def resolve_path(self, path: Sequence[str]) -> Set[SourceColumn]:
+        # 1. qualified: longest-prefix match against aliases / table names
+        for k in range(min(len(path) - 1, 3), 0, -1):
+            qual = ".".join(path[:k]).lower()
+            for rel in self.rels:
+                if qual in rel.names():
+                    return rel.resolve(path[k:])
+        # unnest element referenced by its own alias with no remainder
+        for rel in self.rels:
+            if rel.kind == "unnest":
+                if rel.alias and path[0].lower() == rel.alias.lower():
+                    return rel.resolve(path)
+                if rel.offset_alias and path[0].lower() == rel.offset_alias.lower():
+                    return set()
+        # 2. unqualified
+        name = path[0]
+        claiming = [r for r in self.rels if r.has_column(name) is True]
+        if len(claiming) >= 1:
+            results: Set[SourceColumn] = set()
+            for rel in claiming:
+                results |= rel.resolve(path)
+            return results
+        # script variables shadow unknown-schema columns
+        if name.lower() in self.variables:
+            return set()
+        unknown = [r for r in self.rels if r.has_column(name) is None]
+        if unknown:
+            results = set()
+            for rel in unknown:
+                results |= rel.resolve(path)
+            return results
+        if self.parent is not None:
+            return self.parent.resolve_path(path)
+        return set()
+
+
+# ---------------------------------------------------------------------------
+# Analyzer
+# ---------------------------------------------------------------------------
+
+def _flatten_path(expr: n.Node) -> Optional[List[str]]:
+    if isinstance(expr, n.ColumnRef):
+        return list(expr.path)
+    if isinstance(expr, n.FieldAccess):
+        base = _flatten_path(expr.base)
+        if base is not None:
+            return base + [expr.field_name]
+    return None
+
+
+class _Analyzer:
+    def __init__(self, schema: Optional[Dict[str, List[str]]] = None,
+                 include_indirect: bool = False,
+                 variables: Optional[Set[str]] = None):
+        self.schema = { (k.replace("`", "").lower()): v for k, v in (schema or {}).items() }
+        self.include_indirect = include_indirect
+        self.variables = {v.lower() for v in (variables or set())}
+        self.tables: Set[str] = set()
+        self.cte_names: Set[str] = set()
+        self.usage: Dict[SourceColumn, Set[str]] = {}
+
+    # -------------------------------------------------------------- resolve
+
+    def _collect(self, expr: n.Node, scope: _Scope, out: Set[SourceColumn]):
+        path = _flatten_path(expr)
+        if path is not None:
+            out |= scope.resolve_path(path)
+            return
+        if isinstance(expr, (n.ScalarSubquery, n.ArraySubquery, n.ExistsSubquery)):
+            derived = self.analyze_query(expr.query, {}, scope)
+            for o in derived.outputs.values():
+                out |= o.sources
+            for s in derived.star_sources:
+                out.add(s)
+            return
+        if isinstance(expr, n.Star):
+            return
+        if isinstance(expr, n.Node):
+            for child in expr.children():
+                self._collect(child, scope, out)
+
+    def _tag(self, srcs: Set[SourceColumn], context: str):
+        for s in srcs:
+            self.usage.setdefault(s, set()).add(context)
+
+    def _use(self, expr: n.Node, scope: _Scope, context: str) -> Set[SourceColumn]:
+        """Collect sources of *expr* and record which clause used them."""
+        srcs: Set[SourceColumn] = set()
+        self._collect(expr, scope, srcs)
+        self._tag(srcs, context)
+        return srcs
+
+    def _classify(self, expr: n.Node, sources: Set[SourceColumn]) -> str:
+        has_window = False
+        has_agg = False
+        for node in expr.walk():
+            if isinstance(node, n.FuncCall):
+                if node.over is not None:
+                    has_window = True
+                if is_aggregate(node.name_str) or is_navigation(node.name_str):
+                    if node.over is not None:
+                        has_window = True
+                    elif is_aggregate(node.name_str):
+                        has_agg = True
+        if has_window:
+            return WINDOW
+        if has_agg:
+            return AGGREGATION
+        if _flatten_path(expr) is not None:
+            return IDENTITY
+        if not sources:
+            return CONSTANT
+        return EXPRESSION
+
+    # ---------------------------------------------------------------- query
+
+    def analyze_query(self, query: n.Query, env: Dict[str, _Derived],
+                      parent: Optional[_Scope] = None) -> _Derived:
+        env = dict(env)
+        if query.recursive:
+            for cte in query.ctes:
+                env.setdefault(cte.name.lower(), _Derived({}, [], []))
+        for cte in query.ctes:
+            self.cte_names.add(cte.name)
+            env[cte.name.lower()] = self.analyze_query(cte.query, env, parent)
+        return self._analyze_body(query.body, env, parent, query.order_by)
+
+    def _tag_order_by_outputs(self, derived: _Derived,
+                              order_by: Sequence[n.OrderItem]):
+        """For set ops / nested queries, ORDER BY refers to output names."""
+        for oi in order_by:
+            path = _flatten_path(oi.expr)
+            if path is not None and len(path) == 1:
+                out = derived.get(path[0])
+                if out is not None:
+                    self._tag(out.sources, "ORDER_BY")
+
+    def _analyze_body(self, body: n.Node, env: Dict[str, _Derived],
+                      parent: Optional[_Scope],
+                      order_by: Sequence[n.OrderItem] = ()) -> _Derived:
+        if isinstance(body, n.Query):
+            derived = self.analyze_query(body, env, parent)
+            self._tag_order_by_outputs(derived, order_by)
+            return derived
+        if isinstance(body, n.SetOp):
+            left = self._analyze_body(body.left, env, parent)
+            right = self._analyze_body(body.right, env, parent)
+            outputs: Dict[str, _Out] = {}
+            order: List[str] = []
+            rnames = right.order
+            for idx, key in enumerate(left.order):
+                lo = left.outputs[key]
+                merged = _Out(lo.name, set(lo.sources), lo.transformation,
+                              lo.expression, set(lo.indirect))
+                if idx < len(rnames):
+                    ro = right.outputs[rnames[idx]]
+                    merged.sources |= ro.sources
+                    merged.indirect |= ro.indirect
+                outputs[key] = merged
+                order.append(key)
+            derived = _Derived(outputs, order,
+                               left.star_sources + right.star_sources)
+            self._tag_order_by_outputs(derived, order_by)
+            return derived
+        if isinstance(body, n.Select):
+            return self._analyze_select(body, env, parent, order_by)
+        raise LineageError(f"Unexpected query body: {type(body).__name__}")
+
+    # --------------------------------------------------------------- select
+
+    def _build_scope(self, from_: Optional[n.Node], env: Dict[str, _Derived],
+                     parent: Optional[_Scope],
+                     indirect_exprs: List[Tuple[str, n.Node]]) -> _Scope:
+        rels: List[_Rel] = []
+        scope = _Scope(rels, parent, self.variables)
+
+        def add(item: Optional[n.Node]):
+            if item is None:
+                return
+            if isinstance(item, n.Join):
+                add(item.left)
+                add(item.right)
+                if item.on is not None:
+                    indirect_exprs.append(("JOIN", item.on))
+                for col in item.using:
+                    indirect_exprs.append(("JOIN", n.ColumnRef([col])))
+                return
+            if isinstance(item, n.TableRef):
+                name = item.full_name.replace("`", "")
+                key = name.lower()
+                if len(item.path) == 1 and key in env:
+                    rels.append(_Rel("derived", alias=item.alias or item.path[0],
+                                     derived=env[key]))
+                else:
+                    self.tables.add(name)
+                    rels.append(_Rel("table", alias=item.alias, table_name=name,
+                                     schema_cols=self.schema.get(key)))
+                return
+            if isinstance(item, n.SubqueryRef):
+                derived = self.analyze_query(item.query, env, scope)
+                rels.append(_Rel("derived", alias=item.alias, derived=derived))
+                return
+            if isinstance(item, n.UnnestRef):
+                srcs = self._use(item.expr, scope, "UNNEST")
+                rels.append(_Rel("unnest", alias=item.alias,
+                                 unnest_sources=srcs,
+                                 offset_alias=item.offset_alias))
+                return
+            if isinstance(item, (n.PivotRef, n.UnpivotRef)):
+                inner_indirect: List[Tuple[str, n.Node]] = []
+                inner_scope = self._build_scope(item.input, env, parent, inner_indirect)
+                stars: List[SourceColumn] = []
+                for rel in inner_scope.rels:
+                    if rel.kind == "table":
+                        if rel.schema_cols:
+                            stars.extend(SourceColumn(rel.table_name, c)
+                                         for c in rel.schema_cols)
+                        else:
+                            stars.append(SourceColumn(rel.table_name, "*"))
+                    elif rel.kind == "derived":
+                        stars.extend(rel.derived.star_sources)
+                        for o in rel.derived.outputs.values():
+                            stars.extend(o.sources)
+                rels.append(_Rel("derived", alias=item.alias,
+                                 derived=_Derived({}, [], stars)))
+                if isinstance(item, n.PivotRef):
+                    indirect_exprs.append(("PIVOT", item.for_col))
+                return
+            if isinstance(item, n.TableFuncRef):
+                name = ".".join(item.name).replace("`", "")
+                self.tables.add(name + "()")
+                rels.append(_Rel("table", alias=item.alias,
+                                 table_name=name + "()"))
+                for a in item.args:
+                    indirect_exprs.append(("TABLE_FUNCTION", a))
+                return
+            raise LineageError(f"Unexpected FROM item: {type(item).__name__}")
+
+        add(from_)
+        return scope
+
+    def _expand_rel(self, rel: _Rel, star: n.Star, add_out, star_sources: List[SourceColumn]):
+        except_ = {e.lower() for e in star.except_}
+        replace = {r.name.lower(): r for r in star.replace}
+        if rel.kind == "table":
+            if rel.schema_cols is not None:
+                for col in rel.schema_cols:
+                    if col.lower() in except_:
+                        continue
+                    add_out(col, {SourceColumn(rel.table_name, col)}, IDENTITY, col)
+            else:
+                src = SourceColumn(rel.table_name, "*")
+                star_sources.append(src)
+                add_out("*", {src}, STAR, star.sql())
+        elif rel.kind == "derived":
+            for key in rel.derived.order:
+                o = rel.derived.outputs[key]
+                if key in except_:
+                    continue
+                add_out(o.name, set(o.sources), IDENTITY, o.name, set(o.indirect))
+            for s in rel.derived.star_sources:
+                star_sources.append(s)
+                add_out("*", {s}, STAR, star.sql())
+        elif rel.kind == "unnest":
+            name = rel.alias or "f0_"
+            if name.lower() not in except_:
+                add_out(name, set(rel.unnest_sources), IDENTITY, name)
+            if rel.offset_alias and rel.offset_alias.lower() not in except_:
+                add_out(rel.offset_alias, set(), OFFSET, "OFFSET")
+        # apply REPLACE overrides afterwards (handled by caller via replace map)
+        return replace
+
+    def _analyze_select(self, select: n.Select, env: Dict[str, _Derived],
+                        parent: Optional[_Scope],
+                        order_by: Sequence[n.OrderItem] = ()) -> _Derived:
+        indirect_exprs: List[Tuple[str, n.Node]] = []
+        scope = self._build_scope(select.from_, env, parent, indirect_exprs)
+
+        for ctx, clause in (("WHERE", select.where), ("HAVING", select.having),
+                            ("QUALIFY", select.qualify)):
+            if clause is not None:
+                indirect_exprs.append((ctx, clause))
+        if select.group_by is not None:
+            for e in select.group_by.exprs:
+                indirect_exprs.append(("GROUP_BY", e))
+            for s in select.group_by.sets:
+                for e in s:
+                    indirect_exprs.append(("GROUP_BY", e))
+
+        indirect: Set[SourceColumn] = set()
+        for ctx, e in indirect_exprs:
+            srcs = self._use(e, scope, ctx)
+            if self.include_indirect:
+                indirect |= srcs
+
+        outputs: Dict[str, _Out] = {}
+        order: List[str] = []
+        star_sources: List[SourceColumn] = []
+        anon = [0]
+
+        def add_out(name: str, sources: Set[SourceColumn], transformation: str,
+                    expression: str, extra_indirect: Optional[Set[SourceColumn]] = None):
+            key = name.lower()
+            if key in outputs and key != "*":
+                key = f"{key}#{len(order)}"
+            self._tag(sources, "SELECT")
+            o = _Out(name, sources, transformation, expression,
+                     set(indirect) | (extra_indirect or set()))
+            outputs[key] = o
+            order.append(key)
+
+        for i, item in enumerate(select.items):
+            if isinstance(item, n.Star):
+                target_rels = scope.rels
+                if item.prefix:
+                    qual = ".".join(item.prefix).lower()
+                    target_rels = [r for r in scope.rels if qual in r.names()]
+                    if not target_rels:
+                        # t.struct_col.* — expand a struct column
+                        srcs = scope.resolve_path(item.prefix)
+                        add_out(item.prefix[-1], srcs, EXPRESSION, item.sql())
+                        continue
+                replace_map = {r.name.lower(): r for r in item.replace}
+                for rel in target_rels:
+                    self._expand_rel(rel, item, add_out, star_sources)
+                for rname, ritem in replace_map.items():
+                    srcs: Set[SourceColumn] = set()
+                    self._collect(ritem.expr, scope, srcs)
+                    key = rname
+                    if key in outputs:
+                        o = outputs[key]
+                        o.sources = srcs
+                        o.transformation = self._classify(ritem.expr, srcs)
+                        o.expression = ritem.expr.sql()
+                    else:
+                        add_out(ritem.name, srcs,
+                                self._classify(ritem.expr, srcs), ritem.expr.sql())
+                continue
+
+            expr = item.expr
+            srcs = set()
+            self._collect(expr, scope, srcs)
+            name = item.alias
+            if not name:
+                path = _flatten_path(expr)
+                if path is not None:
+                    name = path[-1]
+                elif isinstance(expr, n.FuncCall):
+                    anon[0] += 1
+                    name = f"f{anon[0] - 1}_"
+                else:
+                    anon[0] += 1
+                    name = f"f{anon[0] - 1}_"
+            add_out(name, srcs, self._classify(expr, srcs), expr.sql())
+
+        # ORDER BY: prefer select-list aliases, else resolve in FROM scope
+        for oi in order_by:
+            path = _flatten_path(oi.expr)
+            if path is not None and len(path) == 1 and path[0].lower() in outputs:
+                self._tag(outputs[path[0].lower()].sources, "ORDER_BY")
+            else:
+                self._use(oi.expr, scope, "ORDER_BY")
+
+        return _Derived(outputs, order, star_sources)
+
+    # ----------------------------------------------------------- statements
+
+    def analyze_statement(self, stmt: n.Node) -> LineageResult:
+        if isinstance(stmt, n.Query):
+            derived = self.analyze_query(stmt, {})
+            return self._result(None, derived)
+        if isinstance(stmt, n.CreateTableAsSelect):
+            if stmt.query is None:
+                raise LineageError("CREATE statement has no AS SELECT query")
+            derived = self.analyze_query(stmt.query, {})
+            declared = [c.name for c in stmt.columns]
+            if declared:
+                derived = self._rename_positionally(derived, declared)
+            return self._result(stmt.full_name, derived)
+        if isinstance(stmt, n.InsertStmt):
+            if stmt.query is not None:
+                derived = self.analyze_query(stmt.query, {})
+            else:
+                outputs: Dict[str, _Out] = {}
+                order: List[str] = []
+                for row in stmt.values[:1]:
+                    for idx, v in enumerate(row):
+                        name = stmt.columns[idx] if idx < len(stmt.columns) else f"f{idx}_"
+                        srcs: Set[SourceColumn] = set()
+                        outputs[name.lower()] = _Out(
+                            name, srcs, CONSTANT, v.sql())
+                        order.append(name.lower())
+                derived = _Derived(outputs, order, [])
+            if stmt.columns:
+                derived = self._rename_positionally(derived, stmt.columns)
+            return self._result(stmt.full_name, derived)
+        if isinstance(stmt, n.UpdateStmt):
+            return self._analyze_update(stmt)
+        if isinstance(stmt, n.MergeStmt):
+            return self._analyze_merge(stmt)
+        if isinstance(stmt, n.DeleteStmt):
+            name = stmt.full_name.replace("`", "")
+            self.tables.add(name)
+            if stmt.where is not None:
+                rel = _Rel("table", alias=stmt.alias, table_name=name,
+                           schema_cols=self.schema.get(name.lower()))
+                self._use(stmt.where, _Scope([rel], variables=self.variables),
+                          "WHERE")
+            return LineageResult(name, [], sorted(self.tables),
+                                 ctes=sorted(self.cte_names),
+                                 usage=self._usage_list())
+        if isinstance(stmt, n.CreateFunction):
+            return self._analyze_create_function(stmt)
+        if isinstance(stmt, _SCRIPT_NODES):
+            results = _script_results(
+                [stmt], self.schema, self.include_indirect, self.variables)
+            if len(results) == 1:
+                return results[0]
+            if not results:
+                return LineageResult(None, [], [])
+            raise LineageError(
+                f"Script contains {len(results)} lineage-bearing statements; "
+                "use extract_script_lineage() to get one result per statement")
+        if isinstance(stmt, _NO_LINEAGE_NODES):
+            return LineageResult(None, [], [])
+        raise LineageError(
+            f"Lineage not supported for {type(stmt).__name__}")
+
+    def _analyze_create_function(self, stmt: n.CreateFunction) -> LineageResult:
+        params = {p.name.lower() for p in stmt.params}
+        inner = _Analyzer(self.schema, self.include_indirect,
+                          self.variables | params)
+        if stmt.body_query is not None:
+            derived = inner.analyze_query(stmt.body_query, {})
+            result = inner._result(stmt.full_name, derived)
+            return result
+        if stmt.body_expr is not None:
+            scope = _Scope([], variables=inner.variables)
+            srcs = inner._use(stmt.body_expr, scope, "SELECT")
+            col = ColumnLineage("<return>", stmt.body_expr.sql(),
+                                inner._classify(stmt.body_expr, srcs),
+                                sorted(srcs))
+            return LineageResult(stmt.full_name, [col], sorted(inner.tables),
+                                 usage=inner._usage_list())
+        return LineageResult(stmt.full_name, [], [])
+
+    def _analyze_update(self, stmt: n.UpdateStmt) -> LineageResult:
+        target = stmt.full_name.replace("`", "")
+        self.tables.add(target)
+        rels = [_Rel("table", alias=stmt.alias, table_name=target,
+                     schema_cols=self.schema.get(target.lower()))]
+        indirect_exprs: List[Tuple[str, n.Node]] = []
+        if stmt.from_ is not None:
+            from_scope = self._build_scope(stmt.from_, {}, None, indirect_exprs)
+            rels.extend(from_scope.rels)
+        scope = _Scope(rels, variables=self.variables)
+        indirect: Set[SourceColumn] = set()
+        if stmt.where is not None:
+            indirect_exprs.append(("WHERE", stmt.where))
+        for ctx, e in indirect_exprs:
+            srcs = self._use(e, scope, ctx)
+            if self.include_indirect:
+                indirect |= srcs
+        cols: List[ColumnLineage] = []
+        for a in stmt.assignments:
+            srcs = self._use(a.value, scope, "SET")
+            cols.append(ColumnLineage(
+                ".".join(a.target), a.value.sql(),
+                self._classify(a.value, srcs),
+                sorted(srcs), sorted(indirect)))
+        return LineageResult(target, cols, sorted(self.tables),
+                             ctes=sorted(self.cte_names),
+                             usage=self._usage_list())
+
+    def _analyze_merge(self, stmt: n.MergeStmt) -> LineageResult:
+        target = stmt.full_name.replace("`", "")
+        self.tables.add(target)
+        target_rel = _Rel("table", alias=stmt.alias, table_name=target,
+                          schema_cols=self.schema.get(target.lower()))
+        indirect_exprs: List[Tuple[str, n.Node]] = [("JOIN", stmt.on)]
+        src_scope = self._build_scope(stmt.source, {}, None, indirect_exprs)
+        scope = _Scope([target_rel] + src_scope.rels, variables=self.variables)
+        indirect: Set[SourceColumn] = set()
+        for ctx, e in indirect_exprs:
+            srcs = self._use(e, scope, ctx)
+            if self.include_indirect:
+                indirect |= srcs
+        cols: List[ColumnLineage] = []
+        for when in stmt.whens:
+            if when.condition is not None:
+                srcs = self._use(when.condition, scope, "WHERE")
+                if self.include_indirect:
+                    indirect |= srcs
+            if when.action_kind == "UPDATE":
+                for a in when.assignments:
+                    srcs = self._use(a.value, scope, "SET")
+                    cols.append(ColumnLineage(
+                        ".".join(a.target), a.value.sql(),
+                        self._classify(a.value, srcs),
+                        sorted(srcs), sorted(indirect)))
+            elif when.action_kind == "INSERT":
+                for idx, v in enumerate(when.values):
+                    name = when.columns[idx] if idx < len(when.columns) else f"f{idx}_"
+                    srcs = self._use(v, scope, "INSERT")
+                    cols.append(ColumnLineage(
+                        name, v.sql(), self._classify(v, srcs),
+                        sorted(srcs), sorted(indirect)))
+            elif when.action_kind == "INSERT_ROW":
+                for rel in src_scope.rels:
+                    if rel.kind == "table":
+                        if rel.schema_cols:
+                            for c in rel.schema_cols:
+                                cols.append(ColumnLineage(
+                                    c, c, IDENTITY,
+                                    [SourceColumn(rel.table_name, c)],
+                                    sorted(indirect)))
+                        else:
+                            cols.append(ColumnLineage(
+                                "*", "INSERT ROW", STAR,
+                                [SourceColumn(rel.table_name, "*")],
+                                sorted(indirect)))
+                    elif rel.kind == "derived":
+                        for key in rel.derived.order:
+                            o = rel.derived.outputs[key]
+                            cols.append(ColumnLineage(
+                                o.name, o.expression, o.transformation,
+                                sorted(o.sources), sorted(indirect)))
+        return LineageResult(target, cols, sorted(self.tables),
+                             ctes=sorted(self.cte_names),
+                             usage=self._usage_list())
+
+    # -------------------------------------------------------------- helpers
+
+    def _usage_list(self) -> List[ColumnUsage]:
+        return [
+            ColumnUsage(s.table, s.column, sorted(ctxs))
+            for s, ctxs in sorted(self.usage.items())
+        ]
+
+    @staticmethod
+    def _rename_positionally(derived: _Derived, names: List[str]) -> _Derived:
+        outputs: Dict[str, _Out] = {}
+        order: List[str] = []
+        for idx, key in enumerate(derived.order):
+            o = derived.outputs[key]
+            new_name = names[idx] if idx < len(names) else o.name
+            o.name = new_name
+            outputs[new_name.lower()] = o
+            order.append(new_name.lower())
+        return _Derived(outputs, order, derived.star_sources)
+
+    def _result(self, target: Optional[str], derived: _Derived) -> LineageResult:
+        cols = []
+        for key in derived.order:
+            o = derived.outputs[key]
+            cols.append(ColumnLineage(
+                o.name, o.expression, o.transformation,
+                sorted(o.sources), sorted(o.indirect)))
+        return LineageResult(
+            target.replace("`", "") if target else None,
+            cols, sorted(self.tables),
+            ctes=sorted(self.cte_names),
+            usage=self._usage_list())
+
+
+def extract_lineage(
+    sql_or_ast: Union[str, n.Node],
+    schema: Optional[Dict[str, List[str]]] = None,
+    include_indirect: bool = False,
+    variables: Optional[Sequence[str]] = None,
+) -> LineageResult:
+    """Extract column-level lineage from a BigQuery SQL statement.
+
+    Args:
+        sql_or_ast: SQL text or an already-parsed statement AST.
+        schema: optional mapping of fully-qualified table name to its column
+            names. Enables ``SELECT *`` expansion and unqualified-column
+            disambiguation across joins.
+        include_indirect: also report columns that influence each output
+            indirectly (WHERE / JOIN ON / GROUP BY / HAVING / QUALIFY).
+        variables: script variable names to exclude from column resolution.
+
+    Returns:
+        A :class:`LineageResult` with one :class:`ColumnLineage` per output
+        attribute, each holding its :class:`SourceColumn` ancestors.
+    """
+    stmt = parse_one(sql_or_ast) if isinstance(sql_or_ast, str) else sql_or_ast
+    vars_ = set(variables or ())
+    return _Analyzer(schema, include_indirect, vars_).analyze_statement(stmt)
+
+
+# ---------------------------------------------------------------------------
+# Script (procedural language) lineage
+# ---------------------------------------------------------------------------
+
+_SCRIPT_NODES = (
+    n.ScriptBlock, n.IfStmt, n.LoopStmt, n.WhileStmt, n.RepeatStmt,
+    n.ForInStmt, n.CreateProcedure,
+)
+_NO_LINEAGE_NODES = (
+    n.DeclareStmt, n.SetStmt, n.BreakContinueStmt, n.CallStmt, n.ReturnStmt,
+    n.RaiseStmt, n.ExecuteImmediate, n.AssertStmt, n.TransactionStmt,
+    n.TruncateStmt, n.DropStmt,
+)
+_LINEAGE_NODES = (
+    n.Query, n.CreateTableAsSelect, n.InsertStmt, n.UpdateStmt, n.DeleteStmt,
+    n.MergeStmt, n.CreateFunction,
+)
+
+
+def _collect_script_variables(stmts: List[n.Node]) -> Set[str]:
+    variables: Set[str] = set()
+    for stmt in stmts:
+        for node in stmt.walk():
+            if isinstance(node, n.DeclareStmt):
+                variables |= {name.lower() for name in node.names}
+            elif isinstance(node, n.ForInStmt):
+                variables.add(node.var.lower())
+            elif isinstance(node, n.ExecuteImmediate):
+                variables |= {name.lower() for name in node.into}
+    return variables
+
+
+def _script_results(stmts: List[n.Node], schema, include_indirect,
+                    extra_vars: Set[str]) -> List[LineageResult]:
+    variables = _collect_script_variables(stmts) | set(extra_vars)
+    results: List[LineageResult] = []
+
+    def visit(stmt: n.Node):
+        if isinstance(stmt, _LINEAGE_NODES):
+            analyzer = _Analyzer(schema, include_indirect, variables)
+            results.append(analyzer.analyze_statement(stmt))
+        elif isinstance(stmt, n.ScriptBlock):
+            for s in stmt.statements:
+                visit(s)
+            for s in stmt.exception_statements:
+                visit(s)
+        elif isinstance(stmt, n.IfStmt):
+            for b in stmt.branches:
+                for s in b.statements:
+                    visit(s)
+            for s in stmt.else_statements:
+                visit(s)
+        elif isinstance(stmt, (n.LoopStmt, n.WhileStmt, n.RepeatStmt)):
+            for s in stmt.statements:
+                visit(s)
+        elif isinstance(stmt, n.ForInStmt):
+            visit(stmt.query)  # the driving query reads tables too
+            for s in stmt.statements:
+                visit(s)
+        elif isinstance(stmt, n.CreateProcedure):
+            visit(stmt.body)
+        # DECLARE/SET/CALL/etc. carry no static table lineage
+
+    for s in stmts:
+        visit(s)
+    return results
+
+
+def extract_script_lineage(
+    sql_or_ast: Union[str, n.Node, List[n.Node]],
+    schema: Optional[Dict[str, List[str]]] = None,
+    include_indirect: bool = False,
+) -> List[LineageResult]:
+    """Extract lineage from a BigQuery script, procedure, or multi-statement
+    SQL string.
+
+    Walks BEGIN blocks, IF/ELSEIF branches, LOOP/WHILE/REPEAT/FOR bodies and
+    exception handlers, returning one :class:`LineageResult` per
+    lineage-bearing statement found (queries, CTAS, INSERT, UPDATE, DELETE,
+    MERGE, CREATE FUNCTION). ``DECLARE``d variables, ``FOR`` loop variables
+    and ``EXECUTE IMMEDIATE ... INTO`` targets are automatically excluded
+    from column resolution.
+    """
+    if isinstance(sql_or_ast, str):
+        stmts = parse(sql_or_ast)
+    elif isinstance(sql_or_ast, list):
+        stmts = sql_or_ast
+    else:
+        stmts = [sql_or_ast]
+    return _script_results(stmts, schema, include_indirect, set())