keygraph 0.2.0__py3-none-any.whl

keygraph/model.py

@@ -0,0 +1,149 @@
+"""Normalised schema model shared by every adapter."""
+
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass, field
+
+
+@dataclass
+class Column:
+    name: str
+    type: str = ""
+    nullable: bool = True
+    position: int = 0
+    key: str = ""  # PRI / UNI / MUL as reported by MySQL COLUMNS, or ""
+    default: str | None = None
+    extra: str = ""  # auto_increment, on update ..., VIRTUAL GENERATED
+    charset: str = ""
+    collation: str = ""
+    comment: str = ""
+    values: list[str] = field(default_factory=list)  # enum / set members
+    generated: str = ""  # generation expression
+
+
+@dataclass
+class Index:
+    name: str
+    columns: list[str]
+    unique: bool = False
+    primary: bool = False
+    approximate: bool = False  # derived from COLUMN_KEY only; composition unknown
+
+
+@dataclass
+class ForeignKey:
+    name: str
+    columns: list[str]
+    ref_schema: str
+    ref_table: str
+    ref_columns: list[str]
+
+
+@dataclass
+class Table:
+    schema: str
+    name: str
+    columns: list[Column] = field(default_factory=list)
+    indexes: list[Index] = field(default_factory=list)
+    foreign_keys: list[ForeignKey] = field(default_factory=list)
+    info: dict = field(default_factory=dict)  # engine, rows_estimate, data_bytes, partitioning, comment, checks, sorting_key ...
+
+    @property
+    def id(self) -> str:
+        return f"{self.schema}.{self.name}"
+
+    @property
+    def primary_key(self) -> list[str]:
+        for ix in self.indexes:
+            if ix.primary:
+                return list(ix.columns)
+        return []
+
+    @property
+    def unique_keys(self) -> list[list[str]]:
+        return [list(ix.columns) for ix in self.indexes if ix.unique and not ix.primary]
+
+    def column(self, name: str) -> Column | None:
+        for c in self.columns:
+            if c.name == name:
+                return c
+        return None
+
+
+@dataclass
+class Schema:
+    tables: dict[str, Table] = field(default_factory=dict)
+
+    def add(self, table: Table) -> None:
+        self.tables[table.id] = table
+
+    def get_or_create(self, schema: str, name: str) -> Table:
+        tid = f"{schema}.{name}"
+        if tid not in self.tables:
+            self.tables[tid] = Table(schema=schema, name=name)
+        return self.tables[tid]
+
+
+_INT = {"tinyint", "smallint", "mediumint", "int", "integer", "bigint", "serial", "bigserial",
+        "smallserial", "int2", "int4", "int8", "int1", "int16", "int32", "int64", "uint8",
+        "uint16", "uint32", "uint64", "number"}
+_DEC = {"decimal", "numeric", "money", "dec", "fixed"}
+_FLT = {"float", "double", "real", "float4", "float8", "float32", "float64"}
+_STR = {"char", "varchar", "text", "tinytext", "mediumtext", "longtext", "character", "nvarchar",
+        "nchar", "string", "citext", "enum", "enum8", "enum16", "set", "uuid", "ntext", "fixedstring", "varchar2"}
+_DATE = {"date", "datetime", "timestamp", "timestamptz", "time", "year", "datetime2", "date32"}
+_BIN = {"blob", "tinyblob", "mediumblob", "longblob", "binary", "varbinary", "bytea"}
+_BOOL = {"bool", "boolean", "bit"}
+
+
+def type_family(type_str: str) -> str:
+    """Coarse type family: int, decimal, float, string, datetime, json, binary, bool or unknown.
+
+    Signedness and width are ignored on purpose: a signed key and an unsigned key still join.
+    """
+    t = (type_str or "").lower().strip()
+    if not t:
+        return "unknown"
+    while re.match(r"^(nullable|lowcardinality)\(", t):
+        t = re.sub(r"^(nullable|lowcardinality)\(", "", t).rstrip(")")
+    first = re.split(r"[\s(]", t, maxsplit=1)[0]
+    if first == "double" or first == "character" and "varying" in t:
+        return "float" if first == "double" else "string"
+    if first in _INT or first.startswith("uint") or (first.startswith("int") and first[3:].isdigit()):
+        return "int"
+    if first in _DEC:
+        return "decimal"
+    if first in _FLT:
+        return "float"
+    if first in _STR:
+        return "string"
+    if first in _DATE or first.startswith("timestamp"):
+        return "datetime"
+    if first in {"json", "jsonb"}:
+        return "json"
+    if first in _BIN:
+        return "binary"
+    if first in _BOOL:
+        return "bool"
+    return "unknown"
+
+
+def type_length(type_str: str) -> int | None:
+    m = re.search(r"\((\d+)\)", type_str or "")
+    return int(m.group(1)) if m else None
+
+
+def parse_enum_members(type_str: str) -> list[str]:
+    """Members of MySQL enum('a','b') / set(...), ClickHouse Enum8('a' = 1, 'b' = 2) (possibly wrapped in
+    Nullable/LowCardinality). Returns [] for other types. ClickHouse members keep their numeric code as 'a=1'."""
+    t = (type_str or "").strip()
+    m = re.search(r"\b(enum|set|enum8|enum16)\s*\((.*)\)", t, re.I | re.S)
+    if not m:
+        return []
+    body = m.group(2)
+    out = []
+    for lit in re.finditer(r"'((?:[^'\\]|\\.|'')*)'(?:\s*=\s*(-?\d+))?", body):
+        name = lit.group(1).replace("''", "'").replace("\\'", "'")
+        out.append(f"{name}={lit.group(2)}" if lit.group(2) is not None and m.group(1).lower().startswith("enum") and m.group(1).lower() != "enum" else name)
+    return out

keygraph/notes.py

@@ -0,0 +1,114 @@
+"""Generic importer: one markdown note per table -> table notes and column semantics.
+
+Recognised (case-insensitive) heading words: grain|meaning -> grain; gotcha|caveat|pitfall|warning -> pitfalls;
+column(s) -> bullets/rows like "- `col` - meaning" or "| col | meaning |"; code|status|enum -> "- `value` = label"
+(optionally under a "### column" sub-heading); a verified section is read but does not mark the note verified (only a person or a checked run should). A "**Summary:**" line sets the
+summary. Front matter keys: table (or the file stem), and the key named by schema_key for the schema.
+"""
+
+from __future__ import annotations
+
+import re
+from pathlib import Path
+
+
+def parse_frontmatter(text: str) -> tuple[dict, str]:
+    m = re.match(r"^---\n(.*?)\n---\n", text, re.S)
+    if not m:
+        return {}, text
+    meta = {}
+    for line in m.group(1).splitlines():
+        if ":" in line:
+            k, v = line.split(":", 1)
+            meta[k.strip()] = v.strip().strip("[]")
+    return meta, text[m.end():]
+
+
+def _sections(body: str) -> list[tuple[int, str, list[str]]]:
+    out: list[tuple[int, str, list[str]]] = []
+    for line in body.splitlines():
+        h = re.match(r"^(#{1,4})\s+(.*)$", line)
+        if h:
+            out.append((len(h.group(1)), h.group(2).strip(), []))
+        elif out:
+            out[-1][2].append(line)
+    return out
+
+
+_BUL = re.compile(r"^\s*[-*]\s+`?([A-Za-z0-9_.\-]+)`?\s*(?:[:=—–-]+|\s)\s*(.+)$")
+
+
+def _text(lines: list[str]) -> str:
+    return re.sub(r"\s+", " ", " ".join(l.strip() for l in lines if l.strip() and not l.strip().startswith("_None"))).strip()
+
+
+def parse_note(path: str | Path, schema_key: str = "schema", default_schema: str = "main") -> dict | None:
+    p = Path(path)
+    meta, body = parse_frontmatter(p.read_text(encoding="utf-8"))
+    table = meta.get("table") or p.stem
+    schema = meta.get(schema_key) or default_schema
+    if meta.get("kind") not in (None, "table"):
+        return None
+    note: dict = {"table": f"{schema}.{table}", "columns": [], "source": str(p), "verified": False}
+    sm = re.search(r"\*\*Summary:\*\*\s*(.+)", body)
+    if sm:
+        note["summary"] = sm.group(1).strip()
+    pit: list[str] = []
+    cur_code_col = None
+    for level, title, lines in _sections(body):
+        tl = title.lower()
+        if re.search(r"grain|meaning", tl) and "column" not in tl:
+            note["grain"] = _text(lines)
+        elif re.search(r"gotcha|caveat|pitfall|warning", tl):
+            pit.extend(re.sub(r"^\s*[-*]\s+", "", l).strip() for l in lines if l.strip())
+        elif re.search(r"verified", tl):
+            note["verified_facts"] = _text(lines)  # kept out of the verified flag: a note is a claim, not a check
+        elif re.search(r"column", tl):
+            for l in lines:
+                m = _BUL.match(l)
+                r = re.match(r"^\|\s*`?([A-Za-z0-9_]+)`?\s*\|\s*([^|]+)\|", l)
+                if (m or r) and not (r and set(r.group(1)) <= {"-"}):
+                    c, mean = (m or r).group(1), (m or r).group(2).strip()
+                    note["columns"].append({"column": c, "meaning": mean})
+        elif re.search(r"code|status|enum", tl):
+            cur_code_col = None
+            for l in lines:
+                m = _BUL.match(l)
+                if m:
+                    cur_code_col = cur_code_col or re.sub(r"[^\w]", "", title.split()[-1])
+                    note.setdefault("_codes", {}).setdefault(cur_code_col, {})[m.group(1)] = m.group(2).strip()
+        elif level >= 3 and note.get("_codes") is not None:
+            pass
+    if pit:
+        note["pitfalls"] = " | ".join(pit)
+    note["codes"] = note.pop("_codes", {})
+    return note
+
+
+def notes_to_context(root: str | Path, schema_key: str = "schema", default_schema: str = "main",
+                     skip_names: tuple[str, ...] = ("memory.md", "readme.md")) -> dict:
+    """Convert a folder of per-table markdown notes into a context structure (tables + columns).
+
+    Entries carry source = the note path; they are always unverified.
+    Column meanings and codes are always unverified: a note bullet is a claim, not a check.
+    """
+    out: dict = {"tables": [], "columns": []}
+    for f in sorted(Path(root).rglob("*.md")):
+        if f.name.lower() in skip_names:
+            continue
+        n = parse_note(f, schema_key, default_schema)
+        if not n:
+            continue
+        cols, codes = n.pop("columns"), n.pop("codes", {})
+        n.pop("verified_facts", None)
+        out["tables"].append({k: v for k, v in n.items() if v not in ("", None)})
+        seen = set()
+        for c in cols:
+            out["columns"].append({"table": n["table"], "column": c["column"], "meaning": c["meaning"],
+                                   **({"codes": codes[c["column"]]} if c["column"] in codes else {}),
+                                   "source": str(f), "verified": False})
+            seen.add(c["column"])
+        for col, cd in codes.items():
+            if col not in seen:
+                out["columns"].append({"table": n["table"], "column": col, "codes": cd, "source": str(f), "verified": False})
+    return out

keygraph/overrides.py

@@ -0,0 +1,112 @@
+"""User overrides: confirm / reject / add edges, table aliases, role aliases, prefixes."""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass, field
+from pathlib import Path
+
+
+@dataclass
+class Overrides:
+    aliases: dict[str, str] = field(default_factory=dict)
+    role_aliases: dict[str, list[str]] = field(default_factory=dict)
+    prefixes: list[str] = field(default_factory=list)
+    confirm: list[dict] = field(default_factory=list)
+    reject: list[dict] = field(default_factory=list)
+    add: list[dict] = field(default_factory=list)
+
+    @classmethod
+    def load(cls, path: str | Path) -> "Overrides":
+        p = Path(path)
+        text = p.read_text(encoding="utf-8")
+        if p.suffix.lower() in (".yaml", ".yml"):
+            try:
+                import yaml
+            except ImportError as exc:  # pragma: no cover
+                raise SystemExit("YAML needs PyYAML: pip install 'keygraph[yaml]' (or use a .json file)") from exc
+            raw = yaml.safe_load(text) or {}
+        else:
+            raw = json.loads(text)
+        ra = {k: ([v] if isinstance(v, str) else list(v)) for k, v in (raw.get("role_aliases") or {}).items()}
+        return cls(aliases=dict(raw.get("aliases") or {}), role_aliases=ra, prefixes=list(raw.get("prefixes") or []),
+                   confirm=list(raw.get("confirm") or []), reject=list(raw.get("reject") or []),
+                   add=list(raw.get("add") or []))
+
+    def configure(self, opts) -> None:
+        opts.aliases.update(self.aliases)
+        opts.role_aliases.update(self.role_aliases)
+        opts.prefixes.extend(p for p in self.prefixes if p not in opts.prefixes)
+
+
+def parse_ref(ref: str) -> tuple[str, list[str]]:
+    """'shop.orders.customer_id' or 'orders.a,b' -> (table ref, columns)."""
+    head, cols = ref.rsplit(".", 1)
+    return head, [c.strip() for c in cols.split(",") if c.strip()]
+
+
+def _resolve(nodes: dict, head: str) -> str | None:
+    if head in nodes:
+        return head
+    hits = [n for n in nodes if n.split(".", 1)[1] == head]
+    return hits[0] if len(hits) == 1 else None
+
+
+def apply_overrides(data: dict, ov: Overrides) -> None:
+    nodes = {n["id"]: n for n in data["nodes"]}
+    applied = {"confirmed": 0, "rejected": 0, "added": 0}
+    notes = data["warnings"]
+
+    def locate(item):
+        fh, fc = parse_ref(item["from"])
+        ft = _resolve(nodes, fh)
+        if item.get("to"):
+            th, tc = parse_ref(item["to"])
+            tt = _resolve(nodes, th)
+        else:
+            tt, tc = None, []
+        return ft, fc, tt, tc
+
+    for item in ov.reject:
+        ft, fc, tt, tc = locate(item)
+        before = len(data["edges"])
+        data["edges"] = [e for e in data["edges"] if not (
+            e["kind"] == "inferred" and e["from"] == ft and e["from_columns"] == fc
+            and (tt is None or (e["to"] == tt and (not tc or e["to_columns"] == tc))))]
+        applied["rejected"] += before - len(data["edges"])
+        if ft is None:
+            notes.append(f"override reject: unknown table in {item['from']}")
+    for kind, items, conf in (("confirmed", ov.confirm, 0.98), ("added", ov.add, 0.9)):
+        for item in items:
+            ft, fc, tt, tc = locate(item)
+            if not (ft and tt):
+                notes.append(f"override {kind}: unknown table in {item}")
+                continue
+            have = {c["name"] for c in nodes[ft]["columns"]}, {c["name"] for c in nodes[tt]["columns"]}
+            if not (set(fc) <= have[0] and set(tc) <= have[1]) or not tc:
+                notes.append(f"override {kind}: unknown column in {item}")
+                continue
+            note = item.get("note", "")
+            match = [e for e in data["edges"] if e["from"] == ft and e["from_columns"] == fc and e["kind"] == "inferred"]
+            hit = [e for e in match if e["to"] == tt and e["to_columns"] == tc]
+            if kind == "confirmed" and hit:
+                data["edges"] = [e for e in data["edges"] if e not in match or e is hit[0]]
+                e = hit[0]
+                e["status"], e["confidence"], e["ambiguous"] = "confirmed", conf, False
+                e["evidence"].append("confirmed by user override" + (f": {note}" if note else ""))
+            else:
+                if kind == "confirmed":
+                    data["edges"] = [e for e in data["edges"] if e not in match]
+                data["edges"] = [e for e in data["edges"] if not (
+                    e["from"] == ft and e["from_columns"] == fc and e["to"] == tt and e["to_columns"] == tc)]
+                fk = nodes[ft]
+                from .infer import cardinality
+                from .model import Table, Index
+                t = Table(ft.split(".")[0], fk["table"], [], [Index(i["name"], i["columns"], i["unique"], i["primary"]) for i in fk["indexes"]])
+                data["edges"].append({
+                    "kind": "inferred", "status": kind, "from": ft, "from_columns": fc, "to": tt, "to_columns": tc,
+                    "confidence": conf, "cardinality": cardinality(t, fc), "ambiguous": False,
+                    "evidence": [f"{'confirmed' if kind == 'confirmed' else 'added'} by user override" + (f": {note}" if note else "")],
+                })
+            applied[kind] += 1
+    data["overrides"] = applied

keygraph/probe.py

@@ -0,0 +1,119 @@
+"""Read-only value-overlap probe for inferred edges. Refuses anything that is not a single SELECT."""
+
+from __future__ import annotations
+
+import re
+from decimal import Decimal
+from typing import Callable
+
+from .graph import Graph
+
+_BAD = re.compile(r"\b(insert|update|delete|replace|drop|alter|create|truncate|grant|revoke|rename|call|handler|load|"
+                  r"lock|unlock|merge|exec|execute|set|use|copy|vacuum|attach|pragma)\b|\binto\s+(outfile|dumpfile)\b|"
+                  r"\bfor\s+update\b", re.I)
+
+
+class UnsafeSQL(ValueError):
+    pass
+
+
+def assert_select(sql: str) -> str:
+    s = re.sub(r"/\*.*?\*/", " ", sql, flags=re.S)
+    s = re.sub(r"--[^\n]*", " ", s).strip()
+    if s.endswith(";"):
+        s = s[:-1].strip()
+    if not s:
+        raise UnsafeSQL("empty statement")
+    if ";" in s:
+        raise UnsafeSQL("only one statement is allowed")
+    if not re.match(r"select\b", s, re.I):
+        raise UnsafeSQL("only SELECT statements are allowed")
+    stripped = re.sub(r"'(?:[^']|'')*'", "''", s)
+    m = _BAD.search(stripped)
+    if m:
+        raise UnsafeSQL(f"forbidden keyword in query: {m.group(0)}")
+    return s
+
+
+class SelectOnlyRunner:
+    def __init__(self, execute: Callable[[str], list]):
+        self._execute = execute
+
+    def run(self, sql: str) -> list:
+        return list(self._execute(assert_select(sql)))
+
+    @classmethod
+    def from_url(cls, url: str) -> "SelectOnlyRunner":
+        from sqlalchemy import create_engine, text
+        eng = create_engine(url)
+
+        def ex(sql):
+            with eng.connect() as c:
+                return [tuple(r) for r in c.execute(text(sql))]
+        return cls(ex)
+
+    @classmethod
+    def from_dbapi(cls, conn) -> "SelectOnlyRunner":
+        def ex(sql):
+            cur = conn.cursor()
+            try:
+                cur.execute(sql)
+                return cur.fetchall()
+            finally:
+                cur.close()
+        return cls(ex)
+
+
+def _lit(v) -> str | None:
+    if isinstance(v, bool):
+        return None
+    if isinstance(v, int):
+        return str(v)
+    if isinstance(v, Decimal) and v == v.to_integral_value():
+        return str(int(v))
+    if isinstance(v, str) and "\\" not in v and "\x00" not in v:
+        return "'" + v.replace("'", "''") + "'"
+    return None
+
+
+def probe_graph(graph: Graph, runner: SelectOnlyRunner, sample: int = 500, quote: str = '"',
+                max_edges: int = 200, min_conf: float = 0.0, statuses=("inferred",), profiles: dict | None = None) -> dict:
+    q = lambda name: f"{quote}{name}{quote}"
+    qt = lambda tid: ".".join(q(p) for p in tid.split(".", 1)) if not tid.startswith("main.") else q(tid.split(".", 1)[1])
+    res = {"probed": 0, "skipped": 0, "errors": []}
+    for e in graph.edges:
+        if e["status"] not in statuses or e["confidence"] < min_conf or res["probed"] >= max_edges:
+            continue
+        if len(e["from_columns"]) != 1:
+            res["skipped"] += 1
+            continue
+        fc, tc = e["from_columns"][0], e["to_columns"][0]
+        pf, pt = (profiles or {}).get(f"{e['from']}.{fc}", {}).get("numeric"), (profiles or {}).get(f"{e['to']}.{tc}", {}).get("numeric")
+        if pf and pt and (pf["min"] > pt["max"] or pf["max"] < pt["min"]):
+            e["evidence"].append(f"profile: value ranges do not overlap ({pf['min']:g}..{pf['max']:g} vs {pt['min']:g}..{pt['max']:g})")
+            e["confidence"] = round(e["confidence"] * 0.3, 2)
+            e["status"] = "probed"
+            res["probed"] += 1
+            continue
+        try:
+            vals = [r[0] for r in runner.run(
+                f"SELECT {q(fc)} FROM {qt(e['from'])} WHERE {q(fc)} IS NOT NULL LIMIT {int(sample)}")]
+            lits = sorted({l for l in map(_lit, vals) if l is not None})
+            if not lits:
+                res["skipped"] += 1
+                continue
+            found = runner.run(f"SELECT {q(tc)} FROM {qt(e['to'])} WHERE {q(tc)} IN ({','.join(lits)}) LIMIT {len(lits)}")
+        except Exception as exc:  # noqa: BLE001
+            res["errors"].append(f"{e['id']}: {exc}")
+            continue
+        rate = len({r[0] for r in found}) / len(lits)
+        e["match_rate"] = round(rate, 3)
+        e["evidence"].append(f"probe: {rate:.0%} of {len(lits)} sampled distinct values of {e['from']}.{fc} exist in {e['to']}.{tc}")
+        c = e["confidence"]
+        c = min(0.99, c + 0.25) if rate >= 0.95 else min(0.99, c + 0.10) if rate >= 0.8 else c if rate >= 0.5 else \
+            min(c, 0.1) if rate == 0 else c * 0.5
+        e["confidence"] = round(c, 2)
+        e["status"] = "probed"
+        res["probed"] += 1
+    graph.refresh_stats()
+    return res