@1aboveio/skills 1.0.1 → 1.0.3

package/skills/app-dev/data-quality/scripts/data_quality.py

@@ -0,0 +1,405 @@
+#!/usr/bin/env python3
+"""Data quality engine: check a published table's data against its data-catalog definitions.
+
+Commands:
+  plan     write the read-only SELECT checks for documented tables (JSON)
+  judge    turn query results into findings
+  run      plan, execute through a runner, and judge in one step
+  version  print the engine identity
+
+Definitions come from the repo's data-catalog engine (scripts/data-catalog/data_catalog.py),
+so a table is checkable once its script carries CATALOG_TABLE / CATALOG_FIELDS or a
+`-- @field` block. Every query is a single SELECT; nothing is ever written.
+
+Runner contract (run --runner CMD): CMD reads one SQL statement on stdin and prints the
+result rows on stdout as a JSON array of objects, or as CSV with a header row.
+
+Standard library only.
+"""
+from __future__ import annotations
+
+import argparse
+import csv
+import hashlib
+import importlib.util
+import io
+import json
+import re
+import shlex
+import sqlite3
+import subprocess
+import sys
+from dataclasses import dataclass, field
+from pathlib import Path
+
+ENGINE_VERSION = "1.0.0"
+CATALOG_ENGINE = "scripts/data-catalog/data_catalog.py"
+NUMERIC_UNITS = {"count", "usd", "cny", "eur", "gbp", "hkd", "amount", "ratio", "percent", "seconds", "days", "score"}
+NON_NEGATIVE_UNITS = {"count", "seconds", "days"}
+RANGES = {"ratio": (0.0, 1.0), "percent": (0.0, 100.0)}
+DOCUMENTED_CONSTANT_RE = re.compile(r"始终|恒为|目前.{0,8}(为空|不填)|always|constant|never (filled|set)|not filled", re.I)
+PROVENANCE_RE = re.compile(r"(^|_)(version|sha\d*|hash|checksum|source_table|compute_time|etl_time)$", re.I)
+FLAG_RE = re.compile(r"(_missing|_flag)$|^(is|has)_", re.I)   # healthy when constant at 0 / N
+CLEAR_FLAG = {"0", "0.0", "false", "n", "no"}
+PROFILE_BATCH = 40  # columns per profile query: each COUNT(DISTINCT) is a stage, and engines cap stages
+MOSTLY_EMPTY = 0.01
+CONSTANT_MIN_ROWS = 30  # below this, one distinct value says little
+NULL_KEY = "<null>"
+
+
+class QualityError(Exception):
+    """A usage or setup problem; exit code 2."""
+
+
+# ---------------------------------------------------------------------------
+# Definitions, through the repo's data-catalog engine
+# ---------------------------------------------------------------------------
+
+def load_catalog_engine(root: Path, path: str | None):
+    engine = Path(path) if path else root / CATALOG_ENGINE
+    if not engine.is_file():
+        raise QualityError(f"no data-catalog engine at {engine}; set the repo up with the data-catalog skill first "
+                           "(its init copies the engine), or pass --catalog-engine")
+    spec = importlib.util.spec_from_file_location("data_catalog_for_quality", engine)
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[spec.name] = module  # dataclasses resolve string annotations through sys.modules
+    spec.loader.exec_module(module)
+    return module
+
+
+def documented_tables(root: Path, engine_path: str | None, wanted: list) -> list:
+    dc = load_catalog_engine(root, engine_path)
+    try:
+        config = dc.load_config(root)
+        tables, errors, _ = dc.collect(root, config)
+    except dc.CatalogError as exc:
+        raise QualityError(f"data-catalog: {exc}") from exc
+    if errors:
+        raise QualityError("data-catalog metadata has errors; run its generate first:\n  " + "\n  ".join(errors))
+    if wanted:
+        missing = [w for w in wanted if w not in {t.name for t in tables}]
+        if missing:
+            raise QualityError(f"not documented in the catalog: {', '.join(missing)}")
+        tables = [t for t in tables if t.name in wanted]
+    if not tables:
+        raise QualityError("no documented tables to check")
+    return tables
+
+
+# ---------------------------------------------------------------------------
+# Plan: one profile query per table, one value query per enumerated column
+# ---------------------------------------------------------------------------
+
+def quote(name: str, dialect: str) -> str:
+    if not re.fullmatch(r"[A-Za-z_][A-Za-z0-9_]*", name):
+        raise QualityError(f"column name {name!r} cannot be quoted safely")
+    return f"`{name}`" if dialect == "spark" else f'"{name}"'
+
+
+def table_ref(name: str) -> str:
+    if not re.fullmatch(r"[A-Za-z_][A-Za-z0-9_]*(\.[A-Za-z_][A-Za-z0-9_]*){0,2}", name):
+        raise QualityError(f"table name {name!r} is not a plain identifier")
+    return name
+
+
+def allowed(f) -> tuple:
+    """Declared values; empty for engine copies older than `Field.values`."""
+    return tuple(getattr(f, "values", ()) or ())
+
+
+def is_numeric(f) -> bool:
+    return (f.unit or "").strip().lower() in NUMERIC_UNITS
+
+
+def plan_queries(tables: list, where: str | None, dialect: str, skip: set, batch: int = PROFILE_BATCH,
+                 key: list | None = None) -> list:
+    cond = f" WHERE {where}" if where else ""
+    queries = []
+    for t in tables:
+        if key:
+            unknown = [k for k in key if k not in {f.name for f in t.fields}]
+            if unknown:
+                raise QualityError(f"--key column(s) not documented in {t.name}: {', '.join(unknown)}")
+            cols = ", ".join(quote(k, dialect) for k in key)
+            sql = (f"SELECT COUNT(*) AS dup_keys, COALESCE(SUM(c), 0) AS dup_rows FROM "
+                   f"(SELECT COUNT(*) AS c FROM {table_ref(t.name)}{cond} GROUP BY {cols} HAVING COUNT(*) > 1) dup")
+            queries.append({"id": f"key:{t.name}", "table": t.name, "kind": "key", "columns": list(key), "sql": sql})
+        cols = [f for f in t.fields if f.name not in skip]
+        chunks = [cols[i:i + batch] for i in range(0, len(cols), batch)] or [[]]
+        for k, chunk in enumerate(chunks):
+            parts = ["COUNT(*) AS n"]
+            for i, f in enumerate(chunk):
+                q = quote(f.name, dialect)
+                parts += [f"COUNT({q}) AS nn_{i}", f"COUNT(DISTINCT {q}) AS d_{i}"]
+                if is_numeric(f) or FLAG_RE.search(f.name):
+                    parts += [f"MIN({q}) AS mn_{i}", f"MAX({q}) AS mx_{i}"]
+            sql = f"SELECT {', '.join(parts)} FROM {table_ref(t.name)}{cond}"
+            suffix = f"#{k + 1}" if len(chunks) > 1 else ""
+            queries.append({"id": f"profile:{t.name}{suffix}", "table": t.name, "kind": "profile", "part": k,
+                            "columns": [f.name for f in chunk], "sql": sql})
+        for f in cols:
+            if allowed(f):
+                q = quote(f.name, dialect)
+                sql = f"SELECT {q} AS value, COUNT(*) AS n FROM {table_ref(t.name)}{cond} GROUP BY {q}"
+                queries.append({"id": f"values:{t.name}.{f.name}", "table": t.name, "kind": "values",
+                                "column": f.name, "sql": sql})
+    for q in queries:
+        assert_read_only(q["sql"])
+    return queries
+
+
+def assert_read_only(sql: str) -> None:
+    if not re.match(r"\s*SELECT\b", sql, re.I) or ";" in sql:
+        raise QualityError(f"refusing to run a statement that is not a single SELECT: {sql[:80]}")
+
+
+# ---------------------------------------------------------------------------
+# Judge: results -> findings
+# ---------------------------------------------------------------------------
+
+@dataclass
+class Finding:
+    table: str
+    column: str | None
+    kind: str
+    message: str
+    documented: bool = False   # the definition already says so: report as info, not a problem
+
+
+def number(value):
+    if value is None or value == "":
+        return None
+    try:
+        return float(value)
+    except (TypeError, ValueError):
+        return None
+
+
+def where_columns(where: str | None) -> set:
+    """Identifiers a --where condition names: those columns are constant by construction."""
+    if not where:
+        return set()
+    bare = re.sub(r"'(?:[^']|'')*'", " ", where)
+    return {w.strip('`"').lower() for w in re.findall(r'[`"]?[A-Za-z_][A-Za-z0-9_]*[`"]?', bare)}
+
+
+def judge(tables: list, queries: list, results: dict, where: str | None = None,
+          constant_min_rows: int = CONSTANT_MIN_ROWS) -> list:
+    by_table = {t.name: t for t in tables}
+    pinned = where_columns(where)
+    findings = []
+    for q in queries:
+        rows = results.get(q["id"])
+        if rows is None:
+            findings.append(Finding(q["table"], q.get("column"), "not-run", f"no result for {q['id']}"))
+            continue
+        t = by_table[q["table"]]
+        if q["kind"] == "key":
+            row = {k.lower(): v for k, v in (rows[0] if rows else {}).items()}
+            keys = number(row.get("dup_keys")) or 0
+            if keys:
+                findings.append(Finding(t.name, None, "duplicate-key",
+                                        f"{int(keys)} ({', '.join(q['columns'])}) value(s) repeat across "
+                                        f"{int(number(row.get('dup_rows')) or 0)} rows; the grain allows one each"))
+            continue
+        if q["kind"] == "profile":
+            part = judge_profile(t, q["columns"], rows, pinned, constant_min_rows)
+            if q.get("part", 0) > 0:  # every batch sees the same row count; report an empty range once
+                part = [f for f in part if f.kind != "empty-table"]
+            findings += part
+        else:
+            findings += judge_values(t, q["column"], rows)
+    return findings
+
+
+def judge_profile(t, columns: list, rows: list, pinned: set, constant_min_rows: int) -> list:
+    if len(rows) != 1:
+        return [Finding(t.name, None, "not-run", f"profile returned {len(rows)} rows, expected 1")]
+    row = {k.lower(): v for k, v in rows[0].items()}
+    n = number(row.get("n")) or 0
+    if n == 0:
+        return [Finding(t.name, None, "empty-table", "no rows in the checked range; nothing else can be judged")]
+    fields = {f.name: f for f in t.fields}
+    out = []
+    for i, name in enumerate(columns):
+        f = fields[name]
+        text = " ".join([*f.text.values(), f.source])
+        documented = bool(DOCUMENTED_CONSTANT_RE.search(text))
+        filled = number(row.get(f"nn_{i}")) or 0
+        distinct = number(row.get(f"d_{i}")) or 0
+        lo, hi = number(row.get(f"mn_{i}")), number(row.get(f"mx_{i}"))
+        if filled == 0:
+            out.append(Finding(t.name, name, "never-filled", f"empty in all {int(n)} rows", documented))
+            continue
+        if filled / n < MOSTLY_EMPTY:
+            out.append(Finding(t.name, name, "mostly-empty", f"filled in {int(filled)} of {int(n)} rows"))
+        if distinct == 1 and n >= constant_min_rows and name.lower() not in pinned:
+            value = row.get(f"mn_{i}")
+            shown = f" ({value})" if value is not None else ""
+            clear_flag = FLAG_RE.search(name) and str(row.get(f"mn_{i}")).strip().lower() in CLEAR_FLAG
+            # a build's version is the same on every row; a missing-flag at 0 means nothing was missing
+            expected = documented or bool(PROVENANCE_RE.search(name)) or bool(clear_flag)
+            out.append(Finding(t.name, name, "constant", f"one value{shown} in all {int(filled)} filled rows", expected))
+        unit = (f.unit or "").strip().lower()
+        if unit in NON_NEGATIVE_UNITS and lo is not None and lo < 0:
+            out.append(Finding(t.name, name, "negative-value", f"minimum {lo:g}, but the unit is {f.unit}"))
+        if unit in RANGES and lo is not None and hi is not None:
+            low, high = RANGES[unit]
+            if lo < low or hi > high:
+                out.append(Finding(t.name, name, "out-of-range",
+                                   f"values from {lo:g} to {hi:g}, outside {low:g}–{high:g} for unit {f.unit}"))
+    return out
+
+
+def judge_values(t, column: str, rows: list) -> list:
+    f = next(f for f in t.fields if f.name == column)
+    seen = {}
+    for r in rows:
+        r = {k.lower(): v for k, v in r.items()}
+        value = r.get("value")
+        empty = value is None or (value == "" and "" not in allowed(f))
+        key = NULL_KEY if empty else str(value)
+        seen[key] = seen.get(key, 0) + (number(r.get("n")) or 0)
+    declared = set(allowed(f))
+    out = []
+    unseen = [v for v in allowed(f) if v not in seen]
+    if unseen and seen:
+        out.append(Finding(t.name, column, "unseen-value", f"declared but never present: {', '.join(unseen)}"))
+    extra = sorted(k for k in seen if k not in declared and k != NULL_KEY)
+    if extra:
+        counts = ", ".join(f"{k} ({int(seen[k])})" for k in extra[:10])
+        out.append(Finding(t.name, column, "undeclared-value", f"present but not declared: {counts}"))
+    return out
+
+
+# ---------------------------------------------------------------------------
+# Runners
+# ---------------------------------------------------------------------------
+
+def parse_rows(text: str) -> list:
+    text = text.strip()
+    if not text:
+        return []
+    if text[0] == "[":
+        return json.loads(text)
+    if text[0] == "{":  # one object, or JSON lines (trino/presto --output-format JSON)
+        return [json.loads(line) for line in text.splitlines() if line.strip()]
+    return [dict(r) for r in csv.DictReader(io.StringIO(text))]
+
+
+def run_command(cmd: str, sql: str) -> list:
+    proc = subprocess.run(shlex.split(cmd), input=sql, capture_output=True, text=True)
+    if proc.returncode != 0:
+        raise QualityError(f"runner failed ({proc.returncode}) on: {sql[:80]}\n{proc.stderr.strip()[-400:]}")
+    return parse_rows(proc.stdout)
+
+
+def run_sqlite(db: str, sql: str) -> list:
+    con = sqlite3.connect(f"file:{db}?mode=ro", uri=True)  # read-only by construction
+    try:
+        cur = con.execute(sql)
+        names = [d[0] for d in cur.description]
+        return [dict(zip(names, row)) for row in cur.fetchall()]
+    finally:
+        con.close()
+
+
+def execute(queries: list, runner: str | None, sqlite_db: str | None) -> dict:
+    results = {}
+    for q in queries:
+        assert_read_only(q["sql"])
+        results[q["id"]] = run_sqlite(sqlite_db, q["sql"]) if sqlite_db else run_command(runner, q["sql"])
+    return results
+
+
+# ---------------------------------------------------------------------------
+# Report and CLI
+# ---------------------------------------------------------------------------
+
+def report(findings: list, tables: list, where: str | None, as_json: bool) -> int:
+    problems = [f for f in findings if not f.documented]
+    if as_json:
+        print(json.dumps({"engine": ENGINE_VERSION, "where": where, "tables": [t.name for t in tables],
+                          "problems": len(problems),
+                          "findings": [vars(f) for f in findings]}, ensure_ascii=False, indent=2))
+        return 1 if problems else 0
+    scope = f" where {where}" if where else " (all rows)"
+    print(f"data-quality: {len(tables)} table(s){scope}; {len(problems)} problem(s), "
+          f"{len(findings) - len(problems)} already documented")
+    for t in tables:
+        mine = [f for f in findings if f.table == t.name]
+        if not mine:
+            continue
+        print(f"\n{t.name}")
+        for f in mine:
+            tag = " (documented)" if f.documented else ""
+            print(f"  - {f.kind}{tag}: {f.column + ': ' if f.column else ''}{f.message}")
+    return 1 if problems else 0
+
+
+def engine_identity() -> str:
+    return f"{ENGINE_VERSION}+{hashlib.sha256(Path(__file__).read_bytes()).hexdigest()[:12]}"
+
+
+def find_root(start: Path) -> Path:
+    for candidate in [start, *start.parents]:
+        if (candidate / "data-catalog.json").exists() or (candidate / ".git").exists():
+            return candidate
+    return start
+
+
+def main(argv: list | None = None) -> int:
+    parser = argparse.ArgumentParser(prog="data_quality.py", description=__doc__.split("\n\n")[0])
+    parser.add_argument("--root", type=Path, help="repo root (default: nearest dir with data-catalog.json or .git)")
+    parser.add_argument("--catalog-engine", help=f"path to data_catalog.py (default <root>/{CATALOG_ENGINE})")
+    sub = parser.add_subparsers(dest="command", required=True)
+
+    def scoped(p):
+        p.add_argument("--table", action="append", default=[], help="documented table to check; repeatable (default all)")
+        p.add_argument("--where", help="SQL condition that bounds the scan, e.g. \"dt >= '2026-10-01'\"")
+        p.add_argument("--dialect", choices=["ansi", "spark"], default="ansi", help="identifier quoting: ansi \"x\", spark `x`")
+        p.add_argument("--skip", action="append", default=[], help="column to leave out (e.g. arrays); repeatable")
+        p.add_argument("--key", action="append", default=[],
+                       help="grain column; repeat for a composite key: --key dt --key cust_id")
+        p.add_argument("--batch", type=int, default=PROFILE_BATCH,
+                       help=f"columns per profile query (default {PROFILE_BATCH}); lower it if the engine caps stages")
+        p.add_argument("--constant-min-rows", type=int, default=CONSTANT_MIN_ROWS,
+                       help=f"rows needed before a single value counts as constant (default {CONSTANT_MIN_ROWS})")
+
+    p_plan = sub.add_parser("plan", help="print the SELECT checks as JSON")
+    scoped(p_plan)
+    p_judge = sub.add_parser("judge", help="judge results; --results maps query id to rows")
+    scoped(p_judge)
+    p_judge.add_argument("--results", required=True, help="JSON file: {query id: [row objects]}")
+    p_judge.add_argument("--json", action="store_true")
+    p_run = sub.add_parser("run", help="plan, execute, judge")
+    scoped(p_run)
+    group = p_run.add_mutually_exclusive_group(required=True)
+    group.add_argument("--runner", help="command that reads SQL on stdin and prints JSON or CSV rows")
+    group.add_argument("--sqlite", help="SQLite database file (opened read-only)")
+    p_run.add_argument("--json", action="store_true")
+    sub.add_parser("version", help="print the engine identity")
+    args = parser.parse_args(argv)
+    if args.command == "version":
+        print(engine_identity())
+        return 0
+    root = (args.root or find_root(Path.cwd())).resolve()
+    try:
+        tables = documented_tables(root, args.catalog_engine, args.table)
+        queries = plan_queries(tables, args.where, args.dialect, set(args.skip), args.batch, args.key)
+        if args.command == "plan":
+            print(json.dumps(queries, ensure_ascii=False, indent=2))
+            return 0
+        if args.command == "judge":
+            results = json.loads(Path(args.results).read_text(encoding="utf-8"))
+        else:
+            results = execute(queries, args.runner, args.sqlite)
+        found = judge(tables, queries, results, args.where, args.constant_min_rows)
+        return report(found, tables, args.where, args.json)
+    except QualityError as exc:
+        print(f"data-quality: {exc}", file=sys.stderr)
+        return 2
+
+
+if __name__ == "__main__":
+    sys.exit(main())