mirror of
https://github.com/Hestia-Homes/Model.git
synced 2026-06-30 13:10:47 +00:00
fix db issues
This commit is contained in:
parent
cbc14f1ca2
commit
79a76b10a9
4 changed files with 36 additions and 17 deletions
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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}
|
||||
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue