refactor(db)!: finalize pooled WNDB v2 migration

app/native/wndb/core/database.py

@@ -1,10 +1,19 @@
-from collections.abc import Mapping, Sequence
-from typing import Any
+from collections.abc import Iterator, Mapping, Sequence
+from contextlib import contextmanager
+from typing import Any, Callable
 
 from psycopg import sql
 from psycopg.rows import Row, dict_row
 
-from .connection import is_project_transaction_active, project_connection
+from app.infra.db.project_routing import get_project_database_name
+
+from .connection import (
+    is_model_mutation_lock_active,
+    is_project_transaction_active,
+    mark_model_mutation_lock_active,
+    project_connection,
+    project_transaction,
+)
 
 API_ADD = "add"
 API_UPDATE = "update"
@@ -61,6 +70,20 @@ class DatabaseCommand:
         self.sql = statement
         self.changes = changes
 
+
+class MaterializedViewRefreshAfterCommitError(RuntimeError):
+    """Report a failed view refresh without implying that the write rolled back."""
+
+    changes_committed = True
+
+    def __init__(self, project: str) -> None:
+        self.project = project
+        super().__init__(
+            f"Project {project!r} changes were committed, but materialized view "
+            "refresh failed"
+        )
+
+
 QueryParams = Sequence[Any] | Mapping[str, Any]
 
 
@@ -78,6 +101,27 @@ def _execute(cur, query: str, params: QueryParams | None = None):
     return cur.execute(query, params) if params is not None else cur.execute(query)
 
 
+def acquire_model_mutation_lock(conn, name: str) -> None:
+    """Serialize model replacement and ordinary WNDB mutations per database."""
+    if is_model_mutation_lock_active(name):
+        return
+    physical_name = get_project_database_name(name)
+    with conn.cursor() as cur:
+        cur.execute(
+            "select pg_advisory_xact_lock(hashtextextended(%s, 0))",
+            (f"tjwater:wndb:model:{physical_name}",),
+        )
+    mark_model_mutation_lock_active(name)
+
+
+@contextmanager
+def model_mutation_transaction(name: str) -> Iterator[Any]:
+    """Open a project transaction and acquire its model lock before reading."""
+    with project_transaction(name) as conn:
+        acquire_model_mutation_lock(conn, name)
+        yield conn
+
+
 def read(name: str, query: str, params: QueryParams | None = None) -> Row:
     with project_connection(name) as conn, conn.cursor(row_factory=dict_row) as cur:
         _execute(cur, query, params)
@@ -104,8 +148,15 @@ def try_read(
 
 
 def write(name: str, query: str, params: QueryParams | None = None) -> None:
-    with project_connection(name) as conn, conn.cursor() as cur:
-        _execute(cur, query, params)
+    connection_context = (
+        project_connection(name)
+        if is_project_transaction_active(name)
+        else model_mutation_transaction(name)
+    )
+    with connection_context as conn:
+        acquire_model_mutation_lock(conn, name)
+        with conn.cursor() as cur:
+            _execute(cur, query, params)
 
 
 def refresh_materialized_views(name: str, *, concurrently: bool = True) -> None:
@@ -114,6 +165,13 @@ def refresh_materialized_views(name: str, *, concurrently: bool = True) -> None:
         cur.execute("CALL gis.refresh_all_materialized_views(%s)", (concurrently,))
 
 
+def refresh_materialized_views_after_commit(name: str) -> None:
+    try:
+        refresh_materialized_views(name)
+    except Exception as exc:
+        raise MaterializedViewRefreshAfterCommitError(name) from exc
+
+
 _MATERIALIZED_VIEW_SOURCES = (
     "network.nodes",
     "network.junctions",
@@ -135,9 +193,48 @@ def _affects_materialized_views(command: DatabaseCommand) -> bool:
     return any(source in statement for source in _MATERIALIZED_VIEW_SOURCES)
 
 
+_MATERIALIZED_VIEW_ELEMENT_TYPES = frozenset(
+    {
+        "junction",
+        "reservoir",
+        "tank",
+        "pipe",
+        "pump",
+        "valve",
+        "demand",
+        "vertex",
+    }
+)
+
+
+def changes_affect_materialized_views(change_set: ChangeSet) -> bool:
+    """Return whether a dispatched WNDB batch changes a published GIS source."""
+    return any(
+        operation.get("type") in _MATERIALIZED_VIEW_ELEMENT_TYPES
+        for operation in change_set.operations
+    )
+
+
 def execute_command(name: str, command: DatabaseCommand) -> ChangeSet:
     """Apply a model mutation without the removed database undo/redo journal."""
     write(name, command.sql)
     if _affects_materialized_views(command) and not is_project_transaction_active(name):
-        refresh_materialized_views(name)
+        refresh_materialized_views_after_commit(name)
     return ChangeSet.from_list(command.changes)
+
+
+def execute_locked_command(
+    name: str,
+    builder: Callable[[], DatabaseCommand | None],
+) -> ChangeSet:
+    """Build a read-modify-write command only after acquiring the model lock."""
+    nested_transaction = is_project_transaction_active(name)
+    command: DatabaseCommand | None = None
+    with model_mutation_transaction(name):
+        command = builder()
+        if command is None:
+            return ChangeSet()
+        result = execute_command(name, command)
+    if not nested_transaction and _affects_materialized_views(command):
+        refresh_materialized_views_after_commit(name)
+    return result