110 lines
No EOL
4.5 KiB
Python
110 lines
No EOL
4.5 KiB
Python
from uuid import UUID
|
|
from sqlalchemy import select, update, func
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
from datetime import datetime
|
|
from app.models import Job, Task, TaskStatus, JobStatus, TaskDependency
|
|
from app.job_registry import list_tasks_for_job, list_dependencies, routing_key
|
|
from app.rabbitmq import publish_task_message
|
|
from app.config import settings
|
|
|
|
async def create_job(session: AsyncSession, job_type: str, domain: str, payload: dict) -> Job:
|
|
job = Job(job_type=job_type, domain=domain, payload=payload, status=JobStatus.PENDING)
|
|
session.add(job)
|
|
await session.flush()
|
|
|
|
# 1) Tasks anlegen
|
|
name_to_task: dict[str, Task] = {}
|
|
for name in list_tasks_for_job(job_type):
|
|
t = Task(
|
|
job_id=job.id,
|
|
name=name,
|
|
status=TaskStatus.PENDING,
|
|
routing_key=routing_key(job_type, name),
|
|
max_retries=settings.MAX_RETRIES,
|
|
)
|
|
session.add(t)
|
|
await session.flush()
|
|
name_to_task[name] = t
|
|
|
|
# 2) Dependencies anlegen
|
|
for name, task in name_to_task.items():
|
|
for dep_name in list_dependencies(job_type, name):
|
|
dep_task = name_to_task[dep_name]
|
|
session.add(TaskDependency(task_id=task.id, depends_on_task_id=dep_task.id))
|
|
|
|
await session.commit()
|
|
|
|
# 3) Nur Tasks ohne Dependencies veröffentlichen
|
|
for name, task in name_to_task.items():
|
|
deps = list_dependencies(job_type, name)
|
|
if not deps:
|
|
await _publish_task(session, job, task)
|
|
|
|
await refresh_job_status(session, job.id)
|
|
await session.commit()
|
|
return job
|
|
|
|
async def _publish_task(session: AsyncSession, job: Job, task: Task):
|
|
await publish_task_message(
|
|
routing_key=task.routing_key,
|
|
message={
|
|
"job_id": str(job.id),
|
|
"task_id": str(task.id),
|
|
"task_name": task.name,
|
|
"job_type": job.job_type,
|
|
"payload": {"domain": job.domain, **(job.payload or {})},
|
|
},
|
|
)
|
|
await session.execute(update(Task).where(Task.id == task.id).values(status=TaskStatus.QUEUED))
|
|
await session.commit()
|
|
|
|
async def release_dependent_tasks(session: AsyncSession, completed_task_id: UUID):
|
|
# Finde Tasks, die von completed_task_id abhängen
|
|
res = await session.execute(select(TaskDependency.task_id).where(TaskDependency.depends_on_task_id == completed_task_id))
|
|
candidate_ids = [row[0] for row in res.all()]
|
|
if not candidate_ids:
|
|
return
|
|
# Für jeden Kandidaten prüfen, ob ALLE Dependencies SUCCESS sind
|
|
for task_id in candidate_ids:
|
|
# task + job laden
|
|
res_task = await session.execute(select(Task, Job).join(Job, Task.job_id == Job.id).where(Task.id == task_id))
|
|
row = res_task.first()
|
|
if not row:
|
|
continue
|
|
task, job = row
|
|
# Alle deps status
|
|
res_deps = await session.execute(
|
|
select(func.count()).select_from(TaskDependency)
|
|
.join(Task, TaskDependency.depends_on_task_id == Task.id)
|
|
.where(TaskDependency.task_id == task.id, Task.status != TaskStatus.SUCCESS)
|
|
)
|
|
remaining = res_deps.scalar_one()
|
|
if remaining == 0 and task.status == TaskStatus.PENDING:
|
|
await _publish_task(session, job, task)
|
|
|
|
async def get_job(session: AsyncSession, job_id: UUID) -> Job | None:
|
|
res = await session.execute(select(Job).where(Job.id == job_id))
|
|
return res.scalar_one_or_none()
|
|
|
|
async def get_job_with_tasks(session: AsyncSession, job_id: UUID) -> Job | None:
|
|
return await get_job(session, job_id)
|
|
|
|
async def refresh_job_status(session: AsyncSession, job_id: UUID) -> None:
|
|
res = await session.execute(select(Task).where(Task.job_id == job_id))
|
|
tasks = res.scalars().all()
|
|
if not tasks:
|
|
await session.execute(update(Job).where(Job.id == job_id).values(status=JobStatus.PENDING))
|
|
return
|
|
if all(t.status == TaskStatus.SUCCESS for t in tasks):
|
|
new_status = JobStatus.SUCCESS
|
|
elif any(t.status == TaskStatus.FAILED for t in tasks) and not any(
|
|
t.status in (TaskStatus.RUNNING, TaskStatus.QUEUED, TaskStatus.PENDING) for t in tasks
|
|
):
|
|
new_status = JobStatus.FAILED
|
|
elif any(t.status == TaskStatus.RUNNING for t in tasks):
|
|
new_status = JobStatus.RUNNING
|
|
elif any(t.status == TaskStatus.QUEUED for t in tasks):
|
|
new_status = JobStatus.QUEUED
|
|
else:
|
|
new_status = JobStatus.PENDING
|
|
await session.execute(update(Job).where(Job.id == job_id).values(status=new_status, updated_at=datetime.utcnow())) |