Model/repositories/plan/plan_postgres_repository.py
Daniel Roth 4764bc7c15 Batch plan saves reduce RDS CPU during bulk modelling runs 🟪
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-06-29 15:08:47 +00:00

110 lines
4.1 KiB
Python

from __future__ import annotations
from typing import Any
from sqlalchemy import insert as _sa_insert
from sqlmodel import Session, col, update
from domain.modelling.plan import Plan
from infrastructure.postgres.modelling import PlanModel, RecommendationModel
from repositories.plan.plan_repository import PlanRepository, PlanSaveRequest
def _col_values(model: Any, exclude: frozenset[str] = frozenset()) -> dict[str, Any]:
"""Extract column-keyed values from a SQLModel instance for Core INSERT."""
return {
c.name: getattr(model, c.name)
for c in model.__table__.c
if c.name not in exclude
}
class PlanPostgresRepository(PlanRepository):
"""Maps a Plan and its Plan Measures onto the live ``plan`` /
``recommendation`` tables (ADR-0017). Does not commit — the Unit of Work
owns the transaction (ADR-0012).
A re-run INSERTs a fresh Plan rather than deleting the prior one (the cascade
delete was slow); when the new Plan is the default it demotes any prior
default Plan for the same (property_id, scenario_id) to ``is_default=False``,
so readers can select the current Plan via ``is_default=True``."""
def __init__(self, session: Session) -> None:
self._session = session
def save(
self,
plan: Plan,
*,
property_id: int,
scenario_id: int,
portfolio_id: int,
is_default: bool,
) -> int:
return self.save_batch(
[PlanSaveRequest(plan, property_id=property_id, scenario_id=scenario_id, portfolio_id=portfolio_id, is_default=is_default)]
)[0]
def save_batch(self, requests: list[PlanSaveRequest]) -> list[int]:
"""Persist all Plans in three statements regardless of batch size.
1. One demote UPDATE (only when any request has ``is_default=True``).
2. One bulk plan INSERT with RETURNING to capture ids positionally.
3. One bulk recommendation INSERT (skipped when no measures exist).
"""
if not requests:
return []
# Demote prior default Plans for every property in the batch that is
# receiving a new default Plan — one UPDATE for the whole batch.
default_pids = [r.property_id for r in requests if r.is_default]
if default_pids:
# scenario_id is uniform per batch (one scenario per SQS message).
scenario_id = requests[0].scenario_id
self._session.exec( # type: ignore[call-overload]
update(PlanModel)
.where(
col(PlanModel.property_id).in_(default_pids),
col(PlanModel.scenario_id) == scenario_id,
)
.values(is_default=False)
)
# Bulk INSERT all plan rows; capture returned ids positionally.
plan_rows = [
_col_values(
PlanModel.from_domain(
r.plan,
property_id=r.property_id,
scenario_id=r.scenario_id,
portfolio_id=r.portfolio_id,
is_default=r.is_default,
),
exclude=frozenset({"id"}),
)
for r in requests
]
returned = self._session.execute( # type: ignore[deprecated]
_sa_insert(PlanModel).returning(PlanModel.__table__.c["id"]), # type: ignore[attr-defined]
plan_rows,
).all()
plan_ids = [row[0] for row in returned]
# Accumulate recommendation rows across all requests; properties with
# zero measures contribute nothing (no special-casing needed).
rec_rows = [
_col_values(
RecommendationModel.from_domain(
measure, property_id=r.property_id, plan_id=plan_id
),
exclude=frozenset({"id"}),
)
for r, plan_id in zip(requests, plan_ids)
for measure in r.plan.measures
]
if rec_rows:
self._session.execute( # type: ignore[deprecated]
_sa_insert(RecommendationModel), rec_rows
)
return plan_ids