fire-planner/fire_planner/ingest/wealthfolio.py

"""Upsert helper for wealthfolio account snapshots.

The actual read happens in `wealthfolio_pg.py` (against the
`wealthfolio_sync` PG mirror). This module keeps the upsert helper that
both prod and tests use, so callers can:

    rows = await read_account_snapshots_from_pg(wf_session)
    await upsert_snapshots(session, rows)
"""
from __future__ import annotations

from typing import Any

from sqlalchemy.dialects.postgresql import insert as pg_insert
from sqlalchemy.dialects.sqlite import insert as sqlite_insert
from sqlalchemy.ext.asyncio import AsyncSession

from fire_planner.db import AccountSnapshot


def _dialect_insert(session: AsyncSession) -> Any:
    bind = session.get_bind()
    if bind.dialect.name == "sqlite":
        return sqlite_insert
    return pg_insert


async def upsert_snapshots(session: AsyncSession, rows: list[dict[str, Any]]) -> int:
    """Idempotent upsert on `external_id` (one row per account per day)."""
    if not rows:
        return 0
    insert_ = _dialect_insert(session)
    stmt = insert_(AccountSnapshot).values(rows)
    update_cols = {
        "market_value": stmt.excluded.market_value,
        "market_value_gbp": stmt.excluded.market_value_gbp,
        "snapshot_date": stmt.excluded.snapshot_date,
        "account_name": stmt.excluded.account_name,
        "account_type": stmt.excluded.account_type,
        "currency": stmt.excluded.currency,
        "cost_basis_gbp": stmt.excluded.cost_basis_gbp,
    }
    stmt = stmt.on_conflict_do_update(index_elements=["external_id"], set_=update_cols)
    await session.execute(stmt)
    return len(rows)