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

import ast
import json
import sys
from datetime import datetime, timedelta, timezone
from pathlib import Path
from typing import Any, cast

from utilities.logger import setup_logger

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

# Max number of properties to process this run (cost cap).
PROPERTIES_LIMIT: int = 32000

# Number of properties bundled into each SQS message / Lambda invocation.
BATCH_SIZE: int = 50

# Skip properties whose modelling_e2e sub_task completed at or after this time.
# Set to None to disable the filter and process all properties.
COMPLETED_SINCE: datetime | None = datetime(
    2026, 6, 24, 12, 27, 54, 34000, tzinfo=timezone(timedelta(hours=1))
)

# 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
# ---------------------------------------------------------------------------

_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(since: datetime) -> set[int]:
    """Return property IDs with a completed modelling_e2e sub_task on or after *since*."""
    load_env(ENV_PATH)
    engine = build_engine()
    with engine.connect() as conn:
        rows = conn.execute(
            text("""
                SELECT DISTINCT ((st.inputs::jsonb)->>'property_id')::int AS property_id
                FROM sub_task st
                JOIN tasks t ON t.id = st.task_id
                WHERE t.task_source = 'modelling_e2e'
                AND st.status = 'complete'
                AND st.job_completed >= :since
                AND (st.inputs::jsonb) ? 'property_id'
            """),
            {"since": since},
        ).fetchall()
    return {int(r[0]) for r in rows}


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

    completed: set[int] = set()
    if COMPLETED_SINCE is not None:
        completed = _completed_property_ids(COMPLETED_SINCE)
        logger.info(
            f"skipping {len(completed)} properties already completed since {COMPLETED_SINCE}"
        )

    # Filter to pending IDs, keeping postcode grouping intact.
    pending: list[tuple[str, list[int]]] = [
        (pc, [i for i in ids if i not in completed])
        for pc, ids in postcode_map.items()
        if any(i not in completed for i in ids)
    ]

    # Apply PROPERTIES_LIMIT: skip whole postcodes that would exceed the cap.
    selected: list[tuple[str, list[int]]] = []
    property_count = 0
    for postcode, ids in pending:
        if property_count + len(ids) > PROPERTIES_LIMIT:
            continue
        selected.append((postcode, ids))
        property_count += len(ids)

    # Pack postcodes into batches of ~BATCH_SIZE, never splitting a postcode.
    # A postcode larger than BATCH_SIZE becomes its own oversized message.
    batches: list[list[int]] = []
    current: list[int] = []
    for _postcode, ids in selected:
        if current and len(current) + len(ids) > BATCH_SIZE:
            batches.append(current)
            current = list(ids)
        else:
            current.extend(ids)
    if current:
        batches.append(current)

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

    logger.info(
        f"selected {property_count} properties across {len(batches)} batches of ~{BATCH_SIZE} "
        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(batches)} messages "
        f"(portfolio={PORTFOLIO_ID}, scenario={SCENARIO_ID}, "
        f"dry_run={DRY_RUN}, no_solar={NO_SOLAR}) → {sqs_url}"
    )

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

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


main()