fire-planner/alembic/versions/0002_user_scenarios_events_goals.py

"""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")