Проект разделен на два репозитория
This commit is contained in:
mi
@@ -0,0 +1,383 @@
|
||||
import asyncio
|
||||
import uuid
|
||||
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 load_settings, 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.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:
|
||||
session.add(
|
||||
DeliveryOutbox(
|
||||
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)
|
||||
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()
|
||||
|
||||
|
||||
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())
|
||||
Reference in New Issue
Block a user