mi
443 lines
17 KiB
Python
443 lines
17 KiB
Python
import asyncio
|
|
import logging
|
|
import uuid
|
|
from collections.abc import AsyncIterator, Awaitable, Callable
|
|
from contextlib import asynccontextmanager
|
|
from datetime import UTC, datetime, timedelta
|
|
|
|
import httpx
|
|
import redis.asyncio as redis
|
|
import structlog
|
|
from opentelemetry import trace
|
|
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.logging_security import redact_event
|
|
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_active_client_profile,
|
|
load_settings,
|
|
openlines_delivery_payload,
|
|
publish_dialog_status,
|
|
publish_message_status,
|
|
)
|
|
from app.settings import Settings, get_settings
|
|
from app.telemetry import TelemetryRuntime, add_trace_context, init_telemetry, origin_links
|
|
|
|
log = structlog.get_logger()
|
|
|
|
|
|
def configure_logging(level: str, telemetry: TelemetryRuntime | None = None) -> None:
|
|
logging.basicConfig(level=level, format="%(message)s")
|
|
if telemetry and telemetry.logging_handler not in logging.getLogger().handlers:
|
|
telemetry.logging_handler.addFilter(
|
|
lambda record: not record.name.startswith("opentelemetry")
|
|
)
|
|
logging.getLogger().addHandler(telemetry.logging_handler)
|
|
structlog.configure(
|
|
processors=[
|
|
structlog.contextvars.merge_contextvars,
|
|
add_trace_context,
|
|
redact_event,
|
|
structlog.processors.TimeStamper(fmt="iso", utc=True, key="timestamp"),
|
|
structlog.stdlib.add_log_level,
|
|
structlog.processors.JSONRenderer(),
|
|
],
|
|
logger_factory=structlog.stdlib.LoggerFactory(),
|
|
wrapper_class=structlog.stdlib.BoundLogger,
|
|
cache_logger_on_first_use=True,
|
|
)
|
|
|
|
|
|
def run_worker(service_name: str, target: Callable[[], Awaitable[None]]) -> None:
|
|
settings = get_settings()
|
|
telemetry = init_telemetry(service_name)
|
|
configure_logging(settings.log_level, telemetry)
|
|
structlog.contextvars.bind_contextvars(**{"service.name": service_name})
|
|
try:
|
|
asyncio.run(target())
|
|
finally:
|
|
if telemetry:
|
|
telemetry.shutdown()
|
|
|
|
|
|
@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:
|
|
tracer = trace.get_tracer("han.api.delivery-worker")
|
|
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
|
|
with (
|
|
tracer.start_as_current_span(
|
|
"delivery.process",
|
|
links=origin_links(row.traceparent),
|
|
),
|
|
structlog.contextvars.bound_contextvars(
|
|
delivery_outbox_id=str(row.id),
|
|
message_id=str(row.message_id),
|
|
),
|
|
):
|
|
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:
|
|
tracer = trace.get_tracer("han.api.safety-recovery-worker")
|
|
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:
|
|
with tracer.start_as_current_span(
|
|
"safety.recovery.poll",
|
|
links=origin_links(task.traceparent),
|
|
):
|
|
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:
|
|
profile = await load_active_client_profile(session, user.id)
|
|
await ensure_delivery_outbox(
|
|
session,
|
|
message_id=message.id,
|
|
external_chat_id=dialog.id,
|
|
payload_json=openlines_delivery_payload(
|
|
message=message,
|
|
dialog_id=dialog.id,
|
|
user=user,
|
|
profile=profile,
|
|
attachment=attachment,
|
|
),
|
|
next_attempt_at=datetime.now(UTC),
|
|
traceparent=task.traceparent,
|
|
)
|
|
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:
|
|
run_worker("delivery-worker", lambda: loop("delivery"))
|
|
|
|
|
|
def safety_main() -> None:
|
|
run_worker("safety-recovery-worker", lambda: loop("safety"))
|
|
|
|
|
|
def cleanup_main() -> None:
|
|
run_worker("cleanup-worker", lambda: loop("cleanup"))
|
|
|
|
|
|
def notification_expire_main() -> None:
|
|
run_worker("notification-expire-worker", notification_expire_loop)
|
|
|
|
|
|
def notification_draft_cleanup_main() -> None:
|
|
run_worker("notification-draft-cleanup-worker", notification_draft_cleanup_loop)
|