Compressor/models/common/db_parity_check.py

#!/usr/bin/env python3
"""Re-fetch the full hourly table from the DB and compare it to the cached copy.

This honours the "DB full rerun" requirement without overwriting any cached
artifact.  The result is written to
``models/results/verification/db_parity.json``.
"""
from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[2]))

from models.common import db, paths, queries
from models.common.constants import EXPECTED_PER_HOUR


def key(row: dict):
    return (row["批次"], row["小时"])


def normalize(row: dict):
    values = {field: row.get(field) for field in queries.NUMERIC_HOURLY_FIELDS}
    values["运行样本数"] = row.get("运行样本数")
    return values


def main() -> None:
    parser = argparse.ArgumentParser(description="DB 全量小时中间表一致性校验")
    parser.add_argument("--start", default="2025-01-01 00:00:00")
    parser.add_argument("--end", default="2026-09-01 00:00:00")
    parser.add_argument("--chunk-days", type=int, default=14)
    parser.add_argument("--cached", default=str(paths.HOURLY_FULL))
    args = parser.parse_args()

    batch_to_unit = db.load_pks_batches()
    selected = sorted(batch_to_unit)
    configs = db.load_alarm_config({batch: batch_to_unit[batch] for batch in selected})
    start, end = queries.parse_dt(args.start), queries.parse_dt(args.end)
    print(f"从数据库重新聚合 {start} .. {end}，批次 {selected}", flush=True)
    fetched = queries.fetch_all_hourly(start, end, args.chunk_days, batch_to_unit, configs, selected)
    print(f"数据库小时行：{len(fetched):,}", flush=True)
    cached = queries.load_hourly_csv(Path(args.cached))
    print(f"缓存小时行：{len(cached):,}", flush=True)

    fetched_by_key = {key(row): row for row in fetched}
    cached_by_key = {key(row): row for row in cached}
    only_db = sorted(set(fetched_by_key) - set(cached_by_key))
    only_cache = sorted(set(cached_by_key) - set(fetched_by_key))
    mismatched = []
    for item_key in sorted(set(fetched_by_key) & set(cached_by_key)):
        if normalize(fetched_by_key[item_key]) != normalize(cached_by_key[item_key]):
            mismatched.append({"key": item_key, "db": normalize(fetched_by_key[item_key]), "cache": normalize(cached_by_key[item_key])})

    report = {
        "db_rows": len(fetched),
        "cached_rows": len(cached),
        "only_in_db": len(only_db),
        "only_in_cache": len(only_cache),
        "mismatched_rows": len(mismatched),
        "identical": not only_db and not only_cache and not mismatched,
        "only_in_db_examples": only_db[:10],
        "only_in_cache_examples": only_cache[:10],
        "mismatched_examples": mismatched[:5],
    }
    output = paths.RESULTS / "verification"
    output.mkdir(parents=True, exist_ok=True)
    (output / "db_parity.json").write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8")
    print(json.dumps({k: v for k, v in report.items() if not k.endswith("examples")}, ensure_ascii=False, indent=2), flush=True)
    print(f"报告：{output / 'db_parity.json'}", flush=True)


if __name__ == "__main__":
    main()