"""Draft persistence and the send/undo flow (SPEC ยง6.5).""" from __future__ import annotations import uuid from datetime import UTC, datetime, timedelta from sqlalchemy import ColumnElement, exists, select from sqlalchemy.ext.asyncio import AsyncSession from app.config import get_settings from app.models import Account, Draft, Mailbox, Message, Outbox from app.outbox.schemas import DraftIn class DraftError(Exception): """Draft could not be created/updated (e.g. account not owned).""" async def delete_mirrored_draft_row( session: AsyncSession, *, account_id: uuid.UUID, imap_uid: int ) -> None: """Delete the DB row for a draft that was mirrored into the IMAP Drafts folder. The server copy is expunged separately (send_job / remove_imap_draft_job). Without also dropping this row the draft resurfaces in ``list_remote_drafts`` as an "other client" draft once its local ``drafts`` row is gone โ€” until the next Drafts poll would catch the expunge. Used on both send and discard. """ mirror = await session.scalar( select(Message) .join(Mailbox, Mailbox.id == Message.mailbox_id) .where( Message.account_id == account_id, Message.uid == imap_uid, Mailbox.special_use == "drafts", ) ) if mirror is not None: await session.delete(mirror) async def upsert_draft( session: AsyncSession, *, user_id: uuid.UUID, body: DraftIn, draft_id: uuid.UUID | None ) -> Draft: """Create or update a draft owned by ``user_id``.""" account = await session.scalar( select(Account).where(Account.id == body.account_id, Account.user_id == user_id) ) if account is None: raise DraftError("account not found") payload = body.payload.model_dump(mode="json") if draft_id is None: draft = Draft( user_id=user_id, account_id=body.account_id, thread_id=body.thread_id, reply_to_message_id=body.reply_to_message_id, kind=body.kind, payload=payload, ) session.add(draft) else: fetched = await session.scalar( select(Draft).where(Draft.id == draft_id, Draft.user_id == user_id) ) if fetched is None: raise DraftError("draft not found") draft = fetched draft.account_id = body.account_id draft.thread_id = body.thread_id draft.reply_to_message_id = body.reply_to_message_id draft.kind = body.kind draft.payload = payload draft.updated_at = datetime.now(UTC) await session.flush() return draft # A draft with one of these sends is in flight โ€” or waiting in Scheduled for a # retry โ€” so it is not listed (or counted) in Drafts. OUTBOX_ACTIVE = ("undo_window", "sending", "retrying", "failed") def not_being_sent() -> ColumnElement[bool]: """SQL condition: this Draft row has no send in flight or waiting.""" return ~exists().where(Outbox.draft_id == Draft.id, Outbox.state.in_(OUTBOX_ACTIVE)) async def create_outbox(session: AsyncSession, *, user_id: uuid.UUID, draft: Draft) -> Outbox: """Create an outbox row in the undo window for a draft.""" send_after = datetime.now(UTC) + timedelta(seconds=get_settings().undo_seconds) outbox = Outbox(user_id=user_id, draft_id=draft.id, state="undo_window", send_after=send_after) session.add(outbox) await session.flush() return outbox async def cancel_outbox(session: AsyncSession, *, user_id: uuid.UUID, outbox_id: uuid.UUID) -> bool: """Cancel a send if still in the undo window. Returns False if too late/unknown.""" outbox = await session.scalar( select(Outbox).where(Outbox.id == outbox_id, Outbox.user_id == user_id) ) if outbox is None or outbox.state != "undo_window": return False outbox.state = "cancelled" await session.flush() return True