diff --git a/alembic/versions/0002_user_scenarios_events_goals.py b/alembic/versions/0002_user_scenarios_events_goals.py new file mode 100644 index 0000000..fbb5ea3 --- /dev/null +++ b/alembic/versions/0002_user_scenarios_events_goals.py @@ -0,0 +1,116 @@ +"""extend scenario, add life_event + retirement_goal + +Revision ID: 0002 +Revises: 0001 +Create Date: 2026-05-09 00:00:00.000000 + +ProjectionLab parity surface: user-defined scenarios, life-event timeline, +retirement goals. +""" +from collections.abc import Sequence + +import sqlalchemy as sa +from sqlalchemy.dialects import postgresql + +from alembic import op + +revision: str = "0002" +down_revision: str | None = "0001" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + +SCHEMA = "fire_planner" + + +def _jsonb() -> sa.types.TypeEngine[object]: + return postgresql.JSONB().with_variant(sa.JSON(), "sqlite") + + +def upgrade() -> None: + # Scenario: add kind, name, description, parent_scenario_id + with op.batch_alter_table("scenario", schema=SCHEMA) as batch: + batch.add_column( + sa.Column("kind", + sa.Text(), + nullable=False, + server_default=sa.text("'cartesian'"))) + batch.add_column(sa.Column("name", sa.Text(), nullable=True)) + batch.add_column(sa.Column("description", sa.Text(), nullable=True)) + batch.add_column(sa.Column("parent_scenario_id", sa.Integer(), nullable=True)) + op.create_index("idx_scenario_kind", "scenario", ["kind"], schema=SCHEMA) + op.create_index("idx_scenario_parent", + "scenario", ["parent_scenario_id"], + schema=SCHEMA) + + op.create_table( + "life_event", + sa.Column("id", sa.Integer(), primary_key=True, autoincrement=True), + sa.Column("scenario_id", sa.Integer(), nullable=False), + sa.Column("kind", sa.Text(), nullable=False), + sa.Column("name", sa.Text(), nullable=False), + sa.Column("year_start", sa.Integer(), nullable=False), + sa.Column("year_end", sa.Integer(), nullable=True), + sa.Column("delta_gbp_per_year", + sa.Numeric(12, 2), + nullable=False, + server_default=sa.text("0")), + sa.Column("one_time_amount_gbp", sa.Numeric(14, 2), nullable=True), + sa.Column("enabled", + sa.Boolean(), + nullable=False, + server_default=sa.text("true")), + sa.Column("payload", _jsonb(), nullable=True), + sa.Column("created_at", + sa.TIMESTAMP(timezone=True), + nullable=False, + server_default=sa.text("now()")), + schema=SCHEMA, + ) + op.create_index("idx_life_event_scenario", + "life_event", ["scenario_id"], + schema=SCHEMA) + + op.create_table( + "retirement_goal", + sa.Column("id", sa.Integer(), primary_key=True, autoincrement=True), + sa.Column("scenario_id", sa.Integer(), nullable=False), + sa.Column("kind", sa.Text(), nullable=False), + sa.Column("name", sa.Text(), nullable=False), + sa.Column("target_amount_gbp", sa.Numeric(16, 2), nullable=True), + sa.Column("target_year", sa.Integer(), nullable=True), + sa.Column("comparator", + sa.Text(), + nullable=False, + server_default=sa.text("'>='")), + sa.Column("success_threshold", + sa.Numeric(4, 3), + nullable=False, + server_default=sa.text("0.95")), + sa.Column("enabled", + sa.Boolean(), + nullable=False, + server_default=sa.text("true")), + sa.Column("payload", _jsonb(), nullable=True), + sa.Column("created_at", + sa.TIMESTAMP(timezone=True), + nullable=False, + server_default=sa.text("now()")), + schema=SCHEMA, + ) + op.create_index("idx_retirement_goal_scenario", + "retirement_goal", ["scenario_id"], + schema=SCHEMA) + + +def downgrade() -> None: + op.drop_index("idx_retirement_goal_scenario", table_name="retirement_goal", schema=SCHEMA) + op.drop_table("retirement_goal", schema=SCHEMA) + op.drop_index("idx_life_event_scenario", table_name="life_event", schema=SCHEMA) + op.drop_table("life_event", schema=SCHEMA) + op.drop_index("idx_scenario_parent", table_name="scenario", schema=SCHEMA) + op.drop_index("idx_scenario_kind", table_name="scenario", schema=SCHEMA) + with op.batch_alter_table("scenario", schema=SCHEMA) as batch: + batch.drop_column("parent_scenario_id") + batch.drop_column("description") + batch.drop_column("name") + batch.drop_column("kind") diff --git a/fire_planner/db.py b/fire_planner/db.py index 650fe0f..5c4ac13 100644 --- a/fire_planner/db.py +++ b/fire_planner/db.py @@ -3,7 +3,7 @@ from datetime import date, datetime from decimal import Decimal from typing import Any -from sqlalchemy import JSON, TIMESTAMP, Date, Integer, Numeric, String, func, text +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 @@ -45,15 +45,25 @@ class AccountSnapshot(Base): class Scenario(Base): - """A simulation scenario — Cartesian point in (jurisdiction × strategy × - leave_year × glide × spending) space. The Cartesian product is rebuilt - from `scenarios.py` every recompute; rows are upserted on `external_id`. + """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) @@ -156,6 +166,73 @@ class ScenarioSummary(Base): 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) diff --git a/tests/test_db_schema.py b/tests/test_db_schema.py index fcbe980..76ea443 100644 --- a/tests/test_db_schema.py +++ b/tests/test_db_schema.py @@ -7,9 +7,11 @@ from sqlalchemy.ext.asyncio import AsyncSession from fire_planner.db import ( AccountSnapshot, + LifeEvent, McPath, McRun, ProjectionYearly, + RetirementGoal, Scenario, ScenarioSummary, ) @@ -109,3 +111,144 @@ async def test_remaining_tables_smoke(session: AsyncSession) -> None: p90_ending_gbp=Decimal("10000000"), median_lifetime_tax_gbp=Decimal("800000"))) await session.commit() + + +async def test_user_scenario_with_clone(session: AsyncSession) -> None: + """User-defined scenarios can declare a parent (cloned-from) scenario.""" + base = Scenario( + external_id="cyprus-vpw-leave-y3-glide-rising", + kind="cartesian", + jurisdiction="cyprus", + strategy="vpw", + leave_uk_year=3, + glide_path="rising", + spending_gbp=Decimal("100000"), + nw_seed_gbp=Decimal("1500000"), + savings_per_year_gbp=Decimal("0"), + config_json={"horizon_years": 60}, + ) + session.add(base) + await session.commit() + + user = Scenario( + external_id="user-aggressive-fire", + kind="user", + name="Aggressive FIRE", + description="Same as Cyprus baseline but spending £80k", + parent_scenario_id=base.id, + jurisdiction="cyprus", + strategy="vpw", + leave_uk_year=3, + glide_path="rising", + spending_gbp=Decimal("80000"), + nw_seed_gbp=Decimal("1500000"), + savings_per_year_gbp=Decimal("0"), + config_json={"horizon_years": 60, "floor": 60000}, + ) + session.add(user) + await session.commit() + + rows = (await session.execute(select(Scenario).where(Scenario.kind == "user"))).scalars().all() + assert len(rows) == 1 + assert rows[0].name == "Aggressive FIRE" + assert rows[0].parent_scenario_id == base.id + + +async def test_life_event_roundtrip(session: AsyncSession) -> None: + ev = LifeEvent( + scenario_id=1, + kind="retirement", + name="Retire at 50", + year_start=15, + year_end=15, + delta_gbp_per_year=Decimal("0"), + ) + childcare = LifeEvent( + scenario_id=1, + kind="expense_range", + name="Childcare", + year_start=2, + year_end=20, + delta_gbp_per_year=Decimal("-12000"), + ) + inheritance = LifeEvent( + scenario_id=1, + kind="one_time_income", + name="Inheritance", + year_start=22, + year_end=22, + one_time_amount_gbp=Decimal("250000"), + ) + session.add_all([ev, childcare, inheritance]) + await session.commit() + rows = (await session.execute( + select(LifeEvent).where(LifeEvent.scenario_id == 1))).scalars().all() + assert len(rows) == 3 + by_kind = {r.kind: r for r in rows} + assert by_kind["retirement"].year_start == 15 + assert by_kind["expense_range"].delta_gbp_per_year == Decimal("-12000") + assert by_kind["one_time_income"].one_time_amount_gbp == Decimal("250000") + + +async def test_life_event_default_enabled(session: AsyncSession) -> None: + ev = LifeEvent(scenario_id=1, kind="retirement", name="Retire", year_start=10) + session.add(ev) + await session.commit() + fetched = (await session.execute(select(LifeEvent))).scalars().first() + assert fetched is not None + assert fetched.enabled is True + + +async def test_retirement_goal_roundtrip(session: AsyncSession) -> None: + target_nw = RetirementGoal( + scenario_id=1, + kind="target_nw", + name="≥ £2M at age 50", + target_amount_gbp=Decimal("2000000"), + target_year=15, + comparator=">=", + success_threshold=Decimal("0.90"), + ) + never_run_out = RetirementGoal( + scenario_id=1, + kind="never_run_out", + name="Last to age 95", + target_year=65, + comparator=">=", + success_threshold=Decimal("0.95"), + ) + inheritance = RetirementGoal( + scenario_id=1, + kind="inheritance", + name="Leave £500k", + target_amount_gbp=Decimal("500000"), + target_year=65, + comparator=">=", + success_threshold=Decimal("0.50"), + ) + session.add_all([target_nw, never_run_out, inheritance]) + await session.commit() + rows = (await session.execute(select(RetirementGoal))).scalars().all() + assert len(rows) == 3 + kinds = {r.kind for r in rows} + assert kinds == {"target_nw", "never_run_out", "inheritance"} + + +async def test_scenario_kind_default_cartesian(session: AsyncSession) -> None: + """Existing Cartesian scenarios omit `kind` — default applies.""" + scen = Scenario( + external_id="uk-trinity-leave-y0-glide-static", + jurisdiction="uk", + strategy="trinity", + leave_uk_year=0, + glide_path="static", + spending_gbp=Decimal("60000"), + nw_seed_gbp=Decimal("1500000"), + savings_per_year_gbp=Decimal("0"), + config_json={}, + ) + session.add(scen) + await session.commit() + fetched = (await session.execute(select(Scenario))).scalars().first() + assert fetched is not None + assert fetched.kind == "cartesian"