269 lines
9.2 KiB
Python
269 lines
9.2 KiB
Python
from __future__ import annotations
|
|
|
|
from collections.abc import Mapping
|
|
from datetime import UTC, datetime, timedelta
|
|
|
|
from sqlalchemy import select, update
|
|
from sqlalchemy.exc import IntegrityError
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from backup_tool.db.models import Execution, ExecutionEvent, Job
|
|
from backup_tool.security.redaction import redact
|
|
|
|
ACTIVE_STATES = frozenset({"queued", "preparing", "running", "verifying", "cancelling"})
|
|
TERMINAL_STATES = frozenset({"committed", "cancelled", "failed"})
|
|
TRANSIENT_REASONS = frozenset({"worker_lost", "timeout", "unavailable", "transient_io"})
|
|
|
|
_ALLOWED: Mapping[str, frozenset[str]] = {
|
|
"queued": frozenset({"preparing", "cancelled", "failed"}),
|
|
"preparing": frozenset({"running", "cancelling", "failed"}),
|
|
"running": frozenset({"verifying", "cancelling", "failed"}),
|
|
"verifying": frozenset({"committed", "failed"}),
|
|
"cancelling": frozenset({"cancelled", "failed"}),
|
|
"committed": frozenset(),
|
|
"cancelled": frozenset(),
|
|
"failed": frozenset(),
|
|
}
|
|
|
|
|
|
class TransitionError(ValueError):
|
|
code = "invalid_transition"
|
|
|
|
|
|
class EnqueueError(ValueError):
|
|
def __init__(self, code: str, detail: str, active_execution_id: str | None = None):
|
|
super().__init__(detail)
|
|
self.code = code
|
|
self.active_execution_id = active_execution_id
|
|
|
|
|
|
def transition(current: str, target: str) -> str:
|
|
"""Validate and return one durable execution-state transition."""
|
|
if target not in _ALLOWED.get(current, frozenset()):
|
|
raise TransitionError(f"Cannot transition execution from {current!r} to {target!r}.")
|
|
return target
|
|
|
|
|
|
async def enqueue(db: AsyncSession, job_id: str, trigger: str = "manual") -> Execution:
|
|
"""Create exactly one active execution for an enabled active job."""
|
|
job = await db.get(Job, job_id)
|
|
if job is None:
|
|
raise EnqueueError("resource_not_found", "Job was not found.")
|
|
if job.state != "active" or not job.enabled:
|
|
raise EnqueueError("job_disabled", "Job is disabled or archived.")
|
|
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.commit()
|
|
except IntegrityError as error:
|
|
await db.rollback()
|
|
existing = await db.scalar(
|
|
select(Execution.id).where(
|
|
Execution.job_id == job_identifier, Execution.state.in_(ACTIVE_STATES)
|
|
)
|
|
)
|
|
raise EnqueueError(
|
|
"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
|
|
|
|
|
|
async def claim(
|
|
db: AsyncSession, execution_id: str, owner: str, lease_seconds: int = 60
|
|
) -> Execution | None:
|
|
"""Atomically lease a queued execution; a stale lease may be recovered."""
|
|
now = datetime.now(UTC)
|
|
expires = now + timedelta(seconds=lease_seconds)
|
|
result = await db.execute(
|
|
update(Execution)
|
|
.where(
|
|
Execution.id == execution_id,
|
|
Execution.state.in_({"queued", "preparing"}),
|
|
(Execution.lease_expires_at.is_(None)) | (Execution.lease_expires_at < now),
|
|
)
|
|
.values(state="preparing", lease_owner=owner, lease_expires_at=expires, heartbeat_at=now)
|
|
)
|
|
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()
|
|
return execution
|
|
|
|
|
|
async def heartbeat(
|
|
db: AsyncSession, execution_id: str, owner: str, lease_seconds: int = 60
|
|
) -> bool:
|
|
now = datetime.now(UTC)
|
|
result = await db.execute(
|
|
update(Execution)
|
|
.where(
|
|
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))
|
|
)
|
|
await db.commit()
|
|
return getattr(result, "rowcount", 0) == 1
|
|
|
|
|
|
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 == "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
|
|
await db.commit()
|
|
execution = await db.get(Execution, execution_id)
|
|
if execution is not None:
|
|
await record_event(db, execution)
|
|
await db.commit()
|
|
return execution
|
|
|
|
|
|
async def complete_cancellation(db: AsyncSession, execution_id: str, owner: str) -> bool:
|
|
result = await db.execute(
|
|
update(Execution)
|
|
.where(
|
|
Execution.id == execution_id,
|
|
Execution.lease_owner == owner,
|
|
Execution.lease_expires_at >= datetime.now(UTC),
|
|
Execution.state == "cancelling",
|
|
)
|
|
.values(
|
|
state="cancelled",
|
|
completed_at=datetime.now(UTC),
|
|
lease_owner=None,
|
|
lease_expires_at=None,
|
|
)
|
|
)
|
|
await db.commit()
|
|
return getattr(result, "rowcount", 0) == 1
|
|
|
|
|
|
async def retry(db: AsyncSession, execution_id: str) -> Execution | None:
|
|
execution = await db.get(Execution, execution_id)
|
|
if execution is None:
|
|
return None
|
|
if execution.state != "failed" or execution.reason_code not in TRANSIENT_REASONS:
|
|
raise EnqueueError("retry_not_allowed", "Execution failure is not retryable.")
|
|
execution.state = "queued"
|
|
execution.attempt += 1
|
|
execution.lease_owner = None
|
|
execution.lease_expires_at = None
|
|
execution.heartbeat_at = None
|
|
execution.reason_code = None
|
|
execution.operator_message = None
|
|
try:
|
|
await db.commit()
|
|
except IntegrityError as error:
|
|
await db.rollback()
|
|
active_id = await db.scalar(
|
|
select(Execution.id).where(
|
|
Execution.job_id == execution.job_id, Execution.state.in_(ACTIVE_STATES)
|
|
)
|
|
)
|
|
raise EnqueueError(
|
|
"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
|
|
|
|
|
|
def _redact_progress(value: object) -> object:
|
|
"""Apply common redaction plus execution-specific path redaction recursively."""
|
|
if isinstance(value, Mapping):
|
|
safe: dict[str, object] = {}
|
|
for key, item in value.items():
|
|
normalized = str(key).lower()
|
|
if any(marker in normalized for marker in ("path", "secret", "token", "password")):
|
|
safe[str(key)] = "[REDACTED]"
|
|
else:
|
|
safe[str(key)] = _redact_progress(item)
|
|
return safe
|
|
if isinstance(value, list):
|
|
return [_redact_progress(item) for item in value]
|
|
return redact(value)
|
|
|
|
|
|
async def record_event(db: AsyncSession, execution: Execution) -> ExecutionEvent:
|
|
"""Append the public representation as the next durable event revision."""
|
|
execution.revision += 1
|
|
event = ExecutionEvent(
|
|
execution_id=execution.id,
|
|
revision=execution.revision,
|
|
payload=public_event(execution),
|
|
)
|
|
db.add(event)
|
|
await db.flush()
|
|
return event
|
|
|
|
|
|
def public_event(execution: Execution) -> dict[str, object]:
|
|
"""Return redacted progress suitable for polling or SSE."""
|
|
progress = _redact_progress(execution.progress)
|
|
return {
|
|
"id": execution.id,
|
|
"state": execution.state,
|
|
"attempt": execution.attempt,
|
|
"revision": execution.revision,
|
|
"reason_code": execution.reason_code,
|
|
"progress": progress,
|
|
}
|
|
|
|
|
|
async def recover_stale(db: AsyncSession) -> int:
|
|
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",
|
|
)
|
|
)
|
|
await db.commit()
|
|
return (getattr(cancelled, "rowcount", 0) or 0) + (getattr(recovered, "rowcount", 0) or 0)
|