"""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()]