nsot_loader.py

# -*- coding: utf-8 -*-

"""
Loader for Trigger NetDevices using NSoT API.

Right now this loads ALL devices ALL the time, which scales very poorly with
the number of devices and attributes in NSoT.

Note that ``NETDEVICES_SOURCE`` is ignored because the settings from your
``~/.pynsotrc``.

To use this:

1. Ensure that this module is in your ``PYTHONPATH``  and then add it to
``settings.NETDEVICES_LOADERS``, for example::

    NETDEVICES_LOADERS = ('nsot_loader.NsotLoader',)

Other stuff:

- There is little to no error-handling.
- Authentication/credentials defaults to whatever is used by pynsot (e.g.
  (``~/.pynsotrc``)
"""

from __future__ import unicode_literals
import time

try:
    import pynsot
except ImportError:
    PYNSOT_AVAILABLE = False
else:
    PYNSOT_AVAILABLE = True

from trigger.netdevices.loader import BaseLoader
from trigger.exceptions import LoaderFailed
from twisted.python import log


__author__ = 'jathan@gmail.com'
__version__ = '0.4'


# Map NSoT fields to ones that Trigger requires or uses.
TRANSFORM_FIELDS = {
    # 'hostname': 'nodeName',
    'hw_type': 'deviceType',
    'metro': 'site',
    'row': 'coordinate',
    'vendor': 'manufacturer',
}

# Whether to force adminStatus='PRODUCTION'
FORCE_PRODUCTION = True

# Cache timeout in seconds until live results are retrieved.
CACHE_TIMEOUT = 60


def _is_usable():
    """Assert whether this loader can be used."""
    if not PYNSOT_AVAILABLE:
        return False

    # Try to get a client and retrieve sites.
    try:
        api_client = pynsot.client.get_api_client()
        api_client.sites.get()

    # If we error for any reason, this loader no good.
    except:
        return False

    return True


class NsotLoader(BaseLoader):
    """
    Wrapper for loading metadata via NSoT.

    Note that ``NETDEVICES_SOURCE`` is ignored because the settings from your
    ``~/.pynsotrc``.
    """
    is_usable = _is_usable()

    def __init__(self, *args, **kwargs):
        self.cache_last_checked = 0
        self.__dict = {}  # For internal storage of NetDevice objects.
        super(NsotLoader, self).__init__(*args, **kwargs)

    @property
    def _dict(self):
        """Overload NetDevices._dict calls so we can refresh device cache."""
        self.refresh_device_cache()
        return self.__dict

    def get_data(self, url=None, **kwargs):
        """
        Fetch data from the NSoT API.

        Url and kwargs are currently ignored.
        """
        api_client = pynsot.client.get_api_client()
        self.client = api_client.sites(api_client.default_site)

        # Fetch the devices and cache them internally.
        self.refresh_device_cache()

        return self._devices

    def refresh_device_cache(self):
        """Refresh the NSoT device cache."""
        if not self.cache_ok():
            log.msg('Refreshing devices cache from NSoT.')
            self.fetch_devices()

    def fetch_devices(self):
        """Fetch devices from NSoT."""
        log.msg('Retrieving all devices from NSoT.')
        self.raw_devices = self.client.devices.get()
        self._devices = self.transform_devices(self.raw_devices)

        # Swap the dict in place.
        # TODO(jathan): If we encounter any race conditions this will have to
        # be revisited.
        self.__dict = {d.nodeName: d for d in self._devices}

    def transform_devices(self, devices):
        """Call ``self.transform_device()`` on each device in ``devices``."""
        return [self.transform_device(device) for device in devices]

    def transform_device(self, device):
        """Transform ``device`` fields if they are present."""
        device['nodeName'] = device['hostname']
        attributes = device.pop('attributes', {})

        # If this is adminStatus, change the value to something Trigger
        # expects.
        if FORCE_PRODUCTION or attributes.get('monitor') != 'ignored':
            admin_val = 'PRODUCTION'
        else:
            admin_val = 'NON-PRODUCTION'
        device['adminStatus'] = admin_val

        # Fixups
        for key, val in attributes.iteritems():
            # Include mapped keys for Trigger semantics
            mapped_key = TRANSFORM_FIELDS.get(key, None)  # KEY? KEY. KEY!

            # Trigger expects required field values to be uppercase
            if mapped_key is not None:
                device[mapped_key] = val.upper()

            # Trigger also has a baked-in "make" field
            if key == 'model':
                device['make'] = val

            device[key] = val

        from trigger.netdevices import NetDevice
        return NetDevice(data=device)

    def transform_fields(self, devices):
        """Transform the fields if they are present."""
        for device in devices:
            device['nodeName'] = device['hostname']
            attributes = device.pop('attributes')

            # If this is adminStatus, change the value to something Trigger
            # expects.
            if FORCE_PRODUCTION or attributes.get('monitor') != 'ignored':
                admin_val = 'PRODUCTION'
            else:
                admin_val = 'NON-PRODUCTION'
            device['adminStatus'] = admin_val

            # Fixups
            for key, val in attributes.iteritems():
                # Include mapped keys for Trigger semantics
                mapped_key = TRANSFORM_FIELDS.get(key, None)  # KEY? KEY. KEY!

                # Trigger expects required field values to be uppercase
                if mapped_key is not None:
                    device[mapped_key] = val.upper()

                # Trigger also has a baked-in "make" field
                if key == 'model':
                    device['make'] = val

                device[key] = val

        return devices

    def load_data_source(self, url, **kwargs):
        """Load initial data from NSoT."""
        try:
            return self.get_data(url, **kwargs)
        except Exception as err:
            raise LoaderFailed("Tried NSoT; and failed: %r" % (url, err))

    def cache_ok(self):
        """Is the internal NSoT cache ok?"""
        then = self.cache_last_checked
        now = time.time()

        # If cache expired, update timestamp to now.
        if now - then > CACHE_TIMEOUT:
            log.msg('NSoT cache invalid/expired!')
            self.cache_last_checked = now
            return False

        return True

    def find(self, hostname):
        """
        Find a device by hostname.

        :param hostname:
            Device hostname
        """
        log.msg('Retrieving device: %s' % hostname)

        if self.cache_ok() and hostname in self._dict:
            log.msg('Device %s found in cache.' % hostname )
            return self._dict[hostname]
        else:
            log.msg('Device %s NOT found in cache.' % hostname )

        try:
            log.msg('Fetching device from NSoT: %s' % hostname)
            device = self.client.devices(hostname).get()
        except:
            raise KeyError(hostname)
        else:
            device = self.transform_device(device)
            self._dict[device.nodeName] = device
            return device

    def all(self):
        """Return all devices from NSoT."""
        log.msg('Retrieving all devices from NSoT cache.')
        return self.get_data()

    def match(self, **kwargs):
        """
        Perform a set query on devices.

        If ``skip_loader=True`` is provided, any kwargs will be passed onto
        default `~NetDevices.match()`

        Otherwise, ``query`` argument must be provided with an NSoT set query
        string as the argument.
        """
        try:
            query = kwargs['query']
        except KeyError:
            raise NameError('query argument is required.')

        log.msg('Retrieving devices from NSoT by set query: %r' % query)
        devices = self.client.devices.query.get(query=query)
        return self.transform_devices(devices)