Merge pull request #238 from vipyrsec/jobs-cte

Use a CTE for the `/POST jobs` endpoint

src/mainframe/endpoints/job.py

@@ -3,8 +3,8 @@
 
 import structlog
 from fastapi import APIRouter, Depends
-from sqlalchemy import and_, or_, select
-from sqlalchemy.orm import Session, joinedload
+from sqlalchemy import and_, or_, select, update
+from sqlalchemy.orm import Session, joinedload, aliased
 
 from mainframe.constants import mainframe_settings
 from mainframe.database import get_db
@@ -41,33 +41,38 @@ def get_jobs(
     """
 
     with session, session.begin():
-        scans = (
-            session.scalars(
-                select(Scan)
-                .where(
-                    or_(
-                        Scan.status == Status.QUEUED,
-                        and_(
-                            Scan.pending_at
-                            < datetime.now(timezone.utc) - timedelta(seconds=mainframe_settings.job_timeout),
-                            Scan.status == Status.PENDING,
-                        ),
-                    )
+        # Use a CTE to limit the number of rows we fetch
+        cte = (
+            select(Scan)
+            .where(
+                or_(
+                    Scan.status == Status.QUEUED,
+                    and_(
+                        Scan.pending_at
+                        < datetime.now(timezone.utc) - timedelta(seconds=mainframe_settings.job_timeout),
+                        Scan.status == Status.PENDING,
+                    ),
                 )
-                .order_by(Scan.pending_at.nulls_first(), Scan.queued_at)
-                .limit(batch)
-                .options(joinedload(Scan.download_urls))
             )
-            .unique()
-            .all()
+            .order_by(Scan.pending_at.nulls_first(), Scan.queued_at)
+            .limit(batch)
+            .options(joinedload(Scan.download_urls))
+            .with_for_update(skip_locked=True)
+            .cte()
+        )
+
+        scan_cte = aliased(Scan, cte)
+
+        # Uses a Postgres `UPDATE .. FROM`. https://docs.sqlalchemy.org/en/20/tutorial/data_update.html#update-from
+        scans = session.scalars(
+            update(Scan)
+            .where(Scan.scan_id == scan_cte.scan_id)
+            .values(status=Status.PENDING, pending_at=datetime.now(timezone.utc), pending_by=auth.subject)
+            .returning(Scan)
         )
 
         response_body: list[JobResult] = []
         for scan in scans:
-            scan.status = Status.PENDING
-            scan.pending_at = datetime.now(timezone.utc)
-            scan.pending_by = auth.subject
-
             logger.info(
                 "Job given and status set to pending in database",
                 package={