diff --git a/backend/src/backup_tool/execution.py b/backend/src/backup_tool/execution.py index 7dd2d42..49bb2c9 100644 --- a/backend/src/backup_tool/execution.py +++ b/backend/src/backup_tool/execution.py @@ -54,9 +54,9 @@ async def enqueue(db: AsyncSession, job_id: str, trigger: str = "manual") -> Exe job_identifier = job.id execution = Execution(job_id=job_identifier, trigger=trigger, progress={}) db.add(execution) - await db.flush() - await record_event(db, execution) try: + await db.flush() + await record_event(db, execution) await db.commit() except IntegrityError as error: await db.rollback() @@ -69,8 +69,6 @@ async def enqueue(db: AsyncSession, job_id: str, trigger: str = "manual") -> Exe "execution_active", "Job already has an active execution.", existing ) from error await db.refresh(execution) - await record_event(db, execution) - await db.commit() return execution @@ -92,11 +90,10 @@ async def claim( if getattr(result, "rowcount", 0) != 1: await db.rollback() return None - await db.commit() execution = await db.get(Execution, execution_id) if execution is not None: await record_event(db, execution) - await db.commit() + await db.commit() return execution @@ -135,11 +132,10 @@ async def request_cancellation(db: AsyncSession, execution_id: str) -> Execution if getattr(result, "rowcount", 0) != 1: await db.rollback() return None - await db.commit() execution = await db.get(Execution, execution_id) if execution is not None: await record_event(db, execution) - await db.commit() + await db.commit() return execution @@ -159,8 +155,14 @@ async def complete_cancellation(db: AsyncSession, execution_id: str, owner: str) lease_expires_at=None, ) ) + if getattr(result, "rowcount", 0) != 1: + await db.rollback() + return False + execution = await db.get(Execution, execution_id) + if execution is not None: + await record_event(db, execution) await db.commit() - return getattr(result, "rowcount", 0) == 1 + return True async def retry(db: AsyncSession, execution_id: str) -> Execution | None: @@ -176,6 +178,7 @@ async def retry(db: AsyncSession, execution_id: str) -> Execution | None: execution.heartbeat_at = None execution.reason_code = None execution.operator_message = None + await record_event(db, execution) try: await db.commit() except IntegrityError as error: @@ -189,8 +192,6 @@ async def retry(db: AsyncSession, execution_id: str) -> Execution | None: "execution_active", "Job already has an active execution.", active_id ) from error await db.refresh(execution) - await record_event(db, execution) - await db.commit() return execution @@ -237,32 +238,29 @@ def public_event(execution: Execution) -> dict[str, object]: async def recover_stale(db: AsyncSession) -> int: + """Atomically recover expired executions and append one event per mutation.""" now = datetime.now(UTC) - cancelled = await db.execute( - update(Execution) - .where(Execution.state == "cancelling", Execution.lease_expires_at < now) - .values( - state="cancelled", - completed_at=now, - lease_owner=None, - lease_expires_at=None, - heartbeat_at=None, - reason_code="cancellation_requested", - ) - ) - recovered = await db.execute( - update(Execution) - .where( - Execution.state.in_({"preparing", "running", "verifying"}), - Execution.lease_expires_at < now, - ) - .values( - state="queued", - lease_owner=None, - lease_expires_at=None, - heartbeat_at=None, - reason_code="worker_lost", - ) + executions = list( + ( + await db.scalars( + select(Execution).where( + Execution.state.in_({"preparing", "running", "verifying", "cancelling"}), + Execution.lease_expires_at < now, + ) + ) + ).all() ) + for execution in executions: + if execution.state == "cancelling": + execution.state = "cancelled" + execution.completed_at = now + execution.reason_code = "cancellation_requested" + else: + execution.state = "queued" + execution.reason_code = "worker_lost" + execution.lease_owner = None + execution.lease_expires_at = None + execution.heartbeat_at = None + await record_event(db, execution) await db.commit() - return (getattr(cancelled, "rowcount", 0) or 0) + (getattr(recovered, "rowcount", 0) or 0) + return len(executions) diff --git a/backend/src/backup_tool/worker.py b/backend/src/backup_tool/worker.py index ae537e9..5837443 100644 --- a/backend/src/backup_tool/worker.py +++ b/backend/src/backup_tool/worker.py @@ -17,7 +17,14 @@ from sqlalchemy.ext.asyncio import async_sessionmaker from backup_tool.config import Settings from backup_tool.db.engine import create_engine from backup_tool.db.models import Execution -from backup_tool.execution import claim, complete_cancellation, heartbeat, recover_stale, transition +from backup_tool.execution import ( + claim, + complete_cancellation, + heartbeat, + record_event, + recover_stale, + transition, +) class Worker: @@ -51,7 +58,7 @@ class Worker: await complete_cancellation(db, execution.id, self.owner) return True now = datetime.now(UTC) - await db.execute( + result = await db.execute( update(Execution) .where( Execution.id == execution.id, @@ -61,6 +68,9 @@ class Worker: ) .values(state=transition("preparing", "running"), started_at=now) ) + if getattr(result, "rowcount", 0) == 1: + await db.refresh(execution) + await record_event(db, execution) await db.commit() await heartbeat(db, execution.id, self.owner) return True