Improved Scheduler (#611)

- Track `last_fired_at` for schedules.
- Enable automatic backfill at startup.
- Configurable timezones.

Author: kraftp

dbos/_client.py

@@ -16,6 +16,7 @@
     TypeVar,
     Union,
 )
+from zoneinfo import ZoneInfo
 
 import sqlalchemy as sa
 
@@ -715,6 +716,8 @@ def create_schedule(
         schedule: str,
         context: Any = None,
         workflow_class_name: Optional[str] = None,
+        automatic_backfill: bool = False,
+        cron_timezone: Optional[str] = None,
     ) -> None:
         """
         Create a cron schedule that periodically invokes a workflow.
@@ -725,12 +728,19 @@ def create_schedule(
             schedule: A cron expression (supports seconds with 6 fields).
             context: A context object passed as the second argument to every invocation. Defaults to ``None``.
             workflow_class_name: Class name for static class method workflows. Defaults to ``None``.
+            automatic_backfill: If ``True``, on startup the scheduler will automatically backfill missed executions since the last time the schedule fired. Defaults to ``False``.
+            cron_timezone: IANA timezone name (e.g. ``"America/New_York"``) in which to evaluate the cron expression. Defaults to ``None`` (UTC).
 
         Raises:
             DBOSException: If the cron expression is invalid or a schedule with the same name already exists.
         """
         if not croniter.is_valid(schedule, second_at_beginning=True):
             raise DBOSException(f"Invalid cron schedule: '{schedule}'")
+        if cron_timezone is not None:
+            try:
+                ZoneInfo(cron_timezone)
+            except (KeyError, Exception):
+                raise DBOSException(f"Invalid timezone: '{cron_timezone}'")
         self._sys_db.create_schedule(
             WorkflowSchedule(
                 schedule_id=generate_uuid(),
@@ -740,6 +750,9 @@ def create_schedule(
                 schedule=schedule,
                 status="ACTIVE",
                 context=self._sys_db.serializer.serialize(context),
+                last_fired_at=None,
+                automatic_backfill=automatic_backfill,
+                cron_timezone=cron_timezone,
             )
         )
 
@@ -788,6 +801,8 @@ async def create_schedule_async(
         schedule: str,
         context: Any = None,
         workflow_class_name: Optional[str] = None,
+        automatic_backfill: bool = False,
+        cron_timezone: Optional[str] = None,
     ) -> None:
         """Async version of :meth:`create_schedule`."""
         await asyncio.to_thread(
@@ -797,6 +812,8 @@ async def create_schedule_async(
             schedule=schedule,
             context=context,
             workflow_class_name=workflow_class_name,
+            automatic_backfill=automatic_backfill,
+            cron_timezone=cron_timezone,
         )
 
     async def list_schedules_async(
@@ -844,7 +861,19 @@ def apply_schedules(
             DBOSException: If a cron expression is invalid
         """
         to_apply: List[WorkflowSchedule] = []
-        for entry in schedules:
+        for i, entry in enumerate(schedules):
+            if "schedule_name" not in entry:
+                raise DBOSException(
+                    f"Schedule entry {i} is missing required field 'schedule_name'"
+                )
+            if "workflow_name" not in entry:
+                raise DBOSException(
+                    f"Schedule entry {i} is missing required field 'workflow_name'"
+                )
+            if "schedule" not in entry:
+                raise DBOSException(
+                    f"Schedule entry {i} is missing required field 'schedule'"
+                )
             cron = entry["schedule"]
             if not croniter.is_valid(cron, second_at_beginning=True):
                 raise DBOSException(f"Invalid cron schedule: '{cron}'")
@@ -857,6 +886,9 @@ def apply_schedules(
                     schedule=cron,
                     status="ACTIVE",
                     context=self._sys_db.serializer.serialize(entry["context"]),
+                    last_fired_at=None,
+                    automatic_backfill=entry.get("automatic_backfill", False),
+                    cron_timezone=entry.get("cron_timezone"),
                 )
             )
         with self._sys_db.engine.begin() as c:

dbos/_conductor/protocol.py

@@ -468,6 +468,9 @@ class ScheduleOutput:
     schedule: str
     status: str
     context: str
+    last_fired_at: Optional[str]
+    automatic_backfill: bool
+    cron_timezone: Optional[str]
 
     @classmethod
     def from_schedule(
@@ -482,6 +485,9 @@ def from_schedule(
             schedule=s["schedule"],
             status=s["status"],
             context=context_str,
+            last_fired_at=s.get("last_fired_at"),
+            automatic_backfill=s.get("automatic_backfill", False),
+            cron_timezone=s.get("cron_timezone"),
         )
 
 

dbos/_dbos.py

@@ -32,6 +32,7 @@
     Union,
     overload,
 )
+from zoneinfo import ZoneInfo
 
 from dbos._conductor.conductor import ConductorWebsocket
 from dbos._debouncer import debouncer_workflow
@@ -288,11 +289,13 @@ def get_internal_queue(self) -> Queue:
         return Queue(INTERNAL_QUEUE_NAME)
 
 
-class ScheduleInput(TypedDict):
+class ScheduleInput(TypedDict, total=False):
     schedule_name: str
     workflow_fn: Callable[[datetime, Any], None]
     schedule: str
     context: Any
+    automatic_backfill: bool
+    cron_timezone: Optional[str]
 
 
 class DBOS:
@@ -2037,6 +2040,8 @@ def create_schedule(
         workflow_fn: ScheduledWorkflow,
         schedule: str,
         context: Any = None,
+        automatic_backfill: bool = False,
+        cron_timezone: Optional[str] = None,
     ) -> None:
         """
         Create a cron schedule that periodically invokes a workflow function.
@@ -2050,6 +2055,8 @@ def create_schedule(
             workflow_fn: The workflow function to invoke. Must accept ``(datetime, context)``.
             schedule: A cron expression (supports seconds with 6 fields).
             context: A context object passed as the second argument to every invocation. Defaults to ``None``.
+            automatic_backfill: If ``True``, on startup the scheduler will automatically backfill missed executions since the last time the schedule fired. Defaults to ``False``.
+            cron_timezone: IANA timezone name (e.g. ``"America/New_York"``) in which to evaluate the cron expression. Defaults to ``None`` (UTC).
 
         Raises:
             DBOSException: If the cron expression is invalid, the workflow is not registered, or a schedule with the same name already exists
@@ -2072,6 +2079,11 @@ def create_schedule(
             if fi and fi.class_info and fi.func_type == DBOSFuncType.Class
             else None
         )
+        if cron_timezone is not None:
+            try:
+                ZoneInfo(cron_timezone)
+            except (KeyError, Exception):
+                raise DBOSException(f"Invalid timezone: '{cron_timezone}'")
         sched = WorkflowSchedule(
             schedule_id=generate_uuid(),
             schedule_name=schedule_name,
@@ -2080,6 +2092,9 @@ def create_schedule(
             schedule=schedule,
             status="ACTIVE",
             context=dbos._sys_db.serializer.serialize(context),
+            last_fired_at=None,
+            automatic_backfill=automatic_backfill,
+            cron_timezone=cron_timezone,
         )
         ctx = snapshot_step_context(reserve_sleep_id=False)
         if ctx and ctx.is_workflow():
@@ -2179,6 +2194,8 @@ async def create_schedule_async(
         workflow_fn: ScheduledWorkflow,
         schedule: str,
         context: Any = None,
+        automatic_backfill: bool = False,
+        cron_timezone: Optional[str] = None,
     ) -> None:
         """Async version of :meth:`create_schedule`."""
         await cls._configure_asyncio_thread_pool()
@@ -2188,6 +2205,8 @@ async def create_schedule_async(
             workflow_fn=workflow_fn,
             schedule=schedule,
             context=context,
+            automatic_backfill=automatic_backfill,
+            cron_timezone=cron_timezone,
         )
 
     @classmethod
@@ -2272,7 +2291,19 @@ def apply_schedules(
             )
         dbos = _get_dbos_instance()
         to_apply: List[WorkflowSchedule] = []
-        for entry in schedules:
+        for i, entry in enumerate(schedules):
+            if "schedule_name" not in entry:
+                raise DBOSException(
+                    f"Schedule entry {i} is missing required field 'schedule_name'"
+                )
+            if "workflow_fn" not in entry:
+                raise DBOSException(
+                    f"Schedule entry {i} is missing required field 'workflow_fn'"
+                )
+            if "schedule" not in entry:
+                raise DBOSException(
+                    f"Schedule entry {i} is missing required field 'schedule'"
+                )
             cron = entry["schedule"]
             if not croniter.is_valid(cron, second_at_beginning=True):
                 raise DBOSException(f"Invalid cron schedule: '{cron}'")
@@ -2301,6 +2332,9 @@ def apply_schedules(
                     schedule=cron,
                     status="ACTIVE",
                     context=dbos._sys_db.serializer.serialize(entry["context"]),
+                    last_fired_at=None,
+                    automatic_backfill=entry.get("automatic_backfill", False),
+                    cron_timezone=entry.get("cron_timezone"),
                 )
             )
         with dbos._sys_db.engine.begin() as c:

dbos/_migration.py

@@ -326,6 +326,7 @@ def get_dbos_migration_thirteen(schema: str) -> str:
 );
 """
 
+
 def get_dbos_migration_fourteen(schema: str) -> str:
     return f"""
 CREATE FUNCTION "{schema}".enqueue_workflow(
@@ -431,6 +432,15 @@ def get_dbos_migration_fourteen(schema: str) -> str:
 $$ LANGUAGE plpgsql;
 """
 
+
+def get_dbos_migration_fifteen(schema: str) -> str:
+    return f"""
+ALTER TABLE "{schema}".workflow_schedules ADD COLUMN "last_fired_at" TEXT DEFAULT NULL;
+ALTER TABLE "{schema}".workflow_schedules ADD COLUMN "automatic_backfill" BOOLEAN NOT NULL DEFAULT FALSE;
+ALTER TABLE "{schema}".workflow_schedules ADD COLUMN "cron_timezone" TEXT DEFAULT NULL;
+"""
+
+
 def get_dbos_migrations(schema: str, use_listen_notify: bool) -> list[str]:
     return [
         get_dbos_migration_one(schema, use_listen_notify),
@@ -447,6 +457,7 @@ def get_dbos_migrations(schema: str, use_listen_notify: bool) -> list[str]:
         get_dbos_migration_twelve(schema),
         get_dbos_migration_thirteen(schema),
         get_dbos_migration_fourteen(schema),
+        get_dbos_migration_fifteen(schema),
     ]
 
 
@@ -613,6 +624,12 @@ def get_sqlite_timestamp_expr() -> str:
     created_at INTEGER NOT NULL DEFAULT {get_sqlite_timestamp_expr()}
 );
 """
+sqlite_migration_fifteen = """
+ALTER TABLE workflow_schedules ADD COLUMN "last_fired_at" TEXT DEFAULT NULL;
+ALTER TABLE workflow_schedules ADD COLUMN "automatic_backfill" BOOLEAN NOT NULL DEFAULT FALSE;
+ALTER TABLE workflow_schedules ADD COLUMN "cron_timezone" TEXT DEFAULT NULL;
+"""
+
 sqlite_migrations = [
     sqlite_migration_one,
     sqlite_migration_two,
@@ -627,4 +644,5 @@ def get_sqlite_timestamp_expr() -> str:
     sqlite_migration_twelve,
     sqlite_migration_thirteen,
     # Note, there is no sqlite version of migration fourteen
+    sqlite_migration_fifteen,
 ]

dbos/_scheduler.py

@@ -3,6 +3,7 @@
 import traceback
 from datetime import datetime, timezone
 from typing import TYPE_CHECKING, Any, Callable, Coroutine, Optional
+from zoneinfo import ZoneInfo
 
 from ._croniter import croniter  # type: ignore
 from ._error import DBOSException
@@ -30,6 +31,8 @@ def __init__(self, schedule: WorkflowSchedule, serializer: Serializer):
         self.cron: str = schedule["schedule"]
         self.serialized_context: str = schedule["context"]
         self.context: Any = serializer.deserialize(self.serialized_context)
+        tz_name = schedule.get("cron_timezone")
+        self.tzinfo = ZoneInfo(tz_name) if tz_name else timezone.utc
         self._stop_event = threading.Event()
         self._thread = threading.Thread(target=self._loop, daemon=True)
         self._thread.start()
@@ -40,7 +43,7 @@ def _loop(self) -> None:
         dbos = _get_dbos_instance()
         try:
             it = croniter(
-                self.cron, datetime.now(timezone.utc), second_at_beginning=True
+                self.cron, datetime.now(self.tzinfo), second_at_beginning=True
             )
         except Exception:
             dbos_logger.error(
@@ -69,6 +72,9 @@ def _loop(self) -> None:
                         self.context,
                         self.class_name,
                     )
+                dbos._sys_db.update_last_fired_at(
+                    self.schedule_name, next_exec_time.isoformat()
+                )
             except Exception:
                 dbos_logger.warning(
                     f"Exception in schedule '{self.schedule_name}': "
@@ -148,7 +154,10 @@ def backfill_schedule(
         raise DBOSException(f"Schedule '{schedule_name}' does not exist")
     context = sys_db.serializer.deserialize(schedule["context"])
     class_name = schedule["workflow_class_name"]
-    it = croniter(schedule["schedule"], start, second_at_beginning=True)
+    tz_name = schedule.get("cron_timezone")
+    tz = ZoneInfo(tz_name) if tz_name else timezone.utc
+    start_in_tz = start.astimezone(tz)
+    it = croniter(schedule["schedule"], start_in_tz, second_at_beginning=True)
     workflow_ids: list[str] = []
     while True:
         next_time = it.get_next(datetime)
@@ -222,6 +231,26 @@ def dynamic_scheduler_loop(
                     schedule_thread.stop()
                     del schedule_threads[schedule_id]
             elif schedule_thread is None:
+                # Automatic backfill: if enabled and last_fired_at is set,
+                # backfill missed executions before starting the thread.
+                if schedule.get("automatic_backfill") and schedule.get("last_fired_at"):
+                    try:
+                        assert schedule["last_fired_at"]
+                        last_fired = datetime.fromisoformat(schedule["last_fired_at"])
+                        now = datetime.now(timezone.utc)
+                        if last_fired < now:
+                            backfill_schedule(
+                                dbos._sys_db,
+                                schedule["schedule_name"],
+                                last_fired,
+                                now,
+                            )
+                    except Exception:
+                        dbos_logger.warning(
+                            f"Exception during automatic backfill for "
+                            f"schedule '{schedule['schedule_name']}': "
+                            f"{traceback.format_exc()}"
+                        )
                 schedule_threads[schedule_id] = _ScheduleThread(
                     schedule, dbos._sys_db.serializer
                 )

dbos/_schemas/system_database.py

@@ -213,6 +213,9 @@ def set_schema(cls, schema_name: Optional[str]) -> None:
         Column("schedule", Text, nullable=False),
         Column("status", Text, nullable=False, server_default="ACTIVE"),
         Column("context", Text, nullable=False),
+        Column("last_fired_at", Text, nullable=True),
+        Column("automatic_backfill", Boolean, nullable=False, server_default="false"),
+        Column("cron_timezone", Text, nullable=True),
     )
 
     application_versions = Table(