mirror of
https://github.com/Hestia-Homes/Model.git
synced 2026-07-12 13:29:04 +00:00
103 lines
3.3 KiB
Python
103 lines
3.3 KiB
Python
"""Enqueue one SQS message per property for the modelling_e2e Lambda.
|
|
|
|
Reads all property IDs for the given portfolio from the DB and sends a batch of
|
|
SQS messages, one per property. The Lambda then processes each message
|
|
independently, enabling concurrent modelling at scale.
|
|
|
|
Edit the CONFIG block below, then run via VSCode Run button or Jupyter.
|
|
AWS creds come from the ambient ~/.aws profile; DB creds from backend/.env.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# CONFIG — edit these before running
|
|
# ---------------------------------------------------------------------------
|
|
PORTFOLIO_ID: int = 785
|
|
SCENARIO_ID: int = 1266
|
|
SQS_URL: str = "https://sqs.eu-west-2.amazonaws.com/ACCOUNT_ID/modelling-e2e-STAGE"
|
|
|
|
# Set to a positive integer to enqueue only the first N properties (trial run).
|
|
LIMIT: int | None = 10
|
|
|
|
# True → Lambda runs the full pipeline but skips all DB writes (safe for testing).
|
|
DRY_RUN: bool = True
|
|
|
|
# True → Lambda skips the Google Solar fetch.
|
|
NO_SOLAR: bool = False
|
|
# ---------------------------------------------------------------------------
|
|
|
|
import json
|
|
import sys
|
|
from pathlib import Path
|
|
from typing import Any, cast
|
|
from uuid import uuid4
|
|
|
|
_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
|
|
|
|
_BATCH_SIZE = 10
|
|
|
|
|
|
def _property_ids(portfolio_id: int, limit: int | None, engine: object) -> list[int]:
|
|
from sqlalchemy.engine import Engine
|
|
|
|
assert isinstance(engine, Engine)
|
|
query = "SELECT id FROM property WHERE portfolio_id = :pid ORDER BY id"
|
|
if limit is not None:
|
|
query += f" LIMIT {int(limit)}"
|
|
with engine.connect() as conn:
|
|
rows = conn.execute(text(query), {"pid": portfolio_id}).fetchall()
|
|
return [int(r[0]) for r in rows]
|
|
|
|
|
|
def _batches(items: list[int], size: int) -> list[list[int]]:
|
|
return [items[i : i + size] for i in range(0, len(items), size)]
|
|
|
|
|
|
def main() -> None:
|
|
load_env(ENV_PATH)
|
|
engine = build_engine()
|
|
|
|
ids = _property_ids(PORTFOLIO_ID, LIMIT, engine)
|
|
if not ids:
|
|
print(f"no properties found for portfolio {PORTFOLIO_ID}")
|
|
return
|
|
|
|
print(
|
|
f"enqueuing {len(ids)} properties "
|
|
f"(portfolio={PORTFOLIO_ID}, scenario={SCENARIO_ID}, "
|
|
f"no_solar={NO_SOLAR}, dry_run={DRY_RUN}) → {SQS_URL}"
|
|
)
|
|
|
|
sqs: Any = cast(Any, boto3.client("sqs")) # pyright: ignore[reportUnknownMemberType]
|
|
sent = 0
|
|
for batch in _batches(ids, _BATCH_SIZE):
|
|
entries = [
|
|
{
|
|
"Id": str(uuid4()).replace("-", "")[:8] + str(i),
|
|
"MessageBody": json.dumps(
|
|
{
|
|
"property_id": pid,
|
|
"portfolio_id": PORTFOLIO_ID,
|
|
"scenario_id": SCENARIO_ID,
|
|
"no_solar": NO_SOLAR,
|
|
"dry_run": DRY_RUN,
|
|
}
|
|
),
|
|
}
|
|
for i, pid in enumerate(batch)
|
|
]
|
|
sqs.send_message_batch(QueueUrl=SQS_URL, Entries=entries)
|
|
sent += len(batch)
|
|
print(f" sent {sent}/{len(ids)}", end="\r")
|
|
|
|
print(f"\ndone — {sent} messages enqueued")
|
|
|
|
|
|
main()
|