diff --git a/backend/src/backup_tool/execution.py b/backend/src/backup_tool/execution.py index 4f4fc0d..70387bb 100644 --- a/backend/src/backup_tool/execution.py +++ b/backend/src/backup_tool/execution.py @@ -8,6 +8,7 @@ from sqlalchemy.exc import IntegrityError from sqlalchemy.ext.asyncio import AsyncSession from backup_tool.db.models import Execution, Job +from backup_tool.security.redaction import redact ACTIVE_STATES = frozenset({"queued", "preparing", "running", "verifying", "cancelling"}) TERMINAL_STATES = frozenset({"committed", "cancelled", "failed"}) @@ -101,6 +102,7 @@ async def heartbeat( Execution.id == execution_id, Execution.lease_owner == owner, Execution.state.in_(ACTIVE_STATES), + Execution.lease_expires_at >= now, ) .values(heartbeat_at=now, lease_expires_at=now + timedelta(seconds=lease_seconds)) ) @@ -109,14 +111,19 @@ async def heartbeat( async def request_cancellation(db: AsyncSession, execution_id: str) -> Execution | None: + # Queued work has no worker and may terminate immediately; leased work asks its owner. + now = datetime.now(UTC) result = await db.execute( update(Execution) - .where( - Execution.id == execution_id, - Execution.state.in_({"queued", "preparing", "running"}), - ) - .values(state="cancelling", reason_code="cancellation_requested") + .where(Execution.id == execution_id, Execution.state == "queued") + .values(state="cancelled", reason_code="cancellation_requested", completed_at=now) ) + if getattr(result, "rowcount", 0) != 1: + result = await db.execute( + update(Execution) + .where(Execution.id == execution_id, Execution.state.in_({"preparing", "running"})) + .values(state="cancelling", reason_code="cancellation_requested") + ) if getattr(result, "rowcount", 0) != 1: await db.rollback() return None @@ -130,6 +137,7 @@ async def complete_cancellation(db: AsyncSession, execution_id: str, owner: str) .where( Execution.id == execution_id, Execution.lease_owner == owner, + Execution.lease_expires_at >= datetime.now(UTC), Execution.state == "cancelling", ) .values( @@ -163,7 +171,7 @@ async def retry(db: AsyncSession, execution_id: str) -> Execution | None: def public_event(execution: Execution) -> dict[str, object]: """Return redacted progress suitable for polling or SSE.""" - progress = {k: v for k, v in execution.progress.items() if k not in {"path", "secret", "token"}} + progress = redact(execution.progress) return { "id": execution.id, "state": execution.state, @@ -177,7 +185,10 @@ async def recover_stale(db: AsyncSession) -> int: now = datetime.now(UTC) result = await db.execute( update(Execution) - .where(Execution.state.in_({"preparing", "running"}), Execution.lease_expires_at < now) + .where( + Execution.state.in_({"preparing", "running", "verifying", "cancelling"}), + Execution.lease_expires_at < now, + ) .values( state="queued", lease_owner=None, diff --git a/tests/integration/test_sources_jobs.py b/tests/integration/test_sources_jobs.py index a2f011e..8b5ce6f 100644 --- a/tests/integration/test_sources_jobs.py +++ b/tests/integration/test_sources_jobs.py @@ -12,7 +12,9 @@ PASSWORD = "correct-horse-battery-staple" async def login(client: httpx.AsyncClient) -> dict[str, str]: - response = await client.post("/api/v2/setup", json={"username": "admin", "password": PASSWORD}) + response = await client.post( + "/api/v2/setup", json={"username": "admin", "password": PASSWORD} + ) assert response.status_code == 201 return {"X-CSRF-Token": client.cookies["backup_tool_csrf"]} @@ -47,7 +49,9 @@ async def test_local_source_probe_archive_and_repository_targeted_job( async with app.state.engine.begin() as connection: await connection.run_sync(Base.metadata.create_all) transport = httpx.ASGITransport(app=app) - async with httpx.AsyncClient(transport=transport, base_url="https://test") as client: + async with httpx.AsyncClient( + transport=transport, base_url="https://test" + ) as client: headers = await login(client) repository = await client.post( "/api/v2/repositories", @@ -95,14 +99,16 @@ async def test_local_source_probe_archive_and_repository_targeted_job( ) assert duplicate.status_code == 409 assert duplicate.json()["code"] == "execution_active" - polled = await client.get(f"/api/v2/executions/{execution.json()['id']}", headers=headers) + polled = await client.get( + f"/api/v2/executions/{execution.json()['id']}", headers=headers + ) assert polled.status_code == 200 assert polled.json()["state"] == "queued" cancellation = await client.post( f"/api/v2/executions/{execution.json()['id']}/cancel", headers=headers ) assert cancellation.status_code == 202 - assert cancellation.json()["state"] == "cancelling" + assert cancellation.json()["state"] == "cancelled" async with app.state.sessions() as db: stored = await db.get(Execution, execution.json()["id"]) assert stored is not None