import asyncio import uuid from collections.abc import AsyncIterator from contextlib import asynccontextmanager from datetime import UTC, datetime, timedelta import httpx import redis.asyncio as redis import structlog from sqlalchemy import delete, select from app.db import ( Database, DeliveryOutbox, Dialog, Message, MessageAttachment, SafetyTask, UserIdentity, ) from app.integrations import ( DependencyFailure, OpenLinesClient, S3Client, SafetyClient, fresh_openlines_payload, ) from app.notification_models import ClientUploadDraft from app.notification_service import expire_notifications from app.realtime import RealtimeFanout from app.services import ( ensure_delivery_outbox, load_settings, publish_dialog_status, publish_message_status, ) from app.settings import Settings, get_settings log = structlog.get_logger() @asynccontextmanager async def worker_http_clients( settings: Settings, ) -> AsyncIterator[tuple[httpx.AsyncClient, httpx.AsyncClient]]: async with ( httpx.AsyncClient() as openlines_http, httpx.AsyncClient(verify=settings.message_safety_ca_file) as safety_http, ): yield openlines_http, safety_http 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.expires_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.poll_location, 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: message.safety_processing_mode = verdict["processing_mode"] message.safety_config_version = verdict["config_version"] message.safety_rules_version = verdict["rules_version"] if attachment and attachment.quarantine_object_key: destination = f"attachments/dialogs/{message.dialog_id}/{attachment.id}" await s3.promote( attachment.quarantine_object_key, destination, version_id=attachment.quarantine_version_id or "", etag=attachment.quarantine_etag or "", ) attachment.storage_bucket = s3.settings.selectel_s3_bucket_attachments attachment.object_key = destination attachment.quarantine_object_key = None attachment.scan_status = ( "bypassed" if verdict["processing_mode"] == "mock" else "clean" ) message.safety_status = "allowed" task.status = "completed" dialog = await session.get(Dialog, message.dialog_id) user = ( await session.get(UserIdentity, dialog.user_id) if dialog else None ) if dialog and user: await ensure_delivery_outbox( session, message_id=message.id, external_chat_id=dialog.id, payload_json={ "message_id": str(message.id), "external_chat_id": str(dialog.id), "occurred_at": message.occurred_at.isoformat(), "user": { "id": str(user.id), "display_name": user.phone_number, }, "message": { "content_kind": message.content_kind, "text": message.text, "files": ( [ { "attachment_id": str(attachment.id), "name": attachment.safe_file_name, "mime_type": attachment.mime_type, "size_bytes": attachment.size_bytes, "_storage_bucket": attachment.storage_bucket, "_object_key": attachment.object_key, } ] if attachment else [] ), }, }, next_attempt_at=datetime.now(UTC), ) elif verdict["_status"] == 403 and message: message.safety_processing_mode = verdict["processing_mode"] message.safety_config_version = verdict["config_version"] message.safety_rules_version = verdict["rules_version"] 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: if verdict["_status"] == 202: task.poll_location = verdict["_location"] task.next_poll_at = datetime.now(UTC) + timedelta(seconds=2) except DependencyFailure as exc: task.attempt_count += 1 terminal = exc.terminal or ( exc.code == "task_not_found" and task.attempt_count >= 2 ) task.status = "terminal_failed" if terminal else "failed" task.last_error_code = exc.code if terminal and message: message.delivery_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) 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: async with worker_http_clients(settings) as (openlines_http, safety_http): safety = SafetyClient(settings, safety_http) openlines = OpenLinesClient(settings, openlines_http) 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 redis_rt.aclose() await db.close() async def notification_expire_loop() -> None: settings = get_settings() db = Database(settings.database_url) try: while True: async with db.sessions() as session: snapshot = await load_settings(session) run_at = snapshot.values["notification.expire_job.run_at"] hour, minute = (int(value) for value in run_at.split(":", 1)) now = datetime.now(UTC) target = now.replace(hour=hour, minute=minute, second=0, microsecond=0) if target <= now: target += timedelta(days=1) await asyncio.sleep((target - now).total_seconds()) async with db.sessions() as session: personal, guest = await expire_notifications(session) log.info( "notification.expired_batch", personal_count=personal, guest_count=guest, ) finally: await db.close() async def notification_draft_cleanup_once(db: Database, s3: S3Client) -> int: async with db.sessions() as session: snapshot = await load_settings(session) cutoff = datetime.now(UTC) - timedelta( days=snapshot.integer("notification.upload_draft.ttl_days") ) rows = ( ( await session.execute( select(ClientUploadDraft) .where(ClientUploadDraft.created_at < cutoff) .with_for_update(skip_locked=True) .limit(100) ) ) .scalars() .all() ) for row in rows: if row.state != "submitted": if row.quarantine_object_key: await s3.delete_quarantine(row.quarantine_object_key) elif row.object_key: await s3.delete(row.storage_bucket, row.object_key) await session.execute( delete(ClientUploadDraft).where(ClientUploadDraft.id == row.id) ) await session.commit() return len(rows) async def notification_draft_cleanup_loop() -> None: settings = get_settings() db = Database(settings.database_url) s3 = S3Client(settings) try: while True: count = await notification_draft_cleanup_once(db, s3) if count < 100: await asyncio.sleep(86400) finally: 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")) def notification_expire_main() -> None: asyncio.run(notification_expire_loop()) def notification_draft_cleanup_main() -> None: asyncio.run(notification_draft_cleanup_loop())