Release Stage 12 database sessions and keep original failures visible

changelog_entry.yaml

@@ -10,3 +10,4 @@
       - Accepted a UK dataset without area codes only when the policyengine.py bundle routes constituency and local-authority regions to another dataset, and kept requiring area codes on every constituency and local-authority run
       - Raised a clear error, never an omission, when a UK constituency or local-authority breakdown is built over a dataset without area codes
       - Bound UK constituency and local-authority weight matrices to the certified data release of the selected dataset instead of the unversioned bucket copy, and stopped runs before they start when the matrix lacks the run year or household dimension
+      - Released each Stage 12 database session after its operation instead of holding pooled sessions through calculations, retried opening each Stage 12 coordinator and worker database connection that the pooler refused without ever retrying a statement or commit, and reported Stage 12 failures to telemetry before recording them in the database so a failed write cannot hide the original error

libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/__init__.py

@@ -1,6 +1,7 @@
 """Direct DML for temporary Stage 12 comparison state."""
 
 from policyengine_stage12_persistence.store import (
+    ConnectionOpener,
     Stage12PersistenceStore,
     create_stage12_engine,
     normalized_database_url,
@@ -11,6 +12,7 @@
 )
 
 __all__ = [
+    "ConnectionOpener",
     "Stage12PersistenceStore",
     "Stage12RuntimeAccessError",
     "create_stage12_engine",

libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/store.py

@@ -2,9 +2,11 @@
 
 from __future__ import annotations
 
+from collections.abc import Callable, Iterator
+from contextlib import contextmanager
 from datetime import datetime
 from enum import StrEnum
-from typing import Any
+from typing import Any, Protocol
 from uuid import UUID
 
 from policyengine_simulation_contract.stage12_execution import (
@@ -17,7 +19,8 @@
 )
 from sqlalchemy import Engine, create_engine, select, update
 from sqlalchemy.dialects.postgresql import insert
-from sqlalchemy.engine import URL, make_url
+from sqlalchemy.engine import URL, Connection, make_url
+from sqlalchemy.pool import NullPool
 
 from policyengine_stage12_persistence.tables import (
     REPORT_IDENTITY_CONSTRAINT,
@@ -115,14 +118,41 @@ def create_stage12_engine(database_url: str) -> Engine:
 
     return create_engine(
         normalized_database_url(database_url),
-        pool_pre_ping=True,
-        pool_recycle=300,
-        pool_size=5,
-        max_overflow=5,
+        # Coordinators and workers spend most of their lifetime calculating
+        # or waiting. Release each database session after its short operation
+        # instead of reserving a pooler slot throughout that idle interval.
+        poolclass=NullPool,
         connect_args={"connect_timeout": 10},
     )
 
 
+class ConnectionOpener(Protocol):
+    """Open the database connection that one store operation uses.
+
+    ``connect`` only opens a connection; it sends no application statement.
+    The operation's statements, and the COMMIT or ROLLBACK that ends it, run
+    after the opener returns. An opener may therefore retry ``connect``
+    without risking a duplicate write, and never sees a later failure.
+    """
+
+    def __call__(
+        self,
+        connect: Callable[[], Connection],
+        *,
+        operation_name: str,
+    ) -> Connection: ...
+
+
+def open_connection_once(
+    connect: Callable[[], Connection],
+    *,
+    operation_name: str,
+) -> Connection:
+    """Open a connection once, so a refused connection fails the operation."""
+
+    return connect()
+
+
 def _database_value(value: Any) -> Any:
     if isinstance(value, StrEnum):
         return value.value
@@ -156,12 +186,31 @@ def __init__(
         database_url: str | None = None,
         *,
         engine: Engine | None = None,
+        open_connection: ConnectionOpener = open_connection_once,
     ) -> None:
         if engine is None:
             if database_url is None:
                 raise ValueError("Stage 12 database URL is required")
             engine = create_stage12_engine(database_url)
         self.engine = engine
+        self._open_connection = open_connection
+
+    @contextmanager
+    def _connection(self, operation_name: str) -> Iterator[Connection]:
+        # Opening the connection is the only step the opener sees. Every
+        # statement, COMMIT and ROLLBACK runs after the opener returns, so
+        # an opener that retries can never repeat work the database received.
+        connection = self._open_connection(
+            self.engine.connect,
+            operation_name=operation_name,
+        )
+        with connection:
+            yield connection
+
+    @contextmanager
+    def _transaction(self, operation_name: str) -> Iterator[Connection]:
+        with self._connection(operation_name) as connection, connection.begin():
+            yield connection
 
     def create_or_resolve_report(
         self,
@@ -174,7 +223,7 @@ def create_or_resolve_report(
             .on_conflict_do_nothing(constraint=REPORT_IDENTITY_CONSTRAINT)
             .returning(*comparison_reports.c)
         )
-        with self.engine.begin() as connection:
+        with self._transaction("create_or_resolve_report") as connection:
             created = connection.execute(statement).mappings().one_or_none()
             if created is not None:
                 return ComparisonReportPersistenceResult(
@@ -203,7 +252,7 @@ def create_or_resolve_report(
             return ComparisonReportPersistenceResult(record=resolved, created=False)
 
     def get_report(self, evaluation_id: UUID) -> ComparisonReportRecord:
-        with self.engine.connect() as connection:
+        with self._connection("get_report") as connection:
             row = (
                 connection.execute(
                     select(comparison_reports).where(
@@ -229,7 +278,7 @@ def list_simulations(
                 comparison_simulations.c.simulation_execution_id,
             )
         )
-        with self.engine.connect() as connection:
+        with self._connection("list_simulations") as connection:
             rows = connection.execute(statement).mappings().all()
         return tuple(_simulation(row) for row in rows)
 
@@ -238,7 +287,7 @@ def replace_report(
         record: ComparisonReportRecord,
     ) -> ComparisonReportRecord:
         values = _record_values(record)
-        with self.engine.begin() as connection:
+        with self._transaction("replace_report") as connection:
             row = (
                 connection.execute(
                     select(comparison_reports)
@@ -280,7 +329,7 @@ def replace_report_result_comparison(
         record: ComparisonReportRecord,
     ) -> ComparisonReportRecord:
         values = _record_values(record)
-        with self.engine.begin() as connection:
+        with self._transaction("replace_report_result_comparison") as connection:
             row = (
                 connection.execute(
                     select(comparison_reports)
@@ -331,7 +380,7 @@ def create_or_resolve_simulation(
             .on_conflict_do_nothing(constraint=SIMULATION_IDENTITY_CONSTRAINT)
             .returning(*comparison_simulations.c)
         )
-        with self.engine.begin() as connection:
+        with self._transaction("create_or_resolve_simulation") as connection:
             parent_row = (
                 connection.execute(
                     select(comparison_reports).where(
@@ -380,7 +429,7 @@ def get_simulation(
         self,
         simulation_execution_id: UUID,
     ) -> ComparisonSimulationRecord:
-        with self.engine.connect() as connection:
+        with self._connection("get_simulation") as connection:
             row = (
                 connection.execute(
                     select(comparison_simulations).where(
@@ -402,7 +451,7 @@ def replace_simulation(
         record: ComparisonSimulationRecord,
     ) -> ComparisonSimulationRecord:
         values = _record_values(record)
-        with self.engine.begin() as connection:
+        with self._transaction("replace_simulation") as connection:
             row = (
                 connection.execute(
                     select(comparison_simulations)
@@ -464,7 +513,7 @@ def attach_simulation_invocation(
             )
             .returning(*comparison_simulations.c)
         )
-        with self.engine.begin() as connection:
+        with self._transaction("attach_simulation_invocation") as connection:
             row = connection.execute(statement).mappings().one_or_none()
             if row is None:
                 row = (