Prune old job executions

repub/job_retention.py

@@ -0,0 +1,220 @@
+from __future__ import annotations
+
+import sys
+from dataclasses import dataclass
+from datetime import UTC, datetime, timedelta
+from pathlib import Path
+from typing import TextIO, cast
+
+from repub.db import get_database_connection, initialize_database
+from repub.jobs import JobArtifacts
+from repub.model import Job, JobExecution, JobExecutionStatus, database
+
+DEFAULT_SUCCESSFUL_EXECUTION_RETENTION_DAYS = 7
+DEFAULT_UNSUCCESSFUL_EXECUTION_RETENTION_DAYS = 90
+UNSUCCESSFUL_EXECUTION_STATUSES = (
+    JobExecutionStatus.FAILED,
+    JobExecutionStatus.CANCELED,
+)
+
+
+@dataclass
+class JobExecutionRetentionResult:
+    log_dir: Path
+    successful_cutoff: datetime
+    unsuccessful_cutoff: datetime
+    dry_run: bool
+    matched_executions: int = 0
+    deleted_executions: int = 0
+    matched_files: int = 0
+    deleted_files: int = 0
+    bytes_deleted: int = 0
+    failures: int = 0
+
+
+@dataclass(frozen=True)
+class _ExecutionRetentionCandidate:
+    execution_id: int
+    job_id: int
+    status: JobExecutionStatus
+    ended_at: datetime
+
+
+def cleanup_job_executions(
+    *,
+    log_dir: str | Path,
+    successful_days: int = DEFAULT_SUCCESSFUL_EXECUTION_RETENTION_DAYS,
+    unsuccessful_days: int = DEFAULT_UNSUCCESSFUL_EXECUTION_RETENTION_DAYS,
+    now: datetime | None = None,
+    dry_run: bool = False,
+    output: TextIO = sys.stdout,
+) -> JobExecutionRetentionResult:
+    if get_database_connection() is None:
+        initialize_database()
+    reference_time = _coerce_datetime(now or datetime.now(UTC))
+    successful_cutoff = reference_time - timedelta(days=successful_days)
+    unsuccessful_cutoff = reference_time - timedelta(days=unsuccessful_days)
+    resolved_log_dir = Path(log_dir).resolve()
+    result = JobExecutionRetentionResult(
+        log_dir=resolved_log_dir,
+        successful_cutoff=successful_cutoff,
+        unsuccessful_cutoff=unsuccessful_cutoff,
+        dry_run=dry_run,
+    )
+
+    candidates = _retention_candidates(
+        successful_cutoff=successful_cutoff,
+        unsuccessful_cutoff=unsuccessful_cutoff,
+    )
+    execution_ids_to_delete: list[int] = []
+    for candidate in candidates:
+        result.matched_executions += 1
+        print(
+            "job retention: matched "
+            f"execution_id={candidate.execution_id} "
+            f"job_id={candidate.job_id} "
+            f"status={candidate.status.name} "
+            f"ended_at={candidate.ended_at.isoformat()}",
+            file=output,
+        )
+        artifacts = JobArtifacts.for_execution(
+            log_dir=resolved_log_dir,
+            job_id=candidate.job_id,
+            execution_id=candidate.execution_id,
+        )
+        artifact_cleanup_succeeded = _cleanup_artifacts(
+            artifacts=artifacts,
+            result=result,
+            dry_run=dry_run,
+            output=output,
+        )
+        if dry_run or not artifact_cleanup_succeeded:
+            continue
+        execution_ids_to_delete.append(candidate.execution_id)
+
+    if execution_ids_to_delete:
+        with database.writer():
+            execution_primary_key = getattr(JobExecution, "_meta").primary_key
+            result.deleted_executions = (
+                JobExecution.delete()
+                .where(execution_primary_key.in_(tuple(execution_ids_to_delete)))
+                .execute()
+            )
+
+    print(
+        "job retention: "
+        f"dry_run={_bool_text(result.dry_run)} "
+        f"successful_cutoff={result.successful_cutoff.isoformat()} "
+        f"unsuccessful_cutoff={result.unsuccessful_cutoff.isoformat()} "
+        f"root={result.log_dir} "
+        f"matched_executions={result.matched_executions} "
+        f"deleted_executions={result.deleted_executions} "
+        f"matched_files={result.matched_files} "
+        f"deleted_files={result.deleted_files} "
+        f"bytes_deleted={result.bytes_deleted} "
+        f"failures={result.failures}",
+        file=output,
+    )
+    return result
+
+
+def _retention_candidates(
+    *, successful_cutoff: datetime, unsuccessful_cutoff: datetime
+) -> tuple[_ExecutionRetentionCandidate, ...]:
+    with database.reader():
+        executions = tuple(
+            JobExecution.select(JobExecution, Job)
+            .join(Job)
+            .where(
+                (
+                    JobExecution.running_status.in_(
+                        (
+                            JobExecutionStatus.SUCCEEDED,
+                            *UNSUCCESSFUL_EXECUTION_STATUSES,
+                        )
+                    )
+                )
+                & (JobExecution.ended_at.is_null(False))
+            )
+        )
+
+    candidates: list[_ExecutionRetentionCandidate] = []
+    for execution in executions:
+        status = JobExecutionStatus(int(execution.running_status))
+        ended_at = _coerce_datetime(cast(datetime | str, execution.ended_at))
+        if status == JobExecutionStatus.SUCCEEDED:
+            if ended_at >= successful_cutoff:
+                continue
+        elif status in UNSUCCESSFUL_EXECUTION_STATUSES:
+            if ended_at >= unsuccessful_cutoff:
+                continue
+        else:
+            continue
+
+        job = cast(Job, execution.job)
+        candidates.append(
+            _ExecutionRetentionCandidate(
+                execution_id=int(execution.get_id()),
+                job_id=int(job.get_id()),
+                status=status,
+                ended_at=ended_at,
+            )
+        )
+    return tuple(candidates)
+
+
+def _cleanup_artifacts(
+    *,
+    artifacts: JobArtifacts,
+    result: JobExecutionRetentionResult,
+    dry_run: bool,
+    output: TextIO,
+) -> bool:
+    succeeded = True
+    for path in artifacts.existing_paths():
+        result.matched_files += 1
+        try:
+            file_size = path.stat().st_size
+        except OSError as error:
+            result.failures += 1
+            succeeded = False
+            print(
+                f"job retention: stat failed path={path} error={error}",
+                file=output,
+            )
+            continue
+
+        if dry_run:
+            continue
+
+        try:
+            path.unlink()
+        except FileNotFoundError:
+            continue
+        except OSError as error:
+            result.failures += 1
+            succeeded = False
+            print(
+                f"job retention: delete failed path={path} error={error}",
+                file=output,
+            )
+            continue
+        result.deleted_files += 1
+        result.bytes_deleted += file_size
+    return succeeded
+
+
+def _coerce_datetime(value: datetime | str) -> datetime:
+    if isinstance(value, datetime):
+        if value.tzinfo is None:
+            return value.replace(tzinfo=UTC)
+        return value.astimezone(UTC)
+
+    parsed = datetime.fromisoformat(value)
+    if parsed.tzinfo is None:
+        return parsed.replace(tzinfo=UTC)
+    return parsed.astimezone(UTC)
+
+
+def _bool_text(value: bool) -> str:
+    return "true" if value else "false"