import os from datetime import date, datetime from decimal import Decimal from typing import Any from sqlalchemy import JSON, TIMESTAMP, Boolean, Date, Integer, Numeric, String, func, text from sqlalchemy.dialects.postgresql import JSONB from sqlalchemy.ext.asyncio import AsyncEngine, async_sessionmaker, create_async_engine from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column SCHEMA_NAME = "fire_planner" class Base(DeclarativeBase): pass # JSONB on Postgres, plain JSON on SQLite — tests use SQLite, prod uses Postgres. JSON_TYPE = JSONB().with_variant(JSON(), "sqlite") class AccountSnapshot(Base): """Daily NW per account from Wealthfolio (filled by ingest). `external_id` is `wealthfolio:{account_id}:{date}` so re-runs on the same day are idempotent — Wealthfolio keeps one snapshot per account per day. """ __tablename__ = "account_snapshot" __table_args__ = {"schema": SCHEMA_NAME} # noqa: RUF012 id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True) external_id: Mapped[str] = mapped_column(String, unique=True, nullable=False) snapshot_date: Mapped[date] = mapped_column(Date, nullable=False, index=True) account_id: Mapped[str] = mapped_column(String, nullable=False, index=True) account_name: Mapped[str] = mapped_column(String, nullable=False) account_type: Mapped[str] = mapped_column(String, nullable=False) currency: Mapped[str] = mapped_column(String(3), nullable=False, server_default="GBP") market_value: Mapped[Decimal] = mapped_column(Numeric(14, 2), nullable=False) market_value_gbp: Mapped[Decimal] = mapped_column(Numeric(14, 2), nullable=False) cost_basis_gbp: Mapped[Decimal | None] = mapped_column(Numeric(14, 2), nullable=True) raw_extraction: Mapped[dict[str, Any] | None] = mapped_column(JSON_TYPE, nullable=True) created_at: Mapped[datetime] = mapped_column(TIMESTAMP(timezone=True), nullable=False, server_default=func.now()) class Scenario(Base): """A simulation scenario. Two kinds: - `kind='cartesian'` — auto-generated from `scenarios.py` Cartesian product; rebuilt every recompute, upserted on `external_id`. - `kind='user'` — user-defined (named, optionally cloned from a base); survives recomputes; `parent_scenario_id` points at the source if any. """ __tablename__ = "scenario" __table_args__ = {"schema": SCHEMA_NAME} # noqa: RUF012 id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True) external_id: Mapped[str] = mapped_column(String, unique=True, nullable=False) kind: Mapped[str] = mapped_column(String(16), nullable=False, server_default=text("'cartesian'")) name: Mapped[str | None] = mapped_column(String, nullable=True) description: Mapped[str | None] = mapped_column(String, nullable=True) parent_scenario_id: Mapped[int | None] = mapped_column(Integer, nullable=True) jurisdiction: Mapped[str] = mapped_column(String(32), nullable=False, index=True) strategy: Mapped[str] = mapped_column(String(32), nullable=False, index=True) leave_uk_year: Mapped[int] = mapped_column(Integer, nullable=False) glide_path: Mapped[str] = mapped_column(String(32), nullable=False) spending_gbp: Mapped[Decimal] = mapped_column(Numeric(12, 2), nullable=False) horizon_years: Mapped[int] = mapped_column(Integer, nullable=False, server_default=text("60")) nw_seed_gbp: Mapped[Decimal] = mapped_column(Numeric(14, 2), nullable=False) savings_per_year_gbp: Mapped[Decimal] = mapped_column(Numeric(12, 2), nullable=False, server_default=text("0")) config_json: Mapped[dict[str, Any]] = mapped_column(JSON_TYPE, nullable=False) created_at: Mapped[datetime] = mapped_column(TIMESTAMP(timezone=True), nullable=False, server_default=func.now()) class McRun(Base): """One MC execution per (scenario, run_at). Stores execution metadata + summary statistics — enough to populate a Grafana cell without touching the per-path tables.""" __tablename__ = "mc_run" __table_args__ = {"schema": SCHEMA_NAME} # noqa: RUF012 id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True) scenario_id: Mapped[int] = mapped_column(Integer, nullable=False, index=True) run_at: Mapped[datetime] = mapped_column(TIMESTAMP(timezone=True), nullable=False, server_default=func.now()) n_paths: Mapped[int] = mapped_column(Integer, nullable=False) seed: Mapped[int] = mapped_column(Integer, nullable=False) success_rate: Mapped[Decimal] = mapped_column(Numeric(6, 4), nullable=False) p10_ending_gbp: Mapped[Decimal] = mapped_column(Numeric(16, 2), nullable=False) p50_ending_gbp: Mapped[Decimal] = mapped_column(Numeric(16, 2), nullable=False) p90_ending_gbp: Mapped[Decimal] = mapped_column(Numeric(16, 2), nullable=False) median_lifetime_tax_gbp: Mapped[Decimal] = mapped_column(Numeric(14, 2), nullable=False) median_years_to_ruin: Mapped[Decimal | None] = mapped_column(Numeric(6, 2), nullable=True) elapsed_seconds: Mapped[Decimal] = mapped_column(Numeric(8, 3), nullable=False) sequence_risk_correlation: Mapped[Decimal | None] = mapped_column(Numeric(6, 4), nullable=True) extra: Mapped[dict[str, Any] | None] = mapped_column(JSON_TYPE, nullable=True) class McPath(Base): """Sparse per-path storage: top decile, bottom decile, and median paths fully stored — enough for a fan chart, not 10k×60 ≈ 600k rows.""" __tablename__ = "mc_path" __table_args__ = {"schema": SCHEMA_NAME} # noqa: RUF012 id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True) mc_run_id: Mapped[int] = mapped_column(Integer, nullable=False, index=True) path_idx: Mapped[int] = mapped_column(Integer, nullable=False) bucket: Mapped[str] = mapped_column(String(16), nullable=False) year_idx: Mapped[int] = mapped_column(Integer, nullable=False) portfolio_gbp: Mapped[Decimal] = mapped_column(Numeric(16, 2), nullable=False) withdrawal_gbp: Mapped[Decimal] = mapped_column(Numeric(12, 2), nullable=False) tax_paid_gbp: Mapped[Decimal] = mapped_column(Numeric(12, 2), nullable=False) real_portfolio_gbp: Mapped[Decimal] = mapped_column(Numeric(16, 2), nullable=False) class ProjectionYearly(Base): """Deterministic point projection per scenario — per-year point estimates that drive fan charts and the per-year Grafana table. One row per (scenario, year).""" __tablename__ = "projection_yearly" __table_args__ = {"schema": SCHEMA_NAME} # noqa: RUF012 id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True) mc_run_id: Mapped[int] = mapped_column(Integer, nullable=False, index=True) year_idx: Mapped[int] = mapped_column(Integer, nullable=False) p10_portfolio_gbp: Mapped[Decimal] = mapped_column(Numeric(16, 2), nullable=False) p25_portfolio_gbp: Mapped[Decimal] = mapped_column(Numeric(16, 2), nullable=False) p50_portfolio_gbp: Mapped[Decimal] = mapped_column(Numeric(16, 2), nullable=False) p75_portfolio_gbp: Mapped[Decimal] = mapped_column(Numeric(16, 2), nullable=False) p90_portfolio_gbp: Mapped[Decimal] = mapped_column(Numeric(16, 2), nullable=False) p50_withdrawal_gbp: Mapped[Decimal] = mapped_column(Numeric(12, 2), nullable=False) p50_tax_gbp: Mapped[Decimal] = mapped_column(Numeric(12, 2), nullable=False) survival_rate: Mapped[Decimal] = mapped_column(Numeric(6, 4), nullable=False) class ScenarioSummary(Base): """Denormalised fast-read for Grafana — one row per (scenario, latest run).""" __tablename__ = "scenario_summary" __table_args__ = {"schema": SCHEMA_NAME} # noqa: RUF012 id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True) scenario_id: Mapped[int] = mapped_column(Integer, unique=True, nullable=False) mc_run_id: Mapped[int] = mapped_column(Integer, nullable=False) jurisdiction: Mapped[str] = mapped_column(String(32), nullable=False, index=True) strategy: Mapped[str] = mapped_column(String(32), nullable=False, index=True) leave_uk_year: Mapped[int] = mapped_column(Integer, nullable=False) glide_path: Mapped[str] = mapped_column(String(32), nullable=False) spending_gbp: Mapped[Decimal] = mapped_column(Numeric(12, 2), nullable=False) success_rate: Mapped[Decimal] = mapped_column(Numeric(6, 4), nullable=False) p10_ending_gbp: Mapped[Decimal] = mapped_column(Numeric(16, 2), nullable=False) p50_ending_gbp: Mapped[Decimal] = mapped_column(Numeric(16, 2), nullable=False) p90_ending_gbp: Mapped[Decimal] = mapped_column(Numeric(16, 2), nullable=False) median_lifetime_tax_gbp: Mapped[Decimal] = mapped_column(Numeric(14, 2), nullable=False) median_years_to_ruin: Mapped[Decimal | None] = mapped_column(Numeric(6, 2), nullable=True) updated_at: Mapped[datetime] = mapped_column(TIMESTAMP(timezone=True), nullable=False, server_default=func.now()) class LifeEvent(Base): """A timed event in a user's plan: retirement, kid born, mortgage payoff, sabbatical, etc. Attached to a scenario. `year_start` and `year_end` are offsets from the scenario start year (year 0 = today). For one-time events, leave `year_end` = `year_start`. `delta_gbp_per_year` is the annual cashflow change while the event is active (negative = expense, positive = income; 0 for events that just mark a milestone like "retire"). Free-form `payload` carries event-kind-specific config that the simulator hasn't yet learned to consume — graceful forward-compat. """ __tablename__ = "life_event" __table_args__ = {"schema": SCHEMA_NAME} # noqa: RUF012 id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True) scenario_id: Mapped[int] = mapped_column(Integer, nullable=False, index=True) kind: Mapped[str] = mapped_column(String(32), nullable=False) name: Mapped[str] = mapped_column(String, nullable=False) year_start: Mapped[int] = mapped_column(Integer, nullable=False) year_end: Mapped[int | None] = mapped_column(Integer, nullable=True) delta_gbp_per_year: Mapped[Decimal] = mapped_column(Numeric(12, 2), nullable=False, server_default=text("0")) one_time_amount_gbp: Mapped[Decimal | None] = mapped_column(Numeric(14, 2), nullable=True) enabled: Mapped[bool] = mapped_column(Boolean, nullable=False, server_default=text("true")) payload: Mapped[dict[str, Any] | None] = mapped_column(JSON_TYPE, nullable=True) created_at: Mapped[datetime] = mapped_column(TIMESTAMP(timezone=True), nullable=False, server_default=func.now()) class RetirementGoal(Base): """A user-defined success criterion for a scenario. Examples: - target_nw: "have ≥£2M real GBP at year 25" → kind=target_nw, target_amount_gbp=2_000_000, target_year=25, comparator='>=' - never_run_out: "never run out before age 95" → kind=never_run_out, target_year=65 (years from start), no amount - inheritance: "leave ≥£500k to heirs" → kind=inheritance, target_amount_gbp=500_000, target_year=horizon, comparator='>=' `success_threshold` is the probability bar that counts as "passing" (e.g. 0.95 = 95% of MC paths must satisfy the comparator). """ __tablename__ = "retirement_goal" __table_args__ = {"schema": SCHEMA_NAME} # noqa: RUF012 id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True) scenario_id: Mapped[int] = mapped_column(Integer, nullable=False, index=True) kind: Mapped[str] = mapped_column(String(32), nullable=False) name: Mapped[str] = mapped_column(String, nullable=False) target_amount_gbp: Mapped[Decimal | None] = mapped_column(Numeric(16, 2), nullable=True) target_year: Mapped[int | None] = mapped_column(Integer, nullable=True) comparator: Mapped[str] = mapped_column(String(4), nullable=False, server_default=text("'>='")) success_threshold: Mapped[Decimal] = mapped_column(Numeric(4, 3), nullable=False, server_default=text("0.95")) enabled: Mapped[bool] = mapped_column(Boolean, nullable=False, server_default=text("true")) payload: Mapped[dict[str, Any] | None] = mapped_column(JSON_TYPE, nullable=True) created_at: Mapped[datetime] = mapped_column(TIMESTAMP(timezone=True), nullable=False, server_default=func.now()) def create_engine_from_env() -> AsyncEngine: url = os.environ["DB_CONNECTION_STRING"] return create_async_engine(url, pool_pre_ping=True) def make_session_factory(engine: AsyncEngine) -> async_sessionmaker[Any]: return async_sessionmaker(engine, expire_on_commit=False)