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 | --super-quick] [--source-from EXPR]

Depth: default compares every column; ``--quick`` drops min/max/len_sum;
``--super-quick`` compares only the row counts (COUNT(*) and COUNT(DISTINCT
merge_key)) — seconds even on multi-million-row tables, and enough to catch
missing or duplicated rows, but blind to changed column values.

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 is_expression(source_name: str) -> bool:
    """True when a columns_json source_name is a SQL expression, not a column.

    Derived columns (e.g. qtnote's merge key ``SUBSTR(GGKEY,1,9)``) are stored
    in columns_json with the expression in ``source_name``. Those must be
    emitted verbatim — quoting them as an identifier yields SQL0206, and
    wrapping them in ``default_expression`` would re-apply a transform the
    module's own SELECT already applied.
    """
    return "(" in (source_name or "")


def source_expression(drv, column: dict) -> str:
    """Source-side SQL for a column: verbatim if derived, else transformed."""
    name = column["source_name"]
    if is_expression(name):
        return name
    return drv.default_expression(column["source_type"], name)


# One name part: "quoted", [bracketed], or bare. A table reference is 1-4 of
# them dot-joined — SQL Server's linked-server form (server.db.schema.table) is
# the reason for 4, GP's CHG.dbo.SOP30200 the reason 2 was never enough.
_PART = r'(?:"[^"]+"|\[[^\]]+\]|[\w#$]+)'
_FROM_RE = re.compile(rf'FROM\s+({_PART}(?:\.{_PART}){{0,3}})', re.IGNORECASE)
_OPENQUERY_RE = re.compile(r'OPENQUERY\s*\(\s*([\w#$]+)\s*,', re.IGNORECASE)


def _base_name(ref: str) -> str:
    """Last part of a dotted table reference, unquoted and lowercased."""
    return ref.split(".")[-1].strip('"').strip("[]").lower()


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 = [m.strip() for m in _FROM_RE.findall(source_query) if "." in m]
    if not matches:
        return None
    dest_base = _base_name(dest_table)
    for m in matches:
        if _base_name(m) == dest_base:
            return m
    return matches[-1]


def detect_openquery_server(source_query: str) -> str | None:
    """Linked-server name if the module reads through ``OPENQUERY(SRV, '...')``.

    The GP modules wrap their whole SELECT in OPENQUERY so the work runs on the
    remote server. The table named inside that string (``CHG.dbo.SOP30200``)
    does not resolve on the local connection, so reconcile has to push its
    aggregate through the same wrapper rather than querying the table directly.
    """
    m = _OPENQUERY_RE.search(source_query)
    return m.group(1) if m else None


def build_openquery_sql(server: str, table: str, exprs: list[str], *,
                        limit: int = 7000) -> str:
    """Aggregate ``exprs`` over ``table`` on linked server ``server``.

    Columns are aliased ``c0..cN`` because OPENQUERY rejects a passthrough
    result set with unnamed columns (msg 8155); the aliases are positional and
    results are still zipped by position, never by name.

    OPENQUERY's passthrough string is capped at 8000 characters and a wide
    module overruns it (sop30200: 291 metrics), so the metric list is split
    into chunks that each fit and CROSS JOINed back into the single row the
    caller expects. `SELECT *` over the join preserves left-to-right order, so
    the columns still line up with the labels. Each chunk costs one extra
    remote scan — hence chunks as large as the cap allows.

    (The uncapped alternative, ``EXEC(@sql) AT server``, needs RPC Out enabled
    on the linked server; GPSERVER has it off.)
    """
    overhead = len("SELECT  FROM ") + len(table)
    chunks: list[list[str]] = []
    cur: list[str] = []
    cur_len = overhead
    for i, e in enumerate(exprs):
        piece = f"{e} AS c{i}"
        add = len(piece) + (2 if cur else 0)
        if cur and cur_len + add > limit:
            chunks.append(cur)
            cur, cur_len = [], overhead
            add = len(piece)
        cur.append(piece)
        cur_len += add
    if cur:
        chunks.append(cur)

    parts = []
    for n, pieces in enumerate(chunks):
        inner = f"SELECT {', '.join(pieces)} FROM {table}"
        escaped = inner.replace("'", "''")
        parts.append(f"OPENQUERY({server}, '{escaped}') q{n}")
    return "SELECT * FROM " + "\n  CROSS JOIN ".join(parts)


def build_metrics(columns, key_names, drv, *, source: bool, quick: bool,
                  super_quick: bool = False):
    """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.

    Three depths: full (every column, incl. min/max/len_sum), ``quick``
    (per-column counts + numeric sums), and ``super_quick`` (the headline row
    counts only — no per-column work, so the DB can often answer from an index).
    """
    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 = col if (source and is_expression(col)) else drv.quote_identifier(col)
        labels.append((f"COUNT(DISTINCT {kn['dest']})", "int"))
        exprs.append(f"COUNT(DISTINCT {q})")

    if super_quick:
        return labels, exprs

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

        col_type = c["source_type"] if source else c["dest_type"]

        if kind == "numeric":
            labels.append((f"{name}: count", "int"))
            exprs.append(f"COUNT({e})")
            labels.append((f"{name}: sum", "num"))
            exprs.append(drv.sum_expression(col_type, 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(drv.sum_expression("int", drv.length_expression(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("--super-quick", action="store_true",
                    help="row counts only: COUNT(*) + COUNT(DISTINCT merge_key). "
                         "Cheapest check — catches missing/duplicated rows, not "
                         "changed column values")
    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

    # --super-quick subsumes --quick; pass both so the depth is unambiguous.
    quick = args.quick or args.super_quick
    labels, src_exprs = build_metrics(columns, key_names, src_drv, source=True,
                                      quick=quick, super_quick=args.super_quick)
    _, dst_exprs = build_metrics(columns, key_names, dst_drv, source=False,
                                 quick=quick, super_quick=args.super_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']}"

    # An explicit --source-from is taken verbatim (the caller may already have
    # written their own OPENQUERY). Otherwise, if the module reads through a
    # linked server, push the aggregate down the same way: computing it here
    # would drag every row across the link, and the inner table name does not
    # resolve locally at all.
    oq_server = None if args.source_from else detect_openquery_server(m["source_query"])
    if oq_server:
        src_sql = build_openquery_sql(oq_server, source_from, src_exprs)

    print(f"module {m['name']}  (id {m['id']})")
    src_label = (f"OPENQUERY({oq_server}) -> {source_from}" if oq_server
                 else source_from)
    print(f"  source {src_conn['name']}  FROM {src_label}")
    print(f"  dest   {dst_conn['name']}  FROM {m['dest_table']}")
    mode = ("  [super-quick]" if args.super_quick
            else "  [quick]" if args.quick else "")
    print(f"  {len(labels)} metrics{mode}\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())