celery/celery/backends/cassandra.py

# -* coding: utf-8 -*-
"""
    celery.backends.cassandra
    ~~~~~~~~~~~~~~~~~~~~~~~~~

    Apache Cassandra result store backend using DataStax driver

"""
from __future__ import absolute_import

import sys
try:  # pragma: no cover
    import cassandra
    import cassandra.cluster
except ImportError:  # pragma: no cover
    cassandra = None   # noqa

from celery import states
from celery.exceptions import ImproperlyConfigured
from celery.utils.log import get_logger
from .base import BaseBackend

__all__ = ['CassandraBackend']

logger = get_logger(__name__)

E_NO_CASSANDRA = """
You need to install the cassandra-driver library to
use the Cassandra backend. See https://github.com/datastax/python-driver
"""

Q_INSERT_RESULT = """
INSERT INTO {table} (
    task_id, status, result, date_done, traceback, children) VALUES (
        %s, %s, %s, %s, %s, %s) {expires};
"""

Q_SELECT_RESULT = """
SELECT status, result, date_done, traceback, children
FROM {table}
WHERE task_id=%s
LIMIT 1
"""

Q_CREATE_RESULT_TABLE = """
CREATE TABLE {table} (
    task_id text,
    status text,
    result blob,
    date_done timestamp,
    traceback blob,
    children blob,
    PRIMARY KEY ((task_id), date_done)
) WITH CLUSTERING ORDER BY (date_done DESC);
"""

Q_EXPIRES = """
    USING TTL {0}
"""

if sys.version_info[0] == 3:
    def buf_t(x):
        return bytes(x, 'utf8')
else:
    buf_t = buffer  # noqa


class CassandraBackend(BaseBackend):
    """Cassandra backend utilizing DataStax driver

    :raises celery.exceptions.ImproperlyConfigured: if
        module :mod:`cassandra` is not available.

    """

    #: List of Cassandra servers with format: ``hostname``.
    servers = None

    supports_autoexpire = True      # autoexpire supported via entry_ttl

    def __init__(self, servers=None, keyspace=None, table=None, entry_ttl=None,
                 port=9042, **kwargs):
        """Initialize Cassandra backend.

        Raises :class:`celery.exceptions.ImproperlyConfigured` if
        the :setting:`cassandra_servers` setting is not set.

        """
        super(CassandraBackend, self).__init__(**kwargs)

        if not cassandra:
            raise ImproperlyConfigured(E_NO_CASSANDRA)

        conf = self.app.conf
        self.servers = (servers or
                        conf.get('cassandra_servers', None))
        self.port = (port or
                     conf.get('cassandra_port', None))
        self.keyspace = (keyspace or
                         conf.get('cassandra_keyspace', None))
        self.table = (table or
                      conf.get('cassandra_table', None))

        if not self.servers or not self.keyspace or not self.table:
            raise ImproperlyConfigured('Cassandra backend not configured.')

        expires = (entry_ttl or conf.get('cassandra_entry_ttl', None))

        self.cqlexpires = (Q_EXPIRES.format(expires)
                           if expires is not None else '')

        read_cons = conf.get('cassandra_read_consistency') or 'LOCAL_QUORUM'
        write_cons = conf.get('cassandra_write_consistency') or 'LOCAL_QUORUM'

        self.read_consistency = getattr(
            cassandra.ConsistencyLevel, read_cons,
            cassandra.ConsistencyLevel.LOCAL_QUORUM,
        )
        self.write_consistency = getattr(
            cassandra.ConsistencyLevel, write_cons,
            cassandra.ConsistencyLevel.LOCAL_QUORUM,
        )

        self._connection = None
        self._session = None
        self._write_stmt = None
        self._read_stmt = None

    def process_cleanup(self):
        if self._connection is not None:
            self._connection = None
        if self._session is not None:
            self._session.shutdown()
            self._session = None

    def _get_connection(self, write=False):
        """Prepare the connection for action

        :param write: bool - are we a writer?

        """
        if self._connection is None:
            try:
                self._connection = cassandra.cluster.Cluster(self.servers,
                                                             port=self.port)
                self._session = self._connection.connect(self.keyspace)

                # We are forced to do concatenation below, as formatting would
                # blow up on superficial %s that will be processed by Cassandra
                self._write_stmt = cassandra.query.SimpleStatement(
                    Q_INSERT_RESULT.format(
                        table=self.table, expires=self.cqlexpires),
                )
                self._write_stmt.consistency_level = self.write_consistency

                self._read_stmt = cassandra.query.SimpleStatement(
                    Q_SELECT_RESULT.format(table=self.table),
                )
                self._read_stmt.consistency_level = self.read_consistency

                if write:
                    # Only possible writers "workers" are allowed to issue
                    # CREATE TABLE. This is to prevent conflicting situations
                    # where both task-creator and task-executor would issue it
                    # at the same time.

                    # Anyway; if you're doing anything critical, you should
                    # have created this table in advance, in which case
                    # this query will be a no-op (AlreadyExists)
                    self._make_stmt = cassandra.query.SimpleStatement(
                        Q_CREATE_RESULT_TABLE.format(table=self.table),
                    )
                    self._make_stmt.consistency_level = self.write_consistency
                    try:
                        self._session.execute(self._make_stmt)
                    except cassandra.AlreadyExists:
                        pass

            except cassandra.OperationTimedOut:
                # a heavily loaded or gone Cassandra cluster failed to respond.
                # leave this class in a consistent state
                self._connection = None
                if self._session is not None:
                    self._session.shutdown()

                raise   # we did fail after all - reraise

    def _store_result(self, task_id, result, status,
                      traceback=None, request=None, **kwargs):
        """Store return value and status of an executed task."""
        self._get_connection(write=True)

        self._session.execute(self._write_stmt, (
            task_id,
            status,
            buf_t(self.encode(result)),
            self.app.now(),
            buf_t(self.encode(traceback)),
            buf_t(self.encode(self.current_task_children(request)))
        ))

    def _get_task_meta_for(self, task_id):
        """Get task metadata for a task by id."""
        self._get_connection()

        res = self._session.execute(self._read_stmt, (task_id, ))
        if not res:
            return {'status': states.PENDING, 'result': None}

        status, result, date_done, traceback, children = res[0]

        return self.meta_from_decoded({
            'task_id': task_id,
            'status': status,
            'result': self.decode(result),
            'date_done': date_done.strftime('%Y-%m-%dT%H:%M:%SZ'),
            'traceback': self.decode(traceback),
            'children': self.decode(children),
        })

    def __reduce__(self, args=(), kwargs={}):
        kwargs.update(
            dict(servers=self.servers,
                 keyspace=self.keyspace,
                 table=self.table))
        return super(CassandraBackend, self).__reduce__(args, kwargs)