nice_consumer.py

import logging
import typing
from traceback import format_exception
from traceback import format_exception_only

logger = logging.getLogger(__name__)

class NiceConsumerMixin:
    async def receive_json(self, content: typing.Dict, **kwargs):
        """
        Called with decoded JSON content.
        """
        request_id = content.pop("request_id", None)
        action = content.pop("action", None)
        if request_id is None or action is None:
            await self.reply(
                action=action,
                errors=["Invalid message content, missing action or request_id."],
                status=500,
                request_id=request_id,
            )
        else:
            await self.handle_action(action, request_id=request_id, **content)

    async def handle_exception(self, exc: Exception, action: str, request_id):
        """
        Handle any exception that occurs, by sending an appropriate message
        """
        # noinspection PyBroadException
        try:
            await super().handle_exception(exc, action, request_id)
        except Exception:
            logger.exception("Unhandled error in handle_action: %(error)s", {"error": exc})
            if settings.DEBUG:
                tb = format_exception(type(exc), exc, exc.__traceback__)
            else:
                tb = format_exception_only(type(exc), exc)
            await self.reply(
                action=action,
                errors=[tb],
                status=500,
                request_id=request_id,
            )


class NiceConsumer(NiceConsumerMixin, AsyncAPIConsumer):
    pass