latest updates
This commit is contained in:
Tanmay Karande
6
scheduler/__init__.py
Normal file
6
scheduler/__init__.py
Normal file
@@ -0,0 +1,6 @@
|
||||
# Aetheel Scheduler
|
||||
# Persistent cron-based task scheduling.
|
||||
|
||||
from scheduler.scheduler import Scheduler
|
||||
|
||||
__all__ = ["Scheduler"]
|
||||
275
scheduler/scheduler.py
Normal file
275
scheduler/scheduler.py
Normal file
@@ -0,0 +1,275 @@
|
||||
"""
|
||||
Aetheel Scheduler
|
||||
=================
|
||||
APScheduler-based task scheduler with SQLite persistence.
|
||||
|
||||
Supports:
|
||||
- One-shot delayed jobs (replaces threading.Timer reminders)
|
||||
- Recurring cron jobs (cron expressions)
|
||||
- Persistent storage (jobs survive restarts)
|
||||
- Callback-based execution (fires handlers with job context)
|
||||
|
||||
Usage:
|
||||
from scheduler import Scheduler
|
||||
|
||||
scheduler = Scheduler(callback=my_handler)
|
||||
scheduler.start()
|
||||
|
||||
# One-shot reminder
|
||||
scheduler.add_once(
|
||||
delay_minutes=5,
|
||||
prompt="Time to stretch!",
|
||||
channel_id="C123",
|
||||
channel_type="slack",
|
||||
)
|
||||
|
||||
# Recurring cron job
|
||||
scheduler.add_cron(
|
||||
cron_expr="0 9 * * *",
|
||||
prompt="Good morning! Here's your daily summary.",
|
||||
channel_id="C123",
|
||||
channel_type="slack",
|
||||
)
|
||||
"""
|
||||
|
||||
import logging
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Callable
|
||||
|
||||
from apscheduler.schedulers.background import BackgroundScheduler
|
||||
from apscheduler.triggers.cron import CronTrigger
|
||||
from apscheduler.triggers.date import DateTrigger
|
||||
|
||||
from scheduler.store import JobStore, ScheduledJob
|
||||
|
||||
logger = logging.getLogger("aetheel.scheduler")
|
||||
|
||||
# Callback type: receives the ScheduledJob when it fires
|
||||
JobCallback = Callable[[ScheduledJob], None]
|
||||
|
||||
|
||||
class Scheduler:
|
||||
"""
|
||||
APScheduler-based task scheduler.
|
||||
|
||||
Wraps BackgroundScheduler with SQLite persistence. When a job fires,
|
||||
it calls the registered callback with the ScheduledJob details. The
|
||||
callback is responsible for routing the job's prompt to the AI handler
|
||||
and sending the response to the right channel.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
callback: JobCallback,
|
||||
db_path: str | None = None,
|
||||
):
|
||||
self._callback = callback
|
||||
self._store = JobStore(db_path=db_path)
|
||||
self._scheduler = BackgroundScheduler(
|
||||
daemon=True,
|
||||
job_defaults={"misfire_grace_time": 60},
|
||||
)
|
||||
self._running = False
|
||||
|
||||
@property
|
||||
def running(self) -> bool:
|
||||
return self._running
|
||||
|
||||
def start(self) -> None:
|
||||
"""Start the scheduler and restore persisted jobs."""
|
||||
if self._running:
|
||||
return
|
||||
|
||||
self._scheduler.start()
|
||||
self._running = True
|
||||
|
||||
# Restore recurring jobs from the database
|
||||
restored = 0
|
||||
for job in self._store.list_recurring():
|
||||
try:
|
||||
self._register_cron_job(job)
|
||||
restored += 1
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to restore job {job.id}: {e}")
|
||||
|
||||
logger.info(f"Scheduler started (restored {restored} recurring jobs)")
|
||||
|
||||
def stop(self) -> None:
|
||||
"""Shut down the scheduler."""
|
||||
if self._running:
|
||||
self._scheduler.shutdown(wait=False)
|
||||
self._running = False
|
||||
logger.info("Scheduler stopped")
|
||||
|
||||
# -------------------------------------------------------------------
|
||||
# Public API: Add jobs
|
||||
# -------------------------------------------------------------------
|
||||
|
||||
def add_once(
|
||||
self,
|
||||
*,
|
||||
delay_minutes: int,
|
||||
prompt: str,
|
||||
channel_id: str,
|
||||
channel_type: str = "slack",
|
||||
thread_id: str | None = None,
|
||||
user_name: str | None = None,
|
||||
) -> str:
|
||||
"""
|
||||
Schedule a one-shot job to fire after a delay.
|
||||
|
||||
Replaces the old threading.Timer-based _schedule_reminder().
|
||||
Returns the job ID.
|
||||
"""
|
||||
job_id = JobStore.new_id()
|
||||
run_at = datetime.now(timezone.utc) + timedelta(minutes=delay_minutes)
|
||||
|
||||
job = ScheduledJob(
|
||||
id=job_id,
|
||||
cron_expr=None,
|
||||
prompt=prompt,
|
||||
channel_id=channel_id,
|
||||
channel_type=channel_type,
|
||||
thread_id=thread_id,
|
||||
user_name=user_name,
|
||||
created_at=datetime.now(timezone.utc).isoformat(),
|
||||
next_run=run_at.isoformat(),
|
||||
)
|
||||
|
||||
# Persist
|
||||
self._store.add(job)
|
||||
|
||||
# Schedule
|
||||
self._scheduler.add_job(
|
||||
self._fire_job,
|
||||
trigger=DateTrigger(run_date=run_at),
|
||||
args=[job_id],
|
||||
id=f"once-{job_id}",
|
||||
)
|
||||
|
||||
logger.info(
|
||||
f"⏰ One-shot scheduled: '{prompt[:50]}' in {delay_minutes} min "
|
||||
f"(id={job_id}, channel={channel_type}/{channel_id})"
|
||||
)
|
||||
return job_id
|
||||
|
||||
def add_cron(
|
||||
self,
|
||||
*,
|
||||
cron_expr: str,
|
||||
prompt: str,
|
||||
channel_id: str,
|
||||
channel_type: str = "slack",
|
||||
thread_id: str | None = None,
|
||||
user_name: str | None = None,
|
||||
) -> str:
|
||||
"""
|
||||
Schedule a recurring cron job.
|
||||
|
||||
Args:
|
||||
cron_expr: Standard cron expression (5 fields: min hour day month weekday)
|
||||
|
||||
Returns the job ID.
|
||||
"""
|
||||
job_id = JobStore.new_id()
|
||||
|
||||
job = ScheduledJob(
|
||||
id=job_id,
|
||||
cron_expr=cron_expr,
|
||||
prompt=prompt,
|
||||
channel_id=channel_id,
|
||||
channel_type=channel_type,
|
||||
thread_id=thread_id,
|
||||
user_name=user_name,
|
||||
created_at=datetime.now(timezone.utc).isoformat(),
|
||||
)
|
||||
|
||||
# Persist
|
||||
self._store.add(job)
|
||||
|
||||
# Register with APScheduler
|
||||
self._register_cron_job(job)
|
||||
|
||||
logger.info(
|
||||
f"🔄 Cron scheduled: '{prompt[:50]}' ({cron_expr}) "
|
||||
f"(id={job_id}, channel={channel_type}/{channel_id})"
|
||||
)
|
||||
return job_id
|
||||
|
||||
# -------------------------------------------------------------------
|
||||
# Public API: Manage jobs
|
||||
# -------------------------------------------------------------------
|
||||
|
||||
def remove(self, job_id: str) -> bool:
|
||||
"""Remove a job by ID. Returns True if found and removed."""
|
||||
# Remove from APScheduler
|
||||
for prefix in ("once-", "cron-"):
|
||||
try:
|
||||
self._scheduler.remove_job(f"{prefix}{job_id}")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Remove from store
|
||||
return self._store.remove(job_id)
|
||||
|
||||
def list_jobs(self) -> list[ScheduledJob]:
|
||||
"""List all scheduled jobs."""
|
||||
return self._store.list_all()
|
||||
|
||||
def list_recurring(self) -> list[ScheduledJob]:
|
||||
"""List only recurring cron jobs."""
|
||||
return self._store.list_recurring()
|
||||
|
||||
# -------------------------------------------------------------------
|
||||
# Internal
|
||||
# -------------------------------------------------------------------
|
||||
|
||||
def _register_cron_job(self, job: ScheduledJob) -> None:
|
||||
"""Register a cron job with APScheduler."""
|
||||
if not job.cron_expr:
|
||||
return
|
||||
|
||||
# Parse cron expression (5 fields: min hour day month weekday)
|
||||
parts = job.cron_expr.strip().split()
|
||||
if len(parts) != 5:
|
||||
raise ValueError(
|
||||
f"Invalid cron expression: '{job.cron_expr}' — "
|
||||
"expected 5 fields (minute hour day month weekday)"
|
||||
)
|
||||
|
||||
trigger = CronTrigger(
|
||||
minute=parts[0],
|
||||
hour=parts[1],
|
||||
day=parts[2],
|
||||
month=parts[3],
|
||||
day_of_week=parts[4],
|
||||
)
|
||||
|
||||
self._scheduler.add_job(
|
||||
self._fire_job,
|
||||
trigger=trigger,
|
||||
args=[job.id],
|
||||
id=f"cron-{job.id}",
|
||||
replace_existing=True,
|
||||
)
|
||||
|
||||
def _fire_job(self, job_id: str) -> None:
|
||||
"""Called by APScheduler when a job triggers."""
|
||||
job = self._store.get(job_id)
|
||||
if not job:
|
||||
logger.warning(f"Job {job_id} not found in store — skipping")
|
||||
return
|
||||
|
||||
logger.info(
|
||||
f"🔔 Job fired: {job_id} (cron={job.cron_expr}, "
|
||||
f"prompt='{job.prompt[:50]}')"
|
||||
)
|
||||
|
||||
try:
|
||||
self._callback(job)
|
||||
except Exception as e:
|
||||
logger.error(f"Job callback failed for {job_id}: {e}", exc_info=True)
|
||||
|
||||
# Clean up one-shot jobs after firing
|
||||
if not job.is_recurring:
|
||||
self._store.remove(job_id)
|
||||
167
scheduler/store.py
Normal file
167
scheduler/store.py
Normal file
@@ -0,0 +1,167 @@
|
||||
"""
|
||||
Aetheel Scheduler — SQLite Job Store
|
||||
=====================================
|
||||
Persistent storage for scheduled jobs.
|
||||
|
||||
Schema:
|
||||
jobs(id, cron_expr, prompt, channel_id, channel_type, created_at, next_run)
|
||||
"""
|
||||
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import sqlite3
|
||||
import uuid
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timezone
|
||||
|
||||
logger = logging.getLogger("aetheel.scheduler.store")
|
||||
|
||||
|
||||
@dataclass
|
||||
class ScheduledJob:
|
||||
"""A persisted scheduled job."""
|
||||
|
||||
id: str
|
||||
cron_expr: str | None # None for one-shot jobs
|
||||
prompt: str
|
||||
channel_id: str
|
||||
channel_type: str # "slack", "telegram", etc.
|
||||
created_at: str
|
||||
next_run: str | None = None
|
||||
thread_id: str | None = None # for threading context
|
||||
user_name: str | None = None # who created the job
|
||||
|
||||
@property
|
||||
def is_recurring(self) -> bool:
|
||||
return self.cron_expr is not None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# SQLite Store
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
CREATE_TABLE_SQL = """
|
||||
CREATE TABLE IF NOT EXISTS jobs (
|
||||
id TEXT PRIMARY KEY,
|
||||
cron_expr TEXT,
|
||||
prompt TEXT NOT NULL,
|
||||
channel_id TEXT NOT NULL,
|
||||
channel_type TEXT NOT NULL DEFAULT 'slack',
|
||||
thread_id TEXT,
|
||||
user_name TEXT,
|
||||
created_at TEXT NOT NULL,
|
||||
next_run TEXT
|
||||
);
|
||||
"""
|
||||
|
||||
|
||||
class JobStore:
|
||||
"""SQLite-backed persistent job store."""
|
||||
|
||||
def __init__(self, db_path: str | None = None):
|
||||
self._db_path = db_path or os.path.join(
|
||||
os.path.expanduser("~/.aetheel"), "scheduler.db"
|
||||
)
|
||||
# Ensure directory exists
|
||||
os.makedirs(os.path.dirname(self._db_path), exist_ok=True)
|
||||
self._init_db()
|
||||
|
||||
def _init_db(self) -> None:
|
||||
"""Initialize the database schema."""
|
||||
with sqlite3.connect(self._db_path) as conn:
|
||||
conn.execute(CREATE_TABLE_SQL)
|
||||
conn.commit()
|
||||
logger.info(f"Scheduler DB initialized: {self._db_path}")
|
||||
|
||||
def _conn(self) -> sqlite3.Connection:
|
||||
conn = sqlite3.connect(self._db_path)
|
||||
conn.row_factory = sqlite3.Row
|
||||
return conn
|
||||
|
||||
def add(self, job: ScheduledJob) -> str:
|
||||
"""Add a job to the store. Returns the job ID."""
|
||||
with self._conn() as conn:
|
||||
conn.execute(
|
||||
"""
|
||||
INSERT INTO jobs (id, cron_expr, prompt, channel_id, channel_type,
|
||||
thread_id, user_name, created_at, next_run)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
""",
|
||||
(
|
||||
job.id,
|
||||
job.cron_expr,
|
||||
job.prompt,
|
||||
job.channel_id,
|
||||
job.channel_type,
|
||||
job.thread_id,
|
||||
job.user_name,
|
||||
job.created_at,
|
||||
job.next_run,
|
||||
),
|
||||
)
|
||||
conn.commit()
|
||||
logger.info(f"Job added: {job.id} (cron={job.cron_expr})")
|
||||
return job.id
|
||||
|
||||
def remove(self, job_id: str) -> bool:
|
||||
"""Remove a job by ID. Returns True if found and removed."""
|
||||
with self._conn() as conn:
|
||||
cursor = conn.execute("DELETE FROM jobs WHERE id = ?", (job_id,))
|
||||
conn.commit()
|
||||
removed = cursor.rowcount > 0
|
||||
if removed:
|
||||
logger.info(f"Job removed: {job_id}")
|
||||
return removed
|
||||
|
||||
def get(self, job_id: str) -> ScheduledJob | None:
|
||||
"""Get a job by ID."""
|
||||
with self._conn() as conn:
|
||||
row = conn.execute(
|
||||
"SELECT * FROM jobs WHERE id = ?", (job_id,)
|
||||
).fetchone()
|
||||
return self._row_to_job(row) if row else None
|
||||
|
||||
def list_all(self) -> list[ScheduledJob]:
|
||||
"""List all jobs."""
|
||||
with self._conn() as conn:
|
||||
rows = conn.execute(
|
||||
"SELECT * FROM jobs ORDER BY created_at"
|
||||
).fetchall()
|
||||
return [self._row_to_job(row) for row in rows]
|
||||
|
||||
def list_recurring(self) -> list[ScheduledJob]:
|
||||
"""List only recurring (cron) jobs."""
|
||||
with self._conn() as conn:
|
||||
rows = conn.execute(
|
||||
"SELECT * FROM jobs WHERE cron_expr IS NOT NULL ORDER BY created_at"
|
||||
).fetchall()
|
||||
return [self._row_to_job(row) for row in rows]
|
||||
|
||||
def clear_oneshot(self) -> int:
|
||||
"""Remove all one-shot (non-cron) jobs. Returns count removed."""
|
||||
with self._conn() as conn:
|
||||
cursor = conn.execute(
|
||||
"DELETE FROM jobs WHERE cron_expr IS NULL"
|
||||
)
|
||||
conn.commit()
|
||||
return cursor.rowcount
|
||||
|
||||
@staticmethod
|
||||
def _row_to_job(row: sqlite3.Row) -> ScheduledJob:
|
||||
return ScheduledJob(
|
||||
id=row["id"],
|
||||
cron_expr=row["cron_expr"],
|
||||
prompt=row["prompt"],
|
||||
channel_id=row["channel_id"],
|
||||
channel_type=row["channel_type"],
|
||||
thread_id=row["thread_id"],
|
||||
user_name=row["user_name"],
|
||||
created_at=row["created_at"],
|
||||
next_run=row["next_run"],
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def new_id() -> str:
|
||||
"""Generate a short unique job ID."""
|
||||
return uuid.uuid4().hex[:8]
|
||||
Reference in New Issue
Block a user