#!/usr/bin/env python3 """ table_attribute_extractor.py Drop this file into a folder with your Snowflake SQL / Python / PySpark / notebook files and run: python table_attribute_extractor.py It scans every code file in the SAME folder and saves two files there: table_attribute_inventory.csv - File Name | Table Name | Attribute Name | Line Numbers | Times Found | Source Location EVERY table and attribute each script references (read, written, created or partitioned) - one row per file + table + attribute, every line it appears on in one cell. code_snapshot.html - all the code on a single page; every table / attribute that went into the CSV is highlighted, everything else is greyed out What it understands * Snowflake SQL: CTEs, subqueries, UNIONs, MERGE / UPDATE / INSERT / COPY INTO, QUALIFY, FLATTEN, dbt ref()/source(), stored procedure bodies. Columns are traced back to their real base table. * Python / PySpark: SQL inside spark.sql / cursor.execute / pd.read_sql (even built with f-strings or variables), spark.read / write on S3 paths (table name derived from the path), spark.table, saveAsTable, Snowflake / JDBC connectors, Glue catalog, boto3, pandas, temp views, and DataFrame columns used in select / col() / groupBy / withColumn / joins / filters. * .ipynb notebooks (including %%sql cells). Other text files are sniffed for SQL or Python. Needs the sqlglot package; it is installed automatically on first run if missing. Optional: python table_attribute_extractor.py """ from __future__ import annotations import argparse import ast import csv import json import os import re import sys from dataclasses import dataclass, asdict, field from pathlib import Path from typing import Optional try: import sqlglot except ImportError: # first run: install the SQL parser automatically import subprocess print("Installing sqlglot (one-time)...") for extra in ([], ["--user"], ["--break-system-packages"]): if subprocess.call([sys.executable, "-m", "pip", "install", "-q", "sqlglot", *extra]) == 0: break else: sys.exit("Could not install sqlglot automatically. Run: pip install sqlglot") import importlib, site importlib.invalidate_caches() try: sys.path.append(site.getusersitepackages()) except Exception: pass import sqlglot from sqlglot import exp from sqlglot.optimizer.scope import traverse_scope, Scope # ───────────────────────────────────────────────────────────────────────────── # Output record # ───────────────────────────────────────────────────────────────────────────── @dataclass class Row: file: str line: int language: str # sql | python | pyspark | notebook context: str # sql_file | spark_sql | cursor_execute | dataframe_api | s3_io | glue_catalog | ... role: str # source | target | defined | derived | partition | filter database: str = "" schema: str = "" table: str = "" column: str = "" full_name: str = "" # database.schema.table.column (whatever parts are known) physical_location: str = "" # s3://… or @stage/… the table was read from / written to resolution: str = "" # how the column was tied to the table note: str = "" end_line: int = 0 # last line of the statement / call this came from written: str = "" # table name exactly as written in the code (for line matching / highlighting) def key(self): return (self.file, self.line, self.context, self.role, self.database.lower(), self.schema.lower(), self.table.lower(), self.column.lower(), self.physical_location) @dataclass class TableRef: database: str = "" schema: str = "" table: str = "" physical_location: str = "" note: str = "" @property def fqn(self): return ".".join(p for p in (self.database, self.schema, self.table) if p) # ───────────────────────────────────────────────────────────────────────────── # Placeholder handling (f-strings, Jinja/dbt, bind params, ${vars}) # ───────────────────────────────────────────────────────────────────────────── PH_RE = re.compile(r"PH__([A-Za-z0-9_]+?)__PH") def _ph(name: str) -> str: clean = re.sub(r"[^A-Za-z0-9_]", "_", name.strip())[:40] or "var" return f"PH__{clean}__PH" def restore_placeholders(s: str) -> str: return PH_RE.sub(lambda m: "{" + m.group(1) + "}", s or "") def preprocess_sql(sql: str) -> str: """Make templated SQL parseable while keeping table names readable.""" # dbt: {{ ref('model') }} / {{ ref('pkg','model') }} → model sql = re.sub(r"\{\{\s*ref\(\s*(?:'[^']*'\s*,\s*)?'([^']+)'\s*\)\s*\}\}", r"\1", sql) sql = re.sub(r'\{\{\s*ref\(\s*(?:"[^"]*"\s*,\s*)?"([^"]+)"\s*\)\s*\}\}', r"\1", sql) # dbt: {{ source('src','tbl') }} → src.tbl sql = re.sub(r"\{\{\s*source\(\s*['\"]([^'\"]+)['\"]\s*,\s*['\"]([^'\"]+)['\"]\s*\)\s*\}\}", r"\1.\2", sql) sql = re.sub(r"\{\{\s*this\s*\}\}", "this_model", sql) sql = re.sub(r"\{%.*?%\}", " ", sql, flags=re.S) # jinja blocks sql = re.sub(r"\{#.*?#\}", " ", sql, flags=re.S) # jinja comments sql = re.sub(r"\{\{\s*([^}]+?)\s*\}\}", lambda m: _ph(m.group(1)), sql) sql = re.sub(r"\$\{([^}]+)\}", lambda m: _ph(m.group(1)), sql) # ${var} sql = re.sub(r"(?])\?(?=[\s,)]|$)", "'__bind__'", sql) # ? return sql # ───────────────────────────────────────────────────────────────────────────── # S3 / cloud path → table name # ───────────────────────────────────────────────────────────────────────────── PATH_RE = re.compile(r"^(s3a?|s3n|gs|abfss?|wasbs?|hdfs|dbfs|file)://([^/]+)/?(.*)$", re.I) GENERIC_SEGMENTS = { "data", "raw", "processed", "curated", "staging", "stage", "stg", "output", "outputs", "input", "inputs", "parquet", "csv", "json", "orc", "avro", "delta", "latest", "current", "snapshot", "snapshots", "part", "parts", "tables", "table", "export", "exports", "landing", "bronze", "silver", "gold", "tmp", "temp", "files", "warehouse", "_delta_log", "incoming", "archive", "load", "loads", "extract", "extracts", "daily", "hourly", } FILE_EXT_RE = re.compile(r"\.(parquet|csv|tsv|json|jsonl|ndjson|orc|avro|txt|gz|bz2|snappy|zip|xlsx?)$", re.I) DATEISH_RE = re.compile(r"^(\d{4}([-_]?\d{2}([-_]?\d{2})?)?|\d{1,2}|v\d+|\d{8,14}|run[-_]?\d+)$", re.I) def looks_like_path(s: str) -> bool: return bool(s) and (PATH_RE.match(s.strip()) is not None or s.strip().startswith("@")) def table_from_path(path: str) -> tuple[TableRef, list[str]]: """Derive a table name (and partition columns) from an S3/stage path. s3://acme-lake/curated/sales/orders/year=2024/month=05/part-0001.parquet → schema=sales, table=orders, partitions=[year, month] """ p = path.strip().rstrip("/") partitions: list[str] = [] m = PATH_RE.match(p) if m: bucket, key = m.group(2), m.group(3) elif p.startswith("@"): # Snowflake stage: @db.schema.stage/path stage, _, key = p[1:].partition("/") bucket = "@" + stage else: bucket, key = "", p segs = [s for s in key.split("/") if s] kept = [] for s in segs: if "=" in s: # hive partition k=v partitions.append(s.split("=", 1)[0]) continue if "*" in s or "?" in s: continue kept.append(s) if kept and FILE_EXT_RE.search(kept[-1]): # drop file name kept.pop() while kept and (kept[-1].lower() in GENERIC_SEGMENTS or DATEISH_RE.match(kept[-1])): kept.pop() note = "" if not kept: # nothing meaningful left – fall back kept = [s for s in segs if "=" not in s and not FILE_EXT_RE.search(s)] or [bucket] note = "path had no distinctive folder; used best guess" table = kept[-1] schema = "" if len(kept) >= 2: prev = kept[-2] if prev.lower().endswith(".db"): # hive warehouse db_name.db/table schema = prev[:-3] elif prev.lower() not in GENERIC_SEGMENTS and not DATEISH_RE.match(prev): schema = prev table = re.sub(r"\.db$", "", table, flags=re.I) return TableRef(database="", schema=schema, table=table, physical_location=restore_placeholders(path), note=note), partitions # ───────────────────────────────────────────────────────────────────────────── # SQL analysis # ───────────────────────────────────────────────────────────────────────────── def split_sql(script: str) -> list[tuple[str, int]]: """Split a script on ';' respecting quotes, comments and $$ blocks. Returns (stmt, start_line).""" out, buf = [], [] i, n, line, start_line = 0, len(script), 1, 1 in_s = in_d = in_dollar = in_lc = in_bc = False started = False while i < n: c = script[i] nxt = script[i + 1] if i + 1 < n else "" if c == "\n": line += 1 in_lc = False if in_lc: buf.append(c); i += 1; continue if in_bc: buf.append(c) if c == "*" and nxt == "/": buf.append(nxt); i += 2; in_bc = False; continue i += 1; continue if in_dollar: buf.append(c) if c == "$" and nxt == "$": buf.append(nxt); i += 2; in_dollar = False; continue i += 1; continue if in_s: buf.append(c) if c == "'" and nxt == "'": buf.append(nxt); i += 2; continue if c == "\\": buf.append(nxt); i += 2; continue if c == "'": in_s = False i += 1; continue if in_d: buf.append(c) if c == '"': in_d = False i += 1; continue if c == "-" and nxt == "-": in_lc = True; buf.append(c); i += 1; continue if c == "/" and nxt == "*": in_bc = True; buf.append(c); i += 1; continue if c == "$" and nxt == "$": in_dollar = True; buf.append("$$"); i += 2; continue if c == "'": in_s = True elif c == '"': in_d = True if c == ";": stmt = "".join(buf).strip() if stmt: out.append((stmt, start_line)) buf, started = [], False i += 1; continue if not started and not c.isspace(): started, start_line = True, line buf.append(c) i += 1 stmt = "".join(buf).strip() if stmt: out.append((stmt, start_line)) return out def strip_comments(sql: str) -> str: sql = re.sub(r"--[^\n]*", "", sql) return re.sub(r"/\*.*?\*/", "", sql, flags=re.S) def _set_scopes(scope) -> list: """Branch scopes of a UNION/INTERSECT/EXCEPT (attribute name differs across sqlglot versions).""" return list(getattr(scope, "set_operation_scopes", None) or getattr(scope, "union_scopes", None) or []) def table_ref(t: exp.Table) -> TableRef: return TableRef(database=restore_placeholders(t.catalog), schema=restore_placeholders(t.db), table=restore_placeholders(t.name)) class SqlAnalyzer: """Extract tables + columns from one SQL statement and resolve each column to its base table.""" def __init__(self, emit, file, language, context, view_map=None, dialect="snowflake"): self.emit = emit self.file = file self.language = language self.context = context self.view_map = view_map or {} # temp view name (lower) → [TableRef] (from PySpark) self.dialect = dialect # ── public ──────────────────────────────────────────────────────────── def analyze_script(self, script: str, base_line: int = 1, python_cb=None): for stmt, ln in split_sql(script): self.analyze_statement(stmt, base_line + ln - 1, python_cb) def analyze_statement(self, raw: str, line: int, python_cb=None): self.cur_end = line + raw.count("\n") # Stored procedures / scripting blocks: recurse into $$ bodies body = re.search(r"\$\$(.*)\$\$", raw, re.S) if body and re.search(r"\bcreate\b.*\b(procedure|function)\b", raw, re.I | re.S): lang = re.search(r"\blanguage\s+(\w+)", raw, re.I) lang = (lang.group(1).lower() if lang else "sql") body_line = line + raw[:body.start(1)].count("\n") if lang == "sql": inner = re.sub(r"^\s*(declare\b.*?)?\bbegin\b|\bend\s*;?\s*$", "", body.group(1), flags=re.I | re.S) self.analyze_script(inner, body_line, python_cb) elif lang == "python" and python_cb: python_cb(body.group(1), body_line) return if re.match(r"^\s*(use|set|unset|alter\s+session|show|describe|desc|grant|revoke|" r"begin|commit|rollback|call|execute\s+immediate|list|ls|put|get|remove)\b", strip_comments(raw), re.I): self._regex_fallback(raw, line, only_dml=True) return sql = preprocess_sql(raw) try: parsed = [p for p in sqlglot.parse(sql, read=self.dialect) if p is not None] except Exception: parsed = [] if self.dialect != "snowflake": try: parsed = [p for p in sqlglot.parse(sql, read="snowflake") if p is not None] except Exception: parsed = [] if not parsed: self._regex_fallback(raw, line) return for tree in parsed: self._analyze_tree(tree, line, raw) # ── core ────────────────────────────────────────────────────────────── def _row(self, line, role, tref: TableRef, column="", resolution="", note=""): written = tref.physical_location or tref.fqn or tref.table for t in self._apply_view(tref): self._row1(line, role, t, column, resolution, note, written) def _row1(self, line, role, tref: TableRef, column="", resolution="", note="", written=""): col = restore_placeholders(column) full = ".".join(p for p in (tref.fqn, col) if p) self.emit(Row(file=self.file, line=line, language=self.language, context=self.context, role=role, database=tref.database, schema=tref.schema, table=tref.table, column=col, full_name=full, physical_location=tref.physical_location, resolution=resolution, note="; ".join(x for x in (tref.note, note) if x), end_line=getattr(self, "cur_end", line), written=restore_placeholders(written))) def _apply_view(self, t: TableRef) -> list[TableRef]: """A Spark temp view name → the real table(s)/S3 path(s) behind it.""" if not t.database and not t.schema and t.table.lower() in self.view_map: srcs = self.view_map[t.table.lower()] note = f"via temp view '{t.table}'" + (" — view joins several sources; column is in one of them" if len(srcs) > 1 else "") return [TableRef(s.database, s.schema, s.table, s.physical_location, "; ".join(x for x in (s.note, note) if x)) for s in srcs] return [t] def _analyze_tree(self, tree: exp.Expression, line: int, raw: str): cte_names = {c.alias_or_name.lower() for c in tree.find_all(exp.CTE)} target_nodes = set() target: Optional[TableRef] = None handled_cols: set[int] = set() # 1. Target of DDL / DML tgt_expr = None if isinstance(tree, (exp.Insert, exp.Create, exp.Merge, exp.Update, exp.Delete, exp.Copy)): tgt_expr = tree.this if isinstance(tree, exp.Alter): tgt_expr = tree.this if tgt_expr is not None: tbl = tgt_expr if isinstance(tgt_expr, exp.Table) else tgt_expr.find(exp.Table) if tbl is not None: target = table_ref(tbl) target_nodes.add(id(tbl)) kind = type(tree).__name__.upper() if isinstance(tree, exp.Create): kind = f"CREATE {tree.args.get('kind') or ''}".strip() self._row(line, "target", target, note=kind) # explicit column list: INSERT INTO t (a,b) / CREATE TABLE t (a INT, …) if isinstance(tgt_expr, exp.Schema): for e in tgt_expr.expressions: name = e.name if not isinstance(e, exp.ColumnDef) else e.this.name role = "defined" if isinstance(tree, exp.Create) else "target" self._row(line, role, target, name, "explicit column list") for c in e.find_all(exp.Column): handled_cols.add(id(c)) # ALTER TABLE ADD/DROP/RENAME COLUMN if isinstance(tree, exp.Alter): for cd in tree.find_all(exp.ColumnDef): self._row(line, "defined", target, cd.this.name, "alter table") # 2. COPY INTO … FROM @stage / 's3://…' if isinstance(tree, exp.Copy): for f in tree.args.get("files") or []: txt = f.sql(dialect="snowflake").strip("'\"") if looks_like_path(txt) or txt.startswith("@"): tr, parts = table_from_path(txt) self._row(line, "source", tr, resolution="stage/path", note="COPY INTO source") for p in parts: self._row(line, "partition", tr, p, "path partition") for t in f.find_all(exp.Table): target_nodes.add(id(t)) for m in ([] if tree.args.get("files") else re.finditer(r"from\s+(@[\w.$/\-{}=]+|'(?:s3|gcs|azure)://[^']+')", raw, re.I)): txt = m.group(1).strip("'") tr, parts = table_from_path(txt) self._row(line, "source", tr, resolution="stage/path", note="COPY INTO source") # 3. Every physical table referenced as a source alias_map: dict[str, TableRef] = {} for t in tree.find_all(exp.Table): if not t.name: continue tr = table_ref(t) if t.alias: alias_map[t.alias.lower()] = tr alias_map.setdefault(t.name.lower(), tr) if id(t) in target_nodes: continue if t.name.lower() in cte_names and not t.db: continue if t.name.startswith("@") or t.name.lower().startswith("s3"): tr2, parts = table_from_path(t.sql()) self._row(line, "source", tr2, resolution="stage/path") continue self._row(line, "source", tr) # 4. Columns: scope-based resolution through CTEs, subqueries, unions try: scopes = list(traverse_scope(tree)) except Exception: scopes = [] scope_by_expr = {id(sc.expression): sc for sc in scopes} sub_alias: dict[str, Scope] = {} for sq in tree.find_all(exp.Subquery): if sq.alias and id(sq.this) in scope_by_expr: sub_alias[sq.alias.lower()] = scope_by_expr[id(sq.this)] for scope in scopes: for col in scope.columns: if id(col) in handled_cols: continue handled_cols.add(id(col)) results = self._resolve(scope, col.table, col.name, set()) role = "source" for tref, cname, how in results: if tref is None: self._row(line, role, TableRef(table=col.table or ""), col.name, how) else: self._row(line, role, tref, cname, how) # SELECT * / t.* for star in scope.expression.find_all(exp.Star): if star.find_ancestor(exp.Select) is not scope.expression and \ not isinstance(scope.expression, exp.Select): continue parent = star.parent proj = parent if isinstance(parent, exp.Column) else star if proj.parent is not scope.expression: # skip COUNT(*), etc. — only SELECT * / t.* continue qualifier = parent.table if isinstance(parent, exp.Column) else "" for tref in self._star_tables(scope, qualifier, set()): self._row(line, "source", tref, "*", "star (column list needs table schema)") # 5. Anything scope.columns skipped (select-alias refs in QUALIFY/ORDER BY) and anything outside # a SELECT scope: UPDATE SET, MERGE ON/SET/INSERT, DELETE WHERE, etc. for col in tree.find_all(exp.Column): if id(col) in handled_cols or not col.name or isinstance(col.this, exp.Star): continue handled_cols.add(id(col)) q = (col.table or "").lower() enclosing = None node = col.parent while node is not None: if id(node) in scope_by_expr: enclosing = scope_by_expr[id(node)] break node = node.parent if enclosing is not None: for tref, cname, how in self._resolve(enclosing, col.table, col.name, set()): if tref is None: self._row(line, "source", TableRef(table=col.table or ""), col.name, how) else: self._row(line, "source", tref, cname, how) continue if q and q in sub_alias: for tref, cname, how in self._resolve_in_source(sub_alias[q], col.name, set(), 0, "traced through USING subquery"): if tref is not None: self._row(line, "source", tref, cname, how) else: self._row(line, "source", TableRef(table=col.table), col.name, how) continue eq = col.parent is_set_target = isinstance(eq, exp.EQ) and eq.this is col and isinstance(eq.parent, exp.Update) in_merge_insert = isinstance(col.parent, exp.Tuple) and isinstance(col.parent.parent, exp.Insert) \ and col.find_ancestor(exp.Merge) is not None if q and q in alias_map: tref = alias_map[q] role = "target" if (is_set_target and target and tref.fqn == target.fqn) else "source" self._row(line, role, tref, col.name, "alias" if q != tref.table.lower() else "qualified") elif (is_set_target or in_merge_insert) and target: self._row(line, "target", target, col.name, "SET/INSERT target column") elif not q and len({t.fqn for t in alias_map.values()}) == 1: self._row(line, "source", list(alias_map.values())[0], col.name, "only table in statement") else: cands = sorted({t.fqn for t in alias_map.values()}) self._row(line, "source", TableRef(table=col.table or ""), col.name, "unresolved", note=("candidates: " + ", ".join(cands)) if cands else "") # 6. Output columns of CREATE … AS SELECT / INSERT … SELECT / CREATE VIEW if target is not None and isinstance(tree, (exp.Create, exp.Insert)): query = tree.expression if isinstance(query, exp.Query) and not isinstance(tree.this, exp.Schema): sel = query while isinstance(sel, exp.SetOperation): sel = sel.this if isinstance(sel, exp.Select): for proj in sel.expressions: if isinstance(proj, exp.Star) or (isinstance(proj, exp.Column) and isinstance(proj.this, exp.Star)): continue name = proj.alias_or_name if name: self._row(line, "target", target, name, "output of SELECT") # ── column resolution through scopes ───────────────────────────────── def _resolve(self, scope: Scope, qualifier: str, name: str, seen: set, depth: int = 0): """Return list of (TableRef|None, column, resolution).""" if depth > 25: return [(None, name, "unresolved (too deep)")] sources = {k.lower(): v for k, v in scope.sources.items()} if qualifier: src = sources.get(qualifier.lower()) if src is None: return [(None, name, "unresolved (qualifier not a known table/alias — lateral/flatten/outer ref?)")] return self._resolve_in_source(src, name, seen, depth, "qualified") # unqualified: only consider sources actually in this FROM/JOIN (not every CTE defined above) try: selected = {k.lower(): v[1] for k, v in scope.selected_sources.items()} except Exception: selected = {} sources = selected or sources if len(sources) == 1: src = next(iter(sources.values())) return self._resolve_in_source(src, name, seen, depth, "only source in scope") # a derived source that projects this name wins hits = [] for alias, src in sources.items(): if isinstance(src, Scope) and name.lower() in self._projected_names(src) and \ all(h is not src for h in hits): hits.append(src) if len(hits) == 1: return self._resolve_in_source(hits[0], name, seen, depth, "matched CTE/subquery output") # SELECT-list alias referenced in WHERE/QUALIFY/ORDER BY (Snowflake lateral alias) if isinstance(scope.expression, exp.Select): for proj in scope.expression.expressions: if isinstance(proj, exp.Alias) and proj.alias.lower() == name.lower(): out = [] for c in proj.find_all(exp.Column): out += self._resolve(scope, c.table, c.name, seen, depth + 1) return out or [(None, name, "select alias (derived, no source column)")] tables = [s for s in sources.values() if isinstance(s, exp.Table)] cands = sorted({table_ref(t).fqn for t in tables}) if not sources: return [(None, name, "unresolved (no FROM)")] return [(None, name, "ambiguous: " + (", ".join(cands) if cands else "multiple sources"))] def _projected_names(self, scope: Scope) -> set[str]: expr = scope.expression while isinstance(expr, exp.SetOperation): expr = expr.this if isinstance(expr, exp.Select): return {p.alias_or_name.lower() for p in expr.expressions if p.alias_or_name} return set() def _resolve_in_source(self, src, name, seen, depth, how): if isinstance(src, exp.Table): return [(table_ref(src), name, how)] if not isinstance(src, Scope) or src.scope_type.name == "UDTF": node = src.expression if isinstance(src, Scope) else src inner = [c for c in node.find_all(exp.Column)] if isinstance(node, exp.Expression) else [] if inner and isinstance(src, Scope): out = [] for c in inner: for r in self._resolve(src, c.table, c.name, seen, depth + 1): if r[0] is not None: out.append((r[0], r[1], f"FLATTEN/lateral output '{name}' → element of {r[1]}")) return out or [(None, name, "unresolved (lateral/table-function output)")] return [(None, name, "unresolved (function/lateral source)")] key = (id(src), name.lower()) if key in seen: return [] seen = seen | {key} # UNION: trace into every branch if _set_scopes(src): out = [] for us in _set_scopes(src): out += self._resolve_in_source(us, name, seen, depth + 1, "traced through UNION") return out expr = src.expression if not isinstance(expr, exp.Select): return [(None, name, "unresolved")] for proj in expr.expressions: if proj.alias_or_name and proj.alias_or_name.lower() == name.lower() and \ not isinstance(proj, exp.Star): cols = list(proj.find_all(exp.Column)) cols = [c for c in cols if not isinstance(c.this, exp.Star)] if not cols: return [(None, name, "derived in CTE/subquery (literal/function, no source column)")] out = [] for c in cols: for r in self._resolve(src, c.table, c.name, seen, depth + 1): tag = "traced through CTE/subquery" + (f" (renamed from {c.name})" if c.name.lower() != name.lower() else "") out.append((r[0], r[1], tag if r[0] is not None else r[2])) return out # not explicitly projected → maybe SELECT * passes it through if any(isinstance(p, exp.Star) or (isinstance(p, exp.Column) and isinstance(p.this, exp.Star)) for p in expr.expressions): return [(r[0], r[1], "traced through SELECT *" if r[0] is not None else r[2]) for r in self._resolve(src, "", name, seen, depth + 1)] return [(None, name, "unresolved (not projected by CTE/subquery)")] def _star_tables(self, scope: Scope, qualifier: str, seen: set) -> list[TableRef]: if id(scope) in seen: return [] seen = seen | {id(scope)} srcs = scope.sources if qualifier: srcs = {k: v for k, v in srcs.items() if k.lower() == qualifier.lower()} out = [] for s in srcs.values(): if isinstance(s, exp.Table): out.append(table_ref(s)) elif isinstance(s, Scope): for us in (_set_scopes(s) or [s]): if isinstance(us.expression, exp.Select) and any( isinstance(p, exp.Star) for p in us.expression.expressions): out += self._star_tables(us, "", seen) return out # ── last resort for statements sqlglot can't parse ─────────────────── def _regex_fallback(self, raw: str, line: int, only_dml: bool = False): sql = strip_comments(preprocess_sql(raw)) ident = r'((?:"[^"]+"|[\w$]+)(?:\.(?:"[^"]+"|[\w$]+)){0,2})' pats = [(r"\b(?:from|join)\s+" + ident, "source"), (r"\b(?:insert\s+(?:overwrite\s+)?into|merge\s+into|update|create\s+(?:or\s+replace\s+)?" r"(?:temporary\s+|transient\s+)?(?:table|view)|truncate\s+table|delete\s+from)\s+" r"(?:if\s+not\s+exists\s+)?" + ident, "target")] if only_dml: pats = [] for pat, role in pats: for m in re.finditer(pat, sql, re.I): parts = [p.strip('"') for p in m.group(1).split(".")] if parts[-1].lower() in {"select", "lateral", "table", "values", "("}: continue parts = [""] * (3 - len(parts)) + parts self._row(line, role, TableRef(*map(restore_placeholders, parts)), resolution="regex fallback (SQL did not parse)") # ───────────────────────────────────────────────────────────────────────────── # Python / PySpark analysis # ───────────────────────────────────────────────────────────────────────────── SQL_START_RE = re.compile(r"^\s*(\(|--[^\n]*\n\s*)*\s*(select|with|insert|merge|create|update|delete|copy\s+into|" r"replace\s+into|unload|truncate)\b", re.I) SQL_CALLS = {"sql", "execute", "executemany", "execute_string", "read_sql", "read_sql_query", "read_sql_table", "run_query", "query", "execute_async", "run", "fetch_pandas_all"} READ_METHODS = {"parquet", "csv", "json", "orc", "avro", "load", "text", "delta", "table", "jdbc", "xml"} WRITE_METHODS = {"parquet", "csv", "json", "orc", "avro", "save", "text", "saveAsTable", "insertInto", "delta", "jdbc", "toTable", "start"} COL_LIST_METHODS = {"select", "groupBy", "groupby", "orderBy", "sort", "drop", "dropDuplicates", "drop_duplicates", "cube", "rollup", "partitionBy", "sortWithinPartitions", "repartition", "dropna", "fillna", "describe", "summary", "sortBy", "bucketBy", "clusterBy"} COL_FUNCS = {"col", "column", "Column"} class PyAnalyzer: def __init__(self, emit, file, language="python", spark_dialect="spark", sql_dialect="snowflake"): self.emit = emit self.file = file self.language = language self.spark_dialect = spark_dialect self.sql_dialect = sql_dialect self.env: dict[str, str] = {} # variable → string value self.env_line: dict[str, int] = {} # variable → line where its string was defined self.dicts: dict[str, dict] = {} # variable → {key: str} for options(**opts) self.lists: dict[str, list] = {} # variable → [str, ...] e.g. cols = ["a", "b"] self.dfs: dict[str, list[TableRef]] = {} # dataframe var → source tables self.views: dict[str, list[TableRef]] = {} # temp view → source tables self.sql_seen: set = set() self.pending_sql: list[tuple[str, str, int]] = [] # (var, sql, line) literals not yet used self.is_spark = False # ── entry ──────────────────────────────────────────────────────────── def analyze(self, code: str, base_line: int = 1): # notebooks magics (%sql, !pip) aren't Python – blank them; then blank any line that won't parse lines = ["" if l.lstrip().startswith(("%", "!")) else l for l in code.splitlines()] tree, skipped = None, [] for _ in range(200): try: tree = ast.parse("\n".join(lines)) break except SyntaxError as e: ln = (e.lineno or 1) - 1 if ln >= len(lines) or ln < 0: break skipped.append(ln + base_line) lines[ln] = re.sub(r"\S.*", "pass", lines[ln]) if not lines[ln].strip().startswith("pass") \ else "" if tree is None: self.emit(Row(self.file, base_line, self.language, "parse_error", "", note="Python could not be parsed")) return if skipped: self.emit(Row(self.file, skipped[0], self.language, "parse_error", "", note=f"skipped unparseable Python line(s): {sorted(set(skipped))[:10]}")) code = "\n".join(lines) self.base = base_line - 1 self.is_spark = bool(re.search(r"pyspark|SparkSession|spark\.read|glueContext|\.createOrReplaceTempView", code)) if self.is_spark: self.language = "pyspark" if self.language == "python" else self.language stmts = sorted((n for n in ast.walk(tree) if isinstance(n, ast.stmt)), key=lambda n: (n.lineno, n.col_offset)) done_calls: set[int] = set() for st in stmts: self._statement(st, done_calls) # SQL strings that were built but never passed to an execute call we recognised for var, sql, ln in self.pending_sql: if (sql, ln) not in self.sql_seen: self._run_sql(sql, ln, "sql_string_literal") def L(self, node) -> int: return self.base + getattr(node, "lineno", 1) # ── string evaluation (constants, f-strings, concat, .format, os.path.join) ── def s(self, node) -> Optional[str]: if node is None: return None if isinstance(node, ast.Constant) and isinstance(node.value, str): return node.value if isinstance(node, ast.Name): return self.env.get(node.id) if isinstance(node, ast.JoinedStr): out = [] for v in node.values: if isinstance(v, ast.Constant): out.append(str(v.value)) elif isinstance(v, ast.FormattedValue): val = self.s(v.value) out.append(val if val is not None else "{" + ast.unparse(v.value) + "}") return "".join(out) if isinstance(node, ast.BinOp) and isinstance(node.op, ast.Add): a, b = self.s(node.left), self.s(node.right) if a is None and b is None: return None return (a if a is not None else "{" + ast.unparse(node.left) + "}") + \ (b if b is not None else "{" + ast.unparse(node.right) + "}") if isinstance(node, ast.BinOp) and isinstance(node.op, ast.Mod): return self.s(node.left) if isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute): if node.func.attr == "format": base = self.s(node.func.value) if base is not None: kw = {k.arg: self.s(k.value) for k in node.keywords if k.arg} pos = [self.s(a) for a in node.args] idx = iter(range(1000)) def rep(m): key = m.group(1) if key == "": i = next(idx) return pos[i] if i < len(pos) and pos[i] is not None else "{arg%d}" % i if key.isdigit(): i = int(key) return pos[i] if i < len(pos) and pos[i] is not None else "{arg%d}" % i return kw.get(key) if kw.get(key) is not None else "{" + key + "}" return re.sub(r"\{([^{}:!]*)(?:[:!][^{}]*)?\}", rep, base) if node.func.attr == "join": sep = self.s(node.func.value) if sep is not None and node.args and isinstance(node.args[0], (ast.List, ast.Tuple)): parts = [self.s(e) or "{" + ast.unparse(e) + "}" for e in node.args[0].elts] return sep.join(parts) if node.func.attr in ("strip", "lower", "upper", "replace") and node.func.attr != "replace": return self.s(node.func.value) if ast.unparse(node.func) in ("os.path.join", "posixpath.join", "path.join"): parts = [self.s(a) or "{" + ast.unparse(a) + "}" for a in node.args] return "/".join(p.strip("/") if i else p.rstrip("/") for i, p in enumerate(parts)) if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "str" and node.args: return self.s(node.args[0]) return None # ── statement handling ─────────────────────────────────────────────── def _statement(self, st, done_calls): targets, value = [], None if isinstance(st, ast.Assign): targets, value = st.targets, st.value elif isinstance(st, ast.AnnAssign) and st.value is not None: targets, value = [st.target], st.value elif isinstance(st, ast.AugAssign) and isinstance(st.target, ast.Name): add = self.s(st.value) if st.target.id in self.env and add is not None: self.env[st.target.id] += add return elif isinstance(st, ast.Expr): value = st.value elif isinstance(st, ast.Return): value = st.value else: # walk calls inside compound statement headers (if/for/with test expressions) for fld in ("test", "iter", "items"): v = getattr(st, fld, None) if v is not None: for n in ([v] if not isinstance(v, list) else v): self._expr(n, done_calls) return if value is None: return sval = self.s(value) names = [t.id for t in targets if isinstance(t, ast.Name)] names += [t.attr for t in targets if isinstance(t, ast.Attribute)] # self.query = "..." if isinstance(value, (ast.List, ast.Tuple)) and names: items = [self.s(e) for e in value.elts] if items and all(isinstance(x, str) for x in items): for nm in names: self.lists[nm] = items if isinstance(value, ast.Dict) and names: d = {(self.s(k) or ""): (self.s(v) or "") for k, v in zip(value.keys, value.values) if k is not None} for nm in names: self.dicts[nm] = d if sval is not None and names: for nm in names: self.env[nm] = sval self.env_line[nm] = self.L(st) if SQL_START_RE.match(sval): self.pending_sql.append((names[0], sval, self.L(st))) return result = self._expr(value, done_calls) if result is not None: for nm in names: self.dfs[nm] = result # tuple unpacking from dict/list of paths: leave alone def _expr(self, node, done_calls) -> Optional[list[TableRef]]: """Analyse an expression; return the DataFrame lineage it produces (if any).""" result = None outer_calls = [n for n in ast.walk(node) if isinstance(n, ast.Call)] # mark inner links of method chains so each chain is processed once from its outermost call inner = set() for c in outer_calls: f = c.func while isinstance(f, ast.Attribute): if isinstance(f.value, ast.Call): inner.add(id(f.value)) f = f.value.func else: f = f.value for c in outer_calls: if id(c) in inner or id(c) in done_calls: continue done_calls.add(id(c)) r = self._chain(c, done_calls) if c is node: result = r if isinstance(node, ast.Name) and node.id in self.dfs: result = self.dfs[node.id] return result # ── method chain: root.m1(a).m2(b).m3(c) ─────────────────────────── def _chain(self, call: ast.Call, done_calls) -> Optional[list[TableRef]]: prev = getattr(self, "cur_end", 0) self.cur_end = self.base + (getattr(call, "end_lineno", None) or call.lineno) try: return self._chain_inner(call, done_calls) finally: self.cur_end = prev def _chain_inner(self, call: ast.Call, done_calls) -> Optional[list[TableRef]]: links = [] # [(method_name, Call)], innermost first f, cur = call.func, call while True: if isinstance(f, ast.Attribute): links.append((f.attr, cur)) # cur is None for properties like .write / .read / .na if isinstance(f.value, ast.Call): cur = f.value f = cur.func continue if isinstance(f.value, ast.Attribute): cur = None f = f.value continue root = f.value break elif isinstance(f, ast.Name): links.append((f.id, cur)) root = None break else: root = None break links.reverse() root_txt = ast.unparse(root) if root is not None else "" names = [m for m, _ in links] line = self.L(call) # lineage carried by the root lineage: Optional[list[TableRef]] = None base = root while isinstance(base, ast.Attribute): # df.write / df.na / df.write.format(...) base = base.value if isinstance(base, ast.Name) and base.id in self.dfs: lineage = list(self.dfs[base.id]) elif isinstance(root, ast.Name) and root.id in self.dfs: lineage = list(self.dfs[root.id]) elif isinstance(root, ast.Subscript) and isinstance(root.value, ast.Name) and root.value.id in self.dfs: lineage = list(self.dfs[root.value.id]) fmt = None options: dict[str, str] = {} is_read = any(n in ("read", "readStream") for n in names) or "read" in root_txt.split(".") is_write = any(n in ("write", "writeStream", "writeTo") for n in names) or root_txt.endswith(".write") for idx, (m, c) in enumerate(links): if c is None: # plain property in the chain (df.write, spark.read) continue args = c.args kws = {k.arg: k.value for k in c.keywords if k.arg} a0 = self.s(args[0]) if args else None # ---------- reading ---------- if m == "format": fmt = (a0 or "").lower() elif m in ("option", "options"): if m == "option" and len(args) >= 2: options[(self.s(args[0]) or "").lower()] = self.s(args[1]) or "" for k, v in kws.items(): options[k.lower()] = self.s(v) or "" for k in c.keywords: if k.arg is None and isinstance(k.value, ast.Name) and k.value.id in self.dicts: options.update({kk.lower(): vv for kk, vv in self.dicts[k.value.id].items()}) if args and isinstance(args[0], ast.Dict): for k, v in zip(args[0].keys, args[0].values): options[(self.s(k) or "").lower()] = self.s(v) or "" elif m == "schema" and is_read and a0: tr = TableRef(table="(schema for read)") for part in a0.split(","): nm = part.strip().split()[0] if part.strip() else "" if nm: self._emit_col(line, "dataframe_api", "defined", lineage or [tr], nm, "read schema DDL") elif is_read and m in READ_METHODS and not is_write: lineage = self._read(line, m, fmt, options, args, kws) or lineage elif m in ("sql",) and args: sql = self.s(args[0]) if sql is not None: self._run_sql(sql, line, "spark_sql" if self.is_spark else "sql_call") lineage = self._sql_lineage(sql) elif m in SQL_CALLS and args and m not in ("run",): sql = self.s(args[0]) if sql is None and "sql" in kws: sql = self.s(kws["sql"]) if sql is not None and SQL_START_RE.match(sql): ctx = {"execute": "cursor_execute", "executemany": "cursor_execute", "read_sql": "pandas_read_sql", "read_sql_query": "pandas_read_sql"}.get(m, "sql_call") self._run_sql(sql, line, ctx) lineage = self._sql_lineage(sql) elif m == "read_sql_table" and sql: lineage = [self._table_name(sql, kws.get("schema"))] self._emit_tbl(line, "pandas_read_sql", "source", lineage[0]) elif m == "table" and args and not is_write: # spark.table("db.t") / session.table name = self.s(args[0]) if name: tr = self._spark_table(name) self._emit_tbl(line, "spark_table", "source", tr) lineage = [tr] # AWS Glue elif m == "from_catalog": db = self.s(kws.get("database")) or "" tb = self.s(kws.get("table_name")) or "" if tb: tr = TableRef(schema=db, table=tb, note="Glue Data Catalog") self._emit_tbl(line, "glue_catalog", "target" if "write" in " ".join(names) + root_txt else "source", tr) lineage = [tr] elif m == "from_options" and kws.get("connection_options") is not None: co = kws["connection_options"] paths = [] if isinstance(co, ast.Dict): for k, v in zip(co.keys, co.values): key = self.s(k) if key == "paths" and isinstance(v, (ast.List, ast.Tuple)): paths += [self.s(e) for e in v.elts] elif key in ("path", "dbtable"): paths.append(self.s(v)) role = "target" if any("write" in n for n in names) or "write" in root_txt else "source" lin = [] for p in filter(None, paths): lin.append(self._path_ref(line, p, role, "glue_s3")) lineage = lin or lineage elif m in ("toDF",) and lineage and args: for a in args: nm = self.s(a) if nm: self._emit_col(line, "dataframe_api", "derived", lineage, nm, "toDF rename") # pandas / awswrangler / boto3 elif m in ("read_parquet", "read_csv", "read_json", "read_orc", "read_excel", "read_feather", "read_table", "read_fwf", "read_delta") and (args or "path" in kws): p = self.s(args[0]) if args else self.s(kws.get("path")) if p: lineage = [self._path_ref(line, p, "source", "s3_io" if looks_like_path(p) else "file_io")] for key in ("columns", "usecols"): if isinstance(kws.get(key), (ast.List, ast.Tuple)): for e in kws[key].elts: nm = self.s(e) if nm: self._emit_col(line, "dataframe_api", "source", lineage, nm, "read columns=") elif m in ("to_parquet", "to_csv", "to_json", "to_sql", "to_excel", "to_orc"): p = self.s(args[0]) if args else self.s(kws.get("path") or kws.get("name")) if p: if m == "to_sql": tr = self._table_name(p, kws.get("schema")) self._emit_tbl(line, "pandas_write", "target", tr) else: self._path_ref(line, p, "target", "s3_io" if looks_like_path(p) else "file_io") elif m in ("get_object", "download_file", "download_fileobj", "put_object", "upload_file", "upload_fileobj", "copy_object", "list_objects_v2", "Object", "select_object_content"): bucket = self.s(kws.get("Bucket")) or (self.s(args[0]) if m in ("Object", "download_file", "download_fileobj") and args else None) key = self.s(kws.get("Key") or kws.get("Prefix")) if m == "download_file" and len(args) >= 2: key = key or self.s(args[1]) if m == "upload_file" and len(args) >= 3: bucket, key = bucket or self.s(args[1]), key or self.s(args[2]) if m == "Object" and len(args) >= 2: key = key or self.s(args[1]) if bucket or key: p = f"s3://{bucket or '{bucket}'}/{key or ''}" role = "target" if m.startswith(("put", "upload", "copy")) else "source" lineage = [self._path_ref(line, p, role, "boto3_s3")] if m == "select_object_content" and "Expression" in kws: sql = self.s(kws["Expression"]) if sql: self._run_sql(sql, line, "s3_select") # snowflake connector / snowpark write elif m == "write_pandas" and len(args) >= 3: tn = self.s(args[2]) or self.s(kws.get("table_name")) if tn: tr = TableRef(database=self.s(kws.get("database")) or "", schema=self.s(kws.get("schema")) or "", table=tn) self._emit_tbl(line, "snowflake_write_pandas", "target", tr) # ---------- writing ---------- elif is_write and m in WRITE_METHODS: if m in ("saveAsTable", "insertInto", "toTable"): nm = self.s(args[0]) if args else self.s(kws.get("tableName") or kws.get("name")) if nm: self._emit_tbl(line, "spark_write", "target", self._spark_table(nm)) else: p = (self.s(args[0]) if args else None) or self.s(kws.get("path")) or options.get("path") dbt = options.get("dbtable") if dbt: self._emit_tbl(line, "spark_write", "target", self._spark_table(dbt)) elif p: self._path_ref(line, p, "target", "spark_write") elif m == "partitionBy" and is_write: for a in args: nm = self.s(a) if nm: self._emit_col(line, "dataframe_api", "partition", lineage or [], nm, "write partitionBy") continue elif m in ("createOrReplaceTempView", "createTempView", "createGlobalTempView", "createOrReplaceGlobalTempView", "registerTempTable"): vn = self.s(args[0]) if args else None if vn: self.views[vn.lower()] = lineage or [TableRef(table=f"(unknown source of view {vn})")] self.views[("global_temp." + vn).lower()] = self.views[vn.lower()] continue # ---------- DataFrame column operations ---------- if m == "join" and args: other = args[0] other_lin = self._lineage_of(other) on = kws.get("on") if "on" in kws else (args[1] if len(args) > 1 else None) combined = (lineage or []) + (other_lin or []) if on is not None: for nm in self._str_list(on): self._emit_col(line, "dataframe_api", "source", combined, nm, "join key (present in both sides)") lineage = combined or lineage elif m == "union" or m == "unionByName" or m == "unionAll": if args: lineage = (lineage or []) + (self._lineage_of(args[0]) or []) elif m in COL_LIST_METHODS and lineage is not None and m != "partitionBy": for a in list(args) + [k.value for k in c.keywords if k.arg in ("subset", "cols", "by")]: for nm in self._str_list(a): self._emit_col(line, "dataframe_api", "source", lineage, nm, f".{m}()") elif m == "selectExpr" and lineage is not None: for a in args: ex = self.s(a) if ex: self._sql_expr_cols(ex, line, lineage, f".selectExpr()") elif m in ("filter", "where") and lineage is not None and args: ex = self.s(args[0]) if ex: self._sql_expr_cols(ex, line, lineage, f".{m}('…')") elif m == "withColumn" and len(args) >= 1: nm = self.s(args[0]) if nm: self._emit_col(line, "dataframe_api", "derived", lineage or [], nm, ".withColumn() new/overwritten column") elif m == "withColumnRenamed" and len(args) >= 2 and lineage is not None: old, new = self.s(args[0]), self.s(args[1]) if old: self._emit_col(line, "dataframe_api", "source", lineage, old, f"renamed to {new}") if new: self._emit_col(line, "dataframe_api", "derived", lineage, new, f"renamed from {old}") elif m in ("withColumnsRenamed", "withColumns") and args and isinstance(args[0], ast.Dict): for k in args[0].keys: nm = self.s(k) if nm: self._emit_col(line, "dataframe_api", "derived" if m == "withColumns" else "source", lineage or [], nm, f".{m}()") elif m in ("agg",) and lineage is not None: for a in args: if isinstance(a, ast.Dict): for k in a.keys: nm = self.s(k) if nm: self._emit_col(line, "dataframe_api", "source", lineage, nm, ".agg({…})") elif m == "alias" and len(links) > 1 and lineage is not None: pass # col("x") / F.col("x") / df["x"] / df.x anywhere inside this call's arguments if lineage is not None or m in COL_FUNCS: for a in list(args) + [k.value for k in c.keywords]: self._inner_cols(a, line, lineage or [], done_calls) return lineage # ── helpers for the chain walker ───────────────────────────────────── def _lineage_of(self, node): if isinstance(node, ast.Name): return self.dfs.get(node.id) if isinstance(node, ast.Call): return self._chain(node, set()) return None def _str_list(self, node) -> list[str]: if node is None: return [] s = self.s(node) if s is not None: return [s] if isinstance(node, (ast.List, ast.Tuple, ast.Set)): out = [] for e in node.elts: out += self._str_list(e) return out if isinstance(node, ast.Starred): return self._str_list(node.value) if isinstance(node, ast.Name) and node.id in self.lists: return list(self.lists[node.id]) if isinstance(node, ast.Call): fn = node.func.attr if isinstance(node.func, ast.Attribute) else getattr(node.func, "id", "") if fn in COL_FUNCS and node.args: return self._str_list(node.args[0]) if fn in ("asc", "desc") and isinstance(node.func, ast.Attribute): return self._str_list(node.func.value) if not node.args else self._str_list(node.args[0]) return [] def _inner_cols(self, node, line, lineage, done_calls): for n in ast.walk(node): if isinstance(n, ast.Call): fn = n.func.attr if isinstance(n.func, ast.Attribute) else getattr(n.func, "id", "") done_calls.add(id(n)) if fn == "alias" and isinstance(n.func, ast.Attribute) and n.args and \ not (isinstance(n.func.value, ast.Name) and n.func.value.id in self.dfs): nm = self.s(n.args[0]) if nm: self._emit_col(line, "dataframe_api", "derived", lineage, nm, ".alias() output column") continue if fn not in COL_FUNCS and fn not in NON_COL_FUNCS and isinstance(n.func, ast.Attribute) \ and isinstance(n.func.value, ast.Name) and n.func.value.id in ("F", "functions", "f", "sf", "fn") \ and n.args and isinstance(n.args[0], ast.Constant) and isinstance(n.args[0].value, str) \ and re.fullmatch(r"[A-Za-z_][\w.]*", n.args[0].value): self._emit_col(line, "dataframe_api", "source", lineage, n.args[0].value.rpartition(".")[2], f"F.{fn}('{n.args[0].value}')") continue if fn in COL_FUNCS and n.args: nm = self.s(n.args[0]) if nm and nm != "*": tbl_hint, _, col = nm.rpartition(".") self._emit_col(line, "dataframe_api", "source", lineage, col or nm, f"col('{nm}')" + (" (qualified by df alias)" if tbl_hint else "")) elif fn == "expr" and n.args: ex = self.s(n.args[0]) if ex: self._sql_expr_cols(ex, line, lineage, "F.expr()") elif isinstance(n, ast.Subscript) and isinstance(n.value, ast.Name) and n.value.id in self.dfs: nm = self.s(n.slice) if nm: self._emit_col(line, "dataframe_api", "source", self.dfs[n.value.id], nm, f"{n.value.id}['{nm}']") elif isinstance(n, ast.Attribute) and isinstance(n.value, ast.Name) and n.value.id in self.dfs \ and not n.attr.startswith("_") and n.attr not in DF_METHODS: self._emit_col(line, "dataframe_api", "source", self.dfs[n.value.id], n.attr, f"{n.value.id}.{n.attr}") def _sql_expr_cols(self, ex: str, line, lineage, how): try: tree = sqlglot.parse_one(preprocess_sql(ex), read=self.spark_dialect) except Exception: return alias = tree.alias if isinstance(tree, exp.Alias) else "" for c in tree.find_all(exp.Column): if c.name and c.name != "*": self._emit_col(line, "dataframe_api", "source", lineage, c.name, how) if alias: self._emit_col(line, "dataframe_api", "derived", lineage, alias, how + " AS alias") def _read(self, line, method, fmt, options, args, kws) -> list[TableRef]: fmt = fmt or method out = [] if fmt in ("jdbc", "snowflake", "net.snowflake.spark.snowflake", "redshift", "io.github.spark_redshift_community.spark.redshift") \ or "dbtable" in options or "query" in options: if options.get("query"): self._run_sql(options["query"], line, f"{fmt}_query") return self._sql_lineage(options["query"]) dbt = options.get("dbtable") or (self.s(kws.get("table")) if kws.get("table") else None) \ or (self.s(args[1]) if method == "jdbc" and len(args) > 1 else None) if dbt: if SQL_START_RE.match(dbt.strip("() ")): self._run_sql(dbt.strip("() ").rsplit(")", 1)[0] if dbt.strip().startswith("(") else dbt, line, f"{fmt}_query") return self._sql_lineage(dbt) tr = TableRef(database=options.get("sfdatabase", ""), schema=options.get("sfschema", ""), table="", note=f"{fmt} read") parts = dbt.split(".") tr.table = parts[-1] if len(parts) >= 2: tr.schema = parts[-2] if len(parts) >= 3: tr.database = parts[-3] self._emit_tbl(line, f"{fmt}_read", "source", tr) return [tr] if method == "table" and args: nm = self.s(args[0]) if nm: tr = self._spark_table(nm) self._emit_tbl(line, "spark_table", "source", tr) return [tr] paths = [] for a in args: if isinstance(a, (ast.List, ast.Tuple)): paths += [self.s(e) for e in a.elts] else: paths.append(self.s(a)) paths.append(self.s(kws.get("path")) if "path" in kws else None) paths.append(options.get("path")) for p in filter(None, paths): out.append(self._path_ref(line, p, "source", "spark_read")) if not out and args: out.append(TableRef(table="{" + ast.unparse(args[0]) + "}", note="path not statically known")) self._emit_tbl(line, "spark_read", "source", out[-1]) return out def _spark_table(self, name: str) -> TableRef: if name.lower() in self.views and len(self.views[name.lower()]) == 1: v = self.views[name.lower()][0] return TableRef(v.database, v.schema, v.table, v.physical_location, f"via temp view '{name}'") parts = name.replace("`", "").split(".") parts = [""] * (3 - len(parts)) + parts[-3:] return TableRef(*parts) def _table_name(self, name, schema_node=None) -> TableRef: tr = self._spark_table(name) if schema_node is not None and not tr.schema: tr.schema = self.s(schema_node) or "" return tr def _path_ref(self, line, path, role, ctx) -> TableRef: tr, parts = table_from_path(path) self._emit_tbl(line, ctx, role, tr, resolution="derived from path") for p in parts: self._emit_col(line, ctx, "partition", [tr], p, "hive partition in path") return tr def _sql_lineage(self, sql: str) -> list[TableRef]: """Source tables of a SQL string (used as the DataFrame's lineage).""" try: tree = sqlglot.parse_one(preprocess_sql(sql), read=self.spark_dialect) except Exception: try: tree = sqlglot.parse_one(preprocess_sql(sql), read="snowflake") except Exception: return [] ctes = {c.alias_or_name.lower() for c in tree.find_all(exp.CTE)} out, seen = [], set() for t in tree.find_all(exp.Table): if t.name.lower() in ctes and not t.db: continue name = ".".join(p for p in (t.catalog, t.db, t.name) if p) refs = self.views.get(name.lower()) or [self._spark_table(name)] # expands view chains for tr in refs: tr = TableRef(*(restore_placeholders(x) for x in (tr.database, tr.schema, tr.table)), physical_location=tr.physical_location, note=tr.note) k = (tr.fqn.lower(), tr.physical_location) if k not in seen: seen.add(k) out.append(tr) return out def _run_sql(self, sql: str, line: int, ctx: str): for v, s_, ln in self.pending_sql: if s_ == sql: line = ln if (sql, line) in self.sql_seen: return self.sql_seen.add((sql, line)) for i, (v, s, ln) in enumerate(self.pending_sql): if s == sql: self.sql_seen.add((s, ln)) dialect = self.spark_dialect if ctx.startswith("spark") else self.sql_dialect SqlAnalyzer(self.emit, self.file, self.language, ctx, self.views, dialect).analyze_script(sql, line) def _emit_tbl(self, line, ctx, role, tr: TableRef, resolution=""): self.emit(Row(self.file, line, self.language, ctx, role, tr.database, tr.schema, tr.table, "", tr.fqn, tr.physical_location, resolution, tr.note, end_line=max(line, getattr(self, "cur_end", 0) or line), written=tr.physical_location or tr.fqn)) def _emit_col(self, line, ctx, role, lineage, col, how): if not lineage: self.emit(Row(self.file, line, self.language, ctx, role, column=col, full_name=col, resolution=how, note="DataFrame source unknown", end_line=max(line, getattr(self, "cur_end", 0) or line))) return amb = len({t.fqn for t in lineage}) > 1 for tr in lineage: self.emit(Row(self.file, line, self.language, ctx, role, tr.database, tr.schema, tr.table, col, ".".join(p for p in (tr.fqn, col) if p), tr.physical_location, how + (" — ambiguous: DataFrame has multiple sources" if amb else ""), tr.note, end_line=max(line, getattr(self, "cur_end", 0) or line), written=tr.physical_location or tr.fqn)) NON_COL_FUNCS = {"lit", "typedLit", "expr", "concat_ws", "format_string", "when", "otherwise", "broadcast", "udf", "pandas_udf", "current_date", "current_timestamp", "regexp_replace_pattern", "array_repeat", "create_map", "struct_lit"} DF_METHODS = {"select", "filter", "where", "join", "groupBy", "agg", "withColumn", "drop", "alias", "show", "count", "collect", "write", "cache", "persist", "columns", "schema", "dtypes", "orderBy", "union", "unionByName", "distinct", "limit", "head", "toPandas", "rdd", "na", "printSchema", "createOrReplaceTempView", "withColumnRenamed", "selectExpr", "sort", "explain", "take", "first", "describe", "iloc", "loc", "shape", "empty", "values", "index", "dropDuplicates"} # ───────────────────────────────────────────────────────────────────────────── # File discovery + dispatch # ───────────────────────────────────────────────────────────────────────────── OUTPUT_CSV = "table_attribute_inventory.csv" OUTPUT_HTML = "code_snapshot.html" SKIP_EXT = {".csv", ".html", ".htm", ".md", ".pdf", ".png", ".jpg", ".jpeg", ".gif", ".zip", ".gz", ".xlsx", ".xls", ".docx", ".pptx", ".pyc", ".exe", ".dll", ".so", ".parquet", ".avro", ".orc", ".log", ".lock", ".ini", ".cfg", ".toml", ".env"} def sniff_language(text: str) -> Optional[str]: """For files without a .sql/.py extension: decide whether they hold Python or SQL.""" if "\x00" in text[:4000]: return None try: ast.parse(text) if re.search(r"^\s*(import|from\s+\S+\s+import|def |class )", text, re.M): return "python" except SyntaxError: pass if re.search(r"\b(select|insert\s+into|merge\s+into|create\s+(or\s+replace\s+)?(table|view)|copy\s+into)\b", text, re.I): return "sql" return None class Extractor: def __init__(self, dialect="snowflake", spark_dialect="spark"): self.rows: list[Row] = [] self._keys = set() self.dialect = dialect self.spark_dialect = spark_dialect self.errors: list[str] = [] self.sources: dict[str, str] = {} # file key (file or file#cellN) → its text self.order: list[str] = [] # file keys in display order def emit(self, r: Row): if not (r.table or r.column or r.physical_location): return k = r.key() if k in self._keys: return self._keys.add(k) self.rows.append(r) def run_folder(self, folder: Path, skip: set[str]): for f in sorted(folder.iterdir(), key=lambda x: x.name.lower()): if not f.is_file() or f.name.startswith(".") or f.name in skip: continue if f.suffix.lower() in SKIP_EXT or f.stat().st_size > 25_000_000: continue try: self.file(f) except Exception as e: # never let one file kill the run self.errors.append(f"{f.name}: {type(e).__name__}: {e}") return self.rows def file(self, f: Path): text = f.read_text(encoding="utf-8", errors="replace") name = f.name ext = f.suffix.lower() lang = {"sql": "sql", ".sql": "sql", ".py": "python", ".ipynb": "notebook"}.get(ext) if lang is None: lang = sniff_language(text) if lang is None: return if lang == "sql": self._register(name, text) SqlAnalyzer(self.emit, name, "sql", "sql_file", dialect=self.dialect).analyze_script( text, 1, python_cb=lambda code, ln: PyAnalyzer(self.emit, name, "python", self.spark_dialect, self.dialect).analyze(code, ln)) elif lang == "python": self._register(name, text) PyAnalyzer(self.emit, name, "python", self.spark_dialect, self.dialect).analyze(text) elif lang == "notebook": nb = json.loads(text) pa = PyAnalyzer(self.emit, name, "notebook", self.spark_dialect, self.dialect) for i, cell in enumerate(nb.get("cells", [])): if cell.get("cell_type") != "code": continue src = cell.get("source", []) src = "".join(src) if isinstance(src, list) else str(src) key = f"{name} (cell {i + 1})" self._register(key, src) if src.lstrip().startswith("%%sql"): body = src.split("\n", 1)[1] if "\n" in src else "" SqlAnalyzer(self.emit, key, "notebook", "sql_cell", pa.views, self.spark_dialect).analyze_script(body, 2) continue pa.file = key pa.analyze(src, 1) def _register(self, key, text): self.sources[key] = text self.order.append(key) # ───────────────────────────────────────────────────────────────────────────── # Exact line numbers: find where each table / attribute actually appears # ───────────────────────────────────────────────────────────────────────────── def _token_re(tok: str) -> Optional[re.Pattern]: tok = (tok or "").strip() if not tok or tok == "*": return None if looks_like_path(tok): # paths: match literally (placeholders → wildcard) parts = re.split(r"\{[^}]*\}", tok.rstrip("/")) return re.compile(r".{0,80}?".join(re.escape(p) for p in parts if p) or re.escape(tok), re.I) # dotted names may be written with "quotes" / `backticks` around each part parts = [re.escape(x) for x in tok.split(".") if x] q = r'["`\[\]]?' body = (q + r"\s*\.\s*" + q).join(parts) return re.compile(r"(? list[int]: if pattern is None: return [] out = [] for n in range(max(1, start), min(len(lines), max(start, end)) + 1): if pattern.search(lines[n - 1]): out.append(n) return out def _table_aliases(r: Row, lines: list[str], start: int, end: int) -> set[str]: """Names this table goes by inside its statement: its own name plus any alias after it.""" names = {x.lower() for x in (r.table, r.written.split(".")[-1] if r.written else "") if x} toks = [t for t in (r.written, r.table) if t and not looks_like_path(t)] if not toks: return names stop = {"on", "where", "join", "left", "right", "inner", "outer", "full", "cross", "using", "group", "order", "qualify", "having", "union", "set", "when", "limit", "natural", "lateral", "values", "select", "from", "as"} text = "\n".join(lines[max(0, start - 1):end]) for tok in toks: pat = _token_re(tok) for m in re.finditer(pat.pattern + r"\s+(?:as\s+)?([A-Za-z_]\w*)", text, re.I): if m.group(1).lower() not in stop: names.add(m.group(1).lower()) return names def _column_lines(r: Row, lines: list[str], start: int, end: int) -> list[int]: """Lines where the column appears either bare or qualified by THIS table's name/alias.""" pat = _token_re(r.column) if pat is None: return [] # target columns named by a SELECT list are produced by *source* aliases, so don't filter those aliases = set() if (r.role in ("target", "defined") or not r.table) else _table_aliases(r, lines, start, end) qual = re.compile(r"([A-Za-z_]\w*)[\"`\]]?\s*\.\s*[\"`\[]?$") out = [] for n in range(max(1, start), min(len(lines), max(start, end)) + 1): line = lines[n - 1] for m in pat.finditer(line): q = qual.search(line[:m.start()]) if q is None or not aliases or q.group(1).lower() in aliases: out.append(n) break if not out and aliases: # alias couldn't be worked out — better to over-report than lose it r2 = Row(**{**asdict(r), "role": "target"}) return _column_lines(r2, lines, start, end) return out def _quoted_literal_lines(lines: list[str], name: str) -> list[int]: """Lines where the name is written inside a string literal anywhere in the file.""" if not name or name == "*": return [] pat = re.compile(r"""["'][^"'\n]*?(? list[int]: """Lines holding an S3/stage path string whose folders include this table name.""" if not table: return [] pat = re.compile(r"""(s3[an]?://|gs://|abfss?://|hdfs://|dbfs:|@|/)[^"'\s]*?/""" + re.escape(table) + r"""(/|["'\s]|$)""", re.I) return [n for n, l in enumerate(lines, 1) if pat.search(l)] def precise_lines(r: Row, lines: list[str]) -> list[int]: start, end = r.line, max(r.end_line or r.line, r.line) if r.column: hits = _column_lines(r, lines, start, end) if r.column == "*": hits = _lines_matching(lines, start, end, re.compile(r"(?:\bselect\s+|\.|,\s*)\*", re.I)) elif not hits: # e.g. cols = ["a", "b"] defined elsewhere hits = _quoted_literal_lines(lines, r.column) hits = [h for h in hits if h <= end] or hits return sorted(set(hits)) or [start] hits: list[int] = [] for tok in (r.written, r.full_name, r.table): # most specific spelling first hits = _lines_matching(lines, start, end, _token_re(tok)) if hits: break extra = [] if r.physical_location: # path built in a variable / f-string elsewhere extra = _path_literal_lines(lines, r.table) elif not hits: # table name held in a variable elsewhere for tok in (r.written, r.full_name): extra = _quoted_literal_lines(lines, tok) if extra: break found = sorted(set(hits) | {x for x in extra if x <= end}) if not hits: found = sorted(set(found) | {start}) # keep the line where it is actually used return found or [start] # ───────────────────────────────────────────────────────────────────────────── # CSV: one row per File + Table + Attribute, every line number on that row # ───────────────────────────────────────────────────────────────────────────── CSV_HEADER = ["File Name", "Table Name", "Attribute Name", "Line Numbers", "Times Found", "Source Location (S3 / Stage)"] def needed_rows(rows: list[Row]) -> list[Row]: """Every table and every attribute the code references - read, written, created, or partitioned. Nothing is filtered out: if a name appears in the script, it has to be in the inventory so it can be mapped/replaced. (Rows are only grouped and de-duplicated, never dropped.) """ return list(rows) def table_display(r: Row) -> str: if r.table: return ".".join(p for p in (r.database, r.schema, r.table) if p) return "(unresolved)" if r.column else "" def build_inventory(ex: Extractor): """Group rows → {(file, table, attribute): info} and collect highlight targets per file/line.""" groups: dict[tuple, dict] = {} highlights: dict[str, dict[int, list[tuple[str, str, str]]]] = {} # file → line → [(token, kind, label)] for r in needed_rows(ex.rows): lines = ex.sources.get(r.file, "").splitlines() tbl = table_display(r) if not tbl: continue found = precise_lines(r, lines) key = (r.file, tbl.upper(), (r.column or "").upper()) g = groups.setdefault(key, {"file": r.file, "table": tbl, "attr": r.column or "", "lines": set(), "loc": set()}) g["lines"].update(found) if r.physical_location: g["loc"].add(r.physical_location) label = tbl + (f".{r.column}" if r.column else "") per_file = highlights.setdefault(r.file, {}) for n in found: if r.column: per_file.setdefault(n, []).append((r.column, "attr", label)) for tok in (r.written, r.full_name if not r.column else "", r.table): if tok: per_file.setdefault(n, []).append((tok, "table", label)) # a table whose attributes are needed always gets its own row, even if the script also writes to it for (f, tkey, akey) in list(groups): if akey and (f, tkey, "") not in groups: lines = ex.sources.get(f, "").splitlines() found = set() for r in ex.rows: if r.file == f and not r.column and table_display(r).upper() == tkey: found.update(precise_lines(r, lines)) g0 = groups[(f, tkey, akey)] groups[(f, tkey, "")] = {"file": f, "table": g0["table"], "attr": "", "lines": found or set(g0["lines"]), "loc": set(g0["loc"])} for n in found: highlights.setdefault(f, {}).setdefault(n, []).append((g0["table"].split(".")[-1], "table", g0["table"])) order = {f: i for i, f in enumerate(ex.order)} rows = sorted(groups.values(), key=lambda g: (order.get(g["file"], 9999), g["table"].upper(), g["attr"] != "", g["attr"].upper())) return rows, highlights def write_csv(rows, path: Path): # utf-8-sig = UTF-8 with BOM: opens cleanly in Google Sheets *and* Excel with open(path, "w", newline="", encoding="utf-8-sig") as fh: w = csv.writer(fh, quoting=csv.QUOTE_MINIMAL) w.writerow(CSV_HEADER) for g in rows: ln = sorted(g["lines"]) w.writerow([g["file"], g["table"], g["attr"], ", ".join(map(str, ln)), len(ln), " | ".join(sorted(g["loc"]))]) # ───────────────────────────────────────────────────────────────────────────── # HTML snapshot: all code on one page, extracted names highlighted # ───────────────────────────────────────────────────────────────────────────── def _esc(t: str) -> str: return t.replace("&", "&").replace("<", "<").replace(">", ">").replace('"', """) def _render_line(text: str, marks: list[tuple[str, str, str]]) -> str: spans = [] # (start, end, kind, label) for tok, kind, label in marks: pat = _token_re(tok) if pat is None: continue for m in pat.finditer(text): spans.append((m.start(), m.end(), kind, label)) if not spans: return _esc(text) # longest-first, drop overlaps; merge labels for identical spans spans.sort(key=lambda x: (-(x[1] - x[0]), x[0])) chosen: list[list] = [] for a, b, kind, label in spans: same = next((c for c in chosen if c[0] == a and c[1] == b), None) if same: if label not in same[3]: same[3].append(label) if kind == "table": same[2] = "table" continue if any(not (b <= c[0] or a >= c[1]) for c in chosen): continue chosen.append([a, b, kind, [label]]) chosen.sort() out, pos = [], 0 for a, b, kind, labels in chosen: out.append(_esc(text[pos:a])) out.append(f'{_esc(text[a:b])}') pos = b out.append(_esc(text[pos:])) return "".join(out) def write_html(ex: Extractor, rows, highlights, path: Path, folder: Path): n_tables = len({g["table"].upper() for g in rows}) n_attrs = sum(1 for g in rows if g["attr"]) parts = [] nav = [] for idx, key in enumerate(ex.order): text = ex.sources[key] marks = highlights.get(key, {}) anchor = f"f{idx}" hit_count = len(marks) nav.append(f'{_esc(key)} {hit_count}') body = [] for n, line in enumerate(text.splitlines() or [""], 1): m = marks.get(n) cls = "hit" if m else "dim" body.append(f'
{n}{_render_line(line, m) if m else _esc(line)}' f'
') parts.append(f'

{_esc(key)}

{"".join(body)}
' f'
') html = f""" Code Snapshot

Code Snapshot

{_esc(str(folder))} · {len(ex.order)} files · {n_tables} tables · {n_attrs} table/attribute pairs · details in {OUTPUT_CSV}
table table / source path attribute attributegrey = no table or attribute on that line
{''.join(parts)} """ path.write_text(html, encoding="utf-8") # ───────────────────────────────────────────────────────────────────────────── # Run # ───────────────────────────────────────────────────────────────────────────── def main(): folder = Path(sys.argv[1]).resolve() if len(sys.argv) > 1 else Path(__file__).resolve().parent me = Path(__file__).name ex = Extractor() ex.run_folder(folder, skip={me, OUTPUT_CSV, OUTPUT_HTML}) rows, highlights = build_inventory(ex) csv_path, html_path = folder / OUTPUT_CSV, folder / OUTPUT_HTML write_csv(rows, csv_path) write_html(ex, rows, highlights, html_path, folder) print(f"Scanned {len(ex.order)} file(s) in {folder}") print(f" {len(rows)} rows -> {csv_path.name}") print(f" code snapshot -> {html_path.name}") for e in ex.errors: print(" could not read", e) if __name__ == "__main__": main()