108 lines
3.9 KiB
Python
108 lines
3.9 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
from collections.abc import Iterator
|
|
from contextlib import contextmanager
|
|
|
|
import pymysql
|
|
from pymysql.cursors import DictCursor
|
|
|
|
from .config import Settings
|
|
from .models import FailureReport
|
|
|
|
|
|
class FailureRepository:
|
|
def __init__(self, settings: Settings) -> None:
|
|
self.settings = settings
|
|
|
|
@contextmanager
|
|
def connection(self) -> Iterator[pymysql.Connection]:
|
|
connection = pymysql.connect(
|
|
host=self.settings.database_host,
|
|
port=self.settings.database_port,
|
|
user=self.settings.database_user,
|
|
password=self.settings.database_password,
|
|
database=self.settings.database_name,
|
|
cursorclass=DictCursor,
|
|
autocommit=False,
|
|
connect_timeout=10,
|
|
read_timeout=30,
|
|
write_timeout=30,
|
|
)
|
|
try:
|
|
yield connection
|
|
finally:
|
|
connection.close()
|
|
|
|
def claim_next(self) -> FailureReport | None:
|
|
with self.connection() as connection, connection.cursor() as cursor:
|
|
try:
|
|
cursor.execute(
|
|
"""
|
|
UPDATE failure_reports
|
|
SET status = 'open', updated_at = NOW()
|
|
WHERE status = 'agent_running'
|
|
AND updated_at < DATE_SUB(NOW(), INTERVAL %s MINUTE)
|
|
""",
|
|
(self.settings.stale_claim_minutes,),
|
|
)
|
|
connection.commit()
|
|
cursor.execute("SET TRANSACTION ISOLATION LEVEL READ COMMITTED")
|
|
cursor.execute("START TRANSACTION")
|
|
except Exception:
|
|
connection.rollback()
|
|
raise
|
|
cursor.execute(
|
|
"""
|
|
SELECT * FROM failure_reports
|
|
WHERE status = 'open'
|
|
ORDER BY id
|
|
LIMIT 1
|
|
FOR UPDATE SKIP LOCKED
|
|
"""
|
|
)
|
|
row = cursor.fetchone()
|
|
if not row:
|
|
connection.commit()
|
|
return None
|
|
cursor.execute(
|
|
"UPDATE failure_reports SET status = 'agent_running', updated_at = NOW() WHERE id = %s",
|
|
(row["id"],),
|
|
)
|
|
connection.commit()
|
|
return FailureReport.from_row(row)
|
|
|
|
def release(self, report_id: int, reason: str) -> None:
|
|
self._finish(report_id, "open", {"last_dry_run": reason})
|
|
|
|
def complete(self, report_id: int, metadata: dict) -> None:
|
|
self._finish(report_id, "resolved", metadata)
|
|
|
|
def fail(self, report_id: int, error: str, metadata: dict | None = None) -> None:
|
|
payload = {"agent_error": error, **(metadata or {})}
|
|
self._finish(report_id, "agent_failed", payload)
|
|
|
|
def _finish(self, report_id: int, status: str, additions: dict) -> None:
|
|
with self.connection() as connection, connection.cursor() as cursor:
|
|
try:
|
|
cursor.execute("SELECT context FROM failure_reports WHERE id = %s", (report_id,))
|
|
row = cursor.fetchone() or {}
|
|
context = row.get("context") or {}
|
|
if isinstance(context, str):
|
|
context = json.loads(context)
|
|
context.update(additions)
|
|
cursor.execute(
|
|
"""
|
|
UPDATE failure_reports
|
|
SET status = %s, context = %s, updated_at = NOW()
|
|
WHERE id = %s
|
|
""",
|
|
(status, json.dumps(context, ensure_ascii=True), report_id),
|
|
)
|
|
connection.commit()
|
|
except Exception as exception:
|
|
connection.rollback()
|
|
raise RuntimeError(
|
|
f"Could not finalize failure report {report_id} as {status}"
|
|
) from exception
|