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

"""Reusable building blocks for list scrapers.

This module establishes the pagination pattern 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.
* :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.

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

from __future__ import annotations

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

from app.session.yunke_client import YunkeClient

__all__ = ["RateLimiter", "paginate_by_page"]

# Informal type aliases for the three 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]"]


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


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