pipekit/reconcile.py

#!/usr/bin/env python3
"""Cross-database reconciliation for a pipekit module.

Compares the live source table against the synced dest table using column-wise
aggregates (COUNT / SUM / MIN / MAX / SUM(LENGTH)). These rely only on
arithmetic and ordering — which DB2 for i and Postgres compute identically — so
no shared hash function or byte-identical row serialization is required.

The source side re-applies the module's per-column transform via the source
driver's ``default_expression`` (RTRIM on char, junk-date -> NULL on dates,
raw on numerics), matching exactly what the sync wrote into the dest, so the
aggregates line up when the row sets agree.

Usage:
    PIPEKIT_SECRETS=/etc/pipekit/secrets.env \\
        .venv/bin/python reconcile.py <module_name> [--quick] [--source-from EXPR]

Exit code is non-zero when any metric diverges.
"""
from __future__ import annotations

import argparse
import json
import re
import sys
from decimal import Decimal, InvalidOperation

from pipekit import repo, jrunner, drivers


def classify(dest_type: str) -> str:
    """Bucket a dest column type into numeric / temporal / text."""
    t = (dest_type or "").strip().lower()
    if t.startswith(("numeric", "decimal", "int", "bigint", "smallint",
                     "double", "real", "float")):
        return "numeric"
    if t.startswith(("date", "time", "timestamp")):
        return "temporal"
    return "text"


def detect_source_from(source_query: str, dest_table: str) -> str | None:
    """Best-effort: find the base source table in the module's SELECT.

    Picks the ``FROM schema.table`` whose table name matches the dest table's
    base name (the modules here name the dest after the source); falls back to
    the last FROM match. Returns None if nothing looks like a table.
    """
    matches = re.findall(r'FROM\s+("?[\w#]+"?\.?"?[\w#]*"?)', source_query,
                         flags=re.IGNORECASE)
    matches = [m.strip() for m in matches if "." in m]
    if not matches:
        return None
    dest_base = dest_table.split(".")[-1].strip('"').lower()
    for m in matches:
        if m.split(".")[-1].strip('"').lower() == dest_base:
            return m
    return matches[-1]


def build_metrics(columns, key_names, drv, *, source: bool, quick: bool):
    """Return (list of (label, kind) metric descriptors, list of SQL exprs).

    Descriptors and exprs are positionally aligned so the two sides zip up.
    ``source`` selects which column name + transform to use.
    """
    labels: list[tuple[str, str]] = []
    exprs: list[str] = []

    labels.append(("COUNT(*)", "int"))
    exprs.append("COUNT(*)")

    for kn in key_names:
        col = kn["source"] if source else kn["dest"]
        q = drv.quote_identifier(col)
        labels.append((f"COUNT(DISTINCT {kn['dest']})", "int"))
        exprs.append(f"COUNT(DISTINCT {q})")

    for c in columns:
        kind = classify(c["dest_type"])
        if source:
            e = drv.default_expression(c["source_type"], c["source_name"])
        else:
            e = drv.quote_identifier(c["dest_name"])
        name = c["dest_name"]

        if kind == "numeric":
            labels.append((f"{name}: count", "int"))
            exprs.append(f"COUNT({e})")
            labels.append((f"{name}: sum", "num"))
            exprs.append(f"SUM({e})")
            if not quick:
                labels.append((f"{name}: min", "str"))
                exprs.append(f"MIN({e})")
                labels.append((f"{name}: max", "str"))
                exprs.append(f"MAX({e})")
        elif kind == "temporal":
            labels.append((f"{name}: count", "int"))
            exprs.append(f"COUNT({e})")
            if not quick:
                labels.append((f"{name}: min", "str"))
                exprs.append(f"MIN({e})")
                labels.append((f"{name}: max", "str"))
                exprs.append(f"MAX({e})")
        else:  # text — avoid MIN/MAX (EBCDIC vs ASCII collation differs)
            labels.append((f"{name}: count", "int"))
            exprs.append(f"COUNT({e})")
            if not quick:
                labels.append((f"{name}: len_sum", "int"))
                exprs.append(f"SUM(LENGTH({e}))")

    return labels, exprs


def values_equal(kind: str, a: str, b: str) -> bool:
    a = "" if a is None else str(a).strip()
    b = "" if b is None else str(b).strip()
    if kind in ("num", "int"):
        try:
            da = Decimal(a) if a != "" else Decimal(0)
            db = Decimal(b) if b != "" else Decimal(0)
            return da == db
        except (InvalidOperation, ValueError):
            return a == b
    return a == b


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("module", help="module name (e.g. ocri)")
    ap.add_argument("--quick", action="store_true",
                    help="counts + numeric sums only (skip min/max/len_sum)")
    ap.add_argument("--source-from",
                    help="override the source FROM target "
                         "(e.g. a schema.table or OPENQUERY(...) t)")
    ap.add_argument("--timeout", type=int, default=1800,
                    help="per-side jrunner timeout seconds (default 1800)")
    ap.add_argument("--show-matches", action="store_true",
                    help="print matching metrics too, not just mismatches")
    args = ap.parse_args()

    m = repo.get_module_by_name(args.module)
    if not m:
        print(f"module not found: {args.module}", file=sys.stderr)
        return 2

    src_conn = repo.get_connection(m["source_connection_id"])
    dst_conn = repo.get_connection(m["dest_connection_id"])
    src_drv = drivers.get_driver(repo.get_driver_row(src_conn["driver_id"])["kind"])
    dst_drv = drivers.get_driver(repo.get_driver_row(dst_conn["driver_id"])["kind"])

    columns = json.loads(m["columns_json"] or "[]")
    if not columns:
        print("module has no columns_json", file=sys.stderr)
        return 2

    # map merge_key (dest names) -> source names via columns_json
    key_names = []
    if m["merge_key"]:
        by_dest = {c["dest_name"]: c for c in columns}
        for k in (x.strip().strip('"') for x in m["merge_key"].split(",") if x.strip()):
            c = by_dest.get(k)
            if c:
                key_names.append({"dest": c["dest_name"], "source": c["source_name"]})

    source_from = args.source_from or detect_source_from(m["source_query"], m["dest_table"])
    if not source_from:
        print("could not detect source table; pass --source-from", file=sys.stderr)
        return 2

    labels, src_exprs = build_metrics(columns, key_names, src_drv,
                                      source=True, quick=args.quick)
    _, dst_exprs = build_metrics(columns, key_names, dst_drv,
                                 source=False, quick=args.quick)

    src_sql = "SELECT\n  " + "\n , ".join(src_exprs) + f"\nFROM {source_from}"
    dst_sql = "SELECT\n  " + "\n , ".join(dst_exprs) + f"\nFROM {m['dest_table']}"

    print(f"module {m['name']}  (id {m['id']})")
    print(f"  source {src_conn['name']}  FROM {source_from}")
    print(f"  dest   {dst_conn['name']}  FROM {m['dest_table']}")
    print(f"  {len(labels)} metrics{'  [quick]' if args.quick else ''}\n"
          f"  running source aggregate ...", flush=True)
    src_res = jrunner.query(src_conn["jdbc_url"], src_conn.get("username"),
                            src_conn.get("password"), src_sql, timeout=args.timeout)
    print("  running dest aggregate ...", flush=True)
    dst_res = jrunner.query(dst_conn["jdbc_url"], dst_conn.get("username"),
                            dst_conn.get("password"), dst_sql, timeout=args.timeout)

    src_row = src_res.rows[0] if src_res.rows else []
    dst_row = dst_res.rows[0] if dst_res.rows else []
    if len(src_row) != len(labels) or len(dst_row) != len(labels):
        print(f"\nunexpected column count "
              f"(labels={len(labels)} src={len(src_row)} dst={len(dst_row)})",
              file=sys.stderr)
        return 2

    # COUNT(*) and the COUNT(DISTINCT key) metrics are the headline row counts;
    # always show them even when they match (they're the primary sanity check).
    n_headline = 1 + len(key_names)

    mismatches = []
    for i, ((label, kind), sv, dv) in enumerate(zip(labels, src_row, dst_row)):
        ok = values_equal(kind, sv, dv)
        if not ok:
            mismatches.append((label, sv, dv))
        if i < n_headline or args.show_matches or not ok:
            mark = "OK " if ok else "XX "
            print(f"  {mark} {label:32} src={sv!s:>22}  dst={dv!s:>22}")
        if i == n_headline - 1:
            print()  # blank line separating row counts from column drift

    print()
    if mismatches:
        print(f"DIVERGED: {len(mismatches)} of {len(labels)} metrics differ")
        return 1
    print(f"IN SYNC: all {len(labels)} metrics match")
    return 0


if __name__ == "__main__":
    sys.exit(main())