501 lines
16 KiB
Python
501 lines
16 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
from typing import Any
|
|
|
|
import sqlalchemy as sa
|
|
from fastapi import APIRouter, Depends, Request
|
|
|
|
from app.api.admin.deps import get_admin_claims
|
|
from app.core.errors import ApiError, ErrorCode
|
|
from app.core.request_id import get_or_create_request_id
|
|
from app.core.time import iso_utc_ms
|
|
from app.db.models.hydration_runs import HydrationRun
|
|
from app.db.models.images import Image
|
|
from app.db.models.jobs import JobRow
|
|
from app.db.session import create_sessionmaker, with_sqlite_busy_retry
|
|
|
|
router = APIRouter()
|
|
|
|
_ALLOWED_RUN_STATUSES = {"pending", "running", "paused", "canceled", "completed", "failed"}
|
|
|
|
|
|
def _parse_json_dict(value: str | None) -> dict[str, Any]:
|
|
raw = str(value or "").strip()
|
|
if not raw:
|
|
return {}
|
|
try:
|
|
loaded = json.loads(raw)
|
|
except Exception:
|
|
return {}
|
|
if isinstance(loaded, dict):
|
|
return loaded
|
|
return {}
|
|
|
|
|
|
def _serialize_job(job: JobRow | None) -> dict[str, Any] | None:
|
|
if job is None:
|
|
return None
|
|
return {
|
|
"id": str(job.id),
|
|
"status": job.status,
|
|
"attempt": int(job.attempt),
|
|
"max_attempts": int(job.max_attempts),
|
|
"run_after": job.run_after,
|
|
"last_error": job.last_error,
|
|
"locked_by": job.locked_by,
|
|
"locked_at": job.locked_at,
|
|
"updated_at": job.updated_at,
|
|
}
|
|
|
|
|
|
def _serialize_run(run: HydrationRun, *, latest_job: JobRow | None = None) -> dict[str, Any]:
|
|
return {
|
|
"id": str(run.id),
|
|
"type": run.type,
|
|
"status": run.status,
|
|
"criteria": _parse_json_dict(run.criteria_json),
|
|
"cursor": _parse_json_dict(run.cursor_json),
|
|
"total": int(run.total) if run.total is not None else None,
|
|
"processed": int(run.processed),
|
|
"success": int(run.success),
|
|
"failed": int(run.failed),
|
|
"started_at": run.started_at,
|
|
"finished_at": run.finished_at,
|
|
"last_error": run.last_error,
|
|
"created_at": run.created_at,
|
|
"updated_at": run.updated_at,
|
|
"latest_job": _serialize_job(latest_job),
|
|
}
|
|
|
|
|
|
async def _latest_jobs_by_run_ids(session, *, run_ids: list[str]) -> dict[str, JobRow]:
|
|
ids = [str(v).strip() for v in run_ids if str(v).strip()]
|
|
if not ids:
|
|
return {}
|
|
|
|
latest_per_ref = (
|
|
sa.select(JobRow.ref_id.label("ref_id"), sa.func.max(JobRow.id).label("max_job_id"))
|
|
.where(JobRow.ref_type == "hydration_run", JobRow.ref_id.in_(ids))
|
|
.group_by(JobRow.ref_id)
|
|
.subquery()
|
|
)
|
|
|
|
rows = (
|
|
await session.execute(
|
|
sa.select(JobRow, latest_per_ref.c.ref_id)
|
|
.join(latest_per_ref, JobRow.id == latest_per_ref.c.max_job_id)
|
|
.order_by(JobRow.id.desc())
|
|
)
|
|
).all()
|
|
|
|
out: dict[str, JobRow] = {}
|
|
for job, ref_id in rows:
|
|
key = str(ref_id or "").strip()
|
|
if not key:
|
|
continue
|
|
out[key] = job
|
|
return out
|
|
|
|
|
|
async def _load_create_json(request: Request) -> dict[str, Any]:
|
|
try:
|
|
data = await request.json()
|
|
except Exception as exc:
|
|
raise ApiError(code=ErrorCode.BAD_REQUEST, message="Invalid JSON body", status_code=400) from exc
|
|
|
|
if not isinstance(data, dict):
|
|
raise ApiError(code=ErrorCode.BAD_REQUEST, message="Invalid JSON body", status_code=400)
|
|
|
|
run_type = str(data.get("type") or "backfill").strip().lower() or "backfill"
|
|
if run_type not in {"backfill", "manual"}:
|
|
raise ApiError(code=ErrorCode.BAD_REQUEST, message="Invalid type", status_code=400)
|
|
|
|
criteria = data.get("criteria")
|
|
if criteria is None:
|
|
criteria = {}
|
|
if not isinstance(criteria, dict):
|
|
raise ApiError(code=ErrorCode.BAD_REQUEST, message="Invalid criteria", status_code=400)
|
|
|
|
return {"type": run_type, "criteria": criteria}
|
|
|
|
|
|
async def _load_manual_job_json(request: Request) -> dict[str, int | None]:
|
|
try:
|
|
data = await request.json()
|
|
except Exception as exc:
|
|
raise ApiError(code=ErrorCode.BAD_REQUEST, message="Invalid JSON body", status_code=400) from exc
|
|
|
|
if not isinstance(data, dict):
|
|
raise ApiError(code=ErrorCode.BAD_REQUEST, message="Invalid JSON body", status_code=400)
|
|
|
|
illust_id_raw = data.get("illust_id")
|
|
image_id_raw = data.get("image_id")
|
|
|
|
if illust_id_raw is None and image_id_raw is None:
|
|
raise ApiError(code=ErrorCode.BAD_REQUEST, message="Either illust_id or image_id is required", status_code=400)
|
|
|
|
if illust_id_raw is not None and image_id_raw is not None:
|
|
raise ApiError(code=ErrorCode.BAD_REQUEST, message="Provide only one of illust_id or image_id", status_code=400)
|
|
|
|
illust_id: int | None = None
|
|
image_id: int | None = None
|
|
if illust_id_raw is not None:
|
|
try:
|
|
illust_id = int(illust_id_raw)
|
|
except Exception as exc:
|
|
raise ApiError(code=ErrorCode.BAD_REQUEST, message="Invalid illust_id", status_code=400) from exc
|
|
if illust_id <= 0:
|
|
raise ApiError(code=ErrorCode.BAD_REQUEST, message="Invalid illust_id", status_code=400)
|
|
|
|
if image_id_raw is not None:
|
|
try:
|
|
image_id = int(image_id_raw)
|
|
except Exception as exc:
|
|
raise ApiError(code=ErrorCode.BAD_REQUEST, message="Invalid image_id", status_code=400) from exc
|
|
if image_id <= 0:
|
|
raise ApiError(code=ErrorCode.BAD_REQUEST, message="Invalid image_id", status_code=400)
|
|
|
|
return {"illust_id": illust_id, "image_id": image_id}
|
|
|
|
|
|
@router.get("/hydration-runs")
|
|
async def list_hydration_runs(
|
|
request: Request,
|
|
limit: int = 20,
|
|
cursor: str | None = None,
|
|
status: str | None = None,
|
|
_claims: dict[str, Any] = Depends(get_admin_claims),
|
|
) -> dict[str, Any]:
|
|
_ = _claims
|
|
if limit < 1 or limit > 200:
|
|
raise ApiError(code=ErrorCode.BAD_REQUEST, message="Unsupported limit", status_code=400)
|
|
|
|
cursor_i: int | None = None
|
|
cursor_raw = str(cursor or "").strip()
|
|
if cursor_raw:
|
|
if not cursor_raw.isdigit():
|
|
raise ApiError(code=ErrorCode.BAD_REQUEST, message="Unsupported cursor", status_code=400)
|
|
cursor_i = int(cursor_raw)
|
|
if cursor_i <= 0:
|
|
raise ApiError(code=ErrorCode.BAD_REQUEST, message="Unsupported cursor", status_code=400)
|
|
|
|
status_norm = str(status or "").strip().lower() or None
|
|
if status_norm is not None and status_norm not in _ALLOWED_RUN_STATUSES:
|
|
raise ApiError(code=ErrorCode.BAD_REQUEST, message="Unsupported status", status_code=400)
|
|
|
|
rid = get_or_create_request_id(request)
|
|
engine = request.app.state.engine
|
|
Session = create_sessionmaker(engine)
|
|
|
|
async with Session() as session:
|
|
stmt = sa.select(HydrationRun).order_by(HydrationRun.id.desc()).limit(limit + 1)
|
|
if cursor_i is not None:
|
|
stmt = stmt.where(HydrationRun.id < cursor_i)
|
|
if status_norm is not None:
|
|
stmt = stmt.where(HydrationRun.status == status_norm)
|
|
|
|
rows = (await session.execute(stmt)).scalars().all()
|
|
current_rows = rows[:limit]
|
|
run_ids = [str(int(r.id)) for r in current_rows]
|
|
jobs_by_run_id = await _latest_jobs_by_run_ids(session, run_ids=run_ids)
|
|
|
|
next_cursor = int(current_rows[-1].id) if len(rows) > limit and current_rows else None
|
|
items = [_serialize_run(row, latest_job=jobs_by_run_id.get(str(int(row.id)))) for row in current_rows]
|
|
|
|
return {
|
|
"ok": True,
|
|
"items": items,
|
|
"next_cursor": str(next_cursor) if next_cursor is not None else "",
|
|
"request_id": rid,
|
|
}
|
|
|
|
|
|
@router.get("/hydration-runs/{run_id}")
|
|
async def get_hydration_run(
|
|
run_id: int,
|
|
request: Request,
|
|
_claims: dict[str, Any] = Depends(get_admin_claims),
|
|
) -> dict[str, Any]:
|
|
_ = _claims
|
|
if run_id <= 0:
|
|
raise ApiError(code=ErrorCode.BAD_REQUEST, message="Invalid hydration_run id", status_code=400)
|
|
|
|
rid = get_or_create_request_id(request)
|
|
engine = request.app.state.engine
|
|
Session = create_sessionmaker(engine)
|
|
|
|
async with Session() as session:
|
|
run = await session.get(HydrationRun, run_id)
|
|
if run is None:
|
|
raise ApiError(code=ErrorCode.NOT_FOUND, message="Hydration run not found", status_code=404)
|
|
|
|
jobs = await _latest_jobs_by_run_ids(session, run_ids=[str(run_id)])
|
|
item = _serialize_run(run, latest_job=jobs.get(str(run_id)))
|
|
|
|
return {"ok": True, "item": item, "request_id": rid}
|
|
|
|
|
|
@router.post("/hydration-runs/manual")
|
|
async def create_manual_hydration_job(
|
|
request: Request,
|
|
_claims: dict[str, Any] = Depends(get_admin_claims),
|
|
) -> dict[str, Any]:
|
|
_ = _claims
|
|
rid = get_or_create_request_id(request)
|
|
now = iso_utc_ms()
|
|
body = await _load_manual_job_json(request)
|
|
|
|
engine = request.app.state.engine
|
|
Session = create_sessionmaker(engine)
|
|
|
|
async def _op() -> dict[str, Any]:
|
|
async with Session() as session:
|
|
illust_id = body.get("illust_id")
|
|
image_id = body.get("image_id")
|
|
if illust_id is None and image_id is not None:
|
|
image = await session.get(Image, int(image_id))
|
|
if image is None:
|
|
raise ApiError(code=ErrorCode.NOT_FOUND, message="Image not found", status_code=404)
|
|
illust_id = int(image.illust_id)
|
|
|
|
if illust_id is None or int(illust_id) <= 0:
|
|
raise ApiError(code=ErrorCode.BAD_REQUEST, message="Invalid illust_id", status_code=400)
|
|
|
|
existing = (
|
|
(
|
|
await session.execute(
|
|
sa.select(JobRow)
|
|
.where(
|
|
JobRow.type == "hydrate_metadata",
|
|
JobRow.ref_type == "manual_hydrate",
|
|
JobRow.ref_id == str(int(illust_id)),
|
|
JobRow.status.in_(["pending", "running"]),
|
|
)
|
|
.order_by(JobRow.id.desc())
|
|
.limit(1)
|
|
)
|
|
)
|
|
.scalars()
|
|
.first()
|
|
)
|
|
|
|
if existing is not None:
|
|
return {
|
|
"ok": True,
|
|
"created": False,
|
|
"job_id": str(int(existing.id)),
|
|
"illust_id": str(int(illust_id)),
|
|
"request_id": rid,
|
|
}
|
|
|
|
job = JobRow(
|
|
type="hydrate_metadata",
|
|
status="pending",
|
|
payload_json=json.dumps({"illust_id": int(illust_id)}, separators=(",", ":"), ensure_ascii=False),
|
|
last_error=None,
|
|
priority=0,
|
|
run_after=None,
|
|
attempt=0,
|
|
max_attempts=3,
|
|
locked_by=None,
|
|
locked_at=None,
|
|
ref_type="manual_hydrate",
|
|
ref_id=str(int(illust_id)),
|
|
updated_at=now,
|
|
)
|
|
session.add(job)
|
|
await session.flush()
|
|
await session.commit()
|
|
|
|
return {
|
|
"ok": True,
|
|
"created": True,
|
|
"job_id": str(int(job.id)),
|
|
"illust_id": str(int(illust_id)),
|
|
"request_id": rid,
|
|
}
|
|
|
|
return await with_sqlite_busy_retry(_op)
|
|
|
|
|
|
@router.post("/hydration-runs")
|
|
async def create_hydration_run(
|
|
request: Request,
|
|
_claims: dict[str, Any] = Depends(get_admin_claims),
|
|
) -> dict[str, Any]:
|
|
_ = _claims
|
|
rid = get_or_create_request_id(request)
|
|
now = iso_utc_ms()
|
|
body = await _load_create_json(request)
|
|
|
|
run_type = str(body["type"])
|
|
criteria = dict(body["criteria"])
|
|
|
|
engine = request.app.state.engine
|
|
Session = create_sessionmaker(engine)
|
|
|
|
async def _op() -> tuple[int, int]:
|
|
async with Session() as session:
|
|
run = HydrationRun(
|
|
type=run_type,
|
|
status="pending",
|
|
criteria_json=json.dumps(criteria, separators=(",", ":"), ensure_ascii=False),
|
|
cursor_json=None,
|
|
total=None,
|
|
processed=0,
|
|
success=0,
|
|
failed=0,
|
|
started_at=None,
|
|
finished_at=None,
|
|
last_error=None,
|
|
updated_at=now,
|
|
)
|
|
session.add(run)
|
|
await session.flush()
|
|
|
|
job = JobRow(
|
|
type="hydrate_metadata",
|
|
status="pending",
|
|
payload_json=json.dumps(
|
|
{"hydration_run_id": int(run.id), "criteria": criteria},
|
|
separators=(",", ":"),
|
|
ensure_ascii=False,
|
|
),
|
|
last_error=None,
|
|
priority=0,
|
|
run_after=None,
|
|
attempt=0,
|
|
max_attempts=3,
|
|
locked_by=None,
|
|
locked_at=None,
|
|
ref_type="hydration_run",
|
|
ref_id=str(int(run.id)),
|
|
updated_at=now,
|
|
)
|
|
session.add(job)
|
|
await session.flush()
|
|
await session.commit()
|
|
return int(run.id), int(job.id)
|
|
|
|
run_id, job_id = await with_sqlite_busy_retry(_op)
|
|
|
|
return {
|
|
"ok": True,
|
|
"hydration_run_id": str(run_id),
|
|
"job_id": str(job_id),
|
|
"request_id": rid,
|
|
}
|
|
|
|
|
|
async def _set_run_and_job_status(
|
|
request: Request,
|
|
*,
|
|
run_id: int,
|
|
target_status: str,
|
|
allowed_from: set[str],
|
|
job_status: str,
|
|
) -> dict[str, Any]:
|
|
if run_id <= 0:
|
|
raise ApiError(code=ErrorCode.BAD_REQUEST, message="Invalid hydration_run id", status_code=400)
|
|
|
|
rid = get_or_create_request_id(request)
|
|
now = iso_utc_ms()
|
|
|
|
engine = request.app.state.engine
|
|
Session = create_sessionmaker(engine)
|
|
|
|
async def _op() -> dict[str, Any]:
|
|
async with Session() as session:
|
|
run = await session.get(HydrationRun, run_id)
|
|
if run is None:
|
|
raise ApiError(code=ErrorCode.NOT_FOUND, message="Hydration run not found", status_code=404)
|
|
|
|
if str(run.status) not in allowed_from:
|
|
raise ApiError(code=ErrorCode.BAD_REQUEST, message="Unsupported status transition", status_code=400)
|
|
|
|
run.status = target_status
|
|
run.updated_at = now
|
|
if target_status in {"canceled", "completed", "failed"} and run.finished_at is None:
|
|
run.finished_at = now
|
|
|
|
job = (
|
|
(
|
|
await session.execute(
|
|
sa.select(JobRow)
|
|
.where(JobRow.ref_type == "hydration_run", JobRow.ref_id == str(run_id))
|
|
.order_by(JobRow.id.desc())
|
|
.limit(1)
|
|
)
|
|
)
|
|
.scalars()
|
|
.first()
|
|
)
|
|
if job is not None:
|
|
job.status = job_status
|
|
if job_status in {"pending", "canceled", "paused", "dlq"}:
|
|
job.run_after = None
|
|
job.locked_by = None
|
|
job.locked_at = None
|
|
job.updated_at = now
|
|
|
|
await session.commit()
|
|
|
|
return {
|
|
"ok": True,
|
|
"hydration_run_id": str(run_id),
|
|
"status": target_status,
|
|
"job_status": job_status if job is not None else "",
|
|
"request_id": rid,
|
|
}
|
|
|
|
return await with_sqlite_busy_retry(_op)
|
|
|
|
|
|
@router.post("/hydration-runs/{run_id}/pause")
|
|
async def pause_hydration_run(
|
|
run_id: int,
|
|
request: Request,
|
|
_claims: dict[str, Any] = Depends(get_admin_claims),
|
|
) -> dict[str, Any]:
|
|
_ = _claims
|
|
return await _set_run_and_job_status(
|
|
request,
|
|
run_id=run_id,
|
|
target_status="paused",
|
|
allowed_from={"pending", "running"},
|
|
job_status="paused",
|
|
)
|
|
|
|
|
|
@router.post("/hydration-runs/{run_id}/resume")
|
|
async def resume_hydration_run(
|
|
run_id: int,
|
|
request: Request,
|
|
_claims: dict[str, Any] = Depends(get_admin_claims),
|
|
) -> dict[str, Any]:
|
|
_ = _claims
|
|
return await _set_run_and_job_status(
|
|
request,
|
|
run_id=run_id,
|
|
target_status="pending",
|
|
allowed_from={"paused"},
|
|
job_status="pending",
|
|
)
|
|
|
|
|
|
@router.post("/hydration-runs/{run_id}/cancel")
|
|
async def cancel_hydration_run(
|
|
run_id: int,
|
|
request: Request,
|
|
_claims: dict[str, Any] = Depends(get_admin_claims),
|
|
) -> dict[str, Any]:
|
|
_ = _claims
|
|
return await _set_run_and_job_status(
|
|
request,
|
|
run_id=run_id,
|
|
target_status="canceled",
|
|
allowed_from={"pending", "running", "paused"},
|
|
job_status="canceled",
|
|
)
|