<gh_stars>10-100
"""Clients manage the main user facing API, and provide functions for sending
events and querying the Riemann server. UDP, TCP and TLS transports are
provided by the :py:mod:`riemann_client.transport` module, and the protocol
buffer objects are provided by the :py:mod:`riemann_client.riemann_pb2` module.
"""

from __future__ import absolute_import

import logging
try:
    from logging import NullHandler
except ImportError:
    # Create a NullHandler class in logging for python 2.6
    class NullHandler(logging.Handler):
        def emit(self, record):
            pass

import socket
try:
    from threading import RLock
    from threading import Timer
except ImportError:
    RLock = None
    Timer = None
import time

import riemann_client.riemann_pb2
import riemann_client.transport

logger = logging.getLogger(__name__)
logger.addHandler(NullHandler())


class Client(object):
    """A client for sending events and querying a Riemann server.

    Two sets of methods are provided - an API dealing directly with protocol
    buffer objects and an extended API that takes and returns dictionaries
    representing events.

    Protocol buffer API:

        - :py:meth:`.send_event`
        - :py:meth:`.send_events`
        - :py:meth:`.send_query`

    Extended API:

        - :py:meth:`.event`
        - :py:meth:`.events`
        - :py:meth:`.query`

    Clients do not directly manage connections to a Riemann server - these are
    managed by :py:class:`riemann_client.transport.Transport` instances, which
    provide methods to read and write messages to the server. Client instances
    can be used as a context manager, and will connect and disconnect the
    transport when entering and exiting the context.

        >>> with Client(transport) as client:
        ...     # Calls transport.connect()
        ...     client.query('true')
        ...     # Calls transport.disconnect()
    """

    def __init__(self, transport=None):
        if transport is None:
            transport = riemann_client.transport.TCPTransport()
        self.transport = transport

    def __enter__(self):
        self.transport.connect()
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self.transport.disconnect()

    @staticmethod
    def create_event(data):
        """Translates a dictionary of event attributes to an Event object

        :param dict data: The attributes to be set on the event
        :returns: A protocol buffer ``Event`` object
        """
        event = riemann_client.riemann_pb2.Event()
        event.host = socket.gethostname()
        event.tags.extend(data.pop('tags', []))

        for key, value in data.pop('attributes', {}).items():
            attribute = event.attributes.add()
            attribute.key, attribute.value = key, value

        for name, value in data.items():
            if value is not None:
                setattr(event, name, value)
        return event

    def send_events(self, events):
        """Sends multiple events to Riemann in a single message

        :param events: A list or iterable of ``Event`` objects
        :returns: The response message from Riemann
        """
        message = riemann_client.riemann_pb2.Msg()
        for event in events:
            message.events.add().MergeFrom(event)
        return self.transport.send(message)

    def send_event(self, event):
        """Sends a single event to Riemann

        :param event: An ``Event`` protocol buffer object
        :returns: The response message from Riemann
        """
        return self.send_events((event,))

    def events(self, *events):
        """Sends multiple events in a single message

        >>> client.events({'service': 'riemann-client', 'state': 'awesome'})

         :param \*events: event dictionaries for :py:func:`create_event`
         :returns: The response message from Riemann
        """
        return self.send_events(self.create_event(e) for e in events)

    def event(self, **data):
        """Sends an event, using keyword arguments to create an Event

        >>> client.event(service='riemann-client', state='awesome')

        :param \*\*data: keyword arguments used for :py:func:`create_event`
        :returns: The response message from Riemann
        """
        return self.send_event(self.create_event(data))

    @staticmethod
    def create_dict(event):
        """Translates an Event object to a dictionary of event attributes

        All attributes are included, so ``create_dict(create_event(input))``
        may return more attributes than were present in the input.

        :param event: A protocol buffer ``Event`` object
        :returns: A dictionary of event attributes
        """

        data = dict()

        for descriptor, value in event.ListFields():
            if descriptor.name == 'tags':
                value = list(value)
            elif descriptor.name == 'attributes':
                value = dict(((a.key, a.value) for a in value))
            data[descriptor.name] = value

        return data

    def send_query(self, query):
        """Sends a query to the Riemann server

        :returns: The response message from Riemann
        """
        message = riemann_client.riemann_pb2.Msg()
        message.query.string = query
        return self.transport.send(message)

    def query(self, query):
        """Sends a query to the Riemann server

        >>> client.query('true')

        :returns: A list of event dictionaries taken from the response
        :raises Exception: if used with a :py:class:`.UDPTransport`
        """
        if isinstance(self.transport, riemann_client.transport.UDPTransport):
            raise Exception('Cannot query the Riemann server over UDP')
        response = self.send_query(query)
        return [self.create_dict(e) for e in response.events]


class QueuedClient(Client):
    """A Riemann client using a queue that can be used to batch send events.

    A message object is used as a queue, with the :py:meth:`.send_event` and
    :py:meth:`.send_events` methods adding new events to the message and the
    :py:meth:`.flush` sending the message.
    """

    def __init__(self, transport=None):
        super(QueuedClient, self).__init__(transport)
        self.clear_queue()

    def flush(self):
        """Sends the waiting message to Riemann

        :returns: The response message from Riemann
        """
        response = self.transport.send(self.queue)
        self.clear_queue()
        return response

    def send_event(self, event):
        """Adds a single event to the queued message

        :returns: None - nothing has been sent to the Riemann server yet
        """
        self.send_events((event,))
        return None

    def send_events(self, events):
        """Adds multiple events to the queued message

        :returns: None - nothing has been sent to the Riemann server yet
        """
        for event in events:
            self.queue.events.add().MergeFrom(event)
        return None

    def clear_queue(self):
        """Resets the message/queue to a blank :py:class:`.Msg` object"""
        self.queue = riemann_client.riemann_pb2.Msg()


if RLock and Timer:  # noqa
    class AutoFlushingQueuedClient(QueuedClient):
        """A Riemann client using a queue and a timer that will automatically
        flush its contents if either:
            - the queue size exceeds :param max_batch_size: or
            - more than :param max_delay: has elapsed since the last flush and
             the queue is non-empty.

        if :param stay_connected: is False, then the transport will be
        disconnected after each flush and reconnected at the beginning of
        the next flush.
        if :param clear_on_fail: is True, then the client will discard its
        buffer after the second retry in the event of a socket error.

        A message object is used as a queue, and the following methods are
        given:
            - :py:meth:`.send_event` - add a new event to the queue
            - :py:meth:`.send_events` add a tuple of new events to the queue
            - :py:meth:`.event` - add a new event to the queue from
               keyword arguments
            - :py:meth:`.events` - add new events to the queue from
               dictionaries
            - :py:meth:`.flush` - manually force flush the queue to the
               transport
        """

        def __init__(self, transport, max_delay=0.5, max_batch_size=100,
                     stay_connected=False, clear_on_fail=False):
            super(AutoFlushingQueuedClient, self).__init__(transport)
            self.stay_connected = stay_connected
            self.clear_on_fail = clear_on_fail
            self.max_delay = max_delay
            self.max_batch_size = max_batch_size
            self.lock = RLock()
            self.event_counter = 0
            self.last_flush = time.time()
            self.timer = None

            # start the timer
            self.start_timer()

        def connect(self):
            """Connect the transport if it is not already connected."""
            if not self.is_connected():
                self.transport.connect()

        def is_connected(self):
            """Check whether the transport is connected."""
            try:
                # this will throw an exception whenever socket isn't connected
                self.transport.socket.type
                return True
            except (AttributeError, RuntimeError, socket.error):
                return False

        def event(self, **data):
            """Enqueues an event, using keyword arguments to create an Event

            >>> client.event(service='riemann-client', state='awesome')

            :param \*\*data: keyword arguments used for :py:func:`create_event`
            """
            self.send_events((self.create_event(data),))

        def events(self, *events):
            """Enqueues multiple events in a single message

            >>> client.events({'service': 'riemann-client',
            >>>                'state': 'awesome'})

             :param \*events: event dictionaries for :py:func:`create_event`
             :returns: The response message from Riemann
            """
            self.send_events(self.create_event(evd) for evd in events)

        def send_events(self, events):
            """Enqueues multiple events

            :param events: A list or iterable of ``Event`` objects
            :returns: The response message from Riemann
            """
            with self.lock:
                for event in events:
                    self.queue.events.add().MergeFrom(event)
                    self.event_counter += 1
                    self.check_for_flush()

        def flush(self):
            """Sends the events in the queue to Riemann in a single protobuf
            message

            :returns: The response message from Riemann
            """
            response = None
            with self.lock:
                if not self.is_connected():
                    self.connect()
                try:
                    response = super(AutoFlushingQueuedClient, self).flush()
                except socket.error:
                    # log and retry
                    logger.warn("Socket error on flushing. "
                                "Attempting reconnect and retry...")
                    try:
                        self.transport.disconnect()
                        self.connect()
                        response = (
                            super(AutoFlushingQueuedClient, self).flush())
                    except:
                        logger.warn("Socket error on flushing "
                                    "second attempt. Batch discarded.")
                        self.transport.disconnect()
                        if self.clear_on_fail:
                            self.clear_queue()
                self.event_counter = 0
                if not self.stay_connected:
                    self.transport.disconnect()
                self.last_flush = time.time()
            self.start_timer()
            return response

        def check_for_flush(self):
            """Checks the conditions for flushing the queue"""
            if (self.event_counter >= self.max_batch_size or
                    (time.time() - self.last_flush) >= self.max_delay):
                self.flush()

        def start_timer(self):
            """Cycle the timer responsible for periodically flushing the queue
            """
            if self.timer:
                self.timer.cancel()
            self.timer = Timer(self.max_delay, self.check_for_flush)
            self.timer.daemon = True
            self.timer.start()

        def stop_timer(self):
            """Stops the current timer

            a :py:meth:`.flush` event will reactviate the timer
            """
            self.timer.cancel()

"""Clients manage the main user facing API, and provide functions for sending
events and querying the Riemann server. UDP, TCP and TLS transports are
provided by the :py:mod:`riemann_client.transport` module, and the protocol
buffer objects are provided by the :py:mod:`riemann_client.riemann_pb2` module.
"""

from __future__ import absolute_import

import logging
try:
    from logging import NullHandler
except ImportError:
    # Create a NullHandler class in logging for python 2.6
    class NullHandler(logging.Handler):
        def emit(self, record):
            pass

import socket
try:
    from threading import RLock
    from threading import Timer
except ImportError:
    RLock = None
    Timer = None
import time

import riemann_client.riemann_pb2
import riemann_client.transport

logger = logging.getLogger(__name__)
logger.addHandler(NullHandler())


class Client(object):
    """A client for sending events and querying a Riemann server.

    Two sets of methods are provided - an API dealing directly with protocol
    buffer objects and an extended API that takes and returns dictionaries
    representing events.

    Protocol buffer API:

        - :py:meth:`.send_event`
        - :py:meth:`.send_events`
        - :py:meth:`.send_query`

    Extended API:

        - :py:meth:`.event`
        - :py:meth:`.events`
        - :py:meth:`.query`

    Clients do not directly manage connections to a Riemann server - these are
    managed by :py:class:`riemann_client.transport.Transport` instances, which
    provide methods to read and write messages to the server. Client instances
    can be used as a context manager, and will connect and disconnect the
    transport when entering and exiting the context.

        >>> with Client(transport) as client:
        ...     # Calls transport.connect()
        ...     client.query('true')
        ...     # Calls transport.disconnect()
    """

    def __init__(self, transport=None):
        if transport is None:
            transport = riemann_client.transport.TCPTransport()
        self.transport = transport

    def __enter__(self):
        self.transport.connect()
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self.transport.disconnect()

    @staticmethod
    def create_event(data):
        """Translates a dictionary of event attributes to an Event object

        :param dict data: The attributes to be set on the event
        :returns: A protocol buffer ``Event`` object
        """
        event = riemann_client.riemann_pb2.Event()
        event.host = socket.gethostname()
        event.tags.extend(data.pop('tags', []))

        for key, value in data.pop('attributes', {}).items():
            attribute = event.attributes.add()
            attribute.key, attribute.value = key, value

        for name, value in data.items():
            if value is not None:
                setattr(event, name, value)
        return event

    def send_events(self, events):
        """Sends multiple events to Riemann in a single message

        :param events: A list or iterable of ``Event`` objects
        :returns: The response message from Riemann
        """
        message = riemann_client.riemann_pb2.Msg()
        for event in events:
            message.events.add().MergeFrom(event)
        return self.transport.send(message)

    def send_event(self, event):
        """Sends a single event to Riemann

        :param event: An ``Event`` protocol buffer object
        :returns: The response message from Riemann
        """
        return self.send_events((event,))

    def events(self, *events):
        """Sends multiple events in a single message

        >>> client.events({'service': 'riemann-client', 'state': 'awesome'})

         :param \*events: event dictionaries for :py:func:`create_event`
         :returns: The response message from Riemann
        """
        return self.send_events(self.create_event(e) for e in events)

    def event(self, **data):
        """Sends an event, using keyword arguments to create an Event

        >>> client.event(service='riemann-client', state='awesome')

        :param \*\*data: keyword arguments used for :py:func:`create_event`
        :returns: The response message from Riemann
        """
        return self.send_event(self.create_event(data))

    @staticmethod
    def create_dict(event):
        """Translates an Event object to a dictionary of event attributes

        All attributes are included, so ``create_dict(create_event(input))``
        may return more attributes than were present in the input.

        :param event: A protocol buffer ``Event`` object
        :returns: A dictionary of event attributes
        """

        data = dict()

        for descriptor, value in event.ListFields():
            if descriptor.name == 'tags':
                value = list(value)
            elif descriptor.name == 'attributes':
                value = dict(((a.key, a.value) for a in value))
            data[descriptor.name] = value

        return data

    def send_query(self, query):
        """Sends a query to the Riemann server

        :returns: The response message from Riemann
        """
        message = riemann_client.riemann_pb2.Msg()
        message.query.string = query
        return self.transport.send(message)

    def query(self, query):
        """Sends a query to the Riemann server

        >>> client.query('true')

        :returns: A list of event dictionaries taken from the response
        :raises Exception: if used with a :py:class:`.UDPTransport`
        """
        if isinstance(self.transport, riemann_client.transport.UDPTransport):
            raise Exception('Cannot query the Riemann server over UDP')
        response = self.send_query(query)
        return [self.create_dict(e) for e in response.events]


class QueuedClient(Client):
    """A Riemann client using a queue that can be used to batch send events.

    A message object is used as a queue, with the :py:meth:`.send_event` and
    :py:meth:`.send_events` methods adding new events to the message and the
    :py:meth:`.flush` sending the message.
    """

    def __init__(self, transport=None):
        super(QueuedClient, self).__init__(transport)
        self.clear_queue()

    def flush(self):
        """Sends the waiting message to Riemann

        :returns: The response message from Riemann
        """
        response = self.transport.send(self.queue)
        self.clear_queue()
        return response

    def send_event(self, event):
        """Adds a single event to the queued message

        :returns: None - nothing has been sent to the Riemann server yet
        """
        self.send_events((event,))
        return None

    def send_events(self, events):
        """Adds multiple events to the queued message

        :returns: None - nothing has been sent to the Riemann server yet
        """
        for event in events:
            self.queue.events.add().MergeFrom(event)
        return None

    def clear_queue(self):
        """Resets the message/queue to a blank :py:class:`.Msg` object"""
        self.queue = riemann_client.riemann_pb2.Msg()


if RLock and Timer:  # noqa
    class AutoFlushingQueuedClient(QueuedClient):
        """A Riemann client using a queue and a timer that will automatically
        flush its contents if either:
            - the queue size exceeds :param max_batch_size: or
            - more than :param max_delay: has elapsed since the last flush and
             the queue is non-empty.

        if :param stay_connected: is False, then the transport will be
        disconnected after each flush and reconnected at the beginning of
        the next flush.
        if :param clear_on_fail: is True, then the client will discard its
        buffer after the second retry in the event of a socket error.

        A message object is used as a queue, and the following methods are
        given:
            - :py:meth:`.send_event` - add a new event to the queue
            - :py:meth:`.send_events` add a tuple of new events to the queue
            - :py:meth:`.event` - add a new event to the queue from
               keyword arguments
            - :py:meth:`.events` - add new events to the queue from
               dictionaries
            - :py:meth:`.flush` - manually force flush the queue to the
               transport
        """

        def __init__(self, transport, max_delay=0.5, max_batch_size=100,
                     stay_connected=False, clear_on_fail=False):
            super(AutoFlushingQueuedClient, self).__init__(transport)
            self.stay_connected = stay_connected
            self.clear_on_fail = clear_on_fail
            self.max_delay = max_delay
            self.max_batch_size = max_batch_size
            self.lock = RLock()
            self.event_counter = 0
            self.last_flush = time.time()
            self.timer = None

            # start the timer
            self.start_timer()

        def connect(self):
            """Connect the transport if it is not already connected."""
            if not self.is_connected():
                self.transport.connect()

        def is_connected(self):
            """Check whether the transport is connected."""
            try:
                # this will throw an exception whenever socket isn't connected
                self.transport.socket.type
                return True
            except (AttributeError, RuntimeError, socket.error):
                return False

        def event(self, **data):
            """Enqueues an event, using keyword arguments to create an Event

            >>> client.event(service='riemann-client', state='awesome')

            :param \*\*data: keyword arguments used for :py:func:`create_event`
            """
            self.send_events((self.create_event(data),))

        def events(self, *events):
            """Enqueues multiple events in a single message

            >>> client.events({'service': 'riemann-client',
            >>>                'state': 'awesome'})

             :param \*events: event dictionaries for :py:func:`create_event`
             :returns: The response message from Riemann
            """
            self.send_events(self.create_event(evd) for evd in events)

        def send_events(self, events):
            """Enqueues multiple events

            :param events: A list or iterable of ``Event`` objects
            :returns: The response message from Riemann
            """
            with self.lock:
                for event in events:
                    self.queue.events.add().MergeFrom(event)
                    self.event_counter += 1
                    self.check_for_flush()

        def flush(self):
            """Sends the events in the queue to Riemann in a single protobuf
            message

            :returns: The response message from Riemann
            """
            response = None
            with self.lock:
                if not self.is_connected():
                    self.connect()
                try:
                    response = super(AutoFlushingQueuedClient, self).flush()
                except socket.error:
                    # log and retry
                    logger.warn("Socket error on flushing. "
                                "Attempting reconnect and retry...")
                    try:
                        self.transport.disconnect()
                        self.connect()
                        response = (
                            super(AutoFlushingQueuedClient, self).flush())
                    except:
                        logger.warn("Socket error on flushing "
                                    "second attempt. Batch discarded.")
                        self.transport.disconnect()
                        if self.clear_on_fail:
                            self.clear_queue()
                self.event_counter = 0
                if not self.stay_connected:
                    self.transport.disconnect()
                self.last_flush = time.time()
            self.start_timer()
            return response

        def check_for_flush(self):
            """Checks the conditions for flushing the queue"""
            if (self.event_counter >= self.max_batch_size or
                    (time.time() - self.last_flush) >= self.max_delay):
                self.flush()

        def start_timer(self):
            """Cycle the timer responsible for periodically flushing the queue
            """
            if self.timer:
                self.timer.cancel()
            self.timer = Timer(self.max_delay, self.check_for_flush)
            self.timer.daemon = True
            self.timer.start()

        def stop_timer(self):
            """Stops the current timer

            a :py:meth:`.flush` event will reactviate the timer
            """
            self.timer.cancel()
Clients manage the main user facing API, and provide functions for sending events and querying the Riemann server. UDP, TCP and TLS transports are provided by the :py:mod:`riemann_client.transport` module, and the protocol buffer objects are provided by the :py:mod:`riemann_client.riemann_pb2` module. # Create a NullHandler class in logging for python 2.6 A client for sending events and querying a Riemann server.      Two sets of methods are provided - an API dealing directly with protocol     buffer objects and an extended API that takes and returns dictionaries     representing events.      Protocol buffer API:          - :py:meth:`.send_event`         - :py:meth:`.send_events`         - :py:meth:`.send_query`      Extended API:          - :py:meth:`.event`         - :py:meth:`.events`         - :py:meth:`.query`      Clients do not directly manage connections to a Riemann server - these are     managed by :py:class:`riemann_client.transport.Transport` instances, which     provide methods to read and write messages to the server. Client instances     can be used as a context manager, and will connect and disconnect the     transport when entering and exiting the context.          >>> with Client(transport) as client:         ...     # Calls transport.connect()         ...     client.query('true')         ...     # Calls transport.disconnect() Translates a dictionary of event attributes to an Event object          :param dict data: The attributes to be set on the event         :returns: A protocol buffer ``Event`` object Sends multiple events to Riemann in a single message          :param events: A list or iterable of ``Event`` objects         :returns: The response message from Riemann Sends a single event to Riemann          :param event: An ``Event`` protocol buffer object         :returns: The response message from Riemann Sends multiple events in a single message          >>> client.events({'service': 'riemann-client', 'state': 'awesome'})           :param \*events: event dictionaries for :py:func:`create_event`          :returns: The response message from Riemann Sends an event, using keyword arguments to create an Event          >>> client.event(service='riemann-client', state='awesome')          :param \*\*data: keyword arguments used for :py:func:`create_event`         :returns: The response message from Riemann Translates an Event object to a dictionary of event attributes          All attributes are included, so ``create_dict(create_event(input))``         may return more attributes than were present in the input.          :param event: A protocol buffer ``Event`` object         :returns: A dictionary of event attributes Sends a query to the Riemann server          :returns: The response message from Riemann Sends a query to the Riemann server          >>> client.query('true')          :returns: A list of event dictionaries taken from the response         :raises Exception: if used with a :py:class:`.UDPTransport` A Riemann client using a queue that can be used to batch send events.      A message object is used as a queue, with the :py:meth:`.send_event` and     :py:meth:`.send_events` methods adding new events to the message and the     :py:meth:`.flush` sending the message. Sends the waiting message to Riemann          :returns: The response message from Riemann Adds a single event to the queued message          :returns: None - nothing has been sent to the Riemann server yet Adds multiple events to the queued message          :returns: None - nothing has been sent to the Riemann server yet Resets the message/queue to a blank :py:class:`.Msg` object # noqa A Riemann client using a queue and a timer that will automatically         flush its contents if either:             - the queue size exceeds :param max_batch_size: or             - more than :param max_delay: has elapsed since the last flush and              the queue is non-empty.          if :param stay_connected: is False, then the transport will be         disconnected after each flush and reconnected at the beginning of         the next flush.         if :param clear_on_fail: is True, then the client will discard its         buffer after the second retry in the event of a socket error.          A message object is used as a queue, and the following methods are         given:             - :py:meth:`.send_event` - add a new event to the queue             - :py:meth:`.send_events` add a tuple of new events to the queue             - :py:meth:`.event` - add a new event to the queue from                keyword arguments             - :py:meth:`.events` - add new events to the queue from                dictionaries             - :py:meth:`.flush` - manually force flush the queue to the                transport # start the timer Connect the transport if it is not already connected. Check whether the transport is connected. # this will throw an exception whenever socket isn't connected Enqueues an event, using keyword arguments to create an Event              >>> client.event(service='riemann-client', state='awesome')              :param \*\*data: keyword arguments used for :py:func:`create_event` Enqueues multiple events in a single message              >>> client.events({'service': 'riemann-client',             >>>                'state': 'awesome'})               :param \*events: event dictionaries for :py:func:`create_event`              :returns: The response message from Riemann Enqueues multiple events              :param events: A list or iterable of ``Event`` objects             :returns: The response message from Riemann Sends the events in the queue to Riemann in a single protobuf             message              :returns: The response message from Riemann # log and retry Checks the conditions for flushing the queue Cycle the timer responsible for periodically flushing the queue Stops the current timer              a :py:meth:`.flush` event will reactviate the timer