Model/scripts/trigger_modelling_e2e_sqs.py

"""Enqueue one SQS message per postcode group for the modelling_e2e Lambda.

Reads postcode → property ID groups from the file produced by
list_properties_by_postcode.py, queries the DB for already-completed
property IDs, then sends one SQS message per postcode batch containing only
the properties that still need processing.

Edit the CONFIG block below, then hit Run.
AWS creds come from the ambient ~/.aws profile.
"""

from __future__ import annotations

from utilities.logger import setup_logger

# ---------------------------------------------------------------------------
# CONFIG — edit these before running
# ---------------------------------------------------------------------------
PORTFOLIO_ID: int = 796
SCENARIO_ID: int = 1268
SQS_QUEUE_NAME: str = "modelling_e2e-queue-dev"

# Max number of properties to process this run (cost cap). Postcodes are added
# whole, smallest-group-first, until adding the next would exceed this — so the
# actual count lands at or just under the limit.
PROPERTIES_LIMIT: int = 5000

# True → Lambda runs the full pipeline but skips all DB writes (safe for testing).
DRY_RUN: bool = False

# True → Lambda skips the Google Solar fetch.
NO_SOLAR: bool = False
# ---------------------------------------------------------------------------

import ast
import json
import sys
from pathlib import Path
from typing import Any, cast

_REPO_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(_REPO_ROOT))

import boto3  # noqa: E402
from sqlalchemy import text  # noqa: E402

from scripts.e2e_common import ENV_PATH, build_engine, load_env  # noqa: E402

logger = setup_logger()

_POSTCODES_FILE = _REPO_ROOT / "scripts" / f"properties_by_postcode_{PORTFOLIO_ID}.txt"


def _load_postcode_map() -> dict[str, list[int]]:
    if not _POSTCODES_FILE.exists():
        raise FileNotFoundError(
            f"{_POSTCODES_FILE} not found — run list_properties_by_postcode.py first"
        )
    result: dict[str, list[int]] = {}
    for line in _POSTCODES_FILE.read_text(encoding="utf-8").splitlines():
        line = line.strip()
        if not line or line.startswith("Total"):
            continue
        postcode_repr, ids_repr = line.split(": ", 1)
        result[ast.literal_eval(postcode_repr)] = ast.literal_eval(ids_repr)
    return result


def _completed_property_ids() -> set[int]:
    """Return all property IDs with a completed modelling_e2e subtask for this
    portfolio + scenario. Single DB round-trip."""
    load_env(ENV_PATH)
    engine = build_engine()
    with engine.connect() as conn:
        rows = conn.execute(
            text("""
                SELECT DISTINCT elem.value::int AS property_id
                FROM sub_task st
                JOIN tasks t ON t.id = st.task_id
                CROSS JOIN jsonb_array_elements_text(
                    (st.inputs::jsonb)->'property_ids'
                ) AS elem(value)
                WHERE t.task_source = 'modelling_e2e'
                AND st.status = 'complete'
                AND ((st.inputs::jsonb)->>'portfolio_id')::int = :portfolio_id
                AND ((st.inputs::jsonb)->>'scenario_id')::int = :scenario_id
            """),
            {
                "portfolio_id": PORTFOLIO_ID,
                "scenario_id": SCENARIO_ID,
            },
        ).fetchall()
    return {int(r[0]) for r in rows}


def main() -> None:
    postcode_map = _load_postcode_map()

    # Pending filter disabled — re-run every property regardless of whether it
    # already has a completed modelling_e2e sub_task for this scenario.
    batches: list[tuple[str, list[int]]] = []
    for postcode, ids in postcode_map.items():
        if ids:
            batches.append((postcode, ids))

    to_process: list[tuple[str, list[int]]] = []
    property_count = 0
    for postcode, ids in batches:
        if property_count + len(ids) > PROPERTIES_LIMIT:
            continue
        to_process.append((postcode, ids))
        property_count += len(ids)

    if not to_process:
        logger.info("Nothing left to process.")
        return

    logger.info(
        f"selected {len(to_process)} postcodes / {property_count} properties "
        f"(limit {PROPERTIES_LIMIT})"
    )

    sqs: Any = cast(
        Any, boto3.client("sqs", region_name="eu-west-2")
    )  # pyright: ignore[reportUnknownMemberType]
    sqs_url: str = sqs.get_queue_url(QueueName=SQS_QUEUE_NAME)["QueueUrl"]

    logger.info(
        f"sending {len(to_process)} messages "
        f"(portfolio={PORTFOLIO_ID}, scenario={SCENARIO_ID}, "
        f"dry_run={DRY_RUN}, no_solar={NO_SOLAR}) → {sqs_url}"
    )

    for postcode, ids in to_process:
        sqs.send_message(
            QueueUrl=sqs_url,
            MessageBody=json.dumps(
                {
                    "property_ids": ids,
                    "portfolio_id": PORTFOLIO_ID,
                    "scenario_id": SCENARIO_ID,
                    "no_solar": NO_SOLAR,
                    "dry_run": DRY_RUN,
                }
            ),
        )
        logger.info(f"  sent {postcode}: {ids}")

    logger.info(f"\ndone — {len(to_process)} messages enqueued")


main()