aurum/backend/app/services/backup_service.py

"""Full-database JSON backup & restore.

Exports every row (accounts, categories, transactions, assets, asset
valuations) as one portable JSON document a user can download from the
browser and re-upload later. Restore fully REPLACES existing data — it's a
snapshot restore, not a merge — so the whole operation runs in one DB
transaction: a corrupt or incompatible file is rejected (referential checks
run first, before any row is touched), and any failure during the swap rolls
the database back to exactly where it was, so a bad file never leaves the
app half-restored.
"""
from datetime import datetime, timezone

from fastapi import HTTPException
from sqlalchemy import delete, select, text
from sqlalchemy.ext.asyncio import AsyncSession

from app.core.config import APP_VERSION
from app.models.account import Account
from app.models.asset import Asset, AssetValuation
from app.models.budget import Budget
from app.models.category import Category
from app.models.goal import Goal, GoalContribution
from app.models.recurring import RecurringTransaction
from app.models.settings import AppSettings
from app.models.transaction import Transaction
from app.schemas.backup import (
    AccountBackup,
    AppSettingsBackup,
    AssetBackup,
    AssetValuationBackup,
    BackupPayload,
    BudgetBackup,
    CategoryBackup,
    GoalBackup,
    GoalContributionBackup,
    RecurringTransactionBackup,
    TransactionBackup,
)

BACKUP_FORMAT_VERSION = 1


async def build_backup(session: AsyncSession) -> BackupPayload:
    accounts = (await session.execute(select(Account))).scalars().all()
    categories = (await session.execute(select(Category))).scalars().all()
    transactions = (await session.execute(select(Transaction))).scalars().all()
    assets = (await session.execute(select(Asset))).scalars().all()
    valuations = (await session.execute(select(AssetValuation))).scalars().all()
    budgets = (await session.execute(select(Budget))).scalars().all()
    goals = (await session.execute(select(Goal))).scalars().all()
    goal_contributions = (await session.execute(select(GoalContribution))).scalars().all()
    recurring_transactions = (await session.execute(select(RecurringTransaction))).scalars().all()
    app_settings = await session.get(AppSettings, 1)

    return BackupPayload(
        aurum_backup_version=BACKUP_FORMAT_VERSION,
        exported_at=datetime.now(timezone.utc),
        app_version=APP_VERSION,
        accounts=[AccountBackup.model_validate(row) for row in accounts],
        categories=[CategoryBackup.model_validate(row) for row in categories],
        transactions=[TransactionBackup.model_validate(row) for row in transactions],
        assets=[AssetBackup.model_validate(row) for row in assets],
        asset_valuations=[AssetValuationBackup.model_validate(row) for row in valuations],
        budgets=[BudgetBackup.model_validate(row) for row in budgets],
        goals=[GoalBackup.model_validate(row) for row in goals],
        goal_contributions=[GoalContributionBackup.model_validate(row) for row in goal_contributions],
        recurring_transactions=[RecurringTransactionBackup.model_validate(row) for row in recurring_transactions],
        app_settings=AppSettingsBackup.model_validate(app_settings) if app_settings else AppSettingsBackup(currency="USD"),
    )


def _validate_references(payload: BackupPayload) -> None:
    account_ids = {row.id for row in payload.accounts}
    category_ids = {row.id for row in payload.categories}
    asset_ids = {row.id for row in payload.assets}

    for c in payload.categories:
        if c.parent_id is not None and c.parent_id not in category_ids:
            raise HTTPException(400, f"Category {c.id} references unknown parent_id {c.parent_id}")

    for t in payload.transactions:
        if t.account_id not in account_ids:
            raise HTTPException(400, f"Transaction {t.id} references unknown account_id {t.account_id}")
        if t.transfer_account_id is not None and t.transfer_account_id not in account_ids:
            raise HTTPException(
                400, f"Transaction {t.id} references unknown transfer_account_id {t.transfer_account_id}"
            )
        if t.category_id is not None and t.category_id not in category_ids:
            raise HTTPException(400, f"Transaction {t.id} references unknown category_id {t.category_id}")

    for v in payload.asset_valuations:
        if v.asset_id not in asset_ids:
            raise HTTPException(400, f"Asset valuation {v.id} references unknown asset_id {v.asset_id}")

    for b in payload.budgets:
        if b.category_id not in category_ids:
            raise HTTPException(400, f"Budget {b.id} references unknown category_id {b.category_id}")

    goal_ids = {row.id for row in payload.goals}
    for c in payload.goal_contributions:
        if c.goal_id not in goal_ids:
            raise HTTPException(400, f"Goal contribution {c.id} references unknown goal_id {c.goal_id}")

    for r in payload.recurring_transactions:
        if r.account_id not in account_ids:
            raise HTTPException(400, f"Recurring transaction {r.id} references unknown account_id {r.account_id}")
        if r.transfer_account_id is not None and r.transfer_account_id not in account_ids:
            raise HTTPException(
                400,
                f"Recurring transaction {r.id} references unknown transfer_account_id {r.transfer_account_id}",
            )
        if r.category_id is not None and r.category_id not in category_ids:
            raise HTTPException(400, f"Recurring transaction {r.id} references unknown category_id {r.category_id}")


async def _reset_sequence(session: AsyncSession, table: str, rows: list) -> None:
    """Bulk-inserting rows with explicit ids doesn't advance the table's
    identity sequence, so the next auto-generated id would collide — bump it
    to max(id) after a restore. `table` is always one of our five hardcoded
    table names, never user input."""
    if not rows:
        return
    max_id = max(row.id for row in rows)
    await session.execute(
        text(f"SELECT setval(pg_get_serial_sequence('{table}', 'id'), :max_id)"), {"max_id": max_id}
    )


async def restore_backup(session: AsyncSession, payload: BackupPayload) -> None:
    if payload.aurum_backup_version != BACKUP_FORMAT_VERSION:
        raise HTTPException(
            400,
            f"Unsupported backup version {payload.aurum_backup_version} "
            f"(this Aurum version supports {BACKUP_FORMAT_VERSION})",
        )

    _validate_references(payload)

    try:
        # Children before parents.
        await session.execute(delete(AssetValuation))
        await session.execute(delete(Budget))
        await session.execute(delete(GoalContribution))
        await session.execute(delete(Goal))
        await session.execute(delete(RecurringTransaction))
        await session.execute(delete(Transaction))
        await session.execute(delete(Asset))
        await session.execute(delete(Category))
        await session.execute(delete(Account))

        # Parents before children. Categories are additionally self-referential
        # (parent_id points at another row in the same table) — sort
        # top-level categories first so a subcategory's FK is never inserted
        # ahead of the row it points to.
        categories_in_order = sorted(payload.categories, key=lambda row: row.parent_id is not None)

        session.add_all(Account(**row.model_dump()) for row in payload.accounts)
        session.add_all(Category(**row.model_dump()) for row in categories_in_order)
        session.add_all(Asset(**row.model_dump()) for row in payload.assets)
        session.add_all(Transaction(**row.model_dump()) for row in payload.transactions)
        session.add_all(AssetValuation(**row.model_dump()) for row in payload.asset_valuations)
        session.add_all(Budget(**row.model_dump()) for row in payload.budgets)
        session.add_all(Goal(**row.model_dump()) for row in payload.goals)
        session.add_all(GoalContribution(**row.model_dump()) for row in payload.goal_contributions)
        session.add_all(RecurringTransaction(**row.model_dump()) for row in payload.recurring_transactions)
        await session.flush()

        await _reset_sequence(session, "accounts", payload.accounts)
        await _reset_sequence(session, "categories", payload.categories)
        await _reset_sequence(session, "assets", payload.assets)
        await _reset_sequence(session, "transactions", payload.transactions)
        await _reset_sequence(session, "asset_valuations", payload.asset_valuations)
        await _reset_sequence(session, "budgets", payload.budgets)
        await _reset_sequence(session, "goals", payload.goals)
        await _reset_sequence(session, "goal_contributions", payload.goal_contributions)
        await _reset_sequence(session, "recurring_transactions", payload.recurring_transactions)

        # Singleton row — updated in place, not deleted/recreated (no
        # sequence to reset, id is always 1).
        app_settings = await session.get(AppSettings, 1)
        if app_settings is None:
            session.add(AppSettings(id=1, **payload.app_settings.model_dump()))
        else:
            for field, value in payload.app_settings.model_dump().items():
                setattr(app_settings, field, value)

        await session.commit()
    except HTTPException:
        await session.rollback()
        raise
    except Exception as exc:
        await session.rollback()
        raise HTTPException(400, f"Restore failed, no changes were made: {exc}") from exc