zl-q
78 lines
2.4 KiB
Python
78 lines
2.4 KiB
Python
from __future__ import annotations
|
|
|
|
from datetime import datetime
|
|
from typing import TYPE_CHECKING
|
|
from uuid import UUID
|
|
|
|
from sqlalchemy import select, update
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from core.db.base_repository import BaseRepository
|
|
from models.agent_chat_session import AgentChatSession, SessionType
|
|
from models.automation_jobs import AutomationJob
|
|
|
|
if TYPE_CHECKING:
|
|
pass
|
|
|
|
|
|
class AutomationJobsRepository(BaseRepository[AutomationJob]):
|
|
def __init__(self, session: AsyncSession) -> None:
|
|
super().__init__(session=session, model=AutomationJob)
|
|
|
|
async def list_due_jobs(
|
|
self,
|
|
*,
|
|
now_utc: datetime,
|
|
limit: int,
|
|
) -> list[AutomationJob]:
|
|
stmt = (
|
|
select(AutomationJob)
|
|
.where(AutomationJob.deleted_at.is_(None))
|
|
.where(AutomationJob.status == "active")
|
|
.where(AutomationJob.next_run_at <= now_utc)
|
|
.order_by(AutomationJob.next_run_at.asc())
|
|
.limit(max(limit, 1))
|
|
)
|
|
rows = (await self._session.execute(stmt)).scalars().all()
|
|
return list(rows)
|
|
|
|
async def update_job_schedule(
|
|
self,
|
|
*,
|
|
job_id: UUID,
|
|
next_run_at: datetime,
|
|
last_run_at: datetime,
|
|
) -> None:
|
|
stmt = (
|
|
update(AutomationJob)
|
|
.where(AutomationJob.id == job_id)
|
|
.where(AutomationJob.deleted_at.is_(None))
|
|
.values(next_run_at=next_run_at, last_run_at=last_run_at)
|
|
)
|
|
await self._session.execute(stmt)
|
|
await self._session.flush()
|
|
|
|
async def get_or_create_chat_session(self, *, owner_id: UUID) -> UUID:
|
|
stmt = (
|
|
select(AgentChatSession.id)
|
|
.where(AgentChatSession.user_id == owner_id)
|
|
.where(AgentChatSession.deleted_at.is_(None))
|
|
.where(AgentChatSession.session_type == SessionType.CHAT)
|
|
.order_by(AgentChatSession.last_activity_at.desc())
|
|
.limit(1)
|
|
)
|
|
existing = (await self._session.execute(stmt)).scalar_one_or_none()
|
|
if existing is not None:
|
|
return existing
|
|
|
|
from uuid import uuid4
|
|
|
|
new_session = AgentChatSession(
|
|
id=uuid4(),
|
|
user_id=owner_id,
|
|
session_type=SessionType.CHAT,
|
|
)
|
|
self._session.add(new_session)
|
|
await self._session.flush()
|
|
return new_session.id
|