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

"""Reusable building blocks for list scrapers.

This module establishes the pagination patterns every list scraper in the
project reuses (see design doc §6.2 拉取层). It is intentionally small:

* :func:`paginate_by_page` — a generic async generator that walks page-numbered
  endpoints until a page comes back empty (the primary stop signal used by the
  yunke CRM), with a defensive over-fetch ceiling keyed off the declared total.
* :func:`paginate_by_cursor` — a generic async generator that walks a
  cursor-paginated endpoint until the fetch callable reports no more pages.
  Used by the WeChat message scraper, whose ``msg/chats`` endpoint pages
  BACKWARD by timestamp cursor (``end = prev_end - 1`` until ``hasLast``
  flips false).
* :class:`RateLimiter` — a tiny asyncio limiter that enforces a minimum
  interval between successive calls. Scrapers that need to be polite to the
  upstream construct one and ``await`` :meth:`acquire` before each call.

``paginate_by_page`` talks to :class:`app.session.yunke_client.YunkeClient`
(which already injects auth headers + handles 401 re-login). It is generic
over the response shape via three callables:

* ``parse_total(resp) -> int | None`` — read the declared row count from the
  envelope (used only for the over-fetch safety; ``None`` disables it).
* ``parse_rows(resp) -> list[dict]`` — extract the list of records.
* ``build_body(page_num, page_size, **extra) -> dict`` — build the request
  body for a given page.

By default the body is sent as JSON. A few yunke endpoints
(e.g. ``msg/chats``, ``wxUser/list``) bind ``@RequestParam`` from
**form-urlencoded** bodies instead — pass ``form=True`` to switch the paginator
to ``data=<body>`` with content-type
``application/x-www-form-urlencoded`` (mirroring the form-encoded
``msg/chats`` request the message scraper already makes directly).

``paginate_by_cursor`` is decoupled from the client entirely: it takes a
single ``fetch(cursor)`` coroutine that performs whatever request the caller
needs and returns ``(rows, next_cursor, has_more)``. The cursor-paginator's
job is just to advance the cursor and stop on ``has_more == False``.
"""

from __future__ import annotations

import asyncio
import random
import time
from collections.abc import AsyncIterator, Awaitable, Callable
from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
    # ``YunkeClient`` is referenced only in the paginator signatures. Kept
    # under TYPE_CHECKING so ``app.session.yunke_client`` may import
    # :class:`RateLimiter` at runtime without forming an import cycle.
    from app.session.yunke_client import YunkeClient

__all__ = [
    "FORM_CONTENT_TYPE",
    "RateLimiter",
    "paginate_by_cursor",
    "paginate_by_page",
]

# Content type for form-encoded bodies. Several yunke endpoints (msg/chats,
# wxUser/list) bind @RequestParam from form data rather than JSON, so the
# paginator + the manual message fetcher both send this content-type.
FORM_CONTENT_TYPE = "application/x-www-form-urlencoded"

# Informal type aliases for the response-handling callables. The yunke
# response shapes are dict-shaped JSON and do not benefit from finer typing.
ParseTotal = Callable[[Any], "int | None"]
ParseRows = Callable[[Any], "list[dict[str, Any]]"]
BuildBody = Callable[..., "dict[str, Any]"]
# fetch(cursor) -> (rows, next_cursor, has_more) — see ``paginate_by_cursor``.
FetchCursor = Callable[[Any], Awaitable["tuple[list[dict[str, Any]], Any, bool]"]]


async def paginate_by_page(
    client: YunkeClient,
    path: str,
    *,
    page_size: int,
    parse_total: ParseTotal,
    parse_rows: ParseRows,
    build_body: BuildBody,
    rate_limiter: RateLimiter | None = None,
    form: bool = False,
    **extra: Any,
) -> AsyncIterator[dict[str, Any]]:
    """Walk a page-numbered list endpoint, yielding one dict per record.

    Pagination strategy
    -------------------
    Start at ``pageNum=1`` and advance while pages are non-empty. The yunke
    CRM signals "no more rows" by returning an empty ``data`` array on the
    page *after* the last real record; that empty page is the **primary stop
    signal**. As a defensive ceiling against a misbehaving upstream that never
    sends an empty page, we also stop once we have yielded strictly more rows
    than the declared ``total`` (strict ``>`` so an honest source yields
    exactly ``total`` rows).

    ``total`` is parsed lazily from the first response and may be ``None`` if
    the envelope does not carry it — in that case only the empty-page signal
    terminates the loop.

    ``rate_limiter`` is an optional :class:`RateLimiter`; when supplied, it is
    acquired before every upstream call.

    ``form`` selects the wire encoding of ``build_body``'s output. The default
    (``False``) sends the body as JSON (``client.post(path, json=body)``) —
    correct for most yunke list endpoints. Pass ``True`` for endpoints that
    bind ``@RequestParam`` from a form-urlencoded body (``msg/chats``,
    ``wxUser/list``): the paginator then sends
    ``client.post(path, data=body, headers={"content-type": ...})`` and lets
    httpx URL-encode the dict. Either way ``build_body`` returns the same
    plain dict of field names → values.
    """
    page_num = 1
    total: int | None = None
    yielded = 0

    while True:
        if rate_limiter is not None:
            await rate_limiter.acquire()

        body = build_body(page_num, page_size, **extra)
        if form:
            resp = await client.post(
                path,
                data=body,
                headers={"content-type": FORM_CONTENT_TYPE},
            )
        else:
            resp = await client.post(path, json=body)

        # Read the declared total once from the first response. The yunke
        # envelopes keep it on every page, but reading it once is enough for
        # the over-fetch ceiling and avoids trusting stale updates.
        if total is None:
            total = parse_total(resp)

        rows = parse_rows(resp)
        if not rows:
            # Empty page == the upstream's "no more rows" signal. This is the
            # normal, primary termination path.
            return

        for row in rows:
            yield row
            yielded += 1

        # Over-fetch safety: a well-behaved source sends an empty page next,
        # but if it never does (or sends duplicates), cap at the declared total.
        # Strict ``>`` so we yield exactly ``total`` when the count is honest.
        if total is not None and yielded > total:
            return

        page_num += 1


async def paginate_by_cursor(
    fetch: FetchCursor,
    *,
    initial_cursor: Any,
    rate_limiter: RateLimiter | None = None,
) -> AsyncIterator[dict[str, Any]]:
    """Walk a cursor-paginated endpoint, yielding rows until no more pages.

    Used by the WeChat message scraper, whose ``msg/chats`` endpoint pages
    BACKWARD by timestamp cursor: each response carries ``data.end`` (the
    oldest timestamp in the batch) and ``data.hasLast`` (true iff there are
    older messages still). The caller advances the cursor to ``end - 1`` and
    repeats until ``hasLast`` is false.

    The contract is intentionally client-agnostic. Each call to
    ``fetch(cursor)`` must return a 3-tuple ``(rows, next_cursor, has_more)``:

    * ``rows``        — list of records from this page (may be empty).
    * ``next_cursor`` — cursor for the next page (e.g. ``prev_end - 1``).
    * ``has_more``    — True if there are more pages to fetch.

    The loop yields every row of every page, then advances ``cursor`` to
    ``next_cursor`` and repeats. It stops after:

    * a page whose ``has_more`` is False (the documented primary stop signal
      for cursor endpoints — e.g. ``hasLast=false``), OR
    * an empty ``rows`` list (a defensive secondary stop, so a misbehaving
      upstream that never flips ``has_more`` to False cannot loop forever).

    ``rate_limiter`` is acquired before each fetch when supplied, so cursor
    scrapers get the same inter-call spacing as page scrapers.
    """
    cursor = initial_cursor
    while True:
        if rate_limiter is not None:
            await rate_limiter.acquire()

        rows, next_cursor, has_more = await fetch(cursor)

        # Defensive stop: an empty page is treated as end-of-data even if the
        # upstream claims has_more. This mirrors paginate_by_page's primary
        # stop signal and bounds the worst case (a stuck cursor that never
        # advances would otherwise loop forever).
        if not rows:
            return

        for row in rows:
            yield row

        if not has_more:
            return

        cursor = next_cursor


class RateLimiter:
    """Minimal asyncio limiter enforcing a minimum interval between calls.

    Two pacing modes are supported:

    * :meth:`acquire` — gap-based spacing. Call it before each upstream
      request; the first call never waits and successive calls wait just long
      enough that the gap between dispatches is at least ``min_interval``
      seconds. Used by the list paginators (:func:`paginate_by_page` /
      :func:`paginate_by_cursor`).

    * :meth:`wait` — fixed per-call throttle. Sleeps ``min_interval`` plus a
      uniform random jitter in ``[0, jitter)`` every time it is awaited. Used
      by :class:`app.session.yunke_client.YunkeClient` to pace every upstream
      request with a conservative, slightly randomized cadence (the jitter
      prevents a perfectly periodic signature that risk control could finger).

    This is deliberately tiny: no tokens, no concurrency cap, no skew handling
    — just a floor on inter-call spacing, which is all the yunke scrapers need
    to stay polite to the CRM. If you need a semaphore or token bucket, that
    belongs in a future helper, not here.

    ``jitter`` defaults to ``0.0`` so existing callers that pass only
    ``min_interval`` keep their exact prior behaviour (gap-based ``acquire``
    never consults ``jitter``).
    """

    def __init__(self, min_interval: float = 0.0, jitter: float = 0.0) -> None:
        self._min_interval = float(min_interval)
        self._jitter = float(jitter)
        self._last: float = 0.0

    async def wait(self) -> None:
        """Sleep ``min_interval`` plus a random ``[0, jitter)`` every call.

        Unlike :meth:`acquire` this is not gap-aware — it is a flat, slightly
        randomized per-call delay, which is what the yunke client wants ahead
        of each upstream request. With ``min_interval == 0`` and
        ``jitter == 0`` it is a no-op.
        """
        delay = self._min_interval
        if self._jitter > 0:
            delay += random.uniform(0.0, self._jitter)
        if delay > 0:
            await asyncio.sleep(delay)

    async def acquire(self) -> None:
        """Block until at least ``min_interval`` has elapsed since the last call."""
        if self._min_interval <= 0:
            self._last = time.monotonic()
            return

        now = time.monotonic()
        wait = self._min_interval - (now - self._last)
        if wait > 0:
            await asyncio.sleep(wait)

        self._last = time.monotonic()