mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-08-01 02:46:02 +00:00
486b51eb51
The once branch of next_run_at returned the run time with the task's local timezone offset attached, while the cron branch normalizes to UTC. The value is persisted into scheduled_tasks.next_run_at, and SQLAlchemy's SQLite dialect discards tzinfo on bind, so a once task declared in a non-UTC timezone fires shifted by the whole offset (e.g. 8 hours late for Asia/Shanghai, hours early for negative offsets - also bypassing the min_once_delay_seconds guard). Postgres timestamptz normalizes on write, which is why only SQLite deployments are affected. Align the once branch with the cron branch by converting to UTC before returning.
60 lines
2.1 KiB
Python
60 lines
2.1 KiB
Python
from __future__ import annotations
|
|
|
|
from datetime import UTC, datetime
|
|
from zoneinfo import ZoneInfo, ZoneInfoNotFoundError
|
|
|
|
from croniter import croniter
|
|
|
|
|
|
def validate_timezone(timezone_name: str) -> str:
|
|
try:
|
|
ZoneInfo(timezone_name)
|
|
except ZoneInfoNotFoundError as exc:
|
|
raise ValueError(f"Unknown timezone: {timezone_name}") from exc
|
|
return timezone_name
|
|
|
|
|
|
def normalize_cron_expression(expr: str) -> str:
|
|
parts = [part for part in expr.split() if part]
|
|
if len(parts) != 5:
|
|
raise ValueError("Cron expression must contain exactly 5 fields")
|
|
return " ".join(parts)
|
|
|
|
|
|
def next_run_at(
|
|
schedule_type: str,
|
|
schedule_spec: dict[str, object],
|
|
timezone_name: str,
|
|
*,
|
|
now: datetime,
|
|
) -> datetime | None:
|
|
validate_timezone(timezone_name)
|
|
if now.tzinfo is None:
|
|
now = now.replace(tzinfo=UTC)
|
|
|
|
if schedule_type == "once":
|
|
run_at_raw = schedule_spec.get("run_at")
|
|
if not isinstance(run_at_raw, str):
|
|
raise ValueError("once schedule requires run_at")
|
|
run_at = datetime.fromisoformat(run_at_raw)
|
|
if run_at.tzinfo is None:
|
|
# A naive run_at means "wall-clock time in the task's declared
|
|
# timezone", matching how cron schedules interpret it.
|
|
run_at = run_at.replace(tzinfo=ZoneInfo(timezone_name))
|
|
# Normalize to UTC like the cron branch: next_run_at is persisted to
|
|
# timezone-discarding columns (SQLite), where a non-UTC offset shifts
|
|
# the effective fire time by the whole offset.
|
|
run_at = run_at.astimezone(UTC)
|
|
return run_at if run_at > now else None
|
|
|
|
if schedule_type == "cron":
|
|
cron_expr = normalize_cron_expression(str(schedule_spec.get("cron", "")))
|
|
zone = ZoneInfo(timezone_name)
|
|
local_now = now.astimezone(zone)
|
|
next_local = croniter(cron_expr, local_now).get_next(datetime)
|
|
if next_local.tzinfo is None:
|
|
next_local = next_local.replace(tzinfo=zone)
|
|
return next_local.astimezone(UTC)
|
|
|
|
raise ValueError(f"Unsupported schedule_type: {schedule_type}")
|