left4me/l4d2web/routes/job_routes.py

60 lines
1.9 KiB
Python

import time
from flask import Blueprint, Response, current_app, request
from sqlalchemy import select
from l4d2web.auth import current_user, require_login
from l4d2web.db import session_scope
from l4d2web.models import Job, JobLog
bp = Blueprint("job", __name__)
TERMINAL_JOB_STATES = {"succeeded", "failed", "cancelled"}
def format_sse_event(seq: int, event: str, data: str) -> str:
lines = [f"id: {seq}", f"event: {event}"]
for line in data.splitlines() or [""]:
lines.append(f"data: {line}")
return "\n".join(lines) + "\n\n"
@bp.get("/jobs/<int:job_id>/stream")
@require_login
def stream_job(job_id: int) -> Response:
user = current_user()
assert user is not None
last_seq = int(request.args.get("last_seq") or request.headers.get("Last-Event-ID") or "0")
limit = int(current_app.config["JOB_LOG_REPLAY_LIMIT"])
poll_seconds = float(current_app.config.get("JOB_WORKER_POLL_SECONDS", 1))
with session_scope() as db:
job = db.scalar(select(Job).where(Job.id == job_id, Job.user_id == user.id))
if job is None:
return Response(status=404)
def generate():
next_seq = last_seq
while True:
with session_scope() as db:
job = db.scalar(select(Job).where(Job.id == job_id))
if job is None:
return
rows = db.scalars(
select(JobLog)
.where(JobLog.job_id == job_id, JobLog.seq > next_seq)
.order_by(JobLog.seq)
.limit(limit)
).all()
terminal = job.state in TERMINAL_JOB_STATES
for row in rows:
next_seq = row.seq
yield format_sse_event(row.seq, row.stream, row.line)
if terminal and len(rows) < limit:
return
time.sleep(poll_seconds)
return Response(generate(), mimetype="text/event-stream")