S054 · Governance & Cost — Lineage, PII, Attribution

"""
governance_demo.py — Toy governance: catalog + column lineage + PII delete.
 
Run:
    pip install networkx
    python governance_demo.py
"""
from __future__ import annotations
 
from dataclasses import dataclass, field
from typing import Iterable
 
import networkx as nx
 
 
# ---------- catalog ----------
@dataclass
class Column:
    name: str
    dtype: str
    tags: set[str] = field(default_factory=set)
 
 
@dataclass
class Table:
    name: str
    owner_team: str
    columns: dict[str, Column]
    storage_bytes: int = 0
 
    def tag_column(self, col: str, *tags: str) -> None:
        self.columns[col].tags.update(tags)
 
 
CATALOG: dict[str, Table] = {}
 
 
def register(tbl: Table) -> Table:
    CATALOG[tbl.name] = tbl
    return tbl
 
 
# ---------- lineage graph ----------
LINEAGE = nx.DiGraph()  # nodes: (table, column); edges: source → derived
 
 
def emit_lineage(src_table: str, src_col: str, dst_table: str, dst_col: str) -> None:
    LINEAGE.add_edge((src_table, src_col), (dst_table, dst_col))
    # Propagate tags automatically (this is the ‘PII flows downstream’ magic).
    src_tags = CATALOG[src_table].columns[src_col].tags
    CATALOG[dst_table].columns[dst_col].tags.update(src_tags)
 
 
# ---------- PII discovery + delete ----------
def columns_with_tag(tag: str) -> list[tuple[str, str]]:
    return [
        (t.name, c.name)
        for t in CATALOG.values()
        for c in t.columns.values()
        if tag in c.tags
    ]
 
 
def gdpr_delete(user_email: str) -> list[str]:
    """Given a user's identifier, walk lineage and return every table needing a delete."""
    affected: set[str] = set()
    # Start from every table containing an ‘email’-tagged column classified as PII.
    for tbl_name, col_name in columns_with_tag("pii.email"):
        # Any downstream node from this (table,col) is affected too.
        node = (tbl_name, col_name)
        affected.add(tbl_name)
        for descendant in nx.descendants(LINEAGE, node):
            affected.add(descendant[0])
    return sorted(affected)
 
 
# ---------- cost attribution ----------
@dataclass
class QueryRun:
    query_tag: str  # "team=ml/pipeline=ltv"
    credits: float
    cost_per_credit: float = 3.0  # USD
 
    @property
    def team(self) -> str:
        return dict(kv.split("=") for kv in self.query_tag.split("/"))["team"]
 
    @property
    def cost_usd(self) -> float:
        return self.credits * self.cost_per_credit
 
 
def cost_by_team(runs: Iterable[QueryRun]) -> dict[str, float]:
    out: dict[str, float] = {}
    for r in runs:
        out[r.team] = out.get(r.team, 0.0) + r.cost_usd
    return dict(sorted(out.items(), key=lambda kv: -kv[1]))
 
 
# ---------- fixture ----------
def build_world() -> None:
    src = register(Table(
        name="raw.customers",
        owner_team="platform",
        columns={
            "id":    Column("id", "int"),
            "email": Column("email", "string", tags={"pii.email"}),
            "name":  Column("name", "string", tags={"pii.name"}),
        },
    ))
    stg = register(Table(
        name="stg_customers",
        owner_team="analytics",
        columns={
            "customer_id":   Column("customer_id", "int"),
            "email_lower":   Column("email_lower", "string"),
            "display_name":  Column("display_name", "string"),
        },
    ))
    fct = register(Table(
        name="fct_customer_orders",
        owner_team="analytics",
        columns={
            "customer_id":   Column("customer_id", "int"),
            "email_lower":   Column("email_lower", "string"),
            "order_count":   Column("order_count", "int"),
        },
    ))
    dash = register(Table(
        name="dash.marketing_kpis",
        owner_team="marketing",
        columns={
            "customer_email": Column("customer_email", "string"),
            "revenue":        Column("revenue", "float"),
        },
    ))
 
    # column-level lineage — this is where governance earns its keep
    emit_lineage("raw.customers", "id",    "stg_customers", "customer_id")
    emit_lineage("raw.customers", "email", "stg_customers", "email_lower")
    emit_lineage("raw.customers", "name",  "stg_customers", "display_name")
    emit_lineage("stg_customers", "customer_id", "fct_customer_orders", "customer_id")
    emit_lineage("stg_customers", "email_lower", "fct_customer_orders", "email_lower")
    emit_lineage("fct_customer_orders", "email_lower", "dash.marketing_kpis", "customer_email")
 
 
def build_query_history() -> list[QueryRun]:
    return [
        QueryRun("team=ml/pipeline=ltv",             credits=1200),
        QueryRun("team=ml/pipeline=fraud",           credits=450),
        QueryRun("team=analytics/pipeline=exec_kpi", credits=280),
        QueryRun("team=marketing/pipeline=daily",    credits=90),
        QueryRun("team=platform/pipeline=metadata",  credits=40),
    ]
 
 
def main() -> None:
    build_world()
 
    print("=== PII columns ===")
    for t, c in columns_with_tag("pii.email"):
        print(f"  {t}.{c}")
 
    print("\n=== GDPR delete plan for alice@example.com ===")
    for t in gdpr_delete("alice@example.com"):
        print(f"  ⇒ scrub in {t} (owner: {CATALOG[t].owner_team})")
 
    print("\n=== Monthly cost by team (USD) ===")
    for team, cost in cost_by_team(build_query_history()).items():
        print(f"  {team:<10} ${cost:>8,.2f}")
 
 
if __name__ == "__main__":
    main()

pip install networkx
python governance_demo.py