Build reusable bot framework
This commit is contained in:
Chelsea Lee
357
core/outbox.py
Normal file
357
core/outbox.py
Normal file
@@ -0,0 +1,357 @@
|
||||
"""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()]
|
||||
Reference in New Issue
Block a user