inkreach-prefect/scripts/api_db_diff.py

"""一次性对账:RIIN/SDS 远程订单量 vs PG Orders,按月 + 按状态(未完结)对比。

口径:
- RIIN:远程 count API(time_type=3 商家推送时间) vs 本地 orderedAt(=pushTime)。
- SDS:远程 statistics(payDate 窗口) vs 本地 orderedAt(=gmtPay),排除"售后"。
运行:uv run python scripts/api_db_diff.py
"""

from __future__ import annotations

import os
from datetime import datetime, timedelta

from dotenv import load_dotenv

load_dotenv(os.path.join(os.path.dirname(__file__), "..", ".env"))

import psycopg2  # noqa: E402

from api.riin.models import OrderPageRequest as RiinReq  # noqa: E402
from api.sds.models import OrderQueryParams  # noqa: E402
from flows.shared import (  # noqa: E402
    create_riin_client, create_sds_client, format_date, start_of_day,
)

RIIN_FINISHED = {"已发货", "已关闭", "已退款"}
SDS_FINISHED = {"已完成", "已取消"}

# count API 返回字段(camelCase 原始 key)→ 状态标签
RIIN_COUNT_MAP = {
    "assignedSuceessed": "已分配",
    "factoryAuditing": "工厂审核中",
    "inProducing": "生产中",
    "shipped": "已发货",
    "closed": "已关闭",
    "refunding": "退款中",
    "refunded": "已退款",
    "rejecting": "拒审中",
    "rejectedToShop": "已退回店铺",
}


def month_windows() -> list[tuple[datetime, datetime, str]]:
    """2025-01 起到当月的连续月窗口。"""
    out = []
    cur = datetime(2025, 1, 1)
    now = datetime.now()
    while cur < now:
        nxt = (cur.replace(day=28) + timedelta(days=4)).replace(day=1)
        out.append((cur, min(nxt, now), cur.strftime("%Y-%m")))
        cur = nxt
    return out


def riin_count(client, start: datetime, end: datetime, retries: int = 4) -> dict[str, int]:
    """直调 count 接口并解析原始 JSON(严格模型在月窗口下会因 dayRangeGroup=null 校验失败)。

    大月份窗口服务端聚合慢,超时按 5/10/20/40s 退避重试。
    """
    import time

    req = RiinReq(time_type=3, start_time=format_date(start), end_time=format_date(end))
    body = client._build_order_page_body(req)
    body.pop("pageIndex", None)
    body.pop("pageSize", None)
    for attempt in range(retries):
        try:
            data = client.post("distribution/order/list/count", body=body)
            break
        except Exception as err:  # noqa: BLE001
            if attempt == retries - 1 or "timed out" not in str(err).lower() and "timeout" not in str(err).lower():
                raise
            wait = 5 * (2 ** attempt)
            print(f"[RIIN count {start:%Y-%m}] 超时,{wait}s 后重试 ({attempt + 1}/{retries})")
            time.sleep(wait)
    raw = data.get("data") or {}
    return {RIIN_COUNT_MAP[k]: v for k, v in raw.items()
            if k in RIIN_COUNT_MAP and v}


def sds_status(client, start: datetime | None, end: datetime | None) -> dict[str, int]:
    params = OrderQueryParams(
        pay_date_begin=format_date(start) if start else None,
        pay_date_end=format_date(end) if end else None,
    )
    stats = client.get_status(params)
    out: dict[str, int] = {}
    for s in stats:
        name = getattr(s, "name", None)
        if not name or name == "售后":
            continue
        out[name] = out.get(name, 0) + (getattr(s, "orderNum", 0) or 0)
    return out


def main() -> None:
    months = month_windows()
    conn = psycopg2.connect(os.environ["DB_DSN"])
    cur = conn.cursor()

    # ---- DB 按月(orderedAt) ----
    cur.execute("""SELECT date_trunc('month',"orderedAt"), 'RIIN', COUNT(DISTINCT "orderNo")
        FROM "Orders" WHERE "systemName"='RIIN' AND "orderedAt" IS NOT NULL GROUP BY 1""")
    db_riin_month = {r[0].strftime("%Y-%m"): r[2] for r in cur.fetchall()}
    cur.execute("""SELECT date_trunc('month',"orderedAt"), COUNT(DISTINCT "orderNo")
        FROM "Orders" WHERE "systemName"='SDS' AND "orderedAt" IS NOT NULL GROUP BY 1""")
    db_sds_month = {r[0].strftime("%Y-%m"): r[1] for r in cur.fetchall()}

    # ---- DB 分状态(全时) ----
    cur.execute("""SELECT "systemStatus", COUNT(DISTINCT "orderNo") FROM "Orders"
        WHERE "systemName"='RIIN' GROUP BY 1""")
    db_riin_status = {r[0] or "(空)": r[1] for r in cur.fetchall()}
    cur.execute("""SELECT "systemStatus", COUNT(DISTINCT "orderNo") FROM "Orders"
        WHERE "systemName"='SDS' GROUP BY 1""")
    db_sds_status = {r[0] or "(空)": r[1] for r in cur.fetchall()}

    riin = create_riin_client()
    sds = create_sds_client()
    try:
        # ---- RIIN 远程按月(含状态明细) ----
        remote_riin_month: dict[str, int] = {}
        riin_status_sum: dict[str, int] = {}
        for start, end, label in months:
            m = riin_count(riin, start, end)
            remote_riin_month[label] = sum(m.values())
            for k, v in m.items():
                riin_status_sum[k] = riin_status_sum.get(k, 0) + v
        # ---- SDS 远程按月(含状态明细) ----
        remote_sds_month: dict[str, int] = {}
        sds_status_sum: dict[str, int] = {}
        for start, end, label in months:
            m = sds_status(sds, start, end)
            remote_sds_month[label] = sum(m.values())
            for k, v in m.items():
                sds_status_sum[k] = sds_status_sum.get(k, 0) + v

        # 全时状态计数 = 按月求和(RIIN 全时窗口 count 接口会超时)
        remote_riin_all = riin_status_sum
        remote_sds_all = sds_status_sum
    finally:
        riin.close()
        sds.close()

    print("=" * 72)
    print("按月对比(RIIN 远程=推送时间窗口计数 / 库=orderedAt;SDS 远程=付款时间 / 库=orderedAt)")
    print(f"{'月份':10}{'RIIN远程':>12}{'RIIN库':>12}{'差':>10}   {'SDS远程':>12}{'SDS库':>12}{'差':>10}")
    for _, _, label in months:
        rr, dr = remote_riin_month.get(label, 0), db_riin_month.get(label, 0)
        rs, ds = remote_sds_month.get(label, 0), db_sds_month.get(label, 0)
        mark = " ←" if abs(rr - dr) > 0 or abs(rs - ds) > 0 else ""
        print(f"{label:10}{rr:>12,}{dr:>12,}{dr - rr:>10,}   {rs:>12,}{ds:>12,}{ds - rs:>10,}{mark}")

    print("=" * 72)
    print("全时状态对比(远程 2025-01-01 ~ 明天 vs 库全量)")
    for title, remote, local in (("RIIN", remote_riin_all, db_riin_status),
                                 ("SDS", remote_sds_all, db_sds_status)):
        print(f"\n[{title}]  远程合计 {sum(remote.values()):,} / 库合计 {sum(local.values()):,}")
        for st in sorted(set(remote) | set(local)):
            r, l = remote.get(st, 0), local.get(st, 0)
            mark = "  ← 差异" if r != l else ""
            print(f"  {st:10} 远程 {r:>10,}   库 {l:>10,}{mark}")

    conn.close()


if __name__ == "__main__":
    main()