Build reusable bot framework
This commit is contained in:
Chelsea Lee
246
modules/reminders/service.py
Normal file
246
modules/reminders/service.py
Normal file
@@ -0,0 +1,246 @@
|
||||
"""Database and scheduling operations for the reminders feature."""
|
||||
|
||||
from datetime import datetime, timedelta, timezone
|
||||
import uuid
|
||||
from zoneinfo import ZoneInfo
|
||||
|
||||
from psycopg2.extras import Json
|
||||
|
||||
from core import jobs, outbox, postgres, users
|
||||
|
||||
|
||||
JOB_TYPE = "reminders.deliver"
|
||||
CHANNEL = "discord_dm"
|
||||
|
||||
|
||||
def _asUtc(value, field="run_at"):
|
||||
if isinstance(value, str):
|
||||
try:
|
||||
value = datetime.fromisoformat(value.replace("Z", "+00:00"))
|
||||
except ValueError as error:
|
||||
raise ValueError(f"{field} must be an ISO-8601 datetime") from error
|
||||
if not isinstance(value, datetime) or value.tzinfo is None:
|
||||
raise ValueError(f"{field} must include a timezone offset")
|
||||
return value.astimezone(timezone.utc)
|
||||
|
||||
|
||||
def normalizeRecurrence(recurrence):
|
||||
if recurrence in (None, False, "none", "once"):
|
||||
return None
|
||||
if isinstance(recurrence, str):
|
||||
recurrence = {"frequency": recurrence}
|
||||
if not isinstance(recurrence, dict):
|
||||
raise ValueError("recurrence must be null or an object")
|
||||
|
||||
frequency = str(recurrence.get("frequency", "")).strip().lower()
|
||||
if frequency not in {"daily", "weekly"}:
|
||||
raise ValueError("recurrence frequency must be daily or weekly")
|
||||
try:
|
||||
interval = int(recurrence.get("interval", 1))
|
||||
except (TypeError, ValueError) as error:
|
||||
raise ValueError("recurrence interval must be a number") from error
|
||||
if interval < 1 or interval > 365:
|
||||
raise ValueError("recurrence interval must be between 1 and 365")
|
||||
return {"frequency": frequency, "interval": interval}
|
||||
|
||||
|
||||
def _jobKey(reminderID, scheduledFor):
|
||||
return f"reminder:{reminderID}:{scheduledFor.isoformat()}"
|
||||
|
||||
|
||||
def createReminder(userUUID, message, runAt, timezoneName, recurrence=None):
|
||||
if not isinstance(message, str) or not message.strip():
|
||||
raise ValueError("reminder message is required")
|
||||
message = message.strip()
|
||||
if len(message) > 1800:
|
||||
raise ValueError("reminder message must be at most 1800 characters")
|
||||
timezoneName = users.normalizeTimezone(timezoneName)
|
||||
runAt = _asUtc(runAt)
|
||||
if runAt <= datetime.now(timezone.utc):
|
||||
raise ValueError("reminder time must be in the future")
|
||||
recurrence = normalizeRecurrence(recurrence)
|
||||
reminderID = str(uuid.uuid4())
|
||||
|
||||
with postgres.get_cursor() as cursor:
|
||||
cursor.execute(
|
||||
"""
|
||||
INSERT INTO reminders (
|
||||
id, user_uuid, message, timezone, recurrence, next_run_at
|
||||
) VALUES (
|
||||
%(id)s, %(user_uuid)s, %(message)s, %(timezone)s,
|
||||
%(recurrence)s, %(next_run_at)s
|
||||
)
|
||||
RETURNING *
|
||||
""",
|
||||
{
|
||||
"id": reminderID,
|
||||
"user_uuid": userUUID,
|
||||
"message": message,
|
||||
"timezone": timezoneName,
|
||||
"recurrence": Json(recurrence) if recurrence else None,
|
||||
"next_run_at": runAt,
|
||||
},
|
||||
)
|
||||
reminder = dict(cursor.fetchone())
|
||||
jobs.create_job(
|
||||
JOB_TYPE,
|
||||
{"reminder_id": reminderID, "scheduled_for": runAt.isoformat()},
|
||||
runAt,
|
||||
user_uuid=userUUID,
|
||||
idempotency_key=_jobKey(reminderID, runAt),
|
||||
cursor=cursor,
|
||||
)
|
||||
return reminder
|
||||
|
||||
|
||||
def listReminders(userUUID, includeFinished=False, limit=50):
|
||||
try:
|
||||
limit = min(max(int(limit), 1), 100)
|
||||
except (TypeError, ValueError):
|
||||
limit = 50
|
||||
statusClause = "" if includeFinished else "AND status = 'active'"
|
||||
return postgres.execute(
|
||||
f"""
|
||||
SELECT * FROM reminders
|
||||
WHERE user_uuid = %(user_uuid)s {statusClause}
|
||||
ORDER BY next_run_at, created_at
|
||||
LIMIT %(limit)s
|
||||
""",
|
||||
{"user_uuid": userUUID, "limit": limit},
|
||||
)
|
||||
|
||||
|
||||
def getReminder(userUUID, reminderID, cursor=None, forUpdate=False):
|
||||
try:
|
||||
uuid.UUID(str(reminderID))
|
||||
except (TypeError, ValueError, AttributeError):
|
||||
return None
|
||||
lock = " FOR UPDATE" if forUpdate else ""
|
||||
query = (
|
||||
"SELECT * FROM reminders WHERE id = %(id)s AND user_uuid = %(user_uuid)s"
|
||||
+ lock
|
||||
)
|
||||
if cursor is not None:
|
||||
cursor.execute(query, {"id": reminderID, "user_uuid": userUUID})
|
||||
record = cursor.fetchone()
|
||||
return dict(record) if record else None
|
||||
rows = postgres.execute(query, {"id": reminderID, "user_uuid": userUUID})
|
||||
return rows[0] if rows else None
|
||||
|
||||
|
||||
def cancelReminder(userUUID, reminderID):
|
||||
with postgres.get_cursor() as cursor:
|
||||
reminder = getReminder(userUUID, reminderID, cursor=cursor, forUpdate=True)
|
||||
if not reminder or reminder["status"] != "active":
|
||||
return None
|
||||
cursor.execute(
|
||||
"""
|
||||
UPDATE reminders
|
||||
SET status = 'cancelled', cancelled_at = NOW(), updated_at = NOW()
|
||||
WHERE id = %(id)s
|
||||
RETURNING *
|
||||
""",
|
||||
{"id": reminderID},
|
||||
)
|
||||
cancelled = dict(cursor.fetchone())
|
||||
cursor.execute(
|
||||
"""
|
||||
UPDATE scheduled_jobs
|
||||
SET status = 'cancelled', cancelled_at = NOW(), updated_at = NOW(),
|
||||
leased_by = NULL, lease_until = NULL
|
||||
WHERE user_uuid = %(user_uuid)s
|
||||
AND job_type = %(job_type)s
|
||||
AND payload->>'reminder_id' = %(reminder_id)s
|
||||
AND status IN ('pending', 'running')
|
||||
""",
|
||||
{
|
||||
"user_uuid": userUUID,
|
||||
"job_type": JOB_TYPE,
|
||||
"reminder_id": reminderID,
|
||||
},
|
||||
)
|
||||
cursor.execute(
|
||||
"""
|
||||
UPDATE outbound_messages
|
||||
SET status = 'cancelled', cancelled_at = NOW(), updated_at = NOW(),
|
||||
leased_by = NULL, lease_until = NULL
|
||||
WHERE user_uuid = %(user_uuid)s
|
||||
AND payload->>'reminder_id' = %(reminder_id)s
|
||||
AND status = 'pending'
|
||||
""",
|
||||
{"user_uuid": userUUID, "reminder_id": reminderID},
|
||||
)
|
||||
return cancelled
|
||||
|
||||
|
||||
def _nextRun(scheduledFor, recurrence, timezoneName, now=None):
|
||||
now = now or datetime.now(timezone.utc)
|
||||
localRun = _asUtc(scheduledFor, "scheduled_for").astimezone(
|
||||
ZoneInfo(timezoneName)
|
||||
)
|
||||
interval = recurrence.get("interval", 1)
|
||||
days = interval if recurrence["frequency"] == "daily" else interval * 7
|
||||
nextLocal = localRun + timedelta(days=days)
|
||||
while nextLocal.astimezone(timezone.utc) <= now:
|
||||
nextLocal += timedelta(days=days)
|
||||
return nextLocal.astimezone(timezone.utc)
|
||||
|
||||
|
||||
def runReminderJob(job, workerID):
|
||||
payload = job.get("payload") or {}
|
||||
reminderID = payload.get("reminder_id")
|
||||
scheduledFor = _asUtc(payload.get("scheduled_for"), "scheduled_for")
|
||||
|
||||
with postgres.get_cursor() as cursor:
|
||||
cursor.execute(
|
||||
"SELECT * FROM reminders WHERE id = %s FOR UPDATE", (reminderID,)
|
||||
)
|
||||
record = cursor.fetchone()
|
||||
reminder = dict(record) if record else None
|
||||
if not reminder or reminder["status"] != "active":
|
||||
return jobs.complete_job(job["id"], workerID, cursor=cursor)
|
||||
|
||||
outbox.enqueue_message(
|
||||
reminder["user_uuid"],
|
||||
CHANNEL,
|
||||
{
|
||||
"content": f"Reminder: {reminder['message']}",
|
||||
"reminder_id": str(reminder["id"]),
|
||||
"scheduled_for": scheduledFor.isoformat(),
|
||||
},
|
||||
idempotency_key=_jobKey(reminder["id"], scheduledFor),
|
||||
cursor=cursor,
|
||||
)
|
||||
|
||||
recurrence = reminder.get("recurrence")
|
||||
if recurrence:
|
||||
nextRun = _nextRun(scheduledFor, recurrence, reminder["timezone"])
|
||||
cursor.execute(
|
||||
"""
|
||||
UPDATE reminders
|
||||
SET next_run_at = %s, updated_at = NOW()
|
||||
WHERE id = %s
|
||||
""",
|
||||
(nextRun, reminder["id"]),
|
||||
)
|
||||
jobs.create_job(
|
||||
JOB_TYPE,
|
||||
{
|
||||
"reminder_id": str(reminder["id"]),
|
||||
"scheduled_for": nextRun.isoformat(),
|
||||
},
|
||||
nextRun,
|
||||
user_uuid=reminder["user_uuid"],
|
||||
idempotency_key=_jobKey(reminder["id"], nextRun),
|
||||
cursor=cursor,
|
||||
)
|
||||
else:
|
||||
cursor.execute(
|
||||
"""
|
||||
UPDATE reminders
|
||||
SET status = 'completed', completed_at = NOW(), updated_at = NOW()
|
||||
WHERE id = %s
|
||||
""",
|
||||
(reminder["id"],),
|
||||
)
|
||||
return jobs.complete_job(job["id"], workerID, cursor=cursor)
|
||||
Reference in New Issue
Block a user