Daniil
135 lines
4.2 KiB
Python
135 lines
4.2 KiB
Python
from __future__ import annotations
|
|
|
|
import uuid
|
|
|
|
from sqlalchemy import Select, select
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from cpv3.modules.jobs.models import Job, JobEvent
|
|
from cpv3.modules.jobs.schemas import (
|
|
JobCreate,
|
|
JobEventCreate,
|
|
JobEventUpdate,
|
|
JobUpdate,
|
|
)
|
|
from cpv3.modules.users.models import User
|
|
|
|
|
|
class JobRepository:
|
|
"""Repository for Job database operations."""
|
|
|
|
def __init__(self, session: AsyncSession) -> None:
|
|
self._session = session
|
|
|
|
async def list_all(self, *, requester: User) -> list[Job]:
|
|
stmt: Select[tuple[Job]] = select(Job).where(Job.is_active.is_(True))
|
|
if not requester.is_staff:
|
|
stmt = stmt.where(Job.user_id == requester.id)
|
|
|
|
result = await self._session.execute(stmt.order_by(Job.created_at.desc()))
|
|
return list(result.scalars().all())
|
|
|
|
async def get_by_id(self, job_id: uuid.UUID) -> Job | None:
|
|
result = await self._session.execute(
|
|
select(Job).where(Job.id == job_id).where(Job.is_active.is_(True))
|
|
)
|
|
return result.scalar_one_or_none()
|
|
|
|
async def list_active_by_type(
|
|
self,
|
|
*,
|
|
requester: User,
|
|
job_type: str,
|
|
project_id: uuid.UUID | None,
|
|
statuses: tuple[str, ...],
|
|
) -> list[Job]:
|
|
stmt: Select[tuple[Job]] = (
|
|
select(Job)
|
|
.where(Job.is_active.is_(True))
|
|
.where(Job.user_id == requester.id)
|
|
.where(Job.job_type == job_type)
|
|
.where(Job.status.in_(statuses))
|
|
.order_by(Job.created_at.desc())
|
|
)
|
|
|
|
if project_id is None:
|
|
stmt = stmt.where(Job.project_id.is_(None))
|
|
else:
|
|
stmt = stmt.where(Job.project_id == project_id)
|
|
|
|
result = await self._session.execute(stmt)
|
|
return list(result.scalars().all())
|
|
|
|
async def create(self, *, requester: User, data: JobCreate) -> Job:
|
|
job = Job(
|
|
user_id=requester.id,
|
|
broker_id=data.broker_id,
|
|
project_id=data.project_id,
|
|
input_data=data.input_data,
|
|
status=data.status,
|
|
job_type=data.job_type,
|
|
)
|
|
|
|
self._session.add(job)
|
|
await self._session.commit()
|
|
await self._session.refresh(job)
|
|
return job
|
|
|
|
async def update(self, job: Job, data: JobUpdate) -> Job:
|
|
for key, value in data.model_dump(exclude_unset=True).items():
|
|
if value is not None:
|
|
setattr(job, key, value)
|
|
|
|
await self._session.commit()
|
|
await self._session.refresh(job)
|
|
return job
|
|
|
|
async def deactivate(self, job: Job) -> None:
|
|
job.is_active = False
|
|
await self._session.commit()
|
|
|
|
|
|
class JobEventRepository:
|
|
"""Repository for JobEvent database operations."""
|
|
|
|
def __init__(self, session: AsyncSession) -> None:
|
|
self._session = session
|
|
|
|
async def list_all(self) -> list[JobEvent]:
|
|
result = await self._session.execute(
|
|
select(JobEvent)
|
|
.where(JobEvent.is_active.is_(True))
|
|
.order_by(JobEvent.created_at.desc())
|
|
)
|
|
return list(result.scalars().all())
|
|
|
|
async def get_by_id(self, event_id: uuid.UUID) -> JobEvent | None:
|
|
result = await self._session.execute(
|
|
select(JobEvent)
|
|
.where(JobEvent.id == event_id)
|
|
.where(JobEvent.is_active.is_(True))
|
|
)
|
|
return result.scalar_one_or_none()
|
|
|
|
async def create(self, data: JobEventCreate) -> JobEvent:
|
|
event = JobEvent(
|
|
job_id=data.job_id, event_type=data.event_type, payload=data.payload
|
|
)
|
|
self._session.add(event)
|
|
await self._session.commit()
|
|
await self._session.refresh(event)
|
|
return event
|
|
|
|
async def update(self, event: JobEvent, data: JobEventUpdate) -> JobEvent:
|
|
for key, value in data.model_dump(exclude_unset=True).items():
|
|
if value is not None:
|
|
setattr(event, key, value)
|
|
|
|
await self._session.commit()
|
|
await self._session.refresh(event)
|
|
return event
|
|
|
|
async def deactivate(self, event: JobEvent) -> None:
|
|
event.is_active = False
|
|
await self._session.commit()
|