sql-lineage-extractor 0.1.1__py3-none-any.whl

sql_lineage/parsing/lineage_extractor.py

@@ -0,0 +1,53 @@
+"""Column-to-column lineage for a single target column, via sqlglot."""
+
+from __future__ import annotations
+
+from sqlglot import exp
+from sqlglot.lineage import Node, lineage
+
+from ..models import ColumnSource
+
+
+def extract_column_lineage(
+    column: str, sql: str, dialect: str
+) -> tuple[str, list[ColumnSource]]:
+    """Return (target expression SQL, sorted upstream (table, column) pairs).
+
+    Walks the sqlglot lineage tree down to its leaves; each leaf backed by a
+    base table becomes one ColumnSource. Results are de-duplicated and sorted so
+    the manifest stays deterministic.
+    """
+    node = lineage(column, sql, dialect=dialect)
+    expression_sql = node.expression.sql(dialect=dialect)
+
+    sources: set[ColumnSource] = set()
+    for leaf in _leaves(node):
+        resolved = _resolve_leaf(leaf)
+        if resolved is not None:
+            sources.add(resolved)
+
+    return expression_sql, sorted(sources, key=lambda s: (s.table, s.column))
+
+
+def _leaves(node: Node):
+    if not node.downstream:
+        yield node
+        return
+    for child in node.downstream:
+        yield from _leaves(child)
+
+
+def _resolve_leaf(leaf: Node) -> ColumnSource | None:
+    source = leaf.source
+    if not isinstance(source, exp.Table):
+        return None
+    table = _qualified_table(source)
+    column = leaf.name.rsplit(".", 1)[-1]
+    if not table or not column:
+        return None
+    return ColumnSource(table=table, column=column)
+
+
+def _qualified_table(table: exp.Table) -> str:
+    parts = [p for p in (table.catalog, table.db, table.name) if p]
+    return ".".join(parts)

sql_lineage/pipeline.py

@@ -0,0 +1,27 @@
+"""Glue the four layers together: source -> parse -> graph -> (nodes, errors)."""
+
+from __future__ import annotations
+
+from .graph import build_graph
+from .models import ObjectNode, ParseError
+from .parsing import parse_unit
+from .sources.base import SQLSource
+
+
+def run_extraction(source: SQLSource) -> tuple[list[ObjectNode], list[ParseError]]:
+    """Extract every unit from ``source`` and resolve the cross-object graph.
+
+    A parse failure on one object is collected as a ParseError and never stops
+    the others.
+    """
+    nodes: list[ObjectNode] = []
+    errors: list[ParseError] = []
+    for unit in source.extract():
+        result = parse_unit(unit)
+        if isinstance(result, ParseError):
+            errors.append(result)
+        else:
+            nodes.append(result)
+    build_graph(nodes)
+    nodes.sort(key=lambda n: n.id)
+    return nodes, errors

sql_lineage/sources/__init__.py

@@ -0,0 +1,13 @@
+from ..config import LineageConfig
+from .base import SQLSource
+from .file_source import FileSQLSource
+from .live_sql_source import LiveSQLEndpointSource
+from .notebook_source import NotebookSQLSource
+
+__all__ = [
+    "SQLSource",
+    "FileSQLSource",
+    "NotebookSQLSource",
+    "LiveSQLEndpointSource",
+    "LineageConfig",
+]

sql_lineage/sources/base.py

@@ -0,0 +1,14 @@
+"""Common source interface. Every source yields SQLUnit objects and nothing else."""
+
+from __future__ import annotations
+
+from typing import Protocol, runtime_checkable
+
+from ..models import SQLUnit
+
+
+@runtime_checkable
+class SQLSource(Protocol):
+    def extract(self) -> list[SQLUnit]:
+        """Return the normalized SQL units this source exposes."""
+        ...

sql_lineage/sources/file_source.py

@@ -0,0 +1,53 @@
+"""FileSQLSource — scan a directory for .sql files, one file == one SQLUnit."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+from ..config import LineageConfig
+from ..models import SQLUnit
+
+
+class FileSQLSource:
+    """Yield one SQLUnit per file matching ``pattern`` under ``root``.
+
+    The ``id`` defaults to the file stem (``escale.sql`` -> ``escale``), but an
+    explicit ``-- lineage:id=`` comment inside the file wins (same convention as
+    notebooks, resolved through the injected :class:`LineageConfig`). Origins are
+    stored as POSIX-style relative paths when possible, for stable, diffable
+    manifests across operating systems.
+    """
+
+    def __init__(
+        self,
+        root: str | Path,
+        dialect: str | None = None,
+        pattern: str = "**/*.sql",
+        *,
+        config: LineageConfig | None = None,
+    ) -> None:
+        self.root = Path(root)
+        self.config = config or LineageConfig()
+        self.dialect = dialect if dialect is not None else self.config.default_dialect
+        self.pattern = pattern
+
+    def extract(self) -> list[SQLUnit]:
+        units: list[SQLUnit] = []
+        for path in sorted(self.root.glob(self.pattern)):
+            if not path.is_file():
+                continue
+            text = path.read_text(encoding="utf-8")
+            try:
+                origine = path.relative_to(self.root).as_posix()
+            except ValueError:
+                origine = path.as_posix()
+            uid = self.config.id_from_comment(text) or path.stem
+            units.append(
+                SQLUnit(
+                    id=uid,
+                    dialecte=self.dialect,
+                    sql_text=text.strip(),
+                    origine=origine,
+                )
+            )
+        return units

sql_lineage/sources/live_sql_source.py

@@ -0,0 +1,81 @@
+"""LiveSQLEndpointSource — read view definitions from a live DB-API connection.
+
+The connection object is *injected*, never built here: authentication (Entra ID
+for Fabric, SQL auth for SQL Server, ...) and driver choice are the caller's
+concern, keeping this library free of any proprietary Fabric SDK dependency.
+
+The default query targets T-SQL system catalogs (``sys.sql_modules``); if that
+fails (permissions, non-SQL-Server engine) it falls back to the ANSI standard
+``information_schema.views``.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from ..config import LineageConfig
+from ..models import SQLUnit
+
+# Full CREATE VIEW ... AS <select> definitions, richest source when available.
+TSQL_QUERY = (
+    "SELECT v.name AS nom_vue, m.definition AS sql_text "
+    "FROM sys.views v "
+    "JOIN sys.sql_modules m ON v.object_id = m.object_id"
+)
+
+# Portable fallback: exposes only the SELECT body, no CREATE VIEW wrapper.
+ANSI_FALLBACK_QUERY = (
+    "SELECT table_name AS nom_vue, view_definition AS sql_text "
+    "FROM information_schema.views"
+)
+
+
+class LiveSQLEndpointSource:
+    def __init__(
+        self,
+        connection: Any,
+        dialect: str | None = None,
+        origine: str = "live",
+        query: str | None = None,
+        fallback_query: str | None = None,
+        *,
+        config: LineageConfig | None = None,
+    ) -> None:
+        self.connection = connection
+        self.config = config or LineageConfig()
+        self.dialect = dialect if dialect is not None else self.config.default_dialect
+        self.origine = origine
+        self.query = query or TSQL_QUERY
+        self.fallback_query = fallback_query or ANSI_FALLBACK_QUERY
+
+    def extract(self) -> list[SQLUnit]:
+        try:
+            rows = self._run(self.query)
+        except Exception:
+            # Any failure on the primary path (missing sys.sql_modules, denied
+            # permissions, unsupported engine) triggers the portable fallback.
+            rows = self._run(self.fallback_query)
+
+        units: list[SQLUnit] = []
+        for name, sql_text in rows:
+            if sql_text is None:
+                continue
+            units.append(
+                SQLUnit(
+                    id=str(name),
+                    dialecte=self.dialect,
+                    sql_text=str(sql_text).strip(),
+                    origine=f"{self.origine}:{name}",
+                )
+            )
+        return units
+
+    def _run(self, query: str) -> list[tuple[Any, Any]]:
+        cursor = self.connection.cursor()
+        try:
+            cursor.execute(query)
+            return [(row[0], row[1]) for row in cursor.fetchall()]
+        finally:
+            close = getattr(cursor, "close", None)
+            if callable(close):
+                close()

sql_lineage/sources/notebook_source.py

@@ -0,0 +1,261 @@
+"""NotebookSQLSource — extract SQL from Jupyter/Fabric ``.ipynb`` notebooks.
+
+Three cell shapes are understood:
+
+* ``%%sql`` magic cells: everything after the magic line is taken verbatim.
+* Python cells calling ``<something>.sql("...")`` with a string literal.
+* Python cells calling ``<something>.sql(f"...")`` with an f-string; each
+  ``{...}`` interpolation is replaced by a neutral token so sqlglot can still
+  parse the statement. The replaced fragments are recorded on the SQLUnit.
+
+Python is parsed with the standard :mod:`ast` module — never with regexes —
+because real notebook code mixes multi-line strings, nested calls and quote
+styles that regexes get wrong.
+
+Every convention (the ``-- lineage:id=`` comment, the ``df_`` variable prefix,
+the neutralization token, the default dialect) comes from an injected
+:class:`~sql_lineage.config.LineageConfig`; nothing is hard-coded here.
+"""
+
+from __future__ import annotations
+
+import ast
+import json
+from pathlib import Path
+from typing import TypeGuard
+
+from ..config import DEFAULT_PARAM_TOKEN, LineageConfig
+from ..models import SQLUnit
+
+# Backward-compatible module-level alias for the default neutralization token.
+# Kept so ``from ...notebook_source import PARAM_TOKEN`` keeps working.
+PARAM_TOKEN = DEFAULT_PARAM_TOKEN
+
+
+class NotebookSQLSource:
+    def __init__(
+        self,
+        root: str | Path,
+        dialect: str | None = None,
+        pattern: str = "**/*.ipynb",
+        *,
+        config: LineageConfig | None = None,
+    ) -> None:
+        self.root = Path(root)
+        self.config = config or LineageConfig()
+        self.dialect = dialect if dialect is not None else self.config.default_dialect
+        self.pattern = pattern
+
+    def extract(self) -> list[SQLUnit]:
+        units: list[SQLUnit] = []
+        for path in sorted(self.root.glob(self.pattern)):
+            if not path.is_file():
+                continue
+            units.extend(self._extract_notebook(path))
+        return units
+
+    # -- per-notebook ---------------------------------------------------------
+
+    def _extract_notebook(self, path: Path) -> list[SQLUnit]:
+        try:
+            origine = path.relative_to(self.root).as_posix()
+        except ValueError:
+            origine = path.as_posix()
+
+        nb = json.loads(path.read_text(encoding="utf-8"))
+        notebook_name = path.stem
+        units: list[SQLUnit] = []
+
+        cell_index = -1
+        for cell in nb.get("cells", []):
+            cell_index += 1
+            if cell.get("cell_type") != "code":
+                continue
+            src = _cell_source(cell)
+            if not src.strip():
+                continue
+
+            magic_sql = _strip_sql_magic(src)
+            if magic_sql is not None:
+                units.append(
+                    self._make_unit(
+                        sql_text=magic_sql,
+                        params=[],
+                        var_name=None,
+                        origine=origine,
+                        notebook_name=notebook_name,
+                        cell_index=cell_index,
+                        call_index=0,
+                    )
+                )
+                continue
+
+            calls = _python_sql_calls(src, self.config.param_token)
+            for call_index, (sql_text, params, var_name) in enumerate(calls):
+                units.append(
+                    self._make_unit(
+                        sql_text=sql_text,
+                        params=params,
+                        var_name=var_name,
+                        origine=origine,
+                        notebook_name=notebook_name,
+                        cell_index=cell_index,
+                        call_index=call_index,
+                    )
+                )
+        return units
+
+    def _make_unit(
+        self,
+        *,
+        sql_text: str,
+        params: list[str],
+        var_name: str | None,
+        origine: str,
+        notebook_name: str,
+        cell_index: int,
+        call_index: int,
+    ) -> SQLUnit:
+        warnings: list[str] = []
+        uid = self.config.id_from_comment(sql_text)
+        if uid is None and var_name:
+            uid = self.config.strip_var_prefix(var_name)
+        if uid is None:
+            uid = f"{notebook_name}#{cell_index}"
+            if call_index:
+                uid = f"{uid}#{call_index}"
+            warnings.append(
+                f"Aucun id explicite trouvé (ni commentaire '-- lineage:id=', ni "
+                f"variable cible) ; repli sur '{uid}'."
+            )
+        return SQLUnit(
+            id=uid,
+            dialecte=self.dialect,
+            sql_text=sql_text.strip(),
+            origine=f"{origine}#cell{cell_index}",
+            parametres_neutralises=params,
+            avertissements=warnings,
+        )
+
+
+# -- module-level helpers -----------------------------------------------------
+
+
+def _cell_source(cell: dict) -> str:
+    src = cell.get("source", "")
+    if isinstance(src, list):
+        return "".join(src)
+    return src
+
+
+def _strip_sql_magic(src: str) -> str | None:
+    """Return the SQL body if ``src`` is a ``%%sql`` cell, else None."""
+    lines = src.splitlines()
+    for i, line in enumerate(lines):
+        if not line.strip():
+            continue
+        if line.strip().startswith("%%sql"):
+            return "\n".join(lines[i + 1 :])
+        return None
+    return None
+
+
+def _python_sql_calls(
+    src: str, param_token: str
+) -> list[tuple[str, list[str], str | None]]:
+    """Find every ``*.sql(<str|fstring>)`` call, returning (sql, params, var)."""
+    try:
+        tree = ast.parse(src)
+    except SyntaxError:
+        return []
+
+    # Map matching Call nodes to the variable they are assigned to (if any).
+    assigned: dict[int, str | None] = {}
+    for node in ast.walk(tree):
+        if isinstance(node, ast.Assign) and _is_sql_call(node.value):
+            assigned[id(node.value)] = _target_name(node.targets)
+
+    calls = [n for n in ast.walk(tree) if _is_sql_call(n)]
+    calls.sort(key=lambda n: (getattr(n, "lineno", 0), getattr(n, "col_offset", 0)))
+
+    results: list[tuple[str, list[str], str | None]] = []
+    for call in calls:
+        extracted = _extract_sql_arg(call.args[0], param_token)
+        if extracted is None:
+            continue
+        sql_text, params = extracted
+        results.append((sql_text, params, assigned.get(id(call))))
+    return results
+
+
+def _is_sql_call(node: ast.AST) -> TypeGuard[ast.Call]:
+    return (
+        isinstance(node, ast.Call)
+        and isinstance(node.func, ast.Attribute)
+        and node.func.attr == "sql"
+        and len(node.args) >= 1
+    )
+
+
+def _target_name(targets: list[ast.expr]) -> str | None:
+    for target in targets:
+        if isinstance(target, ast.Name):
+            return target.id
+    return None
+
+
+def _extract_sql_arg(arg: ast.expr, param_token: str) -> tuple[str, list[str]] | None:
+    """Return (sql_text, neutralized_params) for a string or f-string arg."""
+    if isinstance(arg, ast.Constant) and isinstance(arg.value, str):
+        return arg.value, []
+    if isinstance(arg, ast.JoinedStr):
+        return _neutralize_fstring(arg, param_token)
+    return None
+
+
+def _neutralize_fstring(node: ast.JoinedStr, param_token: str) -> tuple[str, list[str]]:
+    parts: list[str] = []
+    params: list[str] = []
+    values = node.values
+    for i, value in enumerate(values):
+        if isinstance(value, ast.Constant) and isinstance(value.value, str):
+            parts.append(value.value)
+        elif isinstance(value, ast.FormattedValue):
+            try:
+                expr_src = ast.unparse(value.value)
+            except Exception:  # pragma: no cover - very old python fallback
+                expr_src = "?"
+            params.append(expr_src)
+            # An interpolation standing alone on its line is a whole SQL *clause*
+            # (e.g. ``{'' if x is None else 'AND t.c IN ' + x}``), not a value.
+            # Replacing it by the neutral token would leave an unparsable bare
+            # word, so we drop it instead — WHERE/HAVING predicates never feed
+            # column lineage, so removing one is safe.
+            before = "".join(parts)
+            nxt = values[i + 1] if i + 1 < len(values) else None
+            after = (
+                nxt.value
+                if isinstance(nxt, ast.Constant) and isinstance(nxt.value, str)
+                else ""
+            )
+            parts.append("" if _is_statement_fragment(before, after) else param_token)
+    return "".join(parts), params
+
+
+def _is_statement_fragment(before: str, after: str) -> bool:
+    """Whether an interpolation stands alone on its line (a whole SQL clause).
+
+    True when nothing meaningful precedes it on the line (start of string or the
+    preceding text ends the line) and nothing meaningful follows it on the line
+    (end of string, a newline, or a closing paren). Such fragments are dropped
+    rather than tokenized. Value-position interpolations (``>= {x}``, ``'{x}'``,
+    ``IN ({x})``) fail these tests and keep the neutral token.
+    """
+    left_clear = before == "" or before.rstrip(" \t").endswith("\n")
+    stripped_after = after.lstrip(" \t")
+    right_clear = (
+        stripped_after == ""
+        or stripped_after.startswith("\n")
+        or stripped_after.startswith(")")
+    )
+    return left_clear and right_clear