fix: use threads for callbacks and avoid session concurrency

lkstrp · lkstrp · commit 18655aeaf161 · 2026-03-17T11:04:35.000+01:00
diff --git a/src/pypsa_app/backend/api/routes/runs.py b/src/pypsa_app/backend/api/routes/runs.py
@@ -2,6 +2,7 @@
 
 import logging
 import re
+import threading
 import urllib.parse
 import uuid
 from pathlib import PurePosixPath
@@ -33,8 +34,15 @@
     RunSummary,
 )
 from pypsa_app.backend.services.backend_registry import backend_registry
+from pypsa_app.backend.services.callback import (
+    _build_payload,
+    post_callback_sync,
+)
 from pypsa_app.backend.services.run import SnakedispatchClient, SnakedispatchError
-from pypsa_app.backend.services.sync import SYNCED_STATUSES, sync_run_from_job
+from pypsa_app.backend.services.sync import (
+    SYNCED_STATUSES,
+    sync_run_from_job,
+)
 from pypsa_app.backend.settings import settings
 
 router = APIRouter()
@@ -125,7 +133,7 @@ def create_run(
 
     payload = body.model_dump(
         exclude_none=True,
-        exclude={"backend_id", "import_networks", "cache"},
+        exclude={"backend_id", "import_networks", "cache", "callback_url"},
     )
     if body.cache:
         payload["cache_key"] = body.cache.key
@@ -142,6 +150,7 @@ def create_run(
         extra_files=body.extra_files,
         cache=body.cache.model_dump() if body.cache else None,
         import_networks=body.import_networks,
+        callback_url=str(body.callback_url) if body.callback_url else None,
         status=RunStatus(result.get("status", "PENDING")),
     )
     db.add(run)
@@ -281,8 +290,18 @@ def get_run(
         if client:
             try:
                 job = client.get_job(str(run_id))
-                sync_run_from_job(run, job, db)
+                needs_callback = sync_run_from_job(run, job, db)
                 db.commit()
+                if needs_callback and run.callback_url:
+                    # TODO: replace with proper async callback or
+                    # FastAPI BackgroundTasks.
+                    url = str(run.callback_url)
+                    payload = _build_payload(run)
+                    threading.Thread(
+                        target=post_callback_sync,
+                        args=(url, payload),
+                        daemon=True,
+                    ).start()
             except SnakedispatchError:
                 pass
 
@@ -378,13 +397,31 @@ def cancel_run(
     sd_client = _get_client_for_run(run)
     try:
         result = sd_client.cancel_job(str(run_id))
-        sync_run_from_job(run, result, db)
+        needs_callback = sync_run_from_job(run, result, db)
         db.commit()
+        if needs_callback and run.callback_url:
+            # TODO: replace with proper async callback or
+            # FastAPI BackgroundTasks.
+            url = str(run.callback_url)
+            payload = _build_payload(run)
+            threading.Thread(
+                target=post_callback_sync, args=(url, payload), daemon=True
+            ).start()
     except SnakedispatchError as e:
         if e.status_code in (404, 409):
             if run.status not in SYNCED_STATUSES:
                 run.status = RunStatus.CANCELLED
                 db.commit()
+                if run.callback_url:
+                    # TODO: replace with proper async callback or
+                    # FastAPI BackgroundTasks.
+                    url = str(run.callback_url)
+                    payload = _build_payload(run)
+                    threading.Thread(
+                        target=post_callback_sync,
+                        args=(url, payload),
+                        daemon=True,
+                    ).start()
         else:
             raise
 
diff --git a/src/pypsa_app/backend/models.py b/src/pypsa_app/backend/models.py
@@ -294,6 +294,8 @@ class Run(Base):
     extra_files: Mapped[Any | None] = mapped_column(JSON)
     cache: Mapped[Any | None] = mapped_column(JSON)
 
+    callback_url: Mapped[str | None] = mapped_column(String(512))
+
     # Job metadata (synced from Snakedispatch)
     git_ref: Mapped[str | None] = mapped_column(String(255))
     git_sha: Mapped[str | None] = mapped_column(String(40))
diff --git a/src/pypsa_app/backend/schemas/run.py b/src/pypsa_app/backend/schemas/run.py
@@ -2,13 +2,15 @@
 
 import uuid
 from datetime import datetime
+from urllib.parse import urlparse
 
-from pydantic import BaseModel, ConfigDict, Field
+from pydantic import BaseModel, ConfigDict, Field, HttpUrl, field_validator
 
 from pypsa_app.backend.models import RunStatus
 from pypsa_app.backend.schemas.auth import UserPublicResponse
 from pypsa_app.backend.schemas.backend import BackendPublicResponse
 from pypsa_app.backend.schemas.common import PaginationMeta
+from pypsa_app.backend.settings import settings
 
 
 class RunCache(BaseModel):
@@ -46,6 +48,22 @@ class RunCreate(BaseModel):
     cache: RunCache | None = None
     import_networks: list[str] | None = None
     backend_id: uuid.UUID | None = None
+    callback_url: HttpUrl | None = None
+
+    @field_validator("callback_url")
+    @classmethod
+    def _validate_callback_domain(cls, v: HttpUrl | None) -> HttpUrl | None:
+        if v is None:
+            return v
+        allowed = settings.resolved_callback_domains
+        if not allowed:
+            msg = "Callbacks are not enabled on this server"
+            raise ValueError(msg)
+        host = v.host or ""
+        if not any(host == d or host.endswith(f".{d}") for d in allowed):
+            msg = f"callback_url host '{host}' is not in the allowed domains"
+            raise ValueError(msg)
+        return v
 
 
 class RunSummary(BaseModel):
@@ -75,9 +93,18 @@ class RunResponse(RunSummary):
     extra_files: dict[str, str] | None = None
     cache: RunCache | None = None
     import_networks: list[str] | None = None
+    callback_url: str | None = Field(None, validation_alias="callback_url")
     exit_code: int | None = None
     networks: list[RunNetworkSummary] = []
 
+    @field_validator("callback_url", mode="before")
+    @classmethod
+    def _redact_callback_url(cls, v: str | None) -> str | None:
+        if not v:
+            return None
+        parsed = urlparse(v)
+        return f"{parsed.scheme}://{parsed.hostname}/***"
+
 
 class RunListMeta(PaginationMeta):
     """Extended pagination meta with run-specific filter options."""
diff --git a/src/pypsa_app/backend/services/callback.py b/src/pypsa_app/backend/services/callback.py
@@ -0,0 +1,47 @@
+"""Callback helpers for notifying external systems of run status changes."""
+
+import logging
+
+import httpx
+
+from pypsa_app.backend.models import Run
+
+logger = logging.getLogger(__name__)
+
+
+def _build_payload(run: Run) -> dict:
+    return {"run_id": str(run.job_id), "status": run.status.value}
+
+
+def post_callback_sync(url: str, payload: dict) -> None:
+    """POST to a callback URL (blocking)."""
+    try:
+        httpx.post(url, json=payload, timeout=5.0, follow_redirects=False)
+    except Exception:
+        logger.warning(
+            "Callback failed for run %s to %s",
+            payload.get("run_id"),
+            url,
+            exc_info=True,
+        )
+
+
+def fire_callback_sync(run: Run) -> None:
+    """POST to the run's callback URL (blocking)."""
+    if not run.callback_url:
+        return
+    post_callback_sync(str(run.callback_url), _build_payload(run))
+
+
+async def fire_callback_async(url: str, payload: dict) -> None:
+    """POST to a callback URL (async). Used by the background sync loop."""
+    try:
+        async with httpx.AsyncClient() as client:
+            await client.post(url, json=payload, timeout=5.0, follow_redirects=False)
+    except Exception:
+        logger.warning(
+            "Callback failed for run %s to %s",
+            payload.get("run_id"),
+            url,
+            exc_info=True,
+        )
diff --git a/src/pypsa_app/backend/services/sync.py b/src/pypsa_app/backend/services/sync.py
@@ -12,10 +12,14 @@
 from pypsa_app.backend.database import SessionLocal
 from pypsa_app.backend.models import Run, RunStatus
 from pypsa_app.backend.services.backend_registry import backend_registry
+from pypsa_app.backend.services.callback import fire_callback_async
 from pypsa_app.backend.tasks import import_run_outputs_task
 
 logger = logging.getLogger(__name__)
 
+# Hold references to fire-and-forget callback tasks to prevent garbage collection.
+_background_tasks: set[asyncio.Task] = set()
+
 # Statuses where the remote executor is done, no need to sync from Snakedispatch
 SYNCED_STATUSES = {
     RunStatus.UPLOADING,
@@ -41,8 +45,16 @@
 ]
 
 
-def sync_run_from_job(run: Run, job: dict, db: Session) -> None:
-    """Update a Run record from a Snakedispatch response dict."""
+_CALLBACK_STATUSES = SYNCED_STATUSES - {RunStatus.UPLOADING}
+
+
+def sync_run_from_job(run: Run, job: dict, db: Session) -> bool:
+    """Update a Run record from a Snakedispatch response dict.
+
+    Returns:
+        True if a callback should be fired after the transaction commits.
+    """
+    old_status = run.status
     changed = False
     for field in _SYNC_FIELDS:
         new_val = job.get(field)
@@ -65,7 +77,7 @@ def sync_run_from_job(run: Run, job: dict, db: Session) -> None:
             run.status = RunStatus.UPLOADING
             db.flush()
             import_run_outputs_task.apply_async(args=(str(run.job_id),))
-            return
+            return False
         if completed_with_import_pending:
             run.status = RunStatus.COMPLETED
             changed = True
@@ -76,38 +88,65 @@ def sync_run_from_job(run: Run, job: dict, db: Session) -> None:
     if changed:
         db.flush()
 
+    return run.status in _CALLBACK_STATUSES and old_status not in _CALLBACK_STATUSES
+
+
+def sync_non_terminal_runs() -> list[dict]:
+    """Poll all backends and update runs that haven't reached a terminal state.
 
-def sync_non_terminal_runs() -> None:
-    """Poll all backends and update runs that haven't reached a terminal state."""
+    Returns:
+        List of callback dicts ``{"url": ..., "payload": ...}`` to be fired
+        by the async caller after the DB session is closed.
+    """
+    callbacks: list[dict] = []
     db = SessionLocal()
     try:
         non_terminal = db.query(Run).filter(Run.status.notin_(SYNCED_STATUSES)).all()
         if not non_terminal:
-            return
+            return callbacks
 
         for backend_id, client in backend_registry.all_clients().items():
             backend_runs = [r for r in non_terminal if r.backend_id == backend_id]
             if not backend_runs:
                 continue
             try:
                 jobs_by_id = {j["job_id"]: j for j in client.list_jobs()}
+                callback_runs: list[Run] = []
                 for run in backend_runs:
                     job = jobs_by_id.get(str(run.job_id))
-                    if job:
-                        sync_run_from_job(run, job, db)
+                    if job and sync_run_from_job(run, job, db):
+                        callback_runs.append(run)
                 db.commit()
+                callbacks.extend(
+                    {
+                        "url": str(run.callback_url),
+                        "payload": {
+                            "run_id": str(run.job_id),
+                            "status": run.status.value,
+                        },
+                    }
+                    for run in callback_runs
+                    if run.callback_url
+                )
             except Exception:
                 db.rollback()
                 logger.warning("Sync failed for backend %s", backend_id, exc_info=True)
     finally:
         db.close()
+    return callbacks
 
 
 async def run_sync_loop(interval: float = 10.0) -> None:
     """Periodically sync non-terminal runs in a background thread."""
     while True:
         await asyncio.sleep(interval)
         try:
-            await asyncio.to_thread(sync_non_terminal_runs)
+            callbacks = await asyncio.to_thread(sync_non_terminal_runs)
+            for cb in callbacks:
+                task = asyncio.create_task(
+                    fire_callback_async(cb["url"], cb["payload"])
+                )
+                _background_tasks.add(task)
+                task.add_done_callback(_background_tasks.discard)
         except Exception:
             logger.warning("Background run sync failed", exc_info=True)
diff --git a/src/pypsa_app/backend/settings.py b/src/pypsa_app/backend/settings.py
@@ -97,6 +97,26 @@ def networks_path(self) -> Path:
         description="Interval in seconds between background Snakedispatch sync cycles",
         json_schema_extra={"category": "Runs"},
     )
+    callback_url_allowed_domains: str = Field(
+        default="",
+        description=(
+            "Comma-separated list of allowed domains for run callback URLs "
+            "(e.g. hooks.myorg.dev,example.com). "
+            "Callbacks are rejected unless the host matches. "
+            "Empty disables callbacks entirely."
+        ),
+        json_schema_extra={"category": "Runs"},
+    )
+
+    @property
+    def resolved_callback_domains(self) -> list[str]:
+        """Parse CALLBACK_URL_ALLOWED_DOMAINS into a list of domain strings."""
+        if not self.callback_url_allowed_domains:
+            return []
+        return [
+            d.strip() for d in self.callback_url_allowed_domains.split(",") if d.strip()
+        ]
+
     snakedispatch_backends: str | None = Field(
         default=None,
         description=(
diff --git a/src/pypsa_app/backend/tasks.py b/src/pypsa_app/backend/tasks.py
