import logging from datetime import datetime, timezone from dateutil.relativedelta import relativedelta from tortoise.transactions import in_transaction from . import delivery as transport from .models import ScheduledMessage, MessageStatus, Recurrence, JobDelivery logger = logging.getLogger(__name__) RECURRENCE_DELTAS = { Recurrence.DAILY: relativedelta(days=1), Recurrence.WEEKLY: relativedelta(weeks=1), Recurrence.MONTHLY: relativedelta(months=1), } def run_id(msg): # Each recurring occurrence has its own row; edits keep the same identity. return str(msg.id) async def _schedule_next_occurrence(msg, db=None): if db is None: async with in_transaction() as connection: return await _schedule_next_occurrence(msg, connection) delta = RECURRENCE_DELTAS.get(msg.recurrence) if not delta: return series = msg.series_id or msg.id # Serialize series creation/cancellation on the original occurrence. await ( ScheduledMessage.filter(id=series) .using_db(db) .update(updated_at=transport.now()) ) if ( await ScheduledMessage.filter(series_id=series, status=MessageStatus.CANCELLED) .using_db(db) .exists() ): return next_at = msg.scheduled_at + delta now = datetime.now(timezone.utc) while next_at <= now: next_at += delta await ScheduledMessage.get_or_create( previous_occurrence_id=msg.id, defaults=dict( series_id=series, recipient=msg.recipient, channel=msg.channel, content=msg.content, subject=msg.subject, scheduled_at=next_at, recurrence=msg.recurrence, use_agent=msg.use_agent, created_by_id=msg.created_by_id, ), using_db=db, ) async def process_message(msg, *, retry=False): async with in_transaction() as db: # Always lock the series before an occurrence, matching cancellation. await ( ScheduledMessage.filter(id=msg.series_id or msg.id) .using_db(db) .update(updated_at=transport.now()) ) eligible = ScheduledMessage.filter( id=msg.id, status=MessageStatus.FAILED if retry else MessageStatus.PENDING ) if not retry: eligible = eligible.filter(scheduled_at__lte=transport.now()) won = await eligible.using_db(db).update( status=MessageStatus.PROCESSING, series_id=msg.series_id or msg.id ) if not won: return "skipped" msg = await ScheduledMessage.get(id=msg.id).using_db(db) claimed = await transport.claim("scheduled", msg.id, run_id(msg), retry=retry) if not claimed: # Preserve the authoritative delivery state if the job was rescheduled # onto an occurrence already delivered. d = await JobDelivery.get( kind="scheduled", job_id=str(msg.id), run_id=run_id(msg) ) status = ( MessageStatus(d.status) if d.status in {s.value for s in MessageStatus} else MessageStatus.UNCERTAIN ) await ScheduledMessage.filter(id=msg.id).using_db(db).update(status=status) return "skipped" sending = False try: content = msg.content if msg.use_agent: user = await msg.created_by content = await transport.agent_content( user, content, f"scheduled-{claimed.id}", "scheduled" ) # A cancellation which wins before this transaction prevents sending. async with in_transaction() as db: await ( ScheduledMessage.filter(id=msg.series_id or msg.id) .using_db(db) .update(updated_at=transport.now()) ) active = ( await ScheduledMessage.filter( id=msg.id, status=MessageStatus.PROCESSING ) .using_db(db) .update(updated_at=transport.now()) ) if not active: await transport.transition(claimed, "cancelled") return "cancelled" if not await transport.transition(claimed, "sending"): return "skipped" sending = True await transport.send(msg.channel, msg.recipient, content, msg.subject) async with in_transaction() as db: await ( ScheduledMessage.filter(id=msg.series_id or msg.id) .using_db(db) .update(updated_at=transport.now()) ) await transport.transition(claimed, "sent") active = ( await ScheduledMessage.filter( id=msg.id, status=MessageStatus.PROCESSING ) .using_db(db) .update(status=MessageStatus.SENT, error_message=None) ) if active: await _schedule_next_occurrence(msg, db) return "sent" except Exception as exc: status = ( "uncertain" if sending and not isinstance(exc, transport.SafeDeliveryError) else "failed" ) await transport.transition(claimed, status, str(exc)) await ScheduledMessage.filter( id=msg.id, status=MessageStatus.PROCESSING ).update(status=MessageStatus(status), error_message=str(exc)) logger.exception("Scheduled delivery %s ended as %s", msg.id, status) return status async def process_due_messages() -> dict: await transport.expire_claims("scheduled") uncertain = await JobDelivery.filter(kind="scheduled", status="uncertain") for d in uncertain: await ScheduledMessage.filter( id=d.job_id, status=MessageStatus.PROCESSING ).update(status=MessageStatus.UNCERTAIN, error_message=d.error) # Recover recurrence creation if an earlier successful send was committed. # Only rows sent through delivery claims qualify: rows sent before migration # 20 have no successor links, and repairing them would fork each historical # occurrence into its own new series. orphans = await ScheduledMessage.filter( status=MessageStatus.SENT, deleted_at__isnull=True, next_occurrence__id__isnull=True, ).exclude(recurrence=Recurrence.NONE) if orphans: claimed_sends = set( await JobDelivery.filter( kind="scheduled", status="sent", job_id__in=[str(m.id) for m in orphans], ).values_list("job_id", flat=True) ) for msg in orphans: if str(msg.id) in claimed_sends: await _schedule_next_occurrence(msg) due = await ScheduledMessage.filter( status=MessageStatus.PENDING, scheduled_at__lte=transport.now() ) counts = {"sent": 0, "failed": 0, "uncertain": 0} for msg in due: result = await process_message(msg) if result in counts: counts[result] += 1 return counts