"""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 sqlalchemy.orm import selectinload 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.crypto import CryptoHolding, CryptoTransaction from app.models.goal import Goal, GoalContribution from app.models.recurring import RecurringTransaction from app.models.settings import AppSettings from app.models.tag import Tag from app.models.transaction import Transaction from app.schemas.backup import ( AccountBackup, AppSettingsBackup, AssetBackup, AssetValuationBackup, BackupPayload, BudgetBackup, CategoryBackup, CryptoHoldingBackup, CryptoTransactionBackup, GoalBackup, GoalContributionBackup, RecurringTransactionBackup, TagBackup, 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() tags = (await session.execute(select(Tag))).scalars().all() transactions = (await session.execute(select(Transaction).options(selectinload(Transaction.tags)))).scalars().all() assets = (await session.execute(select(Asset))).scalars().all() valuations = (await session.execute(select(AssetValuation))).scalars().all() crypto_holdings = (await session.execute(select(CryptoHolding))).scalars().all() crypto_transactions = (await session.execute(select(CryptoTransaction))).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], tags=[TagBackup.model_validate(row) for row in tags], transactions=[ TransactionBackup.model_validate(row, from_attributes=True).model_copy( update={"tag_ids": [tag.id for tag in row.tags]} ) for row in transactions ], assets=[AssetBackup.model_validate(row) for row in assets], asset_valuations=[AssetValuationBackup.model_validate(row) for row in valuations], crypto_holdings=[CryptoHoldingBackup.model_validate(row) for row in crypto_holdings], crypto_transactions=[CryptoTransactionBackup.model_validate(row) for row in crypto_transactions], 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} tag_ids = {row.id for row in payload.tags} 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 tag_id in t.tag_ids: if tag_id not in tag_ids: raise HTTPException(400, f"Transaction {t.id} references unknown tag_id {tag_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}") crypto_holding_asset_ids = {h.asset_id for h in payload.crypto_holdings} for h in payload.crypto_holdings: if h.asset_id not in asset_ids: raise HTTPException(400, f"Crypto holding {h.asset_id} references unknown asset_id {h.asset_id}") for tx in payload.crypto_transactions: if tx.asset_id not in crypto_holding_asset_ids: raise HTTPException(400, f"Crypto transaction {tx.id} references unknown asset_id {tx.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(CryptoTransaction)) await session.execute(delete(CryptoHolding)) await session.execute(delete(Budget)) await session.execute(delete(GoalContribution)) await session.execute(delete(Goal)) await session.execute(delete(RecurringTransaction)) # Deleting transactions cascades transaction_tags rows (ON DELETE # CASCADE) — tags themselves still need their own delete. await session.execute(delete(Transaction)) await session.execute(delete(Tag)) 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) # Tags and transactions are kept in id-keyed dicts (rather than a # plain add_all) — transaction.tags is a relationship, not a column # in model_dump(), so it has to be wired up from live ORM objects # once everything is flushed and has real identities. tags_by_id = {row.id: Tag(**row.model_dump()) for row in payload.tags} session.add_all(tags_by_id.values()) # tags=[] at construction keeps the collection "loaded" on the # transient object — reassigning it after flush (below) would # otherwise trigger an implicit lazy-load, which async SQLAlchemy # can't do outside an explicit await (MissingGreenlet). transactions_by_id = { row.id: Transaction(**row.model_dump(exclude={"tag_ids"}), tags=[]) for row in payload.transactions } session.add_all(transactions_by_id.values()) session.add_all(AssetValuation(**row.model_dump()) for row in payload.asset_valuations) session.add_all(CryptoHolding(**row.model_dump()) for row in payload.crypto_holdings) session.add_all(CryptoTransaction(**row.model_dump()) for row in payload.crypto_transactions) 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() for row in payload.transactions: if row.tag_ids: transactions_by_id[row.id].tags = [tags_by_id[tag_id] for tag_id in row.tag_ids] if any(row.tag_ids for row in payload.transactions): await session.flush() await _reset_sequence(session, "accounts", payload.accounts) await _reset_sequence(session, "categories", payload.categories) await _reset_sequence(session, "tags", payload.tags) 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, "crypto_transactions", payload.crypto_transactions) 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