diff --git a/applications/modelling_e2e/handler.py b/applications/modelling_e2e/handler.py index de7b5542..a9e351b3 100644 --- a/applications/modelling_e2e/handler.py +++ b/applications/modelling_e2e/handler.py @@ -23,6 +23,7 @@ invocations (ADR-0012). from __future__ import annotations +import dataclasses import io import os from collections.abc import Callable @@ -92,7 +93,11 @@ logger = setup_logger() def _get_engine() -> Engine: global _engine if _engine is None: - _engine = make_engine(PostgresConfig.from_env(dict(os.environ))) + config = PostgresConfig.from_env(dict(os.environ)) + # Reduced pool for Lambda: 32 concurrent containers × 3 connections = 96 max, + # vs the default 3+5=8 which would reach 256+ and exhaust RDS max_connections. + # pool_size=2 covers the simultaneous read_session + UoW session per invocation. + _engine = make_engine(dataclasses.replace(config, pool_size=2, max_overflow=1)) return _engine diff --git a/infrastructure/postgres/engine.py b/infrastructure/postgres/engine.py index 2558532e..cf455219 100644 --- a/infrastructure/postgres/engine.py +++ b/infrastructure/postgres/engine.py @@ -1,20 +1,25 @@ from collections.abc import Iterator from contextlib import contextmanager +from typing import Any, Optional, Type from sqlalchemy.engine import Engine +from sqlalchemy.pool import Pool from sqlmodel import Session, create_engine from infrastructure.postgres.config import PostgresConfig -def make_engine(config: PostgresConfig) -> Engine: - return create_engine( - config.url(), - pool_size=config.pool_size, - max_overflow=config.max_overflow, - pool_pre_ping=config.pool_pre_ping, - pool_recycle=config.pool_recycle, - ) +def make_engine( + config: PostgresConfig, poolclass: Optional[Type[Pool]] = None +) -> Engine: + kwargs: dict[str, Any] = {"pool_pre_ping": config.pool_pre_ping} + if poolclass is None: + kwargs["pool_size"] = config.pool_size + kwargs["max_overflow"] = config.max_overflow + kwargs["pool_recycle"] = config.pool_recycle + else: + kwargs["poolclass"] = poolclass + return create_engine(config.url(), **kwargs) def make_session(engine: Engine) -> Session: diff --git a/scripts/trigger_modelling_e2e_sqs.py b/scripts/trigger_modelling_e2e_sqs.py index 456497ec..f0eaed38 100644 --- a/scripts/trigger_modelling_e2e_sqs.py +++ b/scripts/trigger_modelling_e2e_sqs.py @@ -22,7 +22,7 @@ SQS_QUEUE_NAME: str = "modelling_e2e-queue-dev" # Number of postcodes to process this run (postcodes where all properties are # already completed are skipped and do not count toward this limit). -POSTCODES_LIMIT: int = 100 +POSTCODES_LIMIT: int = 1000 # True → Lambda runs the full pipeline but skips all DB writes (safe for testing). DRY_RUN: bool = False @@ -73,16 +73,21 @@ def _completed_property_ids() -> set[int]: with engine.connect() as conn: rows = conn.execute( text(""" - SELECT DISTINCT (elem.value::int) AS property_id + 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->'property_ids') AS elem(value) + CROSS JOIN jsonb_array_elements_text( + (st.inputs::jsonb)->'property_ids' + ) AS elem(value) WHERE t.task_source = 'modelling_e2e' - AND st.status = 'completed' - AND (st.inputs->>'portfolio_id')::int = :portfolio_id - AND (st.inputs->>'scenario_id')::int = :scenario_id + 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}, + { + "portfolio_id": PORTFOLIO_ID, + "scenario_id": SCENARIO_ID, + }, ).fetchall() return {int(r[0]) for r in rows} diff --git a/utilities/aws_lambda/default_orchestrator.py b/utilities/aws_lambda/default_orchestrator.py index f78886b9..cd4c7d1e 100644 --- a/utilities/aws_lambda/default_orchestrator.py +++ b/utilities/aws_lambda/default_orchestrator.py @@ -2,6 +2,7 @@ import os from collections.abc import Generator from contextlib import contextmanager +from sqlalchemy.pool import NullPool from sqlmodel import Session from infrastructure.postgres.config import PostgresConfig @@ -17,8 +18,11 @@ def default_orchestrator() -> Generator[TaskOrchestrator, None, None]: Connection params come from os.environ via PostgresConfig.from_env. Each handler invocation gets its own session, cleaned up on context exit. + + NullPool is intentional: a new engine is created on every invocation, so + pooling would accumulate idle connections across warm Lambda containers. """ - engine = make_engine(PostgresConfig.from_env(dict(os.environ))) + engine = make_engine(PostgresConfig.from_env(dict(os.environ)), poolclass=NullPool) with Session(engine) as session: yield TaskOrchestrator( task_repo=TaskPostgresRepository(session=session),