247 lines
8.5 KiB
Python
247 lines
8.5 KiB
Python
"""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)
|