bricksentry 0.1.0__py3-none-any.whl

bricksentry/github.py

@@ -0,0 +1,52 @@
+"""Post or update a single bricksentry comment on a GitHub pull request (stdlib only)."""
+
+from __future__ import annotations
+
+import json
+import os
+import urllib.request
+
+from bricksentry.report import MARKER
+
+
+def _req(method: str, url: str, token: str, body: dict | None = None):
+    data = json.dumps(body).encode() if body is not None else None
+    req = urllib.request.Request(url, data=data, method=method)
+    req.add_header("Authorization", f"Bearer {token}")
+    req.add_header("Accept", "application/vnd.github+json")
+    req.add_header("X-GitHub-Api-Version", "2022-11-28")
+    if data:
+        req.add_header("Content-Type", "application/json")
+    with urllib.request.urlopen(req, timeout=30) as r:
+        return json.loads(r.read() or b"null")
+
+
+def upsert_pr_comment(
+    markdown: str,
+    repo: str | None = None,
+    pr: int | str | None = None,
+    token: str | None = None,
+    api: str | None = None,
+) -> str:
+    repo = repo or os.environ["GITHUB_REPOSITORY"]
+    token = token or os.environ["GITHUB_TOKEN"]
+    api = (api or os.environ.get("GITHUB_API_URL") or "https://api.github.com").rstrip("/")
+    if pr is None:
+        event_path = os.environ.get("GITHUB_EVENT_PATH")
+        if event_path and os.path.exists(event_path):
+            with open(event_path) as f:
+                pr = json.load(f).get("pull_request", {}).get("number")
+    if not pr:
+        raise ValueError("Pull request number not found; pass --pr")
+    base = f"{api}/repos/{repo}/issues/{pr}/comments"
+    page = 1
+    while True:
+        comments = _req("GET", f"{base}?per_page=100&page={page}", token) or []
+        for c in comments:
+            if MARKER in (c.get("body") or ""):
+                _req("PATCH", f"{api}/repos/{repo}/issues/comments/{c['id']}", token, {"body": markdown})
+                return c.get("html_url", "")
+        if len(comments) < 100:
+            break
+        page += 1
+    return (_req("POST", base, token, {"body": markdown}) or {}).get("html_url", "")

bricksentry/lineage.py

@@ -0,0 +1,224 @@
+"""Downstream impact ("blast radius").
+
+``UnityCatalogLineage`` reads Databricks lineage system tables, which are filled
+at *runtime* from query plans. That means it sees PySpark notebooks, Jobs,
+Lakeflow pipelines, dashboards and SQL queries, not just SQL you can parse.
+
+``StaticLineage`` takes edges from config, for local work, tests and non-UC estates.
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Any
+
+from bricksentry.backends.base import Backend
+from bricksentry.models import ImpactNode
+
+_ENTITY_KINDS = {
+    "NOTEBOOK": "notebook",
+    "JOB": "job",
+    "PIPELINE": "pipeline",
+    "DLT_PIPELINE": "pipeline",
+    "DASHBOARD": "dashboard",
+    "DASHBOARD_V3": "dashboard",
+    "DBSQL_DASHBOARD": "dashboard",
+    "DBSQL_QUERY": "query",
+    "SQL_QUERY": "query",
+    "GENIE_SPACE": "genie_space",
+}
+
+
+class LineageProvider(ABC):
+    @abstractmethod
+    def downstream(self, table: str, depth: int = 3) -> list[ImpactNode]: ...
+
+    def owner_and_tags(self, table: str) -> tuple[str | None, dict[str, str]]:
+        return None, {}
+
+
+class StaticLineage(LineageProvider):
+    """Lineage from config::
+
+        lineage:
+          type: static
+          edges:
+            main.sales.orders: [main.sales.daily_revenue, "dashboard:Revenue Overview"]
+          owners:
+            main.sales.daily_revenue: finance-team
+          tags:
+            main.sales.daily_revenue: {certified: "true"}
+    """
+
+    def __init__(self, edges: dict[str, list[str]] | None = None, owners: dict | None = None, tags: dict | None = None):
+        self.edges = edges or {}
+        self.owners = owners or {}
+        self.tags = tags or {}
+
+    def downstream(self, table: str, depth: int = 3) -> list[ImpactNode]:
+        out: list[ImpactNode] = []
+        seen = {table}
+        frontier = [table]
+        for d in range(1, depth + 1):
+            nxt = []
+            for src in frontier:
+                for tgt in self.edges.get(src, []):
+                    kind, name = _parse_static(tgt)
+                    if name in seen:
+                        continue
+                    seen.add(name)
+                    out.append(
+                        ImpactNode(
+                            kind=kind,
+                            name=name,
+                            depth=d,
+                            via=src,
+                            owner=self.owners.get(name),
+                            tags={str(k): str(v) for k, v in (self.tags.get(name) or {}).items()},
+                        )
+                    )
+                    if kind == "table":
+                        nxt.append(name)
+            frontier = nxt
+            if not frontier:
+                break
+        return out
+
+    def owner_and_tags(self, table: str) -> tuple[str | None, dict[str, str]]:
+        return self.owners.get(table), {str(k): str(v) for k, v in (self.tags.get(table) or {}).items()}
+
+
+def _parse_static(s: str) -> tuple[str, str]:
+    if ":" in s and not s.startswith(("http:", "https:")):
+        kind, name = s.split(":", 1)
+        return kind.strip().lower(), name.strip()
+    return "table", s
+
+
+class UnityCatalogLineage(LineageProvider):
+    """Breadth-first walk over ``system.access.table_lineage``.
+
+    One query per depth level (not per node), bounded by ``lookback_days``.
+    Table owners come from ``information_schema.tables`` and tags from
+    ``information_schema.table_tags``. All table names are configurable so the
+    same code runs against copies of the system tables in tests.
+    """
+
+    def __init__(
+        self,
+        backend: Backend,
+        lookback_days: int = 90,
+        lineage_table: str = "system.access.table_lineage",
+        tables_table: str = "system.information_schema.tables",
+        tags_table: str = "system.information_schema.table_tags",
+        max_nodes: int = 500,
+    ):
+        self.backend = backend
+        self.lookback_days = lookback_days
+        self.lineage_table = lineage_table
+        self.tables_table = tables_table
+        self.tags_table = tags_table
+        self.max_nodes = max_nodes
+        self._meta_cache: dict[str, tuple[str | None, dict[str, str]]] = {}
+
+    def downstream(self, table: str, depth: int = 3) -> list[ImpactNode]:
+        out: list[ImpactNode] = []
+        seen_tables = {table.lower()}
+        seen_entities: set[tuple[str, str]] = set()
+        frontier = [table]
+        for d in range(1, depth + 1):
+            if not frontier or len(out) >= self.max_nodes:
+                break
+            in_list = ", ".join(_lit(t) for t in frontier)
+            sql = (
+                "SELECT DISTINCT source_table_full_name AS src, target_table_full_name AS tgt, "
+                "entity_type, CAST(entity_id AS STRING) AS entity_id "
+                f"FROM {self.lineage_table} "
+                f"WHERE lower(source_table_full_name) IN ({in_list.lower()}) "
+                f"AND event_date >= {self.backend.days_ago(self.lookback_days)}"
+            )
+            rows = self.backend.query(sql)
+            nxt: list[str] = []
+            for r in rows:
+                src = r.get("src")
+                tgt = r.get("tgt")
+                if tgt and tgt.lower() not in seen_tables:
+                    seen_tables.add(tgt.lower())
+                    nxt.append(tgt)
+                    out.append(ImpactNode(kind="table", name=tgt, depth=d, via=src))
+                et = (r.get("entity_type") or "").upper()
+                eid = r.get("entity_id")
+                if et and eid and (et, eid) not in seen_entities:
+                    seen_entities.add((et, eid))
+                    out.append(
+                        ImpactNode(kind=_ENTITY_KINDS.get(et, et.lower()), name=f"{et.lower()}:{eid}", depth=d, via=src)
+                    )
+            frontier = nxt
+        out = out[: self.max_nodes]
+        self._attach_meta([n for n in out if n.kind == "table"])
+        return out
+
+    def owner_and_tags(self, table: str) -> tuple[str | None, dict[str, str]]:
+        if table.lower() not in self._meta_cache:
+            self._load_meta([table])
+        return self._meta_cache.get(table.lower(), (None, {}))
+
+    def _attach_meta(self, nodes: list[ImpactNode]) -> None:
+        names = [n.name for n in nodes if n.name.lower() not in self._meta_cache]
+        if names:
+            self._load_meta(names)
+        for n in nodes:
+            n.owner, n.tags = self._meta_cache.get(n.name.lower(), (None, {}))
+
+    def _load_meta(self, names: list[str]) -> None:
+        for name in names:
+            self._meta_cache.setdefault(name.lower(), (None, {}))
+        full = "lower(concat_ws('.', table_catalog, table_schema, table_name))"
+        tag_full = "lower(concat_ws('.', catalog_name, schema_name, table_name))"
+        for i in range(0, len(names), 200):
+            batch = names[i : i + 200]
+            in_list = ", ".join(_lit(n.lower()) for n in batch)
+            try:
+                for r in self.backend.query(
+                    f"SELECT {full} AS fqn, table_owner FROM {self.tables_table} WHERE {full} IN ({in_list})"
+                ):
+                    _, tags = self._meta_cache[r["fqn"]]
+                    self._meta_cache[r["fqn"]] = (r.get("table_owner"), tags)
+            except Exception:
+                pass  # missing privileges on information_schema should not break the check
+            try:
+                for r in self.backend.query(
+                    f"SELECT {tag_full} AS fqn, tag_name, tag_value FROM {self.tags_table} "
+                    f"WHERE {tag_full} IN ({in_list})"
+                ):
+                    owner, tags = self._meta_cache[r["fqn"]]
+                    tags = dict(tags)
+                    tags[str(r["tag_name"])] = "" if r.get("tag_value") is None else str(r["tag_value"])
+                    self._meta_cache[r["fqn"]] = (owner, tags)
+            except Exception:
+                pass
+
+
+def _lit(s: str) -> str:
+    s = str(s)
+    if any(ch in s for ch in "'\\\n\r;"):
+        # Table names can come from agents; refuse anything that could break out of a literal.
+        raise ValueError(f"Refusing suspicious table name in lineage lookup: {s!r}")
+    return "'" + s + "'"
+
+
+def make_lineage(cfg: dict[str, Any] | None, backend: Backend) -> LineageProvider:
+    cfg = dict(cfg or {})
+    kind = cfg.get("type") or ("databricks" if backend.name == "databricks" else "static")
+    if kind == "static":
+        return StaticLineage(cfg.get("edges"), cfg.get("owners"), cfg.get("tags"))
+    if kind in ("databricks", "unity_catalog", "uc"):
+        return UnityCatalogLineage(
+            backend,
+            lookback_days=int(cfg.get("lookback_days", 90)),
+            lineage_table=cfg.get("lineage_table", "system.access.table_lineage"),
+            tables_table=cfg.get("tables_table", "system.information_schema.tables"),
+            tags_table=cfg.get("tags_table", "system.information_schema.table_tags"),
+            max_nodes=int(cfg.get("max_nodes", 500)),
+        )
+    raise ValueError(f"Unknown lineage type {kind!r}")

bricksentry/mcp_server.py

@@ -0,0 +1,47 @@
+"""MCP server exposing bricksentry's agent tools.
+
+Run:  ``bricksentry mcp``  (stdio)  or  ``bricksentry mcp --transport streamable-http``
+
+Claude Code:  ``claude mcp add bricksentry -- bricksentry mcp --config /path/bricksentry.yml``
+"""
+
+from __future__ import annotations
+
+from bricksentry import agent_tools
+
+INSTRUCTIONS = """bricksentry checks whether a change to Databricks data is safe.
+Recommended loop when you change a pipeline or table:
+1. whats_downstream(table) to learn the blast radius before editing.
+2. rollback_plan(table) and keep the SQL.
+3. preview_write(sql, target) for any direct write, or assess_change(tables) after
+   your change has built candidate tables.
+4. safe_to_ship(assessment_id) and follow `next_step` exactly.
+Never bypass a needs_human or block verdict."""
+
+
+def build_server():
+    try:  # mcp >= 2
+        from mcp.server.mcpserver import MCPServer as Server
+    except ImportError:
+        try:  # mcp 1.x
+            from mcp.server.fastmcp import FastMCP as Server
+        except ImportError as e:  # pragma: no cover
+            raise ImportError("The MCP server needs `pip install bricksentry[mcp]`") from e
+
+    try:
+        from mcp.types import ToolAnnotations
+    except ImportError:  # pragma: no cover
+        ToolAnnotations = None
+
+    server = Server(name="bricksentry", instructions=INSTRUCTIONS)
+    for fn in agent_tools.TOOLS:
+        ann = None
+        if ToolAnnotations is not None:
+            ro = fn.__name__ in agent_tools.READ_ONLY
+            ann = ToolAnnotations(readOnlyHint=ro, destructiveHint=False, idempotentHint=ro)
+        server.tool(name=fn.__name__, description=fn.__doc__, annotations=ann)(fn)
+    return server
+
+
+def serve(transport: str = "stdio") -> None:
+    build_server().run(transport=transport)

bricksentry/models.py

@@ -0,0 +1,212 @@
+"""Result types. Every result is a plain dataclass with a stable ``to_dict()``,
+so agents can branch on fields and humans get the same data rendered as markdown."""
+
+from __future__ import annotations
+
+from dataclasses import asdict, dataclass, field
+from enum import Enum
+from typing import Any
+
+
+class Verdict(str, Enum):
+    """Ordered from safest to most severe."""
+
+    PASS = "pass"
+    WARN = "warn"
+    NEEDS_HUMAN = "needs_human"
+    BLOCK = "block"
+
+    @property
+    def rank(self) -> int:
+        return [Verdict.PASS, Verdict.WARN, Verdict.NEEDS_HUMAN, Verdict.BLOCK].index(self)
+
+    @property
+    def risk_tier(self) -> str:
+        return {
+            Verdict.PASS: "low",
+            Verdict.WARN: "medium",
+            Verdict.NEEDS_HUMAN: "high",
+            Verdict.BLOCK: "blocked",
+        }[self]
+
+    @staticmethod
+    def worst(verdicts: list["Verdict"]) -> "Verdict":
+        return max(verdicts, key=lambda v: v.rank, default=Verdict.PASS)
+
+
+@dataclass
+class ColumnChange:
+    column: str
+    changed_rows: int
+    base_nulls: int | None = None
+    candidate_nulls: int | None = None
+    base_sum: float | None = None
+    candidate_sum: float | None = None
+    samples: list[dict[str, Any]] = field(default_factory=list)
+
+    @property
+    def sum_change_pct(self) -> float | None:
+        if self.base_sum in (None, 0) or self.candidate_sum is None:
+            return None
+        return (self.candidate_sum - self.base_sum) / abs(self.base_sum) * 100
+
+
+@dataclass
+class SchemaChange:
+    added: dict[str, str] = field(default_factory=dict)
+    removed: dict[str, str] = field(default_factory=dict)
+    type_changed: dict[str, tuple[str, str]] = field(default_factory=dict)
+
+    @property
+    def is_empty(self) -> bool:
+        return not (self.added or self.removed or self.type_changed)
+
+
+@dataclass
+class Segment:
+    column: str
+    value: Any
+    affected_rows: int
+
+
+@dataclass
+class TableDiff:
+    table: str
+    base: str
+    candidate: str
+    keys: list[str]
+    base_rows: int
+    candidate_rows: int
+    added_rows: int
+    removed_rows: int
+    changed_rows: int
+    unchanged_rows: int
+    schema: SchemaChange
+    columns: list[ColumnChange] = field(default_factory=list)
+    duplicate_keys_base: int = 0
+    duplicate_keys_candidate: int = 0
+    segments: list[Segment] = field(default_factory=list)
+    sampled_pct: float | None = None
+    error: str | None = None
+
+    def _pct(self, n: int) -> float:
+        denom = max(self.base_rows, 1)
+        return round(n / denom * 100, 4)
+
+    @property
+    def added_pct(self) -> float:
+        return self._pct(self.added_rows)
+
+    @property
+    def removed_pct(self) -> float:
+        return self._pct(self.removed_rows)
+
+    @property
+    def changed_pct(self) -> float:
+        return self._pct(self.changed_rows)
+
+    @property
+    def row_count_change_pct(self) -> float:
+        return self._pct(self.candidate_rows - self.base_rows)
+
+    @property
+    def is_identical(self) -> bool:
+        return (
+            self.error is None
+            and self.schema.is_empty
+            and self.added_rows == 0
+            and self.removed_rows == 0
+            and self.changed_rows == 0
+            and self.base_rows == self.candidate_rows
+            and self.duplicate_keys_candidate <= self.duplicate_keys_base
+        )
+
+    def to_dict(self) -> dict[str, Any]:
+        d = asdict(self)
+        d["schema"]["type_changed"] = {
+            k: {"base": v[0], "candidate": v[1]} for k, v in self.schema.type_changed.items()
+        }
+        d["added_pct"] = self.added_pct
+        d["removed_pct"] = self.removed_pct
+        d["changed_pct"] = self.changed_pct
+        d["row_count_change_pct"] = self.row_count_change_pct
+        d["is_identical"] = self.is_identical
+        for c, cc in zip(d["columns"], self.columns):
+            c["sum_change_pct"] = cc.sum_change_pct
+        return d
+
+
+@dataclass
+class ImpactNode:
+    """Something downstream of a changed table: a table, job, notebook, pipeline, dashboard, query..."""
+
+    kind: str
+    name: str
+    depth: int
+    via: str
+    owner: str | None = None
+    tags: dict[str, str] = field(default_factory=dict)
+
+    def to_dict(self) -> dict[str, Any]:
+        return asdict(self)
+
+
+@dataclass
+class RuleHit:
+    rule: str
+    verdict: Verdict
+    table: str
+    reason: str
+
+    def to_dict(self) -> dict[str, Any]:
+        return {"rule": self.rule, "verdict": self.verdict.value, "table": self.table, "reason": self.reason}
+
+
+@dataclass
+class TableAssessment:
+    table: str
+    diff: TableDiff | None
+    downstream: list[ImpactNode]
+    owner: str | None
+    tags: dict[str, str]
+    rollback: str | None
+    verdict: Verdict = Verdict.PASS
+    hits: list[RuleHit] = field(default_factory=list)
+    summary: str = ""
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "table": self.table,
+            "verdict": self.verdict.value,
+            "risk_tier": self.verdict.risk_tier,
+            "owner": self.owner,
+            "tags": self.tags,
+            "summary": self.summary,
+            "rule_hits": [h.to_dict() for h in self.hits],
+            "diff": self.diff.to_dict() if self.diff else None,
+            "downstream": [n.to_dict() for n in self.downstream],
+            "rollback": self.rollback,
+        }
+
+
+@dataclass
+class Assessment:
+    id: str
+    tables: list[TableAssessment]
+    verdict: Verdict
+    notify: list[str] = field(default_factory=list)
+
+    @property
+    def hits(self) -> list[RuleHit]:
+        return [h for t in self.tables for h in t.hits]
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "id": self.id,
+            "verdict": self.verdict.value,
+            "risk_tier": self.verdict.risk_tier,
+            "safe_to_ship": self.verdict in (Verdict.PASS, Verdict.WARN),
+            "notify": self.notify,
+            "rule_hits": [h.to_dict() for h in self.hits],
+            "tables": [t.to_dict() for t in self.tables],
+        }