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.

``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 time
from collections.abc import AsyncIterator, Awaitable, Callable
from typing import Any

from app.session.yunke_client import YunkeClient

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

# 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,
    **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.
    """
    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)
        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.

    Call ``await limiter.acquire()`` before each upstream request. The first
    call never waits. Successive calls wait just long enough that the gap
    between the previous call's dispatch and this one is at least
    ``min_interval`` seconds.

    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.
    """

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

    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()