sprockets.mixins.avro-publisher/sprockets/mixins/avro_publisher/mixins.py

import io
import json
import logging

from sprockets.mixins import amqp
from tornado import gen, httpclient
import fastavro

LOGGER = logging.getLogger(__name__)


class PublishingMixin(amqp.PublishingMixin):
    """The request handler will connect to RabbitMQ on the first request,
    blocking until the connection and channel are established. If RabbitMQ
    closes its connection to the app at any point, a connection attempt will
    be made on the next request.

    This class implements a pattern for the use of a single AMQP connection
    to RabbitMQ.
    """
    DATUM_MIME_TYPE = 'application/vnd.apache.avro.datum'
    DEFAULT_SCHEMA_URI_FORMAT = 'http://localhost/avro/%(name)s.avsc'
    DEFAULT_FETCH_RETRY_DELAY = 0.5

    def initialize(self, *args, **kwargs):
        self._schema_fetch_failed = False

        if not hasattr(self, '_http_client'):
            self._http_client = httpclient.AsyncHTTPClient(force_instance=True)

        self._schema_uri_format = self.application.settings.get(
            'avro_schema_uri_format',
            self.DEFAULT_SCHEMA_URI_FORMAT)

        self._fetch_retry_delay = self.application.settings.get(
            'avro_schema_fetch_retry_delay',
            self.DEFAULT_FETCH_RETRY_DELAY)

        if hasattr(super(PublishingMixin, self), 'initialize'):
            super(PublishingMixin, self).initialize(*args, **kwargs)

    @gen.coroutine
    def avro_amqp_publish(self, exchange, routing_key, message_type,
                          data, properties=None, mandatory=False):
        """Publish a message to RabbitMQ, serializing the payload data as an
        Avro datum and creating the AMQP message properties.

        :param str exchange: The exchange to publish the message to.
        :param str routing_key: The routing key to publish the message with.
        :param str message_type: The message type for the Avro schema.
        :param dict data: The message data to serialize.
        :param dict properties: An optional dict of additional properties
                                to append. Will not override mandatory
                                properties:
                                content_type, type
        :param bool mandatory: Whether to instruct the server to return an
                               unqueueable message
                               http://www.rabbitmq.com/amqp-0-9-1-reference.html#basic.publish.mandatory

        :raises: sprockets.mixins.avro_publisher.SchemaFetchError
        """
        # Set mandatory Avro-related properties
        properties = properties or {}
        properties['content_type'] = self.DATUM_MIME_TYPE
        properties['type'] = message_type

        yield self.amqp_publish(exchange, routing_key, data, properties,
                                mandatory)

    @gen.coroutine
    def amqp_publish(self, exchange, routing_key, body, properties,
                     mandatory=False):
        """Publish a message to RabbitMQ, serializing the payload data as an
        Avro datum if the message is to be sent as such.

        :param str exchange: The exchange to publish the message to.
        :param str routing_key: The routing key to publish the message with.
        :param dict body: The message data to serialize.
        :param dict properties: A dict of additional properties
                                to append. If publishing an Avro message, it
                                must contain the Avro message type at 'type'
                                and have a content type of
                                'application/vnd.apache.avro.datum'
        :param bool mandatory: Whether to instruct the server to return an
                               unqueueable message
                               http://www.rabbitmq.com/amqp-0-9-1-reference.html#basic.publish.mandatory

        :raises: sprockets.mixins.avro_publisher.SchemaFetchError
        """
        if (('content_type' in properties and
             properties['content_type']) == self.DATUM_MIME_TYPE):
            avro_schema = yield self._schema(properties['type'])
            body = self._serialize(avro_schema, body)

        yield super(PublishingMixin, self).amqp_publish(
            exchange,
            routing_key,
            body,
            properties,
            mandatory
        )

    @gen.coroutine
    def _schema(self, message_type):
        """Fetch the Avro schema file from application cache or the remote URI.
        If the request for the schema from the remote URI fails, a
        :exc:`sprockets.mixins.avro_publisher.SchemaFetchError` will be raised.

        :param str message_type: The message type for the Avro schema.

        :rtype: str

        :raises: sprockets.mixins.avro_publisher.SchemaFetchError

        """
        if message_type not in self.application.avro_schemas:
            schema = yield self._fetch_schema(message_type)
            self.application.avro_schemas[message_type] = schema
        raise gen.Return(self.application.avro_schemas[message_type])

    @gen.coroutine
    def _fetch_schema(self, message_type):
        """Fetch the Avro schema for the given message type from a remote
        location, returning the schema JSON string.

        If fetching the schema results in an ``tornado.httpclient.HTTPError``,
        it will retry once then raise a SchemaFetchError if the retry fails.

        :param str message_type: The message type for the Avro schema.

        :rtype: str

        :raises: sprockets.mixins.avro_publisher.SchemaFetchError

        """
        url = self._schema_url(message_type)
        LOGGER.debug('Loading schema for %s from %s', message_type, url)
        try:
            response = yield self._http_client.fetch(url)
        except httpclient.HTTPError as error:
            if self._schema_fetch_failed:
                LOGGER.error('Could not fetch Avro schema for %s (%s)',
                             message_type, error)
                raise SchemaFetchError(str(error))
            else:
                self._schema_fetch_failed = True
                yield gen.sleep(self._fetch_retry_delay)
                yield self._fetch_schema(message_type)

        else:
            self._schema_fetch_failed = False
            raise gen.Return(json.loads(response.body.decode('utf-8')))

    def _schema_url(self, message_type):
        """Return the URL for the given message type for retrieving the Avro
        schema from a remote location.

        :param str message_type: The message type for the Avro schema.

        :rtype: str

        """
        return self._schema_uri_format % {'name': message_type}

    @staticmethod
    def _serialize(schema, data):
        """Serialize a data structure into an Avro datum.

        :param dict schema: The parsed Avro schema.
        :param dict data: The value to turn into an Avro datum.

        :rtype: bytes

        """
        stream = io.BytesIO()
        fastavro.schemaless_writer(stream, schema, data)
        return stream.getvalue()


class SchemaFetchError(ValueError):
    """Raised when the Avro schema could not be fetched."""
    pass