path: root/kombu/transport/etcd.py

"""Etcd Transport module for Kombu.

It uses Etcd as a store to transport messages in Queues

It uses python-etcd for talking to Etcd's HTTP API

Features
========
* Type: Virtual
* Supports Direct: *Unreviewed*
* Supports Topic: *Unreviewed*
* Supports Fanout: *Unreviewed*
* Supports Priority: *Unreviewed*
* Supports TTL: *Unreviewed*

Connection String
=================

Connection string has the following format:

.. code-block::

    'etcd'://SERVER:PORT

"""

from __future__ import annotations

import os
import socket
from collections import defaultdict
from contextlib import contextmanager
from queue import Empty

from kombu.exceptions import ChannelError
from kombu.log import get_logger
from kombu.utils.json import dumps, loads
from kombu.utils.objects import cached_property

from . import virtual

try:
    import etcd
except ImportError:
    etcd = None

logger = get_logger('kombu.transport.etcd')

DEFAULT_PORT = 2379
DEFAULT_HOST = 'localhost'


class Channel(virtual.Channel):
    """Etcd Channel class which talks to the Etcd."""

    prefix = 'kombu'
    index = None
    timeout = 10
    session_ttl = 30
    lock_ttl = 10

    def __init__(self, *args, **kwargs):
        if etcd is None:
            raise ImportError('Missing python-etcd library')

        super().__init__(*args, **kwargs)

        port = self.connection.client.port or self.connection.default_port
        host = self.connection.client.hostname or DEFAULT_HOST

        logger.debug('Host: %s Port: %s Timeout: %s', host, port, self.timeout)

        self.queues = defaultdict(dict)

        self.client = etcd.Client(host=host, port=int(port))

    def _key_prefix(self, queue):
        """Create and return the `queue` with the proper prefix.

        Arguments:
            queue (str): The name of the queue.
        """
        return f'{self.prefix}/{queue}'

    @contextmanager
    def _queue_lock(self, queue):
        """Try to acquire a lock on the Queue.

        It does so by creating a object called 'lock' which is locked by the
        current session..

        This way other nodes are not able to write to the lock object which
        means that they have to wait before the lock is released.

        Arguments:
            queue (str): The name of the queue.
        """
        lock = etcd.Lock(self.client, queue)
        lock._uuid = self.lock_value
        logger.debug(f'Acquiring lock {lock.name}')
        lock.acquire(blocking=True, lock_ttl=self.lock_ttl)
        try:
            yield
        finally:
            logger.debug(f'Releasing lock {lock.name}')
            lock.release()

    def _new_queue(self, queue, **_):
        """Create a new `queue` if the `queue` doesn't already exist.

        Arguments:
            queue (str): The name of the queue.
        """
        self.queues[queue] = queue
        with self._queue_lock(queue):
            try:
                return self.client.write(
                    key=self._key_prefix(queue), dir=True, value=None)
            except etcd.EtcdNotFile:
                logger.debug(f'Queue "{queue}" already exists')
                return self.client.read(key=self._key_prefix(queue))

    def _has_queue(self, queue, **kwargs):
        """Verify that queue exists.

        Returns:
            bool: Should return :const:`True` if the queue exists
                or :const:`False` otherwise.
        """
        try:
            self.client.read(self._key_prefix(queue))
            return True
        except etcd.EtcdKeyNotFound:
            return False

    def _delete(self, queue, *args, **_):
        """Delete a `queue`.

        Arguments:
            queue (str): The name of the queue.
        """
        self.queues.pop(queue, None)
        self._purge(queue)

    def _put(self, queue, payload, **_):
        """Put `message` onto `queue`.

        This simply writes a key to the Etcd store

        Arguments:
            queue (str): The name of the queue.
            payload (dict): Message data which will be dumped to etcd.
        """
        with self._queue_lock(queue):
            key = self._key_prefix(queue)
            if not self.client.write(
                    key=key,
                    value=dumps(payload),
                    append=True):
                raise ChannelError(f'Cannot add key {key!r} to etcd')

    def _get(self, queue, timeout=None):
        """Get the first available message from the queue.

        Before it does so it acquires a lock on the store so
        only one node reads at the same time. This is for read consistency

        Arguments:
            queue (str): The name of the queue.
            timeout (int): Optional seconds to wait for a response.
        """
        with self._queue_lock(queue):
            key = self._key_prefix(queue)
            logger.debug('Fetching key %s with index %s', key, self.index)

            try:
                result = self.client.read(
                    key=key, recursive=True,
                    index=self.index, timeout=self.timeout)

                if result is None:
                    raise Empty()

                item = result._children[-1]
                logger.debug('Removing key {}'.format(item['key']))

                msg_content = loads(item['value'])
                self.client.delete(key=item['key'])
                return msg_content
            except (TypeError, IndexError, etcd.EtcdException) as error:
                logger.debug(f'_get failed: {type(error)}:{error}')

            raise Empty()

    def _purge(self, queue):
        """Remove all `message`s from a `queue`.

        Arguments:
            queue (str): The name of the queue.
        """
        with self._queue_lock(queue):
            key = self._key_prefix(queue)
            logger.debug(f'Purging queue at key {key}')
            return self.client.delete(key=key, recursive=True)

    def _size(self, queue):
        """Return the size of the `queue`.

        Arguments:
            queue (str): The name of the queue.
        """
        with self._queue_lock(queue):
            size = 0
            try:
                key = self._key_prefix(queue)
                logger.debug('Fetching key recursively %s with index %s',
                             key, self.index)
                result = self.client.read(
                    key=key, recursive=True,
                    index=self.index)
                size = len(result._children)
            except TypeError:
                pass

            logger.debug('Found %s keys under %s with index %s',
                         size, key, self.index)
            return size

    @cached_property
    def lock_value(self):
        return f'{socket.gethostname()}.{os.getpid()}'


class Transport(virtual.Transport):
    """Etcd storage Transport for Kombu."""

    Channel = Channel

    default_port = DEFAULT_PORT
    driver_type = 'etcd'
    driver_name = 'python-etcd'
    polling_interval = 3

    implements = virtual.Transport.implements.extend(
        exchange_type=frozenset(['direct']))

    if etcd:
        connection_errors = (
            virtual.Transport.connection_errors + (etcd.EtcdException, )
        )

        channel_errors = (
            virtual.Transport.channel_errors + (etcd.EtcdException, )
        )

    def __init__(self, *args, **kwargs):
        """Create a new instance of etcd.Transport."""
        if etcd is None:
            raise ImportError('Missing python-etcd library')

        super().__init__(*args, **kwargs)

    def verify_connection(self, connection):
        """Verify the connection works."""
        port = connection.client.port or self.default_port
        host = connection.client.hostname or DEFAULT_HOST

        logger.debug('Verify Etcd connection to %s:%s', host, port)

        try:
            etcd.Client(host=host, port=int(port))
            return True
        except ValueError:
            pass

        return False

    def driver_version(self):
        """Return the version of the etcd library.

        .. note::
           python-etcd has no __version__. This is a workaround.
        """
        try:
            import pip.commands.freeze
            for x in pip.commands.freeze.freeze():
                if x.startswith('python-etcd'):
                    return x.split('==')[1]
        except (ImportError, IndexError):
            logger.warning('Unable to find the python-etcd version.')
            return 'Unknown'