fix db issues

This commit is contained in:
Daniel Roth 2026-06-23 16:36:59 +00:00
parent cbc14f1ca2
commit 79a76b10a9
4 changed files with 36 additions and 17 deletions

View file

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

View file

@ -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:

View file

@ -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}

View file

@ -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),