LedgerLine/services/audit/consumer.py

import asyncio
import json
import logging

import sqlalchemy as sa
from opentelemetry import propagate, trace
from opentelemetry.instrumentation.utils import suppress_instrumentation
from opentelemetry.trace import SpanKind

from common.db import Session
from common.redis_client import redis_client
from common.settings import settings

logger = logging.getLogger("audit")
tracer = trace.get_tracer(__name__)

GROUP = "audit"
CONSUMER = "audit-1"


async def ensure_group():
    try:
        await redis_client.xgroup_create(settings.stream_name, GROUP, id="$", mkstream=True)
    except Exception as exc:
        if "BUSYGROUP" not in str(exc):
            raise


async def _handle(fields: dict):
    # Store the business fields only, not the propagated trace context.
    trace_fields = propagate.get_global_textmap().fields
    payload = {k: v for k, v in fields.items() if k not in trace_fields}
    async with Session() as session:
        await session.execute(
            sa.text(
                "INSERT INTO audit_events (transaction_id, event_type, payload, recorded_at) "
                "VALUES (:tx_id, 'transaction_processed', CAST(:payload AS jsonb), now())"
            ),
            {"tx_id": fields.get("transaction_id"), "payload": json.dumps(payload)},
        )
        await session.commit()


async def run():
    """Second, independent consumer group on the same stream as `notifier` --
    proves the async fork fans out to genuinely separate consumers, not one
    consumer doing both jobs."""
    await ensure_group()
    logger.info("audit consumer started")
    while True:
        try:
            # The blocking poll would otherwise be a ~2s root trace every loop.
            with suppress_instrumentation():
                resp = await redis_client.xreadgroup(
                    groupname=GROUP,
                    consumername=CONSUMER,
                    streams={settings.stream_name: ">"},
                    count=10,
                    block=2000,
                )
        except asyncio.CancelledError:
            raise
        except Exception:
            logger.exception("audit consumer read failed")
            await asyncio.sleep(1)
            continue

        if not resp:
            continue

        for _, messages in resp:
            for msg_id, fields in messages:
                # Parent is the gateway's PRODUCER span, carried in the message.
                with tracer.start_as_current_span(
                    f"process {settings.stream_name}",
                    context=propagate.extract(fields),
                    kind=SpanKind.CONSUMER,
                    attributes={
                        "messaging.system": "redis",
                        "messaging.destination.name": settings.stream_name,
                        "messaging.operation.type": "process",
                        "messaging.consumer.group.name": GROUP,
                        "messaging.message.id": msg_id,
                    },
                ):
                    try:
                        await _handle(fields)
                    except Exception:
                        logger.exception(f"audit failed to process message {msg_id}")
                    await redis_client.xack(settings.stream_name, GROUP, msg_id)