first
This commit is contained in:
valenti
@@ -0,0 +1,13 @@
|
||||
from .rabbitmq import (
|
||||
get_event_publisher,
|
||||
EventPublisher,
|
||||
publish_domain_event,
|
||||
enqueue_job,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"get_event_publisher",
|
||||
"EventPublisher",
|
||||
"publish_domain_event",
|
||||
"enqueue_job",
|
||||
]
|
||||
@@ -0,0 +1,199 @@
|
||||
"""RabbitMQ producer for asf-core.
|
||||
|
||||
asf-core publishes two classes of messages:
|
||||
1. Domain events → ``asf.events`` topic exchange (e.g. artifact.created)
|
||||
2. Job requests → ``asf.jobs`` direct exchange, routed to the correct
|
||||
logical queue (asf.ai.interactive, asf.indexing, …).
|
||||
|
||||
Core never consumes messages directly — that is the responsibility of
|
||||
asf-worker. Connections are opened lazily and reused.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import AsyncIterator
|
||||
from uuid import UUID
|
||||
|
||||
import aio_pika
|
||||
import aio_pika.abc
|
||||
import structlog
|
||||
|
||||
from ..config import get_settings
|
||||
from asf_domain.base import JobType
|
||||
from asf_events.base import ASFEvent, JobRequestEnvelope
|
||||
|
||||
log = structlog.get_logger(__name__)
|
||||
|
||||
_CHANNEL_POOL: aio_pika.abc.AbstractChannel | None = None
|
||||
_CONN: aio_pika.abc.AbstractConnection | None = None
|
||||
|
||||
|
||||
class EventPublisher:
|
||||
"""Wraps a RabbitMQ channel with helpers for ASF message classes."""
|
||||
|
||||
def __init__(self, channel: aio_pika.abc.AbstractChannel) -> None:
|
||||
self._channel = channel
|
||||
self._settings = get_settings()
|
||||
|
||||
async def declare_exchanges(self) -> None:
|
||||
events_ex = await self._channel.declare_exchange(
|
||||
self._settings.rabbitmq_events_exchange,
|
||||
type=aio_pika.ExchangeType.TOPIC,
|
||||
durable=True,
|
||||
)
|
||||
jobs_ex = await self._channel.declare_exchange(
|
||||
self._settings.rabbitmq_jobs_exchange,
|
||||
type=aio_pika.ExchangeType.DIRECT,
|
||||
durable=True,
|
||||
)
|
||||
|
||||
# ---- Domain event queues (TOPIC bindings on asf.events) ----
|
||||
domain_binds: list[tuple[str, list[str]]] = [
|
||||
("asf.organization.created", ["organization.created"]),
|
||||
("asf.project.created", ["project.created"]),
|
||||
("asf.project.updated", ["project.updated"]),
|
||||
("asf.artifact.created", ["artifact.created"]),
|
||||
("asf.artifact.updated", ["artifact.updated"]),
|
||||
("asf.audit", ["audit.#"]),
|
||||
("asf.ai.execution", ["ai.execution.*"]),
|
||||
]
|
||||
for qname, keys in domain_binds:
|
||||
q = await self._channel.declare_queue(qname, durable=True)
|
||||
for rk in keys:
|
||||
await q.bind(events_ex, routing_key=rk)
|
||||
|
||||
# Dead-letter queue (no default bind — consumed explicitly by worker topology)
|
||||
await self._channel.declare_queue("asf.deadletter", durable=True)
|
||||
|
||||
# ---- Job queues (DIRECT bindings on asf.jobs) ----
|
||||
for qname in self._settings.queues.values():
|
||||
q = await self._channel.declare_queue(qname, durable=True)
|
||||
await q.bind(jobs_ex, routing_key=qname)
|
||||
|
||||
# -------- Domain events --------
|
||||
async def publish_event(self, event: ASFEvent) -> None:
|
||||
body = event.model_dump_json().encode("utf-8")
|
||||
exchange = await self._channel.get_exchange(self._settings.rabbitmq_events_exchange)
|
||||
await exchange.publish(
|
||||
aio_pika.Message(
|
||||
body=body,
|
||||
delivery_mode=aio_pika.DeliveryMode.PERSISTENT,
|
||||
content_type="application/json",
|
||||
message_id=str(event.event_id),
|
||||
timestamp=event.timestamp,
|
||||
headers={
|
||||
"tenant_id": str(event.tenant_id),
|
||||
"event_type": event.event_type,
|
||||
},
|
||||
),
|
||||
routing_key=event.event_type,
|
||||
)
|
||||
await log.adebug(
|
||||
"event_published",
|
||||
event_type=event.event_type,
|
||||
event_id=str(event.event_id),
|
||||
)
|
||||
|
||||
# -------- Job enqueueing --------
|
||||
async def enqueue_job(self, envelope: JobRequestEnvelope) -> None:
|
||||
qname = self._routing_key_for_job_type(envelope.type)
|
||||
body = envelope.model_dump_json().encode("utf-8")
|
||||
exchange = await self._channel.get_exchange(self._settings.rabbitmq_jobs_exchange)
|
||||
await exchange.publish(
|
||||
aio_pika.Message(
|
||||
body=body,
|
||||
delivery_mode=aio_pika.DeliveryMode.PERSISTENT,
|
||||
content_type="application/json",
|
||||
message_id=str(envelope.job_id),
|
||||
priority=self._priority(envelope.priority),
|
||||
headers={
|
||||
"tenant_id": str(envelope.tenant_id),
|
||||
"job_type": envelope.type.value,
|
||||
},
|
||||
),
|
||||
routing_key=qname,
|
||||
)
|
||||
await log.adebug(
|
||||
"job_enqueued",
|
||||
job_id=str(envelope.job_id),
|
||||
queue=qname,
|
||||
)
|
||||
|
||||
def _routing_key_for_job_type(self, t: JobType) -> str:
|
||||
mapping = {
|
||||
JobType.AI_INTERACTIVE: "ai_interactive",
|
||||
JobType.AI_BACKGROUND: "ai_background",
|
||||
JobType.INDEXING: "indexing",
|
||||
JobType.CODE_ANALYSIS: "code_analysis",
|
||||
JobType.TESTING: "testing",
|
||||
JobType.DEVOPS: "devops",
|
||||
JobType.NOTIFICATION: "notifications",
|
||||
JobType.CUSTOM: "notifications",
|
||||
}
|
||||
key = mapping[t]
|
||||
return self._settings.queues[key]
|
||||
|
||||
@staticmethod
|
||||
def _priority(p) -> int:
|
||||
return {
|
||||
"low": 1,
|
||||
"normal": 5,
|
||||
"high": 8,
|
||||
"urgent": 10,
|
||||
}.get(getattr(p, "value", str(p)), 5)
|
||||
|
||||
|
||||
async def _get_channel() -> aio_pika.abc.AbstractChannel:
|
||||
global _CHANNEL_POOL, _CONN
|
||||
if _CHANNEL_POOL is None or _CHANNEL_POOL.is_closed:
|
||||
s = get_settings()
|
||||
_CONN = await aio_pika.connect_robust(s.rabbitmq_uri)
|
||||
_CHANNEL_POOL = await _CONN.channel()
|
||||
await _CHANNEL_POOL.set_qos(prefetch_count=100)
|
||||
return _CHANNEL_POOL
|
||||
|
||||
|
||||
async def get_event_publisher() -> EventPublisher:
|
||||
ch = await _get_channel()
|
||||
pub = EventPublisher(ch)
|
||||
await pub.declare_exchanges()
|
||||
return pub
|
||||
|
||||
|
||||
async def close_rabbitmq() -> None:
|
||||
global _CHANNEL_POOL, _CONN
|
||||
try:
|
||||
if _CHANNEL_POOL is not None:
|
||||
await _CHANNEL_POOL.close()
|
||||
except Exception: # noqa: BLE001 - shutdown best effort
|
||||
pass
|
||||
try:
|
||||
if _CONN is not None:
|
||||
await _CONN.close()
|
||||
except Exception: # noqa: BLE001
|
||||
pass
|
||||
_CHANNEL_POOL = None
|
||||
_CONN = None
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def rabbitmq_lifespan() -> AsyncIterator[None]:
|
||||
await get_event_publisher() # prime connection
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
await close_rabbitmq()
|
||||
|
||||
|
||||
# -------- Convenience helpers (used by services) --------
|
||||
|
||||
async def publish_domain_event(event: ASFEvent) -> None:
|
||||
pub = await get_event_publisher()
|
||||
await pub.publish_event(event)
|
||||
|
||||
|
||||
async def enqueue_job(envelope: JobRequestEnvelope) -> None:
|
||||
pub = await get_event_publisher()
|
||||
await pub.enqueue_job(envelope)
|
||||
Reference in New Issue
Block a user