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, UniqueConstraint, 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) # Spending category for flex-spending classification: # essential — never trimmed by flex rules (mortgage, food, kids) # discretionary — trimmed when portfolio drops vs ATH (travel, dining) # not_spending — informational only (a milestone marker that doesn't # change cashflow, e.g. a kid graduating) # Default is `essential` so existing rows keep their full spending impact. category: Mapped[str] = mapped_column(String(16), nullable=False, server_default=text("'essential'")) 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 IncomeStream(Base): """A typed, recurring source of income — first-class income object. Modelled as a per-scenario row so a user can stack salary, dividends, rental, pensions, RSUs, etc. The simulator routes the after-tax amount through the jurisdiction's tax engine using `tax_treatment` as the bucket hint (income / dividend / cgt / tax_free). `start_year` / `end_year` are offsets from the scenario start year. `growth_pct` is real growth; the simulator applies it geometrically. """ __tablename__ = "income_stream" __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) start_year: Mapped[int] = mapped_column(Integer, nullable=False, server_default=text("0")) end_year: Mapped[int | None] = mapped_column(Integer, nullable=True) amount_gbp_per_year: Mapped[Decimal] = mapped_column(Numeric(12, 2), nullable=False, server_default=text("0")) growth_pct: Mapped[Decimal] = mapped_column(Numeric(6, 4), nullable=False, server_default=text("0")) tax_treatment: Mapped[str] = mapped_column(String(16), nullable=False, server_default=text("'income'")) 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 ColSnapshot(Base): """Cached cost-of-living snapshot per (city_slug, source). Phase 2 of the COL subsystem. Replaces the previous "baseline-only" lookup with cache-then-scrape semantics: service.lookup_city(slug) → check ColSnapshot, return if fresh → else scrape Numbeo, upsert, return → if scrape fails, fall back to baseline.py TTL default = 365 days (`expires_at = fetched_at + interval '365 day'`). The user explicitly asked for 1y on 2026-05-21 — Numbeo data doesn't move fast enough to need monthly refresh, and the API/scraper has rate- limit risk we prefer to amortise. Phase-3 CronJob will run a nightly refresh of stale rows so individual user requests never have to scrape. `(city_slug, source_name)` is unique — we can store multiple sources per city (Numbeo + Expatistan) and reconcile in service.py. """ __tablename__ = "col_snapshot" __table_args__ = {"schema": SCHEMA_NAME} # noqa: RUF012 id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True) city_slug: Mapped[str] = mapped_column(String(64), nullable=False, index=True) city_display: Mapped[str] = mapped_column(String(128), nullable=False) country: Mapped[str] = mapped_column(String(64), nullable=False) source_name: Mapped[str] = mapped_column(String(32), nullable=False) source_url: Mapped[str | None] = mapped_column(String, nullable=True) snapshot_date: Mapped[date] = mapped_column(Date, nullable=False) fetched_at: Mapped[datetime] = mapped_column(TIMESTAMP(timezone=True), nullable=False, server_default=func.now()) expires_at: Mapped[datetime] = mapped_column(TIMESTAMP(timezone=True), nullable=False) total_no_rent_gbp: Mapped[Decimal] = mapped_column(Numeric(12, 2), nullable=False) total_with_rent_gbp: Mapped[Decimal] = mapped_column(Numeric(12, 2), nullable=False) rent_1bed_center_gbp: Mapped[Decimal] = mapped_column(Numeric(12, 2), nullable=False) rent_1bed_outside_gbp: Mapped[Decimal | None] = mapped_column(Numeric(12, 2), nullable=True) raw_currency: Mapped[str] = mapped_column(String(3), nullable=False, server_default="GBP") gbp_per_unit: Mapped[Decimal] = mapped_column(Numeric(12, 8), nullable=False, server_default=text("1")) by_category_json: Mapped[dict[str, Any] | None] = mapped_column(JSON_TYPE, nullable=True) 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()) class FireExample(Base): """One Reddit-sourced FIRE example. `reddit_id` UNIQUE makes re-ingest idempotent. Fields are nullable when the LLM couldn't extract them confidently — never inferred. Currency normalisation (portfolio_gbp / annual_exp_gbp) happens at extraction time using `fire_planner/fx.py`; `raw_currency` is kept for traceability. """ __tablename__ = "fire_example" __table_args__ = {"schema": SCHEMA_NAME} # noqa: RUF012 id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True) reddit_id: Mapped[str] = mapped_column(String(16), unique=True, nullable=False) source_sub: Mapped[str] = mapped_column(String(64), nullable=False) post_url: Mapped[str] = mapped_column(String, nullable=False) post_date: Mapped[date] = mapped_column(Date, nullable=False, index=True) post_title: Mapped[str] = mapped_column(String, nullable=False) country: Mapped[str | None] = mapped_column(String(64), nullable=True, index=True) city: Mapped[str | None] = mapped_column(String(128), nullable=True) portfolio_gbp: Mapped[Decimal | None] = mapped_column(Numeric(14, 2), nullable=True) annual_exp_gbp: Mapped[Decimal | None] = mapped_column(Numeric(12, 2), nullable=True) age: Mapped[int | None] = mapped_column(Integer, nullable=True) family_size: Mapped[int | None] = mapped_column(Integer, nullable=True) fi_status: Mapped[str | None] = mapped_column(String(24), nullable=True, index=True) is_retired: Mapped[bool | None] = mapped_column(Boolean, nullable=True) raw_currency: Mapped[str | None] = mapped_column(String(3), nullable=True) raw_excerpt: Mapped[str | None] = mapped_column(String, nullable=True) llm_model: Mapped[str] = mapped_column(String(64), nullable=False) llm_confidence: Mapped[Decimal | None] = mapped_column(Numeric(3, 2), nullable=True) extracted_at: Mapped[datetime] = mapped_column(TIMESTAMP(timezone=True), nullable=False, server_default=func.now()) ingested_at: Mapped[datetime] = mapped_column(TIMESTAMP(timezone=True), nullable=False, server_default=func.now()) class FireTarget(Base): """Solved FIRE number per (Case × country × with-home) for the countdown. One row per combination; the Grafana countdown reads `target_nw_gbp` for the selected country and diffs it against current liquid net worth. Seeded on LIQUID net worth (the pension joins later as `pension_at_unlock_gbp`) — see ADR-0001. `reached_bar=False` flags a combination that can't hit the bar within the search range (target then holds the search ceiling). """ __tablename__ = "fire_target" __table_args__ = ( UniqueConstraint("case", "country_slug", "with_home", "bar", name="uq_fire_target_case_country_home_bar"), {"schema": SCHEMA_NAME}, ) id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True) case: Mapped[str] = mapped_column(String(16), nullable=False, index=True) country_slug: Mapped[str] = mapped_column(String(64), nullable=False, index=True) country_display: Mapped[str] = mapped_column(String(128), nullable=False) jurisdiction: Mapped[str] = mapped_column(String(32), nullable=False) with_home: Mapped[bool] = mapped_column(Boolean, nullable=False, server_default=text("false")) bar: Mapped[Decimal] = mapped_column(Numeric(4, 3), nullable=False, server_default=text("0.99")) strategy: Mapped[str] = mapped_column(String(32), nullable=False, server_default=text("'guyton_klinger'")) annual_spend_gbp: Mapped[Decimal] = mapped_column(Numeric(12, 2), nullable=False) target_nw_gbp: Mapped[Decimal] = mapped_column(Numeric(16, 2), nullable=False) pension_at_unlock_gbp: Mapped[Decimal] = mapped_column(Numeric(16, 2), nullable=False, server_default=text("0")) success_at_target: Mapped[Decimal] = mapped_column(Numeric(6, 4), nullable=False) reached_bar: Mapped[bool] = mapped_column(Boolean, nullable=False, server_default=text("true")) horizon_years: Mapped[int] = mapped_column(Integer, nullable=False) n_paths: Mapped[int] = mapped_column(Integer, nullable=False) updated_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)