simbarag/blueprints/ynab/sync.py

"""Sync the YNAB budget into the local Postgres mirror.

YNAB has no event feed, but every budget response carries a
``server_knowledge`` counter; passing it back as ``last_knowledge_of_server``
returns only entities changed since then. The first sync downloads the whole
budget, later syncs are deltas, and a full resync runs once a day as a safety
net (deltas can't express things like month rollovers).

Each sync diffs incoming entities against the mirror and appends what changed
to ``ynab_changes``, which is the only history of edits we get.
"""

import asyncio
import logging
import os
from dataclasses import dataclass
from datetime import date, datetime, timedelta, timezone
from decimal import Decimal
from typing import Any, Callable

import httpx
from tortoise.models import Model
from tortoise.transactions import in_transaction

from .models import (
    YnabAccount,
    YnabCategory,
    YnabCategoryGroup,
    YnabChange,
    YnabMonth,
    YnabMonthCategory,
    YnabPayee,
    YnabSubtransaction,
    YnabSyncState,
    YnabTransaction,
)

logger = logging.getLogger(__name__)

YNAB_API_URL = "https://api.ynab.com/v1"
FULL_RESYNC_INTERVAL = timedelta(hours=24)
# Cap on rows passed to a single ``id IN (...)`` lookup
LOOKUP_CHUNK = 2000

_sync_lock = asyncio.Lock()


def ynab_configured() -> bool:
    return bool(os.getenv("YNAB_ACCESS_TOKEN"))


def _budget_id() -> str:
    return os.getenv("YNAB_BUDGET_ID") or "last-used"


def _money(milliunits: int | None) -> Decimal | None:
    return None if milliunits is None else Decimal(milliunits).scaleb(-3)


def _date(value: str | None) -> date | None:
    return date.fromisoformat(value) if value else None


def _datetime(value: str | None) -> datetime | None:
    if not value:
        return None
    return datetime.fromisoformat(value.replace("Z", "+00:00"))


def _norm(value: Any) -> Any:
    """Normalize values so DB rows and freshly mapped rows compare equal."""
    if isinstance(value, datetime) and value.tzinfo is None:
        return value.replace(tzinfo=timezone.utc)
    return value


def _jsonable(value: Any) -> Any:
    if isinstance(value, Decimal):
        return float(value)
    if isinstance(value, (date, datetime)):
        return value.isoformat()
    return value


def _map_account(a: dict) -> dict:
    return {
        "id": a["id"],
        "name": a["name"],
        "type": a["type"],
        "on_budget": a["on_budget"],
        "closed": a["closed"],
        "note": a.get("note"),
        "balance": _money(a["balance"]),
        "cleared_balance": _money(a["cleared_balance"]),
        "uncleared_balance": _money(a["uncleared_balance"]),
        "last_reconciled_at": _datetime(a.get("last_reconciled_at")),
        "deleted": a["deleted"],
    }


def _map_payee(p: dict) -> dict:
    return {
        "id": p["id"],
        "name": p["name"],
        "transfer_account_id": p.get("transfer_account_id"),
        "deleted": p["deleted"],
    }


def _map_category_group(g: dict) -> dict:
    return {
        "id": g["id"],
        "name": g["name"],
        "hidden": g["hidden"],
        "internal": g.get("internal") or False,
        "deleted": g["deleted"],
    }


def _map_category(c: dict) -> dict:
    return {
        "id": c["id"],
        "category_group_id": c["category_group_id"],
        "name": c["name"],
        "hidden": c["hidden"],
        "internal": c.get("internal") or False,
        "note": c.get("note"),
        "goal_type": c.get("goal_type"),
        "goal_target": _money(c.get("goal_target")),
        "goal_target_month": _date(c.get("goal_target_month")),
        "deleted": c["deleted"],
    }


def _map_month(m: dict) -> dict:
    return {
        "id": m["month"],
        "month": _date(m["month"]),
        "note": m.get("note"),
        "income": _money(m["income"]),
        "budgeted": _money(m["budgeted"]),
        "activity": _money(m["activity"]),
        "to_be_budgeted": _money(m["to_be_budgeted"]),
        "age_of_money": m.get("age_of_money"),
        "deleted": m["deleted"],
    }


def _map_month_category(c: dict) -> dict:
    return {
        "id": f"{c['month']}:{c['id']}",
        "month": _date(c["month"]),
        "category_id": c["id"],
        "budgeted": _money(c["budgeted"]),
        "activity": _money(c["activity"]),
        "balance": _money(c["balance"]),
        "deleted": c["deleted"],
    }


def _map_transaction(t: dict) -> dict:
    return {
        "id": t["id"],
        "date": _date(t["date"]),
        "amount": _money(t["amount"]),
        "memo": t.get("memo"),
        "cleared": t["cleared"],
        "approved": t["approved"],
        "flag_color": t.get("flag_color"),
        "account_id": t["account_id"],
        "payee_id": t.get("payee_id"),
        "category_id": t.get("category_id"),
        "transfer_account_id": t.get("transfer_account_id"),
        "import_payee_name": t.get("import_payee_name"),
        "deleted": t["deleted"],
    }


def _map_subtransaction(s: dict) -> dict:
    return {
        "id": s["id"],
        "transaction_id": s["transaction_id"],
        "amount": _money(s["amount"]),
        "memo": s.get("memo"),
        "payee_id": s.get("payee_id"),
        "category_id": s.get("category_id"),
        "transfer_account_id": s.get("transfer_account_id"),
        "deleted": s["deleted"],
    }


@dataclass
class EntitySpec:
    entity_type: str
    model: type[Model]
    mapper: Callable[[dict], dict]
    # Fields worth recording in the change log. Derived fields (balances,
    # activity) change with every transaction and would drown out real edits.
    logged_fields: tuple[str, ...]


def _fields(model: type[Model], exclude: tuple[str, ...] = ()) -> tuple[str, ...]:
    return tuple(f for f in model._meta.fields_map if f != "id" and f not in exclude)


SPECS = {
    "accounts": EntitySpec(
        "account",
        YnabAccount,
        _map_account,
        ("name", "type", "on_budget", "closed", "note", "deleted"),
    ),
    "payees": EntitySpec("payee", YnabPayee, _map_payee, ("name", "deleted")),
    "category_groups": EntitySpec(
        "category_group",
        YnabCategoryGroup,
        _map_category_group,
        ("name", "hidden", "deleted"),
    ),
    "categories": EntitySpec(
        "category", YnabCategory, _map_category, _fields(YnabCategory)
    ),
    "months": EntitySpec("month", YnabMonth, _map_month, ("note",)),
    "month_categories": EntitySpec(
        "month_category", YnabMonthCategory, _map_month_category, ("budgeted",)
    ),
    "transactions": EntitySpec(
        "transaction",
        YnabTransaction,
        _map_transaction,
        _fields(
            YnabTransaction, exclude=("account_name", "payee_name", "category_name")
        ),
    ),
    "subtransactions": EntitySpec(
        "subtransaction",
        YnabSubtransaction,
        _map_subtransaction,
        _fields(YnabSubtransaction),
    ),
}


def _budget_entities(budget: dict) -> dict[str, list[dict]]:
    """Split a budget payload into raw entity lists keyed like ``SPECS``."""
    month_categories = [
        {**c, "month": m["month"]}
        for m in budget.get("months") or []
        for c in m.get("categories") or []
    ]
    return {
        "accounts": budget.get("accounts") or [],
        "payees": budget.get("payees") or [],
        "category_groups": budget.get("category_groups") or [],
        "categories": budget.get("categories") or [],
        "months": budget.get("months") or [],
        "month_categories": month_categories,
        "transactions": budget.get("transactions") or [],
        "subtransactions": budget.get("subtransactions") or [],
    }


class _Labeler:
    """Builds human-readable change log labels from ids."""

    def __init__(self, names: dict[str, str]):
        self.names = names

    def __call__(self, entity_type: str, row: dict) -> str:
        n = self.names.get
        if entity_type == "transaction":
            payee = n(row.get("payee_id") or "") or row.get("import_payee_name") or ""
            return f"{row['date']} {payee} {row['amount']:+.2f}".strip()
        if entity_type == "subtransaction":
            category = n(row.get("category_id") or "", "")
            return f"split line {category} {row['amount']:+.2f}".strip()
        if entity_type == "month_category":
            return f"{n(row['category_id'], row['category_id'])} {row['month']:%Y-%m}"
        if entity_type == "month":
            return row["id"][:7]
        return row.get("name") or row["id"]


async def _load_names(conn) -> dict[str, str]:
    names: dict[str, str] = {}
    for model in (YnabAccount, YnabPayee, YnabCategory, YnabCategoryGroup):
        for row in await model.all().using_db(conn).values("id", "name"):
            names[row["id"]] = row["name"]
    return names


async def _apply(
    spec: EntitySpec,
    raw_rows: list[dict],
    conn,
    server_knowledge: int,
    log_changes: bool,
    labeler: _Labeler,
) -> int:
    """Upsert rows that differ from the mirror; return how many changed."""
    if not raw_rows:
        return 0

    incoming = {row["id"]: row for row in map(spec.mapper, raw_rows)}
    # Only YNAB-sourced columns; denormalized names are filled in afterwards
    sample = next(iter(incoming.values()))
    field_names = tuple(f for f in _fields(spec.model) if f in sample)

    ids = list(incoming)
    existing: dict[str, dict] = {}
    for i in range(0, len(ids), LOOKUP_CHUNK):
        rows = (
            await spec.model.filter(id__in=ids[i : i + LOOKUP_CHUNK])
            .using_db(conn)
            .values("id", *field_names)
        )
        existing.update({row["id"]: row for row in rows})

    to_write: list[Model] = []
    changes: list[YnabChange] = []
    for entity_id, new in incoming.items():
        old = existing.get(entity_id)
        if old is not None and all(_norm(old[f]) == _norm(new[f]) for f in field_names):
            continue
        to_write.append(spec.model(**new))

        if not log_changes or not spec.logged_fields:
            continue
        if old is None:
            # Created and deleted between two syncs: nothing worth logging
            if new.get("deleted"):
                continue
            action = "created"
            diff = {
                f: [None, _jsonable(new[f])]
                for f in spec.logged_fields
                if new[f] is not None
            }
        else:
            diff = {
                f: [_jsonable(old[f]), _jsonable(new[f])]
                for f in spec.logged_fields
                if _norm(old[f]) != _norm(new[f])
            }
            if not diff:
                continue
            action = "deleted" if diff.get("deleted") == [False, True] else "updated"
        changes.append(
            YnabChange(
                server_knowledge=server_knowledge,
                entity_type=spec.entity_type,
                entity_id=entity_id,
                action=action,
                label=labeler(spec.entity_type, new)[:500],
                changes=diff,
            )
        )

    if to_write:
        await spec.model.bulk_create(
            to_write,
            batch_size=500,
            on_conflict=["id"],
            update_fields=list(field_names),
            using_db=conn,
        )
    if changes:
        await YnabChange.bulk_create(changes, batch_size=500, using_db=conn)
    return len(to_write)


DENORMALIZE_SQL = (
    """UPDATE ynab_transactions AS t SET payee_name = p.name FROM ynab_payees p
       WHERE t.payee_id = p.id AND t.payee_name IS DISTINCT FROM p.name""",
    """UPDATE ynab_transactions AS t SET category_name = c.name FROM ynab_categories c
       WHERE t.category_id = c.id AND t.category_name IS DISTINCT FROM c.name""",
    """UPDATE ynab_transactions AS t SET account_name = a.name FROM ynab_accounts a
       WHERE t.account_id = a.id AND t.account_name IS DISTINCT FROM a.name""",
    """UPDATE ynab_transactions SET payee_name = NULL
       WHERE payee_id IS NULL AND payee_name IS NOT NULL""",
    """UPDATE ynab_transactions SET category_name = NULL
       WHERE category_id IS NULL AND category_name IS NOT NULL""",
)


async def apply_budget(
    budget: dict, server_knowledge: int, conn, log_changes: bool
) -> dict[str, int]:
    """Write a (full or delta) budget payload into the mirror."""
    entities = _budget_entities(budget)
    counts: dict[str, int] = {}

    # Reference data first so transaction labels can resolve names
    labeler = _Labeler({})
    for key in ("accounts", "payees", "category_groups", "categories"):
        counts[key] = await _apply(
            SPECS[key], entities[key], conn, server_knowledge, log_changes, labeler
        )
    labeler.names = await _load_names(conn)
    for key in ("months", "month_categories", "transactions", "subtransactions"):
        counts[key] = await _apply(
            SPECS[key], entities[key], conn, server_knowledge, log_changes, labeler
        )

    if any(counts.values()):
        for sql in DENORMALIZE_SQL:
            await conn.execute_query(sql)
    return counts


async def fetch_budget(last_knowledge: int | None) -> dict:
    """GET the budget from YNAB, as a delta when ``last_knowledge`` is set."""
    params = {}
    if last_knowledge is not None:
        params["last_knowledge_of_server"] = last_knowledge
    async with httpx.AsyncClient(timeout=httpx.Timeout(60.0, connect=10.0)) as client:
        response = await client.get(
            f"{YNAB_API_URL}/budgets/{_budget_id()}",
            params=params,
            headers={"Authorization": f"Bearer {os.environ['YNAB_ACCESS_TOKEN']}"},
        )
    response.raise_for_status()
    return response.json()["data"]


async def sync_ynab(full: bool = False) -> dict:
    """Sync YNAB into Postgres. Returns per-entity counts of changed rows.

    Safe to call concurrently: an in-process lock plus a Postgres advisory
    lock serialize syncs across both blue/green containers.
    """
    budget_id = _budget_id()
    async with _sync_lock:
        try:
            async with in_transaction() as conn:
                await conn.execute_query(
                    "SELECT pg_advisory_xact_lock(hashtext('ynab_sync'))"
                )
                state, _ = await YnabSyncState.get_or_create(
                    budget_id=budget_id, using_db=conn
                )
                first_sync = state.server_knowledge is None
                now = datetime.now(timezone.utc)
                full = (
                    full
                    or first_sync
                    or (
                        state.last_full_sync_at is None
                        or now - _norm(state.last_full_sync_at) > FULL_RESYNC_INTERVAL
                    )
                )

                data = await fetch_budget(None if full else state.server_knowledge)
                knowledge = data["server_knowledge"]
                counts = await apply_budget(
                    data["budget"], knowledge, conn, log_changes=not first_sync
                )

                state.budget_name = data["budget"].get("name") or state.budget_name
                state.server_knowledge = knowledge
                state.last_synced_at = now
                if full:
                    state.last_full_sync_at = now
                state.last_error = None
                await state.save(using_db=conn)
        except Exception as e:
            await YnabSyncState.filter(budget_id=budget_id).update(
                last_error=f"{type(e).__name__}: {e}"[:2000]
            )
            raise

    changed = {k: v for k, v in counts.items() if v}
    if changed:
        logger.info(f"YNAB {'full' if full else 'delta'} sync @{knowledge}: {changed}")
    return {"full": full, "server_knowledge": knowledge, "changed": changed}


async def ensure_fresh(max_age_seconds: int = 120) -> YnabSyncState | None:
    """Sync if the mirror is older than ``max_age_seconds``.

    Used by agent tools before answering. A failed sync is logged and the
    (possibly stale) mirror is used anyway; returns None only if the mirror
    has never been populated.
    """
    state = await YnabSyncState.get_or_none(budget_id=_budget_id())
    stale = (
        state is None
        or state.last_synced_at is None
        or datetime.now(timezone.utc) - _norm(state.last_synced_at)
        > timedelta(seconds=max_age_seconds)
    )
    if stale:
        try:
            await sync_ynab()
        except Exception:
            logger.exception("YNAB sync failed; answering from existing mirror")
        state = await YnabSyncState.get_or_none(budget_id=_budget_id())
    if state is None or state.server_knowledge is None:
        return None
    return state


async def ynab_sync_loop() -> None:
    """Background task: delta-sync YNAB every ``YNAB_SYNC_INTERVAL`` seconds."""
    interval = int(os.getenv("YNAB_SYNC_INTERVAL", "900"))
    logger.info(f"YNAB sync loop started (interval={interval}s)")
    while True:
        try:
            await sync_ynab()
        except Exception:
            logger.exception("YNAB sync error")
        await asyncio.sleep(interval)