import asyncio import uuid from datetime import UTC, datetime, timedelta import httpx import redis.asyncio as redis import structlog from sqlalchemy import select from app.db import Database, DeliveryOutbox, Dialog, Message, MessageAttachment, SafetyTask from app.integrations import ( DependencyFailure, OpenLinesClient, S3Client, SafetyClient, fresh_openlines_payload, ) from app.realtime import RealtimeFanout from app.services import publish_dialog_status, publish_message_status from app.settings import Settings, get_settings log = structlog.get_logger() async def delivery_once( db: Database, client: OpenLinesClient, s3: S3Client, fanout: RealtimeFanout, settings: Settings, worker_id: str, batch_size: int = 20, ) -> int: async with db.sessions() as session: rows = ( ( await session.execute( select(DeliveryOutbox) .where( DeliveryOutbox.status.in_(["pending", "retry"]), DeliveryOutbox.next_attempt_at <= datetime.now(UTC), ) .with_for_update(skip_locked=True) .limit(batch_size) ) ) .scalars() .all() ) ids = [row.id for row in rows] for row in rows: row.status = "processing" row.locked_at = datetime.now(UTC) row.locked_by = worker_id await session.commit() for row_id in ids: async with db.sessions() as session: row = await session.get(DeliveryOutbox, row_id, with_for_update=True) if row is None: continue message = None dialog = None try: payload = await fresh_openlines_payload(row.payload_json, s3) await client.send(row.message_id, payload, f"worker-{worker_id}") row.status = "delivered" message = await session.get(Message, row.message_id) if message: message.delivery_status = "delivered" dialog = await session.get(Dialog, message.dialog_id) if dialog: dialog.status = "waiting_for_company" dialog.last_message_at = datetime.now(UTC) except DependencyFailure: row.attempt_count += 1 row.status = "dead_letter" if row.attempt_count >= 12 else "retry" row.next_attempt_at = datetime.now(UTC) + timedelta( seconds=min(3600, 2**row.attempt_count) ) row.last_error_code = "dependency_unavailable" row.locked_at = None row.locked_by = None await session.commit() if message: await publish_message_status(fanout, message, settings) if dialog: await publish_dialog_status(fanout, dialog) return len(ids) async def safety_once( db: Database, safety: SafetyClient, s3: S3Client, fanout: RealtimeFanout, settings: Settings, worker_id: str, batch_size: int = 20, ) -> int: async with db.sessions() as session: rows = ( ( await session.execute( select(SafetyTask) .where( SafetyTask.status.in_(["polling", "failed"]), SafetyTask.next_poll_at <= datetime.now(UTC), SafetyTask.deadline_at > datetime.now(UTC), ) .with_for_update(skip_locked=True) .limit(batch_size) ) ) .scalars() .all() ) ids = [row.id for row in rows] for row in rows: row.locked_at, row.locked_by = datetime.now(UTC), worker_id await session.commit() for task_id in ids: async with db.sessions() as session: task = await session.get(SafetyTask, task_id, with_for_update=True) if task is None: continue message = None try: verdict = await safety.poll(task.task_id, f"worker-{worker_id}") message = await session.get(Message, task.message_id) attachment = ( await session.get(MessageAttachment, task.attachment_id) if task.attachment_id else None ) if verdict["_status"] == 200 and message: if attachment and attachment.quarantine_object_key: destination = f"attachments/dialogs/{message.dialog_id}/{attachment.id}" await s3.promote(attachment.quarantine_object_key, destination) attachment.storage_bucket = s3.settings.selectel_s3_bucket_attachments attachment.object_key = destination attachment.quarantine_object_key = None attachment.scan_status = "clean" message.safety_status = "allowed" task.status = "completed" elif verdict["_status"] == 403 and message: message.text = "" message.safety_status = "blocked" message.delivery_status = "rejected" if attachment and attachment.quarantine_object_key: await s3.delete_quarantine(attachment.quarantine_object_key) attachment.scan_status = "infected" task.status = "completed" else: task.next_poll_at = datetime.now(UTC) + timedelta(seconds=2) except DependencyFailure: task.attempt_count += 1 task.status = "failed" task.next_poll_at = datetime.now(UTC) + timedelta( seconds=min(300, 2**task.attempt_count) ) task.locked_at = None task.locked_by = None await session.commit() if message: await publish_message_status(fanout, message, settings) return len(ids) async def cleanup_once(db: Database, s3: S3Client, batch_size: int = 100) -> int: async with db.sessions() as session: rows = ( ( await session.execute( select(MessageAttachment) .where( MessageAttachment.record_status == "A", MessageAttachment.quarantine_object_key.is_not(None), MessageAttachment.upload_expires_at < datetime.now(UTC), MessageAttachment.message_id.is_(None), ) .with_for_update(skip_locked=True) .limit(batch_size) ) ) .scalars() .all() ) for row in rows: if row.quarantine_object_key: await s3.delete_quarantine(row.quarantine_object_key) row.record_status = "D" row.status_changed_at = datetime.now(UTC) row.status_change_reason = "expired_quarantine_cleanup" await session.commit() return len(rows) async def loop(kind: str) -> None: settings = get_settings() db = Database(settings.database_url) http = httpx.AsyncClient() safety = SafetyClient(settings, http) openlines = OpenLinesClient(settings, http) s3 = S3Client(settings) redis_rt = redis.from_url(settings.redis_realtime_url, decode_responses=True) fanout = RealtimeFanout(redis_rt) worker_id = f"{kind}-{uuid.uuid4()}" try: while True: count = 0 if kind == "delivery": count = await delivery_once(db, openlines, s3, fanout, settings, worker_id) elif kind == "safety": count = await safety_once(db, safety, s3, fanout, settings, worker_id) else: count = await cleanup_once(db, s3) if not count: await asyncio.sleep(settings.worker_poll_interval_sec) finally: await http.aclose() await redis_rt.aclose() await db.close() def delivery_main() -> None: asyncio.run(loop("delivery")) def safety_main() -> None: asyncio.run(loop("safety")) def cleanup_main() -> None: asyncio.run(loop("cleanup"))