yunke/yunke-server/app/sync/watermarks.py

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