diff --git a/yunke-server/app/sync/__init__.py b/yunke-server/app/sync/__init__.py new file mode 100644 index 0000000..d94976b --- /dev/null +++ b/yunke-server/app/sync/__init__.py @@ -0,0 +1,48 @@ +"""Incremental-sync coordination layer (Task 2.7). + +This package ties the scrapers (``app.scraper.*``) and the persistence layer +(``app.storage.repos``) together with **resumable, incremental** sync via the +``SyncState`` watermark table and the ``SyncRun`` audit log. + +Public surface: + +* :mod:`app.sync.watermarks` — SyncState CRUD (per-entity / per-channel / + per-scope cursor get/set + convenience accessors). +* :mod:`app.sync.orchestrator` — the sync brain. One entrypoint per entity + (``sync_wechat_messages`` / ``sync_wechat_friends`` / ...) plus a coordinator + (:func:`app.sync.orchestrator.sync_all_wechat`). Each run is logged as a + ``SyncRun`` row (started → success/failed) so the dashboard / ops view can + report progress and errors. +""" + +from app.sync.orchestrator import ( + SyncResult, + sync_all_wechat, + sync_departments, + sync_wechat_accounts, + sync_wechat_friends, + sync_wechat_groups, + sync_wechat_messages, + sync_wechat_voice_calls, +) +from app.sync.watermarks import ( + get_last_timestamp, + get_last_version, + get_watermark, + set_watermark, +) + +__all__ = [ + "SyncResult", + "get_last_timestamp", + "get_last_version", + "get_watermark", + "set_watermark", + "sync_all_wechat", + "sync_departments", + "sync_wechat_accounts", + "sync_wechat_friends", + "sync_wechat_groups", + "sync_wechat_messages", + "sync_wechat_voice_calls", +] diff --git a/yunke-server/app/sync/orchestrator.py b/yunke-server/app/sync/orchestrator.py new file mode 100644 index 0000000..54b3430 --- /dev/null +++ b/yunke-server/app/sync/orchestrator.py @@ -0,0 +1,475 @@ +"""The sync brain — coordinates scrapers, repos, watermarks, and the SyncRun log. + +Each public ``sync_*`` function follows the same shape: + + read watermark -> scrape -> upsert -> (advance watermark) -> log SyncRun + +Two flavors of "incremental" live here, per Task 2.7: + +* **Time-window incremental** (messages only). The message endpoint pages + backward over a ``(start, end]`` timestamp window, so a real incremental query + is possible. ``sync_wechat_messages`` reads ``last_timestamp`` for + ``(message, wechat, ":")``, queries + ``[last_timestamp, now]`` via :func:`scrape_wechat_messages`, upserts each + record, **best-effort downloads media** for media-bearing messages, and + advances the watermark to ``max(timestamp)``. A re-run with no new messages + therefore re-queries the empty window and fetches nothing — the watermark + advanced past the previously-seen rows. +* **Full re-scrape + idempotent upsert** (every other entity). The list + endpoints (friends / accounts / groups / voice_calls / departments) have no + ``since`` param, so "incremental" here = re-scrape everything and rely on the + repos' idempotent Core upsert (same PK twice -> one refreshed row). The + watermark records ``last_synced_at`` only — no cursor advance, because there + is no cursor to advance. + +Failure policy: a scraper that raises is caught, recorded as a ``SyncRun`` row +with ``status='failed'`` and the error string, and a failed :class:`SyncResult` +is returned — the orchestrator does NOT re-raise, so a single broken entity +does not abort a multi-entity batch. The caller decides whether to surface / +retry. Media-download failures are best-effort: a media error is logged and the +sync continues (the message is still upserted; the watermark still advances). + +Transactions: the orchestrator never commits. Every write is flushed into the +caller's ``AsyncSession``; the caller owns the commit / rollback. This keeps a +multi-entity batch atomic at the caller's discretion and lets a test roll a +failed advance back by simply not committing. +""" + +from __future__ import annotations + +import dataclasses +import datetime as dt +import logging +from collections.abc import AsyncIterator, Callable +from typing import Any + +from sqlalchemy.ext.asyncio import AsyncSession + +from app.db.models import SyncRun +from app.scraper.media import download_message_media +from app.scraper.wechat import ( + scrape_departments, + scrape_wechat_accounts, + scrape_wechat_friends, + scrape_wechat_groups, + scrape_wechat_messages, + scrape_wechat_voice_calls, +) +from app.session.yunke_client import YunkeClient +from app.storage.object_store import ObjectStore +from app.storage.repos import ( + upsert_chat_group, + upsert_customer, + upsert_department, + upsert_messages, + upsert_voice_call, + upsert_wechat_account, +) +from app.sync.watermarks import get_last_timestamp, set_watermark + +logger = logging.getLogger(__name__) + +__all__ = [ + "SyncResult", + "sync_all_wechat", + "sync_departments", + "sync_wechat_accounts", + "sync_wechat_friends", + "sync_wechat_groups", + "sync_wechat_messages", + "sync_wechat_voice_calls", +] + +# --------------------------------------------------------------------------- +# Stable SyncRun / SyncState keys (entity_type / channel labels) +# --------------------------------------------------------------------------- + +# entity_type labels — kept as module constants so the SyncRun / SyncState keys +# are stable strings shared between the orchestrator and any dashboard / ops +# consumer (e.g. Task 6.4 monitoring). +ENTITY_DEPARTMENT = "department" +ENTITY_WECHAT_ACCOUNT = "wechat_account" +ENTITY_WECHAT_FRIEND = "wechat_friend" +ENTITY_WECHAT_GROUP = "wechat_group" +ENTITY_WECHAT_VOICE_CALL = "wechat_voice_call" +ENTITY_MESSAGE = "message" + +CHANNEL_WECHAT = "wechat" + + +# --------------------------------------------------------------------------- +# SyncResult — small per-run outcome returned to the caller +# --------------------------------------------------------------------------- + + +@dataclasses.dataclass(frozen=True) +class SyncResult: + """Outcome of one sync run (one entity / one conversation). + + Mirrors the essential fields of the persisted ``SyncRun`` row so callers can + react to a run without re-querying the DB. ``status`` is ``"success"`` or + ``"failed"``; on failure ``error`` carries the exception text. + """ + + entity: str + status: str + fetched: int + upserted: int + channel: str | None = None + scope_id: str | None = None + error: str | None = None + + +# --------------------------------------------------------------------------- +# Time helpers +# --------------------------------------------------------------------------- + + +def _now() -> dt.datetime: + """Aware UTC now (never naive). Used for SyncRun timestamps.""" + return dt.datetime.now(dt.UTC) + + +def _now_ms() -> int: + """Current time as ms-epoch (the message scraper's timestamp unit).""" + return int(_now().timestamp() * 1000) + + +# --------------------------------------------------------------------------- +# SyncRun bookkeeping +# --------------------------------------------------------------------------- + + +async def _start_run( + session: AsyncSession, + entity_type: str, + *, + channel: str | None = None, + scope_id: str | None = None, +) -> SyncRun: + """Open a ``running`` SyncRun row and flush it (assigning run_id).""" + run = SyncRun( + entity_type=entity_type, + channel=channel, + scope_id=scope_id, + status="running", + started_at=_now(), + ) + session.add(run) + await session.flush() + return run + + +async def _finish_run( + session: AsyncSession, + run: SyncRun, + *, + status: str, + fetched: int, + upserted: int, + error: str | None = None, +) -> None: + """Stamp a SyncRun row with final status / counts / finished_at and flush.""" + run.status = status + run.records_fetched = fetched + run.records_upserted = upserted + run.error = error + run.finished_at = _now() + await session.flush() + + +# --------------------------------------------------------------------------- +# Messages — time-window incremental (the workhorse) +# --------------------------------------------------------------------------- + + +def _message_scope(wechat_id: str, friend_id: str) -> str: + """Scope-id key for a 1:1 (or room) WeChat conversation.""" + return f"{wechat_id}:{friend_id}" + + +async def sync_wechat_messages( + session: AsyncSession, + client: YunkeClient, + store: ObjectStore, + wechat_id: str, + friend_id: str, + *, + types: str = "", +) -> SyncResult: + """Incremental sync of one WeChat conversation's messages. + + Reads ``last_timestamp`` for ``(message, wechat, ":")`` + (default ``0`` on first sync), queries ``[last_timestamp, now]`` via the + scraper, upserts every yielded message, best-effort downloads media for + each, and advances the watermark to ``max(timestamp)``. + + Returns a :class:`SyncResult`. On a scraper failure, records a + ``status='failed'`` SyncRun with the error and returns a failed result + (does NOT re-raise). + """ + scope_id = _message_scope(wechat_id, friend_id) + last_ts = await get_last_timestamp( + session, + ENTITY_MESSAGE, + channel=CHANNEL_WECHAT, + scope_id=scope_id, + default=0, + ) + end = _now_ms() + + run = await _start_run( + session, + ENTITY_MESSAGE, + channel=CHANNEL_WECHAT, + scope_id=scope_id, + ) + + # --- Scrape (may raise) --- + try: + msgs: list[dict[str, Any]] = [ + msg + async for msg in scrape_wechat_messages( + client, wechat_id, friend_id, last_ts, end, types=types + ) + ] + except Exception as exc: + logger.exception( + "scrape_wechat_messages failed for %s", scope_id + ) + error = repr(exc) + await _finish_run( + session, run, status="failed", fetched=0, upserted=0, error=error + ) + return SyncResult( + entity=ENTITY_MESSAGE, + status="failed", + fetched=0, + upserted=0, + channel=CHANNEL_WECHAT, + scope_id=scope_id, + error=error, + ) + + # --- Upsert (idempotent; same msg_id twice -> one refreshed row) --- + upserted = await upsert_messages(session, msgs) if msgs else 0 + + # --- Best-effort media download (never aborts the sync) --- + for msg in msgs: + try: + await download_message_media(client, msg, store, session) + except Exception: + logger.warning( + "media download failed for msg %s; sync continues", + msg.get("msg_id"), + exc_info=True, + ) + + # --- Advance watermark to max(timestamp) seen in this batch --- + timestamps = [m.get("timestamp") for m in msgs if m.get("timestamp") is not None] + if timestamps: + await set_watermark( + session, + ENTITY_MESSAGE, + channel=CHANNEL_WECHAT, + scope_id=scope_id, + last_timestamp=max(timestamps), + ) + + await _finish_run( + session, run, status="success", fetched=len(msgs), upserted=upserted + ) + return SyncResult( + entity=ENTITY_MESSAGE, + status="success", + fetched=len(msgs), + upserted=upserted, + channel=CHANNEL_WECHAT, + scope_id=scope_id, + ) + + +# --------------------------------------------------------------------------- +# Full re-scrape entities — idempotent upsert, no since-param +# --------------------------------------------------------------------------- + + +# Type aliases for readability (a scrape factory takes the client captured at +# call time and returns an async generator of records; an upsert takes the +# session + one record). +_ScrapeFactory = Callable[[], AsyncIterator[dict[str, Any]]] +_UpsertOne = Callable[[AsyncSession, dict[str, Any]], Any] + + +async def _sync_full_rescrape( + session: AsyncSession, + *, + entity_type: str, + channel: str | None, + scope_id: str | None, + scrape: _ScrapeFactory, + upsert: _UpsertOne, +) -> SyncResult: + """Shared body for the "re-scrape all + idempotent upsert" entities. + + These list endpoints have no ``since`` param, so "incremental" = re-scrape + everything and rely on the repos' idempotent Core upsert (same PK twice -> + one refreshed row). The watermark records ``last_synced_at`` only; no + version/timestamp cursor is advanced because there is no upstream window to + close. A scraper failure is recorded as a failed SyncRun and a failed + SyncResult is returned (no re-raise). + """ + run = await _start_run( + session, entity_type, channel=channel, scope_id=scope_id + ) + + try: + recs: list[dict[str, Any]] = [rec async for rec in scrape()] + except Exception as exc: + logger.exception("scrape failed for %s", entity_type) + error = repr(exc) + await _finish_run( + session, run, status="failed", fetched=0, upserted=0, error=error + ) + return SyncResult( + entity=entity_type, + status="failed", + fetched=0, + upserted=0, + channel=channel, + scope_id=scope_id, + error=error, + ) + + for rec in recs: + await upsert(session, rec) + + # Stamp the watermark's last_synced_at (no cursor advance — see docstring). + await set_watermark( + session, entity_type, channel=channel, scope_id=scope_id + ) + + await _finish_run( + session, run, status="success", fetched=len(recs), upserted=len(recs) + ) + return SyncResult( + entity=entity_type, + status="success", + fetched=len(recs), + upserted=len(recs), + channel=channel, + scope_id=scope_id, + ) + + +async def sync_departments( + session: AsyncSession, client: YunkeClient +) -> SyncResult: + """Full re-scrape of the department tree and idempotent upsert.""" + return await _sync_full_rescrape( + session, + entity_type=ENTITY_DEPARTMENT, + channel=None, + scope_id=None, + scrape=lambda: scrape_departments(client), + upsert=upsert_department, + ) + + +async def sync_wechat_accounts( + session: AsyncSession, client: YunkeClient +) -> SyncResult: + """Full re-scrape of all WeChat staff accounts (+ embedded staff rows).""" + return await _sync_full_rescrape( + session, + entity_type=ENTITY_WECHAT_ACCOUNT, + channel=None, + scope_id=None, + scrape=lambda: scrape_wechat_accounts(client), + upsert=upsert_wechat_account, + ) + + +async def sync_wechat_friends( + session: AsyncSession, client: YunkeClient +) -> SyncResult: + """Full re-scrape of all WeChat friends (customers) and idempotent upsert.""" + return await _sync_full_rescrape( + session, + entity_type=ENTITY_WECHAT_FRIEND, + channel=CHANNEL_WECHAT, + scope_id=None, + scrape=lambda: scrape_wechat_friends(client), + upsert=upsert_customer, + ) + + +async def sync_wechat_groups( + session: AsyncSession, client: YunkeClient +) -> SyncResult: + """Full re-scrape of all WeChat chat-groups (+members) and idempotent upsert.""" + return await _sync_full_rescrape( + session, + entity_type=ENTITY_WECHAT_GROUP, + channel=CHANNEL_WECHAT, + scope_id=None, + scrape=lambda: scrape_wechat_groups(client), + upsert=upsert_chat_group, + ) + + +async def sync_wechat_voice_calls( + session: AsyncSession, client: YunkeClient +) -> SyncResult: + """Full re-scrape of all WeChat voice-call records and idempotent upsert.""" + return await _sync_full_rescrape( + session, + entity_type=ENTITY_WECHAT_VOICE_CALL, + channel=None, + scope_id=None, + scrape=lambda: scrape_wechat_voice_calls(client), + upsert=upsert_voice_call, + ) + + +# --------------------------------------------------------------------------- +# Coordinator +# --------------------------------------------------------------------------- + + +async def sync_all_wechat( + session: AsyncSession, + client: YunkeClient, + store: ObjectStore, + *, + conversations: list[tuple[str, str]] | None = None, +) -> list[SyncResult]: + """Run the full WeChat sync sequence: org/entities first, then messages. + + Order: departments → accounts → friends → groups → voice_calls, then the + per-conversation message sync for every ``(wechat_id, friend_id)`` in + ``conversations``. Each step returns its own :class:`SyncResult` so a single + failed entity does NOT abort the batch (the failure is recorded and the + next entity proceeds). Returns the results in execution order. + + ``conversations`` is the explicit "which direct / room chats to pull + messages for" list. Full auto-discovery from the talkers endpoint is a + LATER concern (Task 2.5 wires the scraper; wiring it into the coordinator + is intentionally deferred to keep this entrypoint bounded and testable). + When ``conversations`` is ``None`` / empty, only the five full-re-scrape + entities run. + """ + results: list[SyncResult] = [ + await sync_departments(session, client), + await sync_wechat_accounts(session, client), + await sync_wechat_friends(session, client), + await sync_wechat_groups(session, client), + await sync_wechat_voice_calls(session, client), + ] + for wechat_id, friend_id in conversations or []: + results.append( + await sync_wechat_messages( + session, client, store, wechat_id, friend_id + ) + ) + return results diff --git a/yunke-server/app/sync/watermarks.py b/yunke-server/app/sync/watermarks.py new file mode 100644 index 0000000..e9e0d19 --- /dev/null +++ b/yunke-server/app/sync/watermarks.py @@ -0,0 +1,154 @@ +"""SyncState CRUD — resumable sync watermarks. + +The ``SyncState`` table holds one cursor per ``(entity_type, channel, scope_id)`` +triple (unique index ``ix_sync_state_scope``). The orchestrator reads a cursor +before a sync run to decide the upstream query window, and advances it after a +successful run so a crash + re-run does not re-pull the same rows. + +Two cursors live on one row: + +* ``last_version`` — opaque source update stamp (string). Used by entities that + expose a ``version > watermark`` style incremental query. +* ``last_timestamp`` — ms-epoch int. Used by the message sync (time-window + incremental: query ``[last_timestamp, now]``). + +``last_synced_at`` stamps every successful advance (the dashboard surfaces +"last synced" per entity); ``updated_at`` is the ORM-level row mtime. + +The helpers here are intentionally lightweight (select-then-update ORM style, +not Core upsert). ``SyncState`` is a tiny bookkeeping table — one row per +entity/channel/scope — so the extra round-trip is irrelevant and the partial- +update semantics (advance one cursor without clobbering the other with NULL) +are clearer this way than fighting a Core ``ON CONFLICT`` to be conditional. +""" + +from __future__ import annotations + +import datetime as dt + +import sqlalchemy as sa +from sqlalchemy.ext.asyncio import AsyncSession + +from app.db.models import SyncState + +__all__ = [ + "get_last_timestamp", + "get_last_version", + "get_watermark", + "set_watermark", +] + + +async def get_watermark( + session: AsyncSession, + entity_type: str, + *, + channel: str | None = None, + scope_id: str | None = None, +) -> SyncState | None: + """Return the ``SyncState`` row for ``(entity_type, channel, scope_id)``. + + ``channel`` / ``scope_id`` are optional filters: ``None`` matches the row + whose column IS NULL (the canonical "no sub-scope" cursor for an entity). + Returns ``None`` when no such row exists (first sync of a brand-new scope). + """ + stmt = sa.select(SyncState).where(SyncState.entity_type == entity_type) + if channel is not None: + stmt = stmt.where(SyncState.channel == channel) + else: + stmt = stmt.where(SyncState.channel.is_(None)) + if scope_id is not None: + stmt = stmt.where(SyncState.scope_id == scope_id) + else: + stmt = stmt.where(SyncState.scope_id.is_(None)) + result = await session.execute(stmt) + return result.scalar_one_or_none() + + +async def get_last_timestamp( + session: AsyncSession, + entity_type: str, + *, + channel: str | None = None, + scope_id: str | None = None, + default: int = 0, +) -> int: + """Convenience: return the stored ``last_timestamp`` or ``default``. + + ``default`` (``0``) is returned when no row exists yet OR the row exists but + has no timestamp cursor set. A ``0`` start window means "pull everything + upstream" — the correct behavior for a first sync. + """ + row = await get_watermark( + session, entity_type, channel=channel, scope_id=scope_id + ) + if row is None or row.last_timestamp is None: + return default + return int(row.last_timestamp) + + +async def get_last_version( + session: AsyncSession, + entity_type: str, + *, + channel: str | None = None, + scope_id: str | None = None, +) -> str | None: + """Return the stored ``last_version`` cursor, or ``None`` if not set. + + Distinct from :func:`get_last_timestamp`'s ``default=0`` contract: a version + cursor is a string and has no natural "zero" — ``None`` means "no version + cursor yet" and the caller decides what that implies (typically: full + re-scrape). + """ + row = await get_watermark( + session, entity_type, channel=channel, scope_id=scope_id + ) + if row is None: + return None + return row.last_version + + +async def set_watermark( + session: AsyncSession, + entity_type: str, + *, + channel: str | None = None, + scope_id: str | None = None, + last_version: str | None = None, + last_timestamp: int | None = None, +) -> None: + """Upsert the ``SyncState`` cursor row for ``(entity_type, channel, scope_id)``. + + * If no row exists, inserts a fresh one carrying the supplied cursors and + ``last_synced_at = now``. + * If a row exists, advances **only** the cursors the caller provided + (non-``None``); a ``None`` argument is treated as "do not touch", so a + timestamp-only advance does NOT wipe a previously-set version cursor + (and vice versa). ``last_synced_at`` is always refreshed to now. + + The row is flushed (not committed) so callers own the transaction; a + failing batch can still roll the watermark advance back. + """ + existing = await get_watermark( + session, entity_type, channel=channel, scope_id=scope_id + ) + now = dt.datetime.now(dt.UTC) + if existing is None: + session.add( + SyncState( + entity_type=entity_type, + channel=channel, + scope_id=scope_id, + last_version=last_version, + last_timestamp=last_timestamp, + last_synced_at=now, + ) + ) + else: + if last_version is not None: + existing.last_version = last_version + if last_timestamp is not None: + existing.last_timestamp = last_timestamp + existing.last_synced_at = now + await session.flush() diff --git a/yunke-server/tests/test_orchestrator.py b/yunke-server/tests/test_orchestrator.py new file mode 100644 index 0000000..8ec6c1f --- /dev/null +++ b/yunke-server/tests/test_orchestrator.py @@ -0,0 +1,520 @@ +"""Tests for app.sync.orchestrator — incremental sync coordination. + +The orchestrator is the integration seam that ties scrapers + repos together +with resumable, incremental sync. These tests monkeypatch the scraper + repo +functions (so no network / no real upsert SQL beyond what we assert) and drive +the orchestrator against an in-memory SQLite engine to verify the documented +behaviors: + +* **Messages (time-window incremental)**: first run upserts 2 messages and + advances the watermark to ``max(timestamp)``; the second run's scrape is + called with ``start == last_watermark`` and finds nothing new (watermark + advanced past) — the proof of incremental correctness. +* **Full re-scrape entity (friends)**: scraper yields 3 records; upsert is + called once per record; a ``SyncRun`` row with ``status='success'`` and the + right counts is recorded; the watermark's ``last_synced_at`` is set. +* **Failure policy**: when the scraper raises, the orchestrator records a + ``SyncRun`` row with ``status='failed'`` + the error string and returns a + failed result WITHOUT re-raising (the rest of the batch keeps going). +* **Media coupling**: media download is best-effort — a download that raises + does NOT abort the message sync; the message is still upserted and the + watermark still advances. + +All tests build their own in-memory SQLite engine and never touch production. +""" + +from __future__ import annotations + +import sqlalchemy as sa +from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine + +from app.db.base import Base +from app.db.models import SyncRun, SyncState +from app.sync import orchestrator +from app.sync.watermarks import get_last_timestamp + +T1 = 1_700_000_000_000 +T2 = 1_700_000_001_000 + + +# --------------------------------------------------------------------------- +# Shared helpers +# --------------------------------------------------------------------------- + + +async def _make_engine(): + """Create an in-memory async SQLite engine with all tables created.""" + engine = create_async_engine("sqlite+aiosqlite:///:memory:") + async with engine.begin() as conn: + await conn.run_sync(Base.metadata.create_all) + return engine + + +def _msg(msg_id: str, ts: int, *, channel: str = "wechat") -> dict: + """Build a scraper-shaped message dict (snake_case keys).""" + return { + "msg_id": msg_id, + "channel": channel, + "owner_account_id": "wx_staff", + "talker_id": "friend_a", + "timestamp": ts, + "type": "text", + "type_raw": 1, + "raw_json": {"msgSvrId": msg_id, "timestamp": ts}, + } + + +# =========================================================================== +# Messages — time-window incremental +# =========================================================================== + + +async def test_messages_first_run_upserts_and_advances_watermark(monkeypatch): + """First run: scrape [m1(T1), m2(T2)] -> upsert 2, watermark -> T2.""" + engine = await _make_engine() + try: + SessionLocal = async_sessionmaker(engine, expire_on_commit=False) + + scrape_calls: list[dict] = [] + + async def fake_scrape(client, wechat_id, friend_id, start, end, *, types=""): + scrape_calls.append( + { + "wechat_id": wechat_id, + "friend_id": friend_id, + "start": start, + "end": end, + "types": types, + } + ) + for m in [_msg("m1", T1), _msg("m2", T2)]: + yield m + + upsert_calls: list[list[dict]] = [] + + async def fake_upsert(session, recs): + upsert_calls.append(list(recs)) + return len(recs) + + async def fake_media(client, msg, store, session): + return None + + monkeypatch.setattr(orchestrator, "scrape_wechat_messages", fake_scrape) + monkeypatch.setattr(orchestrator, "upsert_messages", fake_upsert) + monkeypatch.setattr(orchestrator, "download_message_media", fake_media) + + async with SessionLocal() as session: + result = await orchestrator.sync_wechat_messages( + session, client=object(), store=object(), wechat_id="wx_staff", + friend_id="friend_a", + ) + await session.commit() + + # First call always starts from the default watermark 0. + assert scrape_calls[0]["start"] == 0 + assert scrape_calls[0]["wechat_id"] == "wx_staff" + assert scrape_calls[0]["friend_id"] == "friend_a" + + # Result reports the fetched / upserted counts. + assert result.status == "success" + assert result.fetched == 2 + assert result.upserted == 2 + + # upsert called once with both records. + assert len(upsert_calls) == 1 + assert len(upsert_calls[0]) == 2 + + # Watermark advanced to max(timestamp) == T2. + async with SessionLocal() as session: + assert ( + await get_last_timestamp( + session, + "message", + channel="wechat", + scope_id="wx_staff:friend_a", + ) + == T2 + ) + finally: + await engine.dispose() + + +async def test_messages_second_run_scrape_starts_at_watermark_and_fetches_nothing( + monkeypatch, +): + """Two consecutive runs: the second scrape is called with start==T2 and + upserts 0 records — the watermark advanced PAST the previously-seen + messages. This is the core incremental-correctness proof.""" + engine = await _make_engine() + try: + SessionLocal = async_sessionmaker(engine, expire_on_commit=False) + + # A queue of batches: first call yields [m1, m2], second call yields []. + batches: list[list[dict]] = [[_msg("m1", T1), _msg("m2", T2)], []] + scrape_calls: list[int] = [] + + async def fake_scrape(client, wechat_id, friend_id, start, end, *, types=""): + scrape_calls.append(start) + for m in batches.pop(0): + yield m + + upsert_calls: list[list[dict]] = [] + + async def fake_upsert(session, recs): + upsert_calls.append(list(recs)) + return len(recs) + + async def fake_media(client, msg, store, session): + return None + + monkeypatch.setattr(orchestrator, "scrape_wechat_messages", fake_scrape) + monkeypatch.setattr(orchestrator, "upsert_messages", fake_upsert) + monkeypatch.setattr(orchestrator, "download_message_media", fake_media) + + # --- Run 1 --- + async with SessionLocal() as session: + r1 = await orchestrator.sync_wechat_messages( + session, client=object(), store=object(), wechat_id="wx_staff", + friend_id="friend_a", + ) + await session.commit() + assert r1.upserted == 2 + assert scrape_calls[0] == 0 # first run starts at default watermark + + # --- Run 2 (no new messages upstream) --- + async with SessionLocal() as session: + r2 = await orchestrator.sync_wechat_messages( + session, client=object(), store=object(), wechat_id="wx_staff", + friend_id="friend_a", + ) + await session.commit() + + # The second scrape MUST be called with start == T2 (the watermark we + # advanced to in run 1) — proving we did not re-query [0, now] again. + assert scrape_calls[1] == T2 + assert r2.status == "success" + assert r2.fetched == 0 + assert r2.upserted == 0 + # No new messages -> the orchestrator does not even call upsert on run 2 + # (the empty-batch round-trip is skipped). The single recorded upsert + # call is from run 1. + assert len(upsert_calls) == 1 + + # Watermark unchanged (no new max to advance to). + async with SessionLocal() as session: + assert ( + await get_last_timestamp( + session, + "message", + channel="wechat", + scope_id="wx_staff:friend_a", + ) + == T2 + ) + finally: + await engine.dispose() + + +async def test_messages_media_download_failure_does_not_abort_sync(monkeypatch): + """If a media download raises, the message is still upserted and the + watermark still advances — best-effort coupling by design.""" + engine = await _make_engine() + try: + SessionLocal = async_sessionmaker(engine, expire_on_commit=False) + + async def fake_scrape(client, wechat_id, friend_id, start, end, *, types=""): + yield _msg("m1", T1) + yield _msg("m2", T2) + + upserted: list[dict] = [] + + async def fake_upsert(session, recs): + upserted.extend(recs) + return len(recs) + + async def failing_media(client, msg, store, session): + raise RuntimeError("oss exploded") + + monkeypatch.setattr(orchestrator, "scrape_wechat_messages", fake_scrape) + monkeypatch.setattr(orchestrator, "upsert_messages", fake_upsert) + monkeypatch.setattr(orchestrator, "download_message_media", failing_media) + + async with SessionLocal() as session: + result = await orchestrator.sync_wechat_messages( + session, client=object(), store=object(), wechat_id="wx_staff", + friend_id="friend_a", + ) + await session.commit() + + # Media failure did NOT abort: both messages upserted, success recorded. + assert result.status == "success" + assert result.upserted == 2 + assert len(upserted) == 2 + async with SessionLocal() as session: + assert ( + await get_last_timestamp( + session, + "message", + channel="wechat", + scope_id="wx_staff:friend_a", + ) + == T2 + ) + finally: + await engine.dispose() + + +async def test_messages_records_sync_run_success(monkeypatch): + """A successful message run records a SyncRun row with the right shape.""" + engine = await _make_engine() + try: + SessionLocal = async_sessionmaker(engine, expire_on_commit=False) + + async def fake_scrape(client, wechat_id, friend_id, start, end, *, types=""): + yield _msg("m1", T1) + + async def fake_upsert(session, recs): + return len(recs) + + async def fake_media(client, msg, store, session): + return None + + monkeypatch.setattr(orchestrator, "scrape_wechat_messages", fake_scrape) + monkeypatch.setattr(orchestrator, "upsert_messages", fake_upsert) + monkeypatch.setattr(orchestrator, "download_message_media", fake_media) + + async with SessionLocal() as session: + await orchestrator.sync_wechat_messages( + session, client=object(), store=object(), wechat_id="wx_staff", + friend_id="friend_a", + ) + await session.commit() + + async with SessionLocal() as session: + runs = ( + await session.execute( + sa.select(SyncRun).where(SyncRun.entity_type == "message") + ) + ).scalars().all() + assert len(runs) == 1 + run = runs[0] + assert run.status == "success" + assert run.channel == "wechat" + assert run.scope_id == "wx_staff:friend_a" + assert run.records_fetched == 1 + assert run.records_upserted == 1 + assert run.error is None + assert run.started_at is not None + assert run.finished_at is not None + finally: + await engine.dispose() + + +# =========================================================================== +# Full re-scrape entity (friends) — idempotent upsert, no since-param +# =========================================================================== + + +async def test_friends_full_rescrape_upserts_all_records_sync_run_and_watermark( + monkeypatch, +): + """Friends list has no `since` param: 'incremental' = re-scrape all and + rely on upsert idempotency. The orchestrator upserts every record, records + a success SyncRun, and stamps the watermark's last_synced_at.""" + engine = await _make_engine() + try: + SessionLocal = async_sessionmaker(engine, expire_on_commit=False) + + friend_recs = [ + {"friend_id": "f1", "channel": "wechat", "version": "v1"}, + {"friend_id": "f2", "channel": "wechat", "version": "v1"}, + {"friend_id": "f3", "channel": "wechat", "version": "v1"}, + ] + + async def fake_scrape_friends(client): + for r in friend_recs: + yield r + + upserted: list[dict] = [] + + async def fake_upsert_customer(session, rec): + upserted.append(rec) + + monkeypatch.setattr(orchestrator, "scrape_wechat_friends", fake_scrape_friends) + monkeypatch.setattr(orchestrator, "upsert_customer", fake_upsert_customer) + + async with SessionLocal() as session: + result = await orchestrator.sync_wechat_friends(session, client=object()) + await session.commit() + + assert result.status == "success" + assert result.fetched == 3 + assert result.upserted == 3 + assert [r["friend_id"] for r in upserted] == ["f1", "f2", "f3"] + + # SyncRun recorded with success + the right counts. + async with SessionLocal() as session: + run = ( + await session.execute( + sa.select(SyncRun).where(SyncRun.entity_type == "wechat_friend") + ) + ).scalar_one() + assert run.status == "success" + assert run.records_fetched == 3 + assert run.records_upserted == 3 + assert run.error is None + + # Watermark row exists with last_synced_at set (no version/timestamp + # cursor advance — these entities have no since-param). + state = ( + await session.execute( + sa.select(SyncState).where(SyncState.entity_type == "wechat_friend") + ) + ).scalar_one() + assert state.last_synced_at is not None + finally: + await engine.dispose() + + +# =========================================================================== +# Failure policy — scraper raises +# =========================================================================== + + +async def test_failing_scraper_records_failed_sync_run_and_does_not_raise( + monkeypatch, +): + """A scraper that raises must NOT crash the batch. The orchestrator records + a SyncRun with status='failed' + the error string and returns a failed + result.""" + engine = await _make_engine() + try: + SessionLocal = async_sessionmaker(engine, expire_on_commit=False) + + async def fake_scrape_friends(client): + raise RuntimeError("upstream 500") + yield {} # pragma: no cover - make this an async generator + + monkeypatch.setattr(orchestrator, "scrape_wechat_friends", fake_scrape_friends) + + async with SessionLocal() as session: + result = await orchestrator.sync_wechat_friends(session, client=object()) + await session.commit() + + # Did not raise; returned a failed result. + assert result.status == "failed" + assert result.fetched == 0 + assert result.upserted == 0 + assert "upstream 500" in (result.error or "") + + async with SessionLocal() as session: + run = ( + await session.execute( + sa.select(SyncRun).where(SyncRun.entity_type == "wechat_friend") + ) + ).scalar_one() + assert run.status == "failed" + assert run.records_fetched == 0 + assert run.records_upserted == 0 + assert "upstream 500" in (run.error or "") + assert run.finished_at is not None + finally: + await engine.dispose() + + +async def test_failing_message_scraper_records_failed_run(monkeypatch): + """Same failure policy for the message sync path.""" + engine = await _make_engine() + try: + SessionLocal = async_sessionmaker(engine, expire_on_commit=False) + + async def fake_scrape(client, wechat_id, friend_id, start, end, *, types=""): + raise RuntimeError("msg endpoint 500") + yield {} # pragma: no cover - make this an async generator + + async def fake_upsert(session, recs): + return 0 + + async def fake_media(client, msg, store, session): + return None + + monkeypatch.setattr(orchestrator, "scrape_wechat_messages", fake_scrape) + monkeypatch.setattr(orchestrator, "upsert_messages", fake_upsert) + monkeypatch.setattr(orchestrator, "download_message_media", fake_media) + + async with SessionLocal() as session: + result = await orchestrator.sync_wechat_messages( + session, client=object(), store=object(), wechat_id="wx_staff", + friend_id="friend_a", + ) + await session.commit() + + assert result.status == "failed" + assert "msg endpoint 500" in (result.error or "") + + async with SessionLocal() as session: + run = ( + await session.execute( + sa.select(SyncRun).where(SyncRun.entity_type == "message") + ) + ).scalar_one() + assert run.status == "failed" + assert "msg endpoint 500" in (run.error or "") + finally: + await engine.dispose() + + +# =========================================================================== +# sync_all_wechat — coordinator +# =========================================================================== + + +async def test_sync_all_wechat_runs_each_entity_then_messages(monkeypatch): + """The coordinator runs departments → accounts → friends → groups → + voice_calls, then per-conversation messages for an explicit list. + Returns one result per entity/conversation in order.""" + engine = await _make_engine() + try: + SessionLocal = async_sessionmaker(engine, expire_on_commit=False) + + async def empty_async_gen(_client=None): + if False: # pragma: no cover + yield {} + + async def empty_messages(client, wechat_id, friend_id, start, end, *, types=""): + if False: # pragma: no cover + yield {} + + monkeypatch.setattr(orchestrator, "scrape_departments", empty_async_gen) + monkeypatch.setattr(orchestrator, "scrape_wechat_accounts", empty_async_gen) + monkeypatch.setattr(orchestrator, "scrape_wechat_friends", empty_async_gen) + monkeypatch.setattr(orchestrator, "scrape_wechat_groups", empty_async_gen) + monkeypatch.setattr(orchestrator, "scrape_wechat_voice_calls", empty_async_gen) + monkeypatch.setattr(orchestrator, "scrape_wechat_messages", empty_messages) + # Scrapers yield nothing -> upsert / media loops never execute, so the + # real upsert_messages / download_message_media (also imported into the + # orchestrator namespace) are never called. No need to patch them. + + async with SessionLocal() as session: + results = await orchestrator.sync_all_wechat( + session, + client=object(), + store=object(), + conversations=[("wx_staff", "friend_a"), ("wx_staff", "friend_b")], + ) + await session.commit() + + entities = [r.entity for r in results] + # Five full-re-scrape entities in order, then two message runs. + assert entities == [ + "department", + "wechat_account", + "wechat_friend", + "wechat_group", + "wechat_voice_call", + "message", + "message", + ] + assert all(r.status == "success" for r in results) + finally: + await engine.dispose() diff --git a/yunke-server/tests/test_watermarks.py b/yunke-server/tests/test_watermarks.py new file mode 100644 index 0000000..91cdad5 --- /dev/null +++ b/yunke-server/tests/test_watermarks.py @@ -0,0 +1,244 @@ +"""Tests for app.sync.watermarks — SyncState CRUD (watermark get/set). + +Exercises the watermark helpers against an isolated in-memory SQLite engine +(same pattern as ``tests/test_repos.py``). Covers the three guarantees the +orchestrator relies on: + +* ``set_watermark`` then ``get_watermark`` returns the just-written row. +* ``set_watermark`` is an **upsert** keyed on ``(entity_type, channel, scope_id)``: + the same key twice produces exactly ONE row, with the cursor fields refreshed + from the later write. +* The convenience accessors (``get_last_timestamp`` / ``get_last_version``) + return the stored value, or the documented default (``0`` / ``None``) when no + row exists yet (first sync of a brand-new scope). +""" + +from __future__ import annotations + +import sqlalchemy as sa +from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine + +from app.db.base import Base +from app.db.models import SyncState +from app.sync.watermarks import ( + get_last_timestamp, + get_last_version, + get_watermark, + set_watermark, +) + + +async def _make_engine(): + """Create an in-memory async SQLite engine with all tables created.""" + engine = create_async_engine("sqlite+aiosqlite:///:memory:") + async with engine.begin() as conn: + await conn.run_sync(Base.metadata.create_all) + return engine + + +# --------------------------------------------------------------------------- +# set then get +# --------------------------------------------------------------------------- + + +async def test_set_then_get_returns_row_with_values(): + """A fresh set creates a row; get returns it carrying the cursors.""" + engine = await _make_engine() + try: + SessionLocal = async_sessionmaker(engine, expire_on_commit=False) + async with SessionLocal() as session: + await set_watermark( + session, + "message", + channel="wechat", + scope_id="wx_a:friend_b", + last_version="v1", + last_timestamp=1_700_000_000_000, + ) + await session.commit() + + async with SessionLocal() as session: + row = await get_watermark( + session, "message", channel="wechat", scope_id="wx_a:friend_b" + ) + assert row is not None + assert row.entity_type == "message" + assert row.channel == "wechat" + assert row.scope_id == "wx_a:friend_b" + assert row.last_version == "v1" + assert row.last_timestamp == 1_700_000_000_000 + assert row.last_synced_at is not None + finally: + await engine.dispose() + + +# --------------------------------------------------------------------------- +# upsert: same key twice -> 1 row, updated +# --------------------------------------------------------------------------- + + +async def test_set_watermark_is_upsert_one_row_updated(): + """Two sets on the same (entity_type, channel, scope_id) -> exactly one row. + + The later write refreshes ``last_timestamp`` / ``last_version`` / + ``last_synced_at``. This is the resumability guarantee the orchestrator's + incremental loop leans on. + """ + engine = await _make_engine() + try: + SessionLocal = async_sessionmaker(engine, expire_on_commit=False) + async with SessionLocal() as session: + await set_watermark( + session, + "message", + channel="wechat", + scope_id="s1", + last_timestamp=100, + last_version="v1", + ) + await session.commit() + async with SessionLocal() as session: + await set_watermark( + session, + "message", + channel="wechat", + scope_id="s1", + last_timestamp=200, + last_version="v2", + ) + await session.commit() + + async with SessionLocal() as session: + count = ( + await session.execute( + sa.select(sa.func.count()) + .select_from(SyncState) + .where(SyncState.scope_id == "s1") + ) + ).scalar_one() + assert count == 1 + + row = await get_watermark( + session, "message", channel="wechat", scope_id="s1" + ) + assert row is not None + assert row.last_timestamp == 200 + assert row.last_version == "v2" + finally: + await engine.dispose() + + +async def test_set_watermark_partial_advance_does_not_clobber_other_cursor(): + """A timestamp-only advance leaves the previously-set version intact. + + The orchestrator advances one cursor at a time (e.g. the message loop + advances ``last_timestamp``); a partial set must NOT wipe ``last_version`` + with None. + """ + engine = await _make_engine() + try: + SessionLocal = async_sessionmaker(engine, expire_on_commit=False) + async with SessionLocal() as session: + await set_watermark( + session, + "wechat_friend", + channel="wechat", + scope_id=None, + last_version="v1", + ) + await session.commit() + async with SessionLocal() as session: + # Advance only the sync time (no new cursors supplied). + await set_watermark( + session, "wechat_friend", channel="wechat", scope_id=None + ) + await session.commit() + + async with SessionLocal() as session: + row = await get_watermark( + session, "wechat_friend", channel="wechat", scope_id=None + ) + assert row is not None + assert row.last_version == "v1" # preserved, not clobbered to None + finally: + await engine.dispose() + + +# --------------------------------------------------------------------------- +# convenience accessors: defaults when row absent +# --------------------------------------------------------------------------- + + +async def test_get_last_timestamp_returns_default_when_absent(): + """No row yet -> default (0) returned, not None.""" + engine = await _make_engine() + try: + SessionLocal = async_sessionmaker(engine, expire_on_commit=False) + async with SessionLocal() as session: + assert ( + await get_last_timestamp( + session, "message", channel="wechat", scope_id="never" + ) + == 0 + ) + # Custom default is honored. + assert ( + await get_last_timestamp( + session, + "message", + channel="wechat", + scope_id="never", + default=42, + ) + == 42 + ) + finally: + await engine.dispose() + + +async def test_get_last_version_returns_none_when_absent(): + """No row yet -> None (no version cursor yet).""" + engine = await _make_engine() + try: + SessionLocal = async_sessionmaker(engine, expire_on_commit=False) + async with SessionLocal() as session: + assert ( + await get_last_version( + session, "message", channel="wechat", scope_id="never" + ) + is None + ) + finally: + await engine.dispose() + + +async def test_get_last_timestamp_returns_stored_value_after_set(): + """After a set, the accessor returns the stored timestamp (not default).""" + engine = await _make_engine() + try: + SessionLocal = async_sessionmaker(engine, expire_on_commit=False) + async with SessionLocal() as session: + await set_watermark( + session, + "message", + channel="wechat", + scope_id="s2", + last_timestamp=999, + ) + await session.commit() + + async with SessionLocal() as session: + assert ( + await get_last_timestamp( + session, "message", channel="wechat", scope_id="s2" + ) + == 999 + ) + assert ( + await get_last_version( + session, "message", channel="wechat", scope_id="s2" + ) + is None + ) + finally: + await engine.dispose()