Model/scripts/run_pashub_modelling.py

"""Run the modelling_e2e handler locally over a portfolio's properties, using
the STORED EPCs (``refetch_epc=False``) — the correct source for PasHub
site-notes cohorts, where the stored site note is the newest assessment (the
gov-API refetch would bypass it; see issue #1589 / ADR-0001).

Invokes the real Lambda handler in-process (no Docker) with one SQS-shaped event
per batch of property ids. Reads ``backend/.env`` and maps ``DB_*`` -> the
``POSTGRES_*`` vars the handler's ``PostgresConfig.from_env`` expects.

Config via environment variables (all optional):

    RUN_PORTFOLIO   portfolio id            (default 838)
    RUN_SCENARIO    scenario id             (default 1297)
    RUN_PIDS        comma-separated property ids; unset -> every property in the
                    portfolio, ordered by id
    RUN_DRY         "1" dry-run / "0" write to the DB   (default "1")
    RUN_BATCH       properties per handler invocation   (default 50)
    RUN_REFETCH_EPC / RUN_REPREDICT_EPC / RUN_REFETCH_SOLAR
                    "1"/"0" to override the refetch flags (default "0" — use
                    stored EPCs/solar; flip to re-fetch from the gov API/Google)

Examples:
    python scripts/run_pashub_modelling.py                      # dry-run, whole portfolio
    RUN_DRY=0 python scripts/run_pashub_modelling.py            # real writes
    RUN_SCENARIO=1301 RUN_DRY=0 python scripts/run_pashub_modelling.py
    RUN_PIDS=754772,754780 RUN_DRY=0 python scripts/run_pashub_modelling.py
"""

from __future__ import annotations

import json
import os

from sqlalchemy import text

from scripts.e2e_common import build_engine, load_env


def _configure_env() -> None:
    """Load ``backend/.env`` and mirror the FastAPI-layer ``DB_*`` creds onto the
    ``POSTGRES_*`` names the modelling handler reads."""
    load_env()
    for pg, db in (
        ("POSTGRES_HOST", "DB_HOST"),
        ("POSTGRES_PORT", "DB_PORT"),
        ("POSTGRES_USERNAME", "DB_USERNAME"),
        ("POSTGRES_PASSWORD", "DB_PASSWORD"),
        ("POSTGRES_DATABASE", "DB_NAME"),
    ):
        if db in os.environ:
            os.environ.setdefault(pg, os.environ[db])


def _flag(name: str, default: bool) -> bool:
    raw = os.environ.get(name)
    return default if raw is None else raw == "1"


def _portfolio_property_ids(portfolio_id: int) -> list[int]:
    engine = build_engine()
    with engine.begin() as conn:
        conn.execute(text("SET statement_timeout = 120000"))
        rows = conn.execute(
            text("SELECT id FROM property WHERE portfolio_id = :p ORDER BY id"),
            {"p": portfolio_id},
        )
        return [int(r[0]) for r in rows]


def main() -> None:
    _configure_env()
    # Imported after env is configured so the handler's module-level engine binds
    # to the mirrored POSTGRES_* vars.
    from applications.modelling_e2e.handler import handler

    portfolio_id = int(os.environ.get("RUN_PORTFOLIO", "838"))
    scenario_id = int(os.environ.get("RUN_SCENARIO", "1297"))
    dry_run = _flag("RUN_DRY", True)
    batch_size = int(os.environ.get("RUN_BATCH", "50"))
    refetch_epc = _flag("RUN_REFETCH_EPC", False)
    repredict_epc = _flag("RUN_REPREDICT_EPC", False)
    refetch_solar = _flag("RUN_REFETCH_SOLAR", False)

    pids_env = os.environ.get("RUN_PIDS")
    property_ids: list[int] = (
        [int(x) for x in pids_env.split(",") if x.strip()]
        if pids_env
        else _portfolio_property_ids(portfolio_id)
    )
    batches: list[list[int]] = [
        property_ids[i : i + batch_size]
        for i in range(0, len(property_ids), batch_size)
    ]

    print(
        f"portfolio={portfolio_id} scenario={scenario_id} dry_run={dry_run} "
        f"refetch_epc={refetch_epc} repredict_epc={repredict_epc} "
        f"refetch_solar={refetch_solar} properties={len(property_ids)} "
        f"batches={len(batches)}"
    )

    for i, chunk in enumerate(batches):
        body = {
            "property_ids": chunk,
            "portfolio_id": portfolio_id,
            "scenario_id": scenario_id,
            "refetch_epc": refetch_epc,
            "repredict_epc": repredict_epc,
            "refetch_solar": refetch_solar,
            "dry_run": dry_run,
        }
        event = {"Records": [{"body": json.dumps(body)}]}
        print(
            f"\n=== batch {i + 1}/{len(batches)}: {len(chunk)} props "
            f"({chunk[0]}..{chunk[-1]}) ==="
        )
        result: object = handler(event, None)
        print("  ->", json.dumps(result, default=str)[:800])


if __name__ == "__main__":
    main()