yunke/yunke-server/app/scraper/wechat.py

"""WeChat friend-list scraper.

Wraps ``POST /wechatweb/wxFriends/friend_list/es`` and yields normalized dicts
ready to be upserted into :class:`app.db.models.Customer` (persistence is a
LATER task — Task 1.7). Each dict carries:

* snake_case keys matching the ``Customer`` columns (see :func:`_normalize`),
* ``channel == "wechat"``,
* ``raw_json`` — the original upstream record, verbatim, so nothing captured
  is lost even when a column has no target field.

Field names are mapped from the real capture (see Task 1.5 spec). The mapping
is a straight ``camelCase -> snake_case`` translation with a couple of
disambiguations where the source has two fields for one column:

* ``phone``       <- ``phone``
* ``wechat_phone`` <- ``wechatPhone``
* ``dept_name``   <- ``deptName`` (falls back to ``deptRouteName``)
* ``head_image_url`` <- ``friendHeadImageUrl`` (falls back to ``headImg``)
"""

from __future__ import annotations

from collections.abc import AsyncIterator
from typing import Any

from app.scraper.base import paginate_by_page
from app.session.yunke_client import YunkeClient

__all__ = ["scrape_wechat_friends"]

FRIEND_LIST_PATH = "/wechatweb/wxFriends/friend_list/es"


def _parse_total(resp: Any) -> int | None:
    """Read the declared row count from the friend_list envelope.

    Shape: ``data.basicResult.total`` (an int). Returns ``None`` if the path
    is absent so the paginator falls back to the empty-page stop signal.
    """
    try:
        body = resp.json()
    except (AttributeError, ValueError):
        return None
    data = body.get("data") if isinstance(body, dict) else None
    if not isinstance(data, dict):
        return None
    basic = data.get("basicResult")
    if not isinstance(basic, dict):
        return None
    total = basic.get("total")
    return int(total) if isinstance(total, (int, float)) else None


def _parse_rows(resp: Any) -> list[dict[str, Any]]:
    """Extract the list of friend records from the envelope.

    Shape: ``data.basicResult.data`` (a list of records). Missing/non-list
    values collapse to an empty list so the paginator sees "empty page".
    """
    try:
        body = resp.json()
    except (AttributeError, ValueError):
        return []
    data = body.get("data") if isinstance(body, dict) else None
    if not isinstance(data, dict):
        return []
    basic = data.get("basicResult")
    if not isinstance(basic, dict):
        return []
    rows = basic.get("data")
    return rows if isinstance(rows, list) else []


def _build_body(
    page_num: int,
    page_size: int,
    *,
    department_id: str = "",
    user_id: str = "",
) -> dict[str, Any]:
    """Build the friend_list request body for a given page.

    The real client sends ``departmentId`` / ``userId`` filters; we accept
    them as optional kwargs (default empty) so per-staff / per-department
    scrapes are a drop-in without changing the paginator contract.
    """
    body: dict[str, Any] = {"pageNum": page_num, "pageSize": page_size}
    if department_id:
        body["departmentId"] = department_id
    if user_id:
        body["userId"] = user_id
    return body


def _normalize(record: dict[str, Any]) -> dict[str, Any]:
    """Map one upstream friend record to the ``Customer`` column shape.

    Adds ``channel="wechat"`` and preserves the original record under
    ``raw_json``. ``None``-safe: missing fields become ``None`` rather than
    raising, so partial payloads do not break the whole scrape.
    """
    return {
        "friend_id": record.get("friendWxId"),
        "channel": "wechat",
        "owner_account_id": record.get("wxId"),
        "friend_nickname": record.get("friendWxNickname"),
        "friend_remark": record.get("friendWxRemark"),
        "real_name": record.get("realName"),
        "phone": record.get("phone"),
        "wechat_phone": record.get("wechatPhone"),
        "gender": record.get("gender"),
        "region": record.get("region"),
        "description": record.get("description"),
        "sale_name": record.get("saleName"),
        "user_id": record.get("userId"),
        "dept_route": record.get("deptRoute"),
        # Prefer the canonical deptName; fall back to the route-display name.
        "dept_name": record.get("deptName") or record.get("deptRouteName"),
        # Prefer friendHeadImageUrl; fall back to the legacy headImg field.
        "head_image_url": record.get("friendHeadImageUrl") or record.get("headImg"),
        "last_chat_time": record.get("lastChatTime"),
        "send_num": record.get("sendNum"),
        "receive_num": record.get("receiveNum"),
        "message_count": record.get("messageCount"),
        "delete_flag": record.get("deleteFlag"),
        "is_block": record.get("isBlock"),
        "friend_state": record.get("friendState"),
        "version": record.get("version"),
        "raw_json": record,
    }


async def scrape_wechat_friends(
    client: YunkeClient,
    *,
    department_id: str = "",
    user_id: str = "",
    page_size: int = 50,
) -> AsyncIterator[dict[str, Any]]:
    """Yield normalized WeChat friend records from the CRM.

    Walks ``POST /wechatweb/wxFriends/friend_list/es`` page by page until the
    upstream reports no more rows, mapping each record to the ``Customer``
    column shape (snake_case) plus a verbatim ``raw_json`` copy.

    Parameters
    ----------
    client:
        Authenticated :class:`YunkeClient`. Auth headers / 401 re-login are
        handled by the client; this scraper only concerns itself with the
        friend-list endpoint contract.
    department_id, user_id:
        Optional upstream filters (passed through to the request body as
        ``departmentId`` / ``userId``). Empty by default — scrape everything
        the logged-in account can see.
    page_size:
        Page size to request. Defaults to 50 (a sane middle ground between
        request count and per-response payload size for this endpoint).
    """
    async for record in paginate_by_page(
        client,
        FRIEND_LIST_PATH,
        page_size=page_size,
        parse_total=_parse_total,
        parse_rows=_parse_rows,
        build_body=_build_body,
        department_id=department_id,
        user_id=user_id,
    ):
        yield _normalize(record)