Lamont/core/outbox.py

"""Durable, lease-based outbound message operations."""

from contextlib import contextmanager
from datetime import datetime, timezone
import uuid

from psycopg2.extras import Json

from core import postgres


DEFAULT_MAX_ATTEMPTS = 3
DEFAULT_LEASE_SECONDS = 300
DEFAULT_RETRY_SECONDS = 30
MAX_RETRY_SECONDS = 900


@contextmanager
def _using_cursor(cursor=None):
    if cursor is not None:
        yield cursor
        return
    with postgres.get_cursor() as owned_cursor:
        yield owned_cursor


def _timestamp(value, field="timestamp"):
    if isinstance(value, str):
        value = datetime.fromisoformat(value.replace("Z", "+00:00"))
    if not isinstance(value, datetime):
        raise ValueError(f"{field} must be a datetime or ISO-8601 string")
    if value.tzinfo is None:
        value = value.replace(tzinfo=timezone.utc)
    return value.astimezone(timezone.utc)


def _positive(value, field, maximum=None):
    try:
        value = int(value)
    except (TypeError, ValueError) as error:
        raise ValueError(f"{field} must be a whole number") from error
    if value < 1 or (maximum is not None and value > maximum):
        suffix = f" and at most {maximum}" if maximum is not None else ""
        raise ValueError(f"{field} must be at least 1{suffix}")
    return value


def _row(cursor):
    record = cursor.fetchone()
    return dict(record) if record else None


def enqueue_message(
    user_uuid,
    channel,
    payload,
    idempotency_key,
    available_at=None,
    max_attempts=DEFAULT_MAX_ATTEMPTS,
    message_id=None,
    cursor=None,
):
    """Queue a message, returning the existing row for a repeated key."""
    if not user_uuid:
        raise ValueError("user_uuid is required")
    if not isinstance(channel, str) or not channel.strip():
        raise ValueError("channel is required")
    if not idempotency_key:
        raise ValueError("idempotency_key is required")
    if payload is None:
        raise ValueError("payload is required")
    max_attempts = _positive(max_attempts, "max_attempts", 100)
    available_at = available_at or datetime.now(timezone.utc)
    values = {
        "id": str(message_id or uuid.uuid4()),
        "user_uuid": user_uuid,
        "channel": channel.strip(),
        "payload": Json(payload),
        "idempotency_key": idempotency_key,
        "available_at": _timestamp(available_at, "available_at"),
        "max_attempts": max_attempts,
    }

    with _using_cursor(cursor) as active_cursor:
        active_cursor.execute(
            """
            INSERT INTO outbound_messages (
                id, user_uuid, channel, payload, idempotency_key,
                available_at, max_attempts
            ) VALUES (
                %(id)s, %(user_uuid)s, %(channel)s, %(payload)s,
                %(idempotency_key)s, %(available_at)s, %(max_attempts)s
            )
            ON CONFLICT (idempotency_key)
            DO UPDATE SET idempotency_key = EXCLUDED.idempotency_key
            WHERE outbound_messages.user_uuid = EXCLUDED.user_uuid
              AND outbound_messages.channel = EXCLUDED.channel
            RETURNING *
            """,
            values,
        )
        return _row(active_cursor)


def get_message(message_id, cursor=None):
    """Return one outbound message by UUID."""
    try:
        message_id = str(uuid.UUID(str(message_id)))
    except (TypeError, ValueError, AttributeError):
        return None
    with _using_cursor(cursor) as active_cursor:
        active_cursor.execute(
            "SELECT * FROM outbound_messages WHERE id = %s", (message_id,)
        )
        return _row(active_cursor)


def list_messages(user_uuid=None, status=None, channel=None, limit=100, cursor=None):
    """List outbound messages newest first with optional filters."""
    limit = _positive(limit, "limit", 500)
    clauses = []
    params = []
    if user_uuid is not None:
        clauses.append("user_uuid = %s")
        params.append(user_uuid)
    if status is not None:
        clauses.append("status = %s")
        params.append(status)
    if channel is not None:
        clauses.append("channel = %s")
        params.append(channel)
    where = f"WHERE {' AND '.join(clauses)}" if clauses else ""
    params.append(limit)

    with _using_cursor(cursor) as active_cursor:
        active_cursor.execute(
            f"""
            SELECT * FROM outbound_messages
            {where}
            ORDER BY created_at DESC
            LIMIT %s
            """,
            params,
        )
        return [dict(record) for record in active_cursor.fetchall()]


def claim_messages(
    worker_id,
    channel=None,
    limit=10,
    lease_seconds=DEFAULT_LEASE_SECONDS,
    cursor=None,
):
    """Atomically lease deliverable messages while skipping other workers."""
    if not worker_id:
        raise ValueError("worker_id is required")
    limit = _positive(limit, "limit", 100)
    lease_seconds = _positive(lease_seconds, "lease_seconds")

    with _using_cursor(cursor) as active_cursor:
        active_cursor.execute(
            """
            WITH exhausted AS (
                SELECT id
                FROM outbound_messages
                WHERE status = 'delivering'
                  AND (lease_until IS NULL OR lease_until <= NOW())
                  AND attempts >= max_attempts
                FOR UPDATE SKIP LOCKED
            )
            UPDATE outbound_messages AS message
            SET status = 'failed',
                lease_until = NULL,
                leased_by = NULL,
                last_error = COALESCE(last_error, 'lease expired'),
                updated_at = NOW()
            FROM exhausted
            WHERE message.id = exhausted.id
            """
        )
        active_cursor.execute(
            """
            WITH candidates AS (
                SELECT id
                FROM outbound_messages
                WHERE available_at <= NOW()
                  AND attempts < max_attempts
                  AND (
                      status = 'pending'
                      OR (
                          status = 'delivering'
                          AND (lease_until IS NULL OR lease_until <= NOW())
                      )
                  )
                  AND (%(channel)s IS NULL OR channel = %(channel)s)
                ORDER BY available_at, created_at
                FOR UPDATE SKIP LOCKED
                LIMIT %(limit)s
            )
            UPDATE outbound_messages AS message
            SET status = 'delivering',
                attempts = message.attempts + 1,
                leased_by = %(worker_id)s,
                lease_until = NOW() + (%(lease_seconds)s * INTERVAL '1 second'),
                updated_at = NOW()
            FROM candidates
            WHERE message.id = candidates.id
            RETURNING message.*
            """,
            {
                "channel": channel,
                "limit": limit,
                "worker_id": worker_id,
                "lease_seconds": lease_seconds,
            },
        )
        return [dict(record) for record in active_cursor.fetchall()]


def renew_message_lease(
    message_id, worker_id, lease_seconds=DEFAULT_LEASE_SECONDS, cursor=None
):
    """Extend a delivery lease owned by the requesting worker."""
    lease_seconds = _positive(lease_seconds, "lease_seconds")
    with _using_cursor(cursor) as active_cursor:
        active_cursor.execute(
            """
            UPDATE outbound_messages
            SET lease_until = NOW() + (%s * INTERVAL '1 second'), updated_at = NOW()
            WHERE id = %s AND status = 'delivering' AND leased_by = %s
            RETURNING *
            """,
            (lease_seconds, message_id, worker_id),
        )
        return _row(active_cursor)


def mark_delivered(
    message_id, worker_id, external_message_id=None, cursor=None
):
    """Mark a message delivered when its lease is owned by the worker."""
    with _using_cursor(cursor) as active_cursor:
        active_cursor.execute(
            """
            UPDATE outbound_messages
            SET status = 'delivered', delivered_at = NOW(), updated_at = NOW(),
                leased_by = NULL, lease_until = NULL, last_error = NULL,
                external_message_id = %s
            WHERE id = %s AND status = 'delivering' AND leased_by = %s
            RETURNING *
            """,
            (external_message_id, message_id, worker_id),
        )
        return _row(active_cursor)


def retry_message(
    message_id,
    worker_id,
    error,
    retry_seconds=DEFAULT_RETRY_SECONDS,
    max_retry_seconds=MAX_RETRY_SECONDS,
    cursor=None,
):
    """Fail or reschedule an owned delivery with exponential backoff."""
    retry_seconds = _positive(retry_seconds, "retry_seconds")
    max_retry_seconds = _positive(max_retry_seconds, "max_retry_seconds")
    with _using_cursor(cursor) as active_cursor:
        active_cursor.execute(
            """
            SELECT attempts, max_attempts
            FROM outbound_messages
            WHERE id = %s AND status = 'delivering' AND leased_by = %s
            FOR UPDATE
            """,
            (message_id, worker_id),
        )
        current = active_cursor.fetchone()
        if not current:
            return None

        exhausted = current["attempts"] >= current["max_attempts"]
        delay = min(
            max_retry_seconds,
            retry_seconds * (2 ** min(max(current["attempts"] - 1, 0), 30)),
        )
        active_cursor.execute(
            """
            UPDATE outbound_messages
            SET status = %(status)s,
                available_at = CASE
                    WHEN %(exhausted)s THEN available_at
                    ELSE NOW() + (%(delay)s * INTERVAL '1 second')
                END,
                leased_by = NULL,
                lease_until = NULL,
                last_error = %(error)s,
                updated_at = NOW()
            WHERE id = %(message_id)s
            RETURNING *
            """,
            {
                "status": "failed" if exhausted else "pending",
                "exhausted": exhausted,
                "delay": delay,
                "error": str(error)[:4000],
                "message_id": message_id,
            },
        )
        return _row(active_cursor)


def cancel_message(message_id, user_uuid=None, cursor=None):
    """Cancel one undelivered message, optionally enforcing its owner."""
    with _using_cursor(cursor) as active_cursor:
        active_cursor.execute(
            """
            UPDATE outbound_messages
            SET status = 'cancelled', cancelled_at = NOW(), updated_at = NOW(),
                leased_by = NULL, lease_until = NULL
            WHERE id = %(message_id)s
              AND status IN ('pending', 'delivering')
              AND (%(user_uuid)s IS NULL OR user_uuid = %(user_uuid)s)
            RETURNING *
            """,
            {"message_id": message_id, "user_uuid": user_uuid},
        )
        return _row(active_cursor)


def cancel_messages(
    user_uuid=None, channel=None, idempotency_key=None, cursor=None
):
    """Cancel matching undelivered messages; at least one filter is required."""
    filters = {
        "user_uuid": user_uuid,
        "channel": channel,
        "idempotency_key": idempotency_key,
    }
    clauses = [f"{name} = %({name})s" for name, value in filters.items() if value is not None]
    if not clauses:
        raise ValueError("at least one cancellation filter is required")

    with _using_cursor(cursor) as active_cursor:
        active_cursor.execute(
            f"""
            UPDATE outbound_messages
            SET status = 'cancelled', cancelled_at = NOW(), updated_at = NOW(),
                leased_by = NULL, lease_until = NULL
            WHERE status IN ('pending', 'delivering')
              AND {' AND '.join(clauses)}
            RETURNING *
            """,
            filters,
        )
        return [dict(record) for record in active_cursor.fetchall()]