feat: 添加自动化任务(automation_jobs)功能模块

2026-03-24 12:38:11 +08:00
parent f4b7eb7e09
commit 23359c2d01
43 changed files with 4266 additions and 1139 deletions
@@ -1,8 +1,8 @@
 from __future__ import annotations

-from datetime import datetime, timedelta, timezone
+from datetime import datetime, time, timedelta, timezone
 from typing import TYPE_CHECKING
-from uuid import UUID, uuid4
+from uuid import UUID

 from sqlalchemy import func, select, update
 from sqlalchemy.ext.asyncio import AsyncSession
@@ -10,7 +10,7 @@ from zoneinfo import ZoneInfo, ZoneInfoNotFoundError

 from core.db.base_repository import BaseRepository
 from models.agent_chat_session import AgentChatSession, SessionType
-from models.automation_jobs import AutomationJob
+from models.automation_jobs import AutomationJob, AutomationJobStatus, ScheduleType

 if TYPE_CHECKING:
    from v1.automation_jobs.schemas import (
@@ -19,144 +19,10 @@ if TYPE_CHECKING:
    )


-def _compute_next_local_time_utc(
-    *,
-    now_utc: datetime,
-    timezone_name: str,
-    local_hour: int,
-    local_minute: int,
-) -> tuple[datetime, datetime]:
-    try:
-        timezone_obj = ZoneInfo(timezone_name)
-    except ZoneInfoNotFoundError:
-        timezone_obj = ZoneInfo("UTC")
-    local_now = now_utc.astimezone(timezone_obj)
-    today_run_local = local_now.replace(
-        hour=local_hour,
-        minute=local_minute,
-        second=0,
-        microsecond=0,
-    )
-    run_local = (
-        today_run_local
-        if local_now <= today_run_local
-        else today_run_local + timedelta(days=1)
-    )
-    next_local = run_local + timedelta(days=1)
-    return run_local.astimezone(timezone.utc), next_local.astimezone(timezone.utc)
-
-
 class AutomationJobsRepository(BaseRepository[AutomationJob]):
    def __init__(self, session: AsyncSession) -> None:
        super().__init__(session=session, model=AutomationJob)

-    async def list_by_owner(self, owner_id: UUID) -> list[AutomationJob]:
-        stmt = (
-            select(AutomationJob)
-            .where(AutomationJob.owner_id == owner_id)
-            .where(AutomationJob.deleted_at.is_(None))
-            .order_by(AutomationJob.created_at.desc())
-        )
-        rows = (await self._session.execute(stmt)).scalars().all()
-        return list(rows)
-
-    async def get_by_id(self, job_id: UUID) -> AutomationJob | None:  # type: ignore[override]
-        stmt = (
-            select(AutomationJob)
-            .where(AutomationJob.id == job_id)
-            .where(AutomationJob.deleted_at.is_(None))
-        )
-        result = await self._session.execute(stmt)
-        return result.scalar_one_or_none()
-
-    async def count_user_jobs(self, owner_id: UUID) -> int:
-        stmt = (
-            select(func.count(AutomationJob.id))
-            .where(AutomationJob.owner_id == owner_id)
-            .where(AutomationJob.bootstrap_key.is_(None))
-            .where(AutomationJob.deleted_at.is_(None))
-        )
-        result = (await self._session.execute(stmt)).scalar_one()
-        return int(result)
-
-    async def create(
-        self, owner_id: UUID, data: "AutomationJobCreateRequest"
-    ) -> AutomationJob:
-        now_utc = datetime.now(timezone.utc)
-        run_at_dt, next_run_at = _compute_next_local_time_utc(
-            now_utc=now_utc,
-            timezone_name=data.timezone,
-            local_hour=data.run_at.hour,
-            local_minute=data.run_at.minute,
-        )
-        new_job = AutomationJob(
-            id=uuid4(),
-            owner_id=owner_id,
-            created_by=owner_id,
-            bootstrap_key=None,
-            title=data.title,
-            config=data.config.model_dump(mode="json"),
-            schedule_type=data.schedule_type,
-            run_at=run_at_dt,
-            next_run_at=next_run_at,
-            timezone=data.timezone,
-            status=data.status,
-        )
-        self._session.add(new_job)
-        await self._session.flush()
-        return new_job
-
-    async def update(
-        self, job_id: UUID, data: "AutomationJobUpdateRequest"
-    ) -> AutomationJob | None:
-        update_values: dict[str, object] = {}
-        if data.title is not None:
-            update_values["title"] = data.title
-        if data.schedule_type is not None:
-            update_values["schedule_type"] = data.schedule_type
-        if data.run_at is not None:
-            stmt = select(AutomationJob).where(AutomationJob.id == job_id)
-            existing = (await self._session.execute(stmt)).scalar_one_or_none()
-            if existing is None:
-                return None
-            run_at_dt, next_run_at = _compute_next_local_time_utc(
-                now_utc=datetime.now(timezone.utc),
-                timezone_name=data.timezone or existing.timezone,
-                local_hour=data.run_at.hour,
-                local_minute=data.run_at.minute,
-            )
-            update_values["run_at"] = run_at_dt
-            update_values["next_run_at"] = next_run_at
-            update_values["timezone"] = data.timezone or existing.timezone
-        if data.status is not None:
-            update_values["status"] = data.status
-        if data.config is not None:
-            update_values["config"] = data.config.model_dump(mode="json")
-
-        if not update_values:
-            return await self.get_by_id(job_id)
-
-        stmt = (
-            update(AutomationJob)
-            .where(AutomationJob.id == job_id)
-            .where(AutomationJob.deleted_at.is_(None))
-            .values(**update_values)
-            .returning(AutomationJob)
-        )
-        result = await self._session.execute(stmt)
-        await self._session.flush()
-        return result.scalar_one_or_none()
-
-    async def soft_delete(self, job_id: UUID) -> None:
-        stmt = (
-            update(AutomationJob)
-            .where(AutomationJob.id == job_id)
-            .where(AutomationJob.deleted_at.is_(None))
-            .values(deleted_at=datetime.now(timezone.utc))
-        )
-        await self._session.execute(stmt)
-        await self._session.flush()
-
    async def list_due_jobs(
        self,
        *,
@@ -166,7 +32,7 @@ class AutomationJobsRepository(BaseRepository[AutomationJob]):
        stmt = (
            select(AutomationJob)
            .where(AutomationJob.deleted_at.is_(None))
-            .where(AutomationJob.status == "active")
+            .where(AutomationJob.status == AutomationJobStatus.ACTIVE)
            .where(AutomationJob.next_run_at <= now_utc)
            .order_by(AutomationJob.next_run_at.asc())
            .limit(max(limit, 1))
@@ -213,3 +79,160 @@ class AutomationJobsRepository(BaseRepository[AutomationJob]):
        self._session.add(new_session)
        await self._session.flush()
        return new_session.id
+
+    async def list_by_owner(self, owner_id: UUID) -> list[AutomationJob]:
+        stmt = (
+            select(AutomationJob)
+            .where(AutomationJob.owner_id == owner_id)
+            .where(AutomationJob.deleted_at.is_(None))
+            .order_by(AutomationJob.created_at.desc())
+        )
+        rows = (await self._session.execute(stmt)).scalars().all()
+        return list(rows)
+
+    async def count_user_jobs(self, owner_id: UUID) -> int:
+        stmt = (
+            select(func.count())
+            .select_from(AutomationJob)
+            .where(AutomationJob.owner_id == owner_id)
+            .where(AutomationJob.deleted_at.is_(None))
+            .where(AutomationJob.bootstrap_key.is_(None))
+        )
+        result = (await self._session.execute(stmt)).scalar_one()
+        return int(result)
+
+    def _resolve_timezone(self, timezone_str: str) -> ZoneInfo:
+        try:
+            return ZoneInfo(timezone_str)
+        except ZoneInfoNotFoundError:
+            return ZoneInfo("UTC")
+
+    def _compute_initial_next_run_at(
+        self,
+        *,
+        run_at: time,
+        timezone_str: str,
+        now_utc: datetime,
+        schedule_type: ScheduleType,
+    ) -> datetime:
+        tz = self._resolve_timezone(timezone_str)
+        local_now = now_utc.astimezone(tz)
+        run_at_local = datetime.combine(local_now.date(), run_at, tz)
+        if run_at_local.tzinfo is None:
+            run_at_local = run_at_local.replace(tzinfo=tz)
+        next_run_at = run_at_local
+        if next_run_at <= local_now:
+            if schedule_type == ScheduleType.DAILY:
+                next_run_at = next_run_at + timedelta(days=1)
+            else:
+                next_run_at = next_run_at + timedelta(weeks=1)
+        return next_run_at.astimezone(timezone.utc)
+
+    async def create(
+        self,
+        owner_id: UUID,
+        data: AutomationJobCreateRequest,
+    ) -> AutomationJob:
+        now_utc = datetime.now(tz=timezone.utc)
+        timezone_obj = self._resolve_timezone(data.timezone)
+        local_now = now_utc.astimezone(timezone_obj)
+        date_ref = local_now.date()
+        local_dt = datetime.combine(date_ref, data.run_at, timezone_obj)
+        run_at_datetime = local_dt.astimezone(timezone.utc)
+        next_run_at = self._compute_initial_next_run_at(
+            run_at=data.run_at,
+            timezone_str=data.timezone,
+            now_utc=now_utc,
+            schedule_type=data.schedule_type,
+        )
+
+        new_job = AutomationJob(
+            owner_id=owner_id,
+            created_by=owner_id,
+            bootstrap_key=None,
+            title=data.title,
+            schedule_type=data.schedule_type,
+            run_at=run_at_datetime,
+            timezone=data.timezone,
+            status=data.status,
+            config=data.config.model_dump(mode="json"),
+            next_run_at=next_run_at,
+        )
+        self._session.add(new_job)
+        await self._session.flush()
+        return new_job
+
+    async def update(
+        self,
+        job_id: UUID,
+        data: AutomationJobUpdateRequest,
+    ) -> AutomationJob | None:
+        update_values: dict[str, object] = {}
+        existing_job: AutomationJob | None = None
+
+        if data.title is not None:
+            update_values["title"] = data.title
+        if data.schedule_type is not None:
+            update_values["schedule_type"] = data.schedule_type
+
+        should_recompute_schedule = (
+            data.run_at is not None
+            or data.schedule_type is not None
+            or data.timezone is not None
+        )
+        if should_recompute_schedule:
+            now_utc = datetime.now(tz=timezone.utc)
+            if existing_job is None:
+                existing_job = await self.get_by_id(job_id)
+                if existing_job is None:
+                    return None
+
+            effective_timezone = data.timezone or existing_job.timezone
+            effective_timezone_obj = self._resolve_timezone(effective_timezone)
+            effective_schedule_type = data.schedule_type or existing_job.schedule_type
+
+            if data.run_at is not None:
+                effective_run_at = data.run_at
+            else:
+                existing_timezone_obj = self._resolve_timezone(existing_job.timezone)
+                effective_run_at = (
+                    existing_job.run_at.astimezone(existing_timezone_obj)
+                    .time()
+                    .replace(microsecond=0)
+                )
+
+            local_now = now_utc.astimezone(effective_timezone_obj)
+            local_dt = datetime.combine(
+                local_now.date(),
+                effective_run_at,
+                effective_timezone_obj,
+            )
+            update_values["run_at"] = local_dt.astimezone(timezone.utc)
+            update_values["next_run_at"] = self._compute_initial_next_run_at(
+                run_at=effective_run_at,
+                timezone_str=effective_timezone,
+                now_utc=now_utc,
+                schedule_type=effective_schedule_type,
+            )
+        if data.timezone is not None:
+            update_values["timezone"] = data.timezone
+        if data.status is not None:
+            update_values["status"] = data.status
+        if data.config is not None:
+            if existing_job is None:
+                existing_job = await self.get_by_id(job_id)
+                if existing_job is None:
+                    return None
+            merged_config = {
+                **existing_job.config,
+                **data.config.model_dump(mode="json", exclude_unset=True),
+            }
+            update_values["config"] = merged_config
+
+        if not update_values:
+            return await self.get_by_id(job_id)
+
+        return await self.update_by_id(job_id, update_values)
+
+    async def soft_delete(self, job_id: UUID) -> None:
+        await self.soft_delete_by_id(job_id)