inkreach-prefect/flows/shared/clients.py

"""同步共享:客户端工厂、顺序分页(含重试/延迟)、RIIN 商品拉取、SDS 物流渠道字典、日期辅助。

行为对齐 inkreach-bi-auth-mng 的 sync/tasks/clients.ts:
- 顺序分页(页间 500ms),非并发,对外部 API 友好。
- fetch_with_retry:仅对 5xx / 429 / 超时重试,最多 3 次,退避 3s*次数。
"""

from __future__ import annotations

import json
import time
from datetime import datetime, timedelta, timezone
from typing import Any, Callable, TypeVar

from prefect.blocks.system import Secret

from api import RiinDisClient, SdsMngClient
from api.riin.models import OrderPageRequest as RiinReq
from api.sds.models import LogisticsPageRequest, OrderPageRequest as SdsReq
from flows.shared.mappers import map_sds_channel_type

T = TypeVar("T")

PAGE_DELAY_MS = 500
MAX_RETRIES = 3
RETRY_DELAY_MS = 3000
TZ_CN = timezone(timedelta(hours=8))

_RIIN_FINISHED = ["已退款", "已关闭", "已发货"]
_SDS_FINISHED = ["已完成", "已取消"]


# ---------- 客户端工厂 ----------

def create_riin_client() -> RiinDisClient:
    cred = _load_secret("riin-credentials")
    return RiinDisClient(cred["username"], cred["password"])


def create_sds_client() -> SdsMngClient:
    cred = _load_secret("sds-credentials")
    return SdsMngClient(cred["username"], cred["password"])


def _load_secret(name: str) -> dict[str, str]:
    value = Secret.load(name).get()
    if isinstance(value, str):
        value = json.loads(value)
    return value


# ---------- 日期辅助 ----------

def format_date(d: datetime) -> str:
    return d.strftime("%Y-%m-%d %H:%M:%S")


def start_of_day(d: datetime) -> datetime:
    return d.replace(hour=0, minute=0, second=0, microsecond=0)


def floor_to_hour(d: datetime) -> datetime:
    return d.replace(minute=0, second=0, microsecond=0)


# ---------- 重试 ----------

def fetch_with_retry(fn: Callable[[], T], label: str) -> T:
    last_err: Exception | None = None
    for attempt in range(1, MAX_RETRIES + 1):
        try:
            return fn()
        except Exception as err:  # noqa: BLE001
            last_err = err
            status = _status_of(err)
            msg = str(err)
            is_timeout = "timeout" in msg.lower() or "timed out" in msg.lower()
            retryable = status == 429 or (status is not None and status >= 500) or is_timeout
            if attempt < MAX_RETRIES and retryable:
                print(f"[sync] {label} 第 {attempt}/{MAX_RETRIES} 次重试 ({status}: {msg[:120]})")
                time.sleep(RETRY_DELAY_MS * attempt / 1000)
                continue
            raise
    assert last_err is not None
    raise last_err


def _status_of(err: Exception) -> int | None:
    for attr in ("status_code", "statusCode", "status"):
        v = getattr(err, attr, None)
        if isinstance(v, int):
            return v
    resp = getattr(err, "response", None)
    return getattr(resp, "status_code", None) if resp is not None else None


# ---------- 分页 ----------

def fetch_riin_all_pages(client: RiinDisClient, params: dict[str, Any], max_pages: int = 1000) -> list[Any]:
    page_size = params.get("page_size", 200)
    base = {**params, "page_size": page_size}

    def fetch(idx: int):
        return fetch_with_retry(lambda: client.get_orders(RiinReq(**{**base, "page_index": idx})), f"RIIN page {idx}")

    first = fetch(1)
    total = first.data.total
    records = list(first.data.records)
    total_pages = min(-(-total // page_size), max_pages)
    print(f"[RIIN sync] 分页开始: 共 {total} 条 / {total_pages} 页 (page_size={page_size})")
    for page in range(2, total_pages + 1):
        resp = fetch(page)
        records.extend(resp.data.records)
        if page % 10 == 0 or page == total_pages:
            print(f"[RIIN sync] 分页进度: {page}/{total_pages} 页, 累计 {len(records)} 条")
        time.sleep(PAGE_DELAY_MS / 1000)
    if total_pages == max_pages:
        print(f"[RIIN sync] 达到最大页数限制 {max_pages},截断到 {len(records)} 条")
    print(f"[RIIN sync] 分页完成: 共拉取 {len(records)} 条")
    return records


def fetch_sds_all_pages(client: SdsMngClient, params: dict[str, Any]) -> list[Any]:
    page_size = params.get("size", 500)
    base = {**params, "size": page_size}

    def fetch(idx: int):
        return fetch_with_retry(lambda: client.get_orders(SdsReq(**{**base, "page": idx})), f"SDS page {idx}")

    first = fetch(1)
    records = list(first.records)
    total_pages = -(-first.total_count // page_size)
    print(f"[SDS sync] 分页开始: 共 {first.total_count} 条 / {total_pages} 页 (page_size={page_size})")
    for page in range(2, total_pages + 1):
        resp = fetch(page)
        records.extend(resp.records)
        if page % 10 == 0 or page == total_pages:
            print(f"[SDS sync] 分页进度: {page}/{total_pages} 页, 累计 {len(records)} 条")
        time.sleep(PAGE_DELAY_MS / 1000)
    print(f"[SDS sync] 分页完成: 共拉取 {len(records)} 条")
    return records


# ---------- RIIN 商品 ----------

# 批量接口 distribution/order/goods/list:单批 ≤200 个 id,且要求串行调用(分销侧约定)。
# 相比旧逐单接口(软限流:高频返回成功但空 data),调用量从日均 ~1 万次降到 ~50 次。
GOODS_BATCH_SIZE = 200
GOODS_BATCH_INTERVAL_MS = 2000  # 相邻两批的串行间隔


def fetch_goods_for_orders(client: RiinDisClient, order_ids: list[str]) -> dict[str, list[Any]]:
    """按 goods/list 批量接口串行拉取多订单商品,返回 {order_id: [goods...]}。

    单批失败(网络/业务错误)整批重试最多 MAX_RETRIES 次,仍失败则该批订单记空列表,
    由调用方按"无商品订单"计数(与旧逐单路径的容错语义一致)。
    """
    results: dict[str, list[Any]] = {oid: [] for oid in order_ids}
    if not order_ids:
        return results
    total_batches = -(-len(order_ids) // GOODS_BATCH_SIZE)
    print(f"[RIIN sync] 商品批量拉取开始: {len(order_ids)} 个订单 / {total_batches} 批"
          f" (goods/list 串行, 批间 {GOODS_BATCH_INTERVAL_MS}ms)")
    for bi in range(total_batches):
        chunk = order_ids[bi * GOODS_BATCH_SIZE : (bi + 1) * GOODS_BATCH_SIZE]
        items = _fetch_goods_batch_with_retry(client, chunk, bi + 1)
        for g in items:
            if g.order_id in results:
                results[g.order_id].append(g)
        if bi < total_batches - 1:
            time.sleep(GOODS_BATCH_INTERVAL_MS / 1000)
    total_goods = sum(len(v) for v in results.values())
    empty_orders = sum(1 for v in results.values() if not v)
    print(f"[RIIN sync] 商品批量拉取完成: 共 {total_goods} 条商品, {empty_orders} 个订单无商品")
    return results


def _fetch_goods_batch_with_retry(client: RiinDisClient, chunk: list[str], batch_no: int) -> list[Any]:
    last_err: Exception | None = None
    for attempt in range(1, MAX_RETRIES + 1):
        try:
            resp = fetch_with_retry(
                lambda: client.get_order_goods_batch(chunk), f"RIIN goods 批次 {batch_no}"
            )
            if resp.successful is False:
                raise RuntimeError(f"goods/list 业务失败: {resp.message}")
            return list(resp.data)
        except Exception as err:  # noqa: BLE001
            last_err = err
            print(f"[RIIN sync] 商品批次 {batch_no} 第 {attempt}/{MAX_RETRIES} 次失败: {str(err)[:120]}")
            if attempt < MAX_RETRIES:
                time.sleep(RETRY_DELAY_MS * attempt / 1000)
    assert last_err is not None
    print(f"[RIIN sync] 商品批次 {batch_no} 放弃, {len(chunk)} 个订单本轮无商品")
    return []


# ---------- SDS 物流渠道字典 ----------

_channel_cache: dict[str, Any] = {}
_CHANNEL_TTL = 24 * 3600
_CHANNEL_KEY = "all"


def get_channel_type_map(client: SdsMngClient) -> dict[int, str | None]:
    cached = _channel_cache.get(_CHANNEL_KEY)
    if cached and cached["expire_at"] > time.time():
        return cached["data"]
    mapping: dict[int, str | None] = {}
    try:
        areas = client.get_logistics_areas()
        area_ids = [a.id for a in areas.items if a.id is not None]
    except Exception as err:  # noqa: BLE001
        print(f"[SDS sync] 物流区域拉取失败: {err}")
        area_ids = []
    for area_id in area_ids:
        try:
            resp = client.get_logistics_page(LogisticsPageRequest(issuing_bay_area_id=area_id, all=True, size=200))
            for item in resp.items:
                if item.id is None:
                    continue
                raw = getattr(item, "channelType", None) or getattr(item, "channel_type", None)
                label = map_sds_channel_type(raw)
                if label is None:
                    continue
                mapping[int(item.id)] = label
        except Exception as err:  # noqa: BLE001
            print(f"[SDS sync] 区域 {area_id} 物流字典拉取失败: {err}")
    _channel_cache[_CHANNEL_KEY] = {"data": mapping, "expire_at": time.time() + _CHANNEL_TTL}
    return mapping


def refresh_channel_type_map(client: SdsMngClient) -> dict[int, str | None]:
    _channel_cache.pop(_CHANNEL_KEY, None)
    return get_channel_type_map(client)


def resolve_courier_source(
    raw_logistics_id: Any,
    mapping: dict[int, str | None],
    refresh: Callable[[], dict[int, str | None]],
    refreshed_ids: set[int],
) -> str | None:
    if raw_logistics_id is None:
        return None
    try:
        lid = int(raw_logistics_id)
    except (TypeError, ValueError):
        return None
    if lid in mapping:
        return mapping[lid]
    if lid in refreshed_ids:
        return None
    refreshed_ids.add(lid)
    try:
        new_map = refresh()
        mapping.update(new_map)
    except Exception:  # noqa: BLE001
        pass
    return mapping.get(lid)